Mapping functions of 2D numpy arrays

Question

I have a function foo that takes a NxM numpy array as an argument and returns a scalar value. I have a AxNxM numpy array data, over which I'd like to map foo to give me a resultant numpy array of length A.

Curently, I'm doing this:

result = numpy.array([foo(x) for x in data])

It works, but it seems like I'm not taking advantage of the numpy magic (and speed). Is there a better way?

I've looked at numpy.vectorize, and numpy.apply_along_axis, but neither works for a function of 2D arrays.

EDIT: I'm doing boosted regression on 24x24 image patches, so my AxNxM is something like 1000x24x24. What I called foo above applies a Haar-like feature to a patch (so, not terribly computationally intensive).

Answer 1

If NxM is big (say, 100), they the cost of iterating over A will be amortized into basically nothing.

Say the array is 1000 X 100 X 100.

Iterating is O(1000), but the cumulative cost of the inside function is O(1000 X 100 X 100) - 10,000 times slower. (Note, my terminology is a bit wonky, but I do know what I'm talking about)

I'm not sure, but you could try this:

result = numpy.empty(data.shape[0])
for i in range(len(data)):
    result[i] = foo(data[i])

You would save a big of memory allocation on building the list ... but the loop overhead would be greater.

Or you could write a parallel version of the loop, and split it across multiple processes. That could be a lot faster, depending on how intensive foo is (as it would have to offset the data handling).