CODEWITHSUNDEEP
pythonnumpyvectorization
user2124834
user2124834

Reputation:

Python: using np.apply_along_axis with a function returning array

I'm trying to do something like the following

import numpy as np
arr = np.array([0,1,2,3])
fn = lambda x: np.array([x/2,x*2])
res = np.apply_along_axis(fn,0,arr)

That is, I want to use np.apply_along_axis with a function that returns an array rather than a single value. However, this results in the following error.

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3.5/site-packages/numpy/lib/shape_base.py", line 117, in apply_along_axis
    outarr[tuple(i.tolist())] = res
ValueError: could not broadcast input array from shape (2,4) into shape (2)

While I would expect the outcome to be equivalent to:

res = np.array([[0.0,0],[0.5,2],[1.0,4],[1.5,6]])

Note that my actual use case is more complicated than the fn given. I would like to know whether there is a way in general to use np.apply_along_axis with a function that returns an array, or whether there is a different (numpythonic) way to achieve the same thing.

Thanks!

Upvotes: 1

Views: 1038

Answers (2)

John Zwinck
John Zwinck

Reputation: 249163

This sort of thing--applying a regular Python function along an array--is basically the same as a for loop. It's not going to be very fast, even if you use NumPy functions to invoke your function. So just keep it simple:

arr = np.array([0,1,2,3])
fn = lambda x: [x/2,x*2]
res = np.array(list(map(fn, arr)))

Or if you care about performance:

arr = np.array([0,1,2,3])
funcs = [lambda x: x/2.0, lambda x: x*2.0]
res = np.empty((len(arr), len(funcs)))
for i, func in enumerate(funcs):
    res[:,i] = func(arr)

This second way is much faster, so long as you have more rows in the array than the number of functions to apply. This is because you only loop a few times, using NumPy vectorized code to apply each function to the entire input at once, rather than looping over the input rows one by one.

Upvotes: 2

P. Camilleri
P. Camilleri

Reputation: 13218

A workaround I have found is to stack the results, to make your lambda return a 1D array instead of a 2D one:

arr = np.array([0, 1, 2, 3])
fn = lambda x: np.hstack([x/2, x*2])
res = np.apply_along_axis(fn, 0, arr)

and then you can reshape res as you want

Upvotes: 2

Related Questions