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

时间:2016-11-23 19:52:01

标签: python arrays numpy

我有一个模拟模型,它集成了一组变量,其状态由任意数量维度的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.splitnp.nditer来执行此操作?

我想这相当于做:

I, J, K = stacked.shape

result = []

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

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

2 个答案:

答案 0 :(得分:2)

如果我将其重塑为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仅适用于一个轴,但您希望在第一个轴上进行分割 - 因此需要重塑或拉平。

忘了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,但它可以更好地控制添加维度的位置。但要看看它的代码。

答案 1 :(得分:1)

您可以在展平版本上使用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])]