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

Question

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

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

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，只不过它可以更好地控制添加维度的位置。但一定要看看它的代码。