pbreach - 1 year ago 92
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.

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.

Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download