pbreach pbreach - 12 days ago 3
Python Question

Splitting an N dimensional numpy array into multiple 1D arrays

I have a simulation model that integrates a set of variables whose states are represented by numpy arrays of an arbitrary number of dimensions. After the simulation, I now have a list of arrays whose elements represent the variable state at a particular point in time.

In order to output the simulation results I want to split these arrays into multiple 1D arrays where the elements correspond to the same component of the state variable through time. Here is an example of a 2D state variable over a number of time steps.

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)


The output I need to get is:

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


I've tried doing
np.split(stacked, sum(stacked.shape[:-1]), axis=...)
(tried everything for
axis=
) but get the following error:
ValueError: array split does not result in an equal division
. Is there a way to do this using
np.split
or maybe
np.nditer
that will work for the general case?

I guess this would be equivalent to doing:

I, J, K = stacked.shape

result = []

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


Which is also the ordering I'm hoping to get. Easy enough, however I'm hoping there is something in numpy that I can take advantage of for this that will be more general.

Answer

If I reshape it to a 9x3 array, then a simple list() will turn it into a list of 3 element arrays:

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) produces a list of 1x3 arrays.

np.split only works on one axis, but you want to split on the 1st 2 - hence the need for a reshape or ravel.

And forget about nditer. That's a stepping stone to reworking code in cython. It does not help with ordinary iteration - except that when used in ndindex it can streamline your i,j double loop:

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])]

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

With the different state, just stack on a different axis

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 is rather like np.array, except it gives more control over where the dimension is added. But do look at it's code.