将N维numpy数组拆分为多个1D数组

Question

我有一个模拟模型，它集成了一组变量，其状态由任意数量维度的numpy数组表示。 在模拟之后，我现在有一个数组列表，其元素表示特定时间点的变量状态。

为了输出模拟结果，我想将这些数组拆分成多个1D数组，其中元素随时间对应于状态变量的相同组件。 以下是多个时间步长的2D状态变量的示例。

import numpy as np

# Arbitrary state that is constant
arr = np.arange(9).reshape((3, 3))

# State variable through 3 time steps
state = [arr.copy() for _ in range(3)]

# Stack the arrays up to 3d. Axis could be rolled here if it makes it easier.
stacked = np.stack(state)

我需要得到的输出是：

[np.array([0, 0, 0]), np.array([1, 1, 1]), np.array([2, 2, 2]), ...]

我试过做np.split(stacked, sum(stacked.shape[:-1]), axis=...) （尝试了所有的axis= ）但得到以下错误： ValueError: array split does not result in an equal division 。 有没有办法使用np.split或者np.nditer来实现这一点？

我想这相当于做：

I, J, K = stacked.shape

result = []

for i in range(I):
    for j in range(J):
        result.append(stacked[i, j, :])

这也是我希望获得的顺序。 很容易，但我希望有一些numpy的东西，我可以利用这一点，这将是更一般的。

Answer 1

如果我将其重塑为9x3数组，那么一个简单的list()会将其转换为3个元素数组的列表：

In [190]: stacked.reshape(-1,3)
Out[190]: 
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2],
       [3, 3, 3],
       [4, 4, 4],
       [5, 5, 5],
       [6, 6, 6],
       [7, 7, 7],
       [8, 8, 8]])
In [191]: list(stacked.reshape(-1,3))
Out[191]: 
[array([0, 0, 0]),
 array([1, 1, 1]),
 array([2, 2, 2]),
 array([3, 3, 3]),
 array([4, 4, 4]),
 array([5, 5, 5]),
 array([6, 6, 6]),
 array([7, 7, 7]),
 array([8, 8, 8])]

np.split(stacked.reshape(-1,3),9)生成1x3数组的列表。

np.split只能在一个轴上工作，但是你想在第一个2上进行分割 - 因此需要重塑或拉平。

忘了nditer 。 这是在cython中重新编写代码的垫脚石。 它对普通迭代没有帮助 - 除了在ndindex它可以简化你的i,j双循环：

In [196]: [stacked[idx] for idx in np.ndindex(stacked.shape[:2])]
Out[196]: 
[array([0, 0, 0]),
 array([1, 1, 1]),
 array([2, 2, 2]),
 array([3, 3, 3]),
 array([4, 4, 4]),
 array([5, 5, 5]),
 array([6, 6, 6]),
 array([7, 7, 7]),
 array([8, 8, 8])]

======================

在不同的state ，只需堆叠在不同的轴上

In [302]: state
Out[302]: 
[array([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]]), array([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]]), array([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])]
In [303]: np.stack(state,axis=2).reshape(-1,3)
Out[303]: 
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2],
       [3, 3, 3],
       [4, 4, 4],
       [5, 5, 5],
       [6, 6, 6],
       [7, 7, 7],
       [8, 8, 8]])

stack更像是np.array ，除了它可以更好地控制添加维度的位置。 但要看看它的代码。

Answer 2

你可以在扁平版本上使用np.split ，并将其切成适当数量的部分作为1D列表，如下所示 -

np.split(stacked.ravel(),np.prod(stacked.shape[:2]))

样品运行 -

In [406]: stacked
Out[406]: 
array([[[0, 0, 0],
        [1, 1, 1]],

       [[2, 2, 2],
        [3, 3, 3]],

       [[4, 4, 4],
        [5, 5, 5]],

       [[6, 6, 6],
        [7, 7, 7]]])

In [407]: np.split(stacked.ravel(),np.prod(stacked.shape[:2]))
Out[407]: 
[array([0, 0, 0]),
 array([1, 1, 1]),
 array([2, 2, 2]),
 array([3, 3, 3]),
 array([4, 4, 4]),
 array([5, 5, 5]),
 array([6, 6, 6]),
 array([7, 7, 7])]