On May 25, 2015 4:05 AM, "Andrew Nelson" <andyfaff@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