Computational Science Stack Exchange is a question and answer site for scientists using computers to solve scientific problems. It only takes a minute to sign up.
Sign up to join this community
I have a numpy function f that takes arrays as arguments and a 3D array x[a,b,c].
I would like to evaluate the function f along a specific column. A long-winded way could be with comprehensions:
y = [ [ f(x[a][b]) for a in range(len(x)) ] for b in range(len(x[0]))]
y = np.array(y)
Is there a numpy way of doing this with broadcasting?
np.apply_along_axis(f, 2, x).T
This is a simpler comprehension:
numpy.array([[f(row) for row in layer] for layer in x]).T
or more concisely,
numpy.array([map(f,layer) for layer in x]).T
Note that I reproduced the transpose from your sample code. You can avoid the nested comprehension with cryptic reshaping:
numpy.reshape(map(f,x.reshape(numpy.prod(x.shape[:2]), x.shape[2])),x.shape[:2]).T
