Discussion:
[Numpy-discussion] Chaining apply_over_axis for multiple axes.
Andrew Nelson
2015-05-25 11:02:42 UTC
Permalink
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?
Sebastian Berg
2015-05-25 15:21:29 UTC
Permalink
Post by Andrew Nelson
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
arr_out[row] = func(arr[row])
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
# arr.shape=(2, 3, 4)
it = np.nditer(arr, flags=['multi_index'])
it.remove_axis(2)
arr_out[it.multi_index] = func(arr[it.multi_index])
it.iternext()
Just warning that nditer is pretty low level (i.e. can be a bit mind
boggling since it is close to the C-side of things).

Anyway, you can of course do this just iterating the result. Since you
have no buffering, etc. this should work fine. There is also
`np.nesterd_iters` but since I am a bit lazy to look it up, you would
have to actually check some examples for it from the numpy tests to see
how it works probably.

- Sebastian
Post by Andrew Nelson
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?
_______________________________________________
NumPy-Discussion mailing list
http://mail.scipy.org/mailman/listinfo/numpy-discussion
Nathaniel Smith
2015-05-25 15:38:26 UTC
Permalink
Post by Andrew Nelson
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
Post by Andrew Nelson
arr_out[row] = func(arr[row])
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

Loading...