[Numpy-discussion] Chaining apply_over_axis for multiple axes.

[Nathaniel Smith]

On May 25, 2015 4:05 AM, "Andrew Nelson" <andyfaff at gmail.com> wrote:
>
> 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?

The crude but effective way is

tmp_in = arr.reshape((-1, arr.shape[-
1]))
tmp_out = np.empty(tmp_in.shape)
for i in range(tmp_in.shape[0]):
    tmp_out[i, :] = func(tmp_in[i, :])
out = tmp_out.reshape(arr.shape)

This won't produce any unnecessary copies if your input array is contiguous.

This also assumes you want to apply the function on the last axis. If not
you can do something like

arr = arr.swapaxes(axis, -1)
... call the code above ...
out = out.swapaxes(axis, -1)

This will result in an extra copy of the input array though if it's >2d and
the requested axis is not the last one.

-n
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20150525/24c70189/attachment.html>