Разбиение 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 работает только на одной оси, но вы хотите разбить на 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])]