pbreach - 6 months ago 55

Python Question

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=...)`

`axis=`

`ValueError: array split does not result in an equal division`

`np.split`

`np.nditer`

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.