[Numpy-discussion] Chaining apply_over_axis for multiple axes.

Andrew Nelson andyfaff at gmail.com
Mon May 25 07:02:42 EDT 2015


I have a function that operates over a 1D array, to return an array of a
similar size.  To use it in a 2D fashion I would have to do something like
the following:

for row in range(np.size(arr, 0):
    arr_out[row] = func(arr[row])
for col in range(np.size(arr, 1):
    arr_out[:, col] = func(arr[:, col])

I would like to generalise this to N dimensions. Does anyone have any
suggestions of how to achieve this?  Presumably what I need to do is build
an iterator, and then remove an axis:

# arr.shape=(2, 3, 4)
it = np.nditer(arr, flags=['multi_index'])
it.remove_axis(2)
while not it.finished:
    arr_out[it.multi_index] = func(arr[it.multi_index])
    it.iternext()

If I have an array with shape (2, 3, 4) this would allow me to iterate over
the 6 1D arrays that are 4 elements long.  However, how do I then construct
the iterator for the preceding axes?
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20150525/f1e6259b/attachment.html>


More information about the NumPy-Discussion mailing list