Splitting an N dimensional numpy array into multiple 1D arrays

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([0, 0, 0]), np.array([1, 1, 1]), np.array([2, 2, 2]), ...]

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.

hpaulj · Accepted Answer · 2016-11-24 20:18:22Z

3

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.

edited Nov 24, 2016 at 20:18

answered Nov 23, 2016 at 20:31

hpaulj

233k14 gold badges260 silver badges392 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Divakar Over a year ago

Lovely use of list! Looks pretty fast. Just to make it generic, maybe we could use the shape param there : list(stacked.reshape(-1,stacked.shape[-1])).

pbreach Over a year ago

I messed up on the example input array. Should be stacked = np.stack(state). Updated the question so now your answer should be list(stacked.reshape(3, -1).T) I think? I agree with @Divakar, great use of list!

pbreach Over a year ago

Looks great, I've just confirmed that this works for more than 2 dimensions. The general solution would be, list(np.stack(state, axis=state[0].ndim).reshape(-1, len(state)))

Divakar · Accepted Answer · 2016-11-23 20:02:01Z

2

You could use np.split on a flattened version and cut into appropriate number of parts as 1D lists, like so -

np.split(stacked.ravel(),np.prod(stacked.shape[:2]))

Sample run -

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

answered Nov 23, 2016 at 20:02

Divakar

222k19 gold badges273 silver badges374 bronze badges

2 Comments

pbreach Over a year ago

Ah okay, I was wondering why I couldn't get split to work. I'll give this a shot later tonight

pbreach Over a year ago

See edit - I think maybe split could still work but not sure.

Collectives™ on Stack Overflow

Splitting an N dimensional numpy array into multiple 1D arrays

2 Answers 2

3 Comments

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

3 Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related