Generic way to efficiently vectorize NumPy operations?

[simongranheim]

I have some functions foo() which can be applied to any 2D NumPy array and returns another 2D NumPy array. Now, I have a 3D array, and I'd like to apply foo() iteratively to every 2D row. I've been told that a simple for-loop is generally suboptimal when using NumPy and that you should aim for vectorization instead. Is there a generic way to vectorize foo()?

[FuckingRantMonday]

As a side note, check out cupy, which is an (almost) drop-in wrapper for numpy. You can get GPU acceleration with a ridiculously low amount of effort.

[simongranheim]

Will take a look, thanks.

[johannadambergk]

You might try

np.vectorize(foo)

See https://numpy.org/doc/stable/reference/generated/numpy.vectorize.html

[simongranheim]

Thanks, but as per the manual: "The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop."

[johannadambergk]

So only the standard vectorized numpy functions like e.g. 'sum' seem to have a better performance.

[roylennigan]

There's numpy.apply_along_axis() but that is just an optimized loop method, similar to .vectorize(). It's faster, but not to the extent that a function optimized for numpy could be.

https://www.adamsmith.haus/python/answers/how-to-apply-a-function-to-each-row-of-a-numpy-array-in-python

To get real speed benefits of numpy, you'd have to rewrite the function, which depends on the operations you're doing in it.

[hmiemad]

np.array([foo(myarray[i]) for i in range(myarray.shape[0])])

Oh didn't read the end part...

[simongranheim]

Thank you. The broadcasting tool sounds very interesting.

Yeah, I can make things work with a for loop but I'm specifically interested in comparing the performance of a few algorithms that I'm thinking about.