Splitting an N dimensional numpy array into multiple 1D arrays
<p>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.</p>

<p>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.</p>

<pre><code>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)
</code></pre>

<p>The output I need to get is:</p>

<pre><code>[np.array([0, 0, 0]), np.array([1, 1, 1]), np.array([2, 2, 2]), ...]
</code></pre>

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

<p>I guess this would be equivalent to doing:</p>

<pre><code>I, J, K = stacked.shape

result = []

for i in range(I):
    for j in range(J):
        result.append(stacked[i, j, :])
</code></pre>

<p>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.</p>