CODEWITHSUNDEEP
pythonnumpyvectorization
atbug
atbug

Reputation: 838

how to vectorize a function whose argument is a vector in numpy

I have an numpy array R of dimension (n, n, n, 3) and a function f which takes a 1-D vector to a scalar. I need a new array A whose relationship with R is

A[i, j, k] = f(R[i, j, k, :])

How can I do this in numpy without three for statements.

Upvotes: 2

Views: 1951

Answers (1)

user2357112
user2357112

Reputation: 280335

Ideally, you'd do this by changing the implementation of f to use techniques that handle high-dimensional input appropriately. For example, you might change np.sum(whatever) to np.sum(whatever, axis=-1) to get a sum over the last axis instead of the whole array. This would produce the most efficient results, but it might be difficult or impossible, depending on f.

The slower, much easier answer is np.apply_along_axis:

A = np.apply_along_axis(f, -1, R)

This is prettier than 3 for loops, but it probably won't be any more efficient.

Upvotes: 2

Related Questions