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

Question

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

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

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 只能在一个轴上工作，但您想在第 1 个 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])]