tisdag 24 maj 2011

Rearrange multi-dimensional numpy arrays with rollaxis

numpy.rollaxis() is a simple-to-use function for rearranging / reordering multi-dimensional numpy arrays. Here are a couple of examples with a 3-dimensional array:
>>> import numpy as np
>>> a = np.arange(4 * 3 * 2).reshape(4, 3, 2)
>>> b = np.rollaxis(a, 0, 3) # Move first dimension to last position
# (before 'virtual' dimension 3)
>>> b.shape
(3, 2, 4)
>>> c = np.rollaxis(a, 2, 0) # Move last dimension to first position
# (before current dimension 0)
>>> c.shape
(2, 4, 3)

It's really fast, too. In ipython, using the %timeit command:
In [1]: a = arange(400*200*91).reshape(400, 200, 91)

In [2]: timeit rollaxis(a, 0, 2)
100000 loops, best of 3: 2.87 us per loop

With some guidance from a thread at stackoverflow.

Inga kommentarer:

Skicka en kommentar