Is there a better way than using np.vectorize to use function on ndarray of ndarray of different shape?

Question

The following function apply numpy functions to two numpy arrays.

import numpy as np

def my_func(a: np.ndarray, b: np.ndarray) -> float:
    return np.nanmin(a, axis=0) + np.nanmin(b, axis=0)

>>> my_func(np.array([1., 2., np.nan]), np.array([1., np.nan]))
2.0

However what is the best way to apply this same function to an np.array of np.array of different shape ?

a = np.array([np.array([1., 2]), np.array([1, 2., 3, np.nan])], dtype=object)  # First array shape (2,), second (3,)
b = np.array([np.array([1]), np.array([1.5, 2.5, np.nan])], dtype=object)

np.vectorize does work

>>> np.vectorize(my_func)(a, b)
array([2. , 2.5])

but as specified by the vectorize documentation:

The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop.

Is there a more clever solution ? I could use np.pad to have identifical shape but it seems sub-optimal as it requires to pad up to the maximum length of the inside arrays (here 4 for a and 3 for b).

I looked at numba and this stack exchange about performance but I am not sure of the best pratice for such a case.

Thanks !

Not only is `np.vectorize` not a speed tool, it passes the arrays element by element to function. It's useful for a function that only works with scalar inputs, and you are too lazy to rewrite to work with arrays (or some cases can't). — hpaulj, Oct 19 '21 at 00:28
OK, I see you are passing object dtype arrays, so `vectorize` will pass those array elements. The fast numpy methods ("real" vectorization) work with numeric dtype arrays. Object dtype arrays are processed at least speads. There isn't a "fast" way of using those object arrays. — hpaulj, Oct 19 '21 at 00:31
Indeed here I have different shapes so I can't just rewrite the function. — Theo, Oct 19 '21 at 07:29

hpaulj · Accepted Answer · 2021-10-19T00:55:10.117

Your function and arrays:

In [222]: def my_func(a: np.ndarray, b: np.ndarray) -> float:
     ...:     return np.nanmin(a, axis=0) + np.nanmin(b, axis=0)
     ...: 
In [223]: a = np.array([np.array([1., 2]), np.array([1, 2., 3, np.nan])], dtype=object
     ...: )  # First array shape (2,), second (3,)
     ...: b = np.array([np.array([1]), np.array([1.5, 2.5, np.nan])], dtype=object)
In [224]: a
Out[224]: array([array([1., 2.]), array([ 1.,  2.,  3., nan])], dtype=object)
In [225]: b
Out[225]: array([array([1]), array([1.5, 2.5, nan])], dtype=object)

Compare vectorize with a straightforward list comprehension:

In [226]: np.vectorize(my_func)(a, b)
Out[226]: array([2. , 2.5])
In [227]: [my_func(i,j) for i,j in zip(a,b)]
Out[227]: [2.0, 2.5]

and their times:

In [228]: timeit np.vectorize(my_func)(a, b)
157 µs ± 117 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [229]: timeit [my_func(i,j) for i,j in zip(a,b)]
85.9 µs ± 148 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [230]: timeit np.array([my_func(i,j) for i,j in zip(a,b)])
89.7 µs ± 1.03 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

If you are going to work with object arrays, frompyfunc is faster than vectorize:

In [231]: np.frompyfunc(my_func,2,1)(a, b)
Out[231]: array([2.0, 2.5], dtype=object)
In [232]: timeit np.frompyfunc(my_func,2,1)(a, b)
83.2 µs ± 50.1 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

I'm a bit surprised that it's even better than the list comprehension.

frompyfunc (and vectorize) are more useful when the inputs need to 'broadcast' against each other:

In [233]: np.frompyfunc(my_func,2,1)(a[:,None], b)
Out[233]: 
array([[2.0, 2.5],
       [2.0, 2.5]], dtype=object)

I'm not a numba expert, but I suspect it doesn't handle object dtype arrays, or it it does it doesn't improve speed much. Remember, object dtype means the elements are object references, just like in lists.

I get better times by using otypes and taking the function creation out of the timing loop:

In [235]: %%timeit f=np.vectorize(my_func, otypes=[float])
     ...: f(a, b)
     ...: 
     ...: 
95.5 µs ± 316 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [236]: %%timeit f=np.frompyfunc(my_func,2,1)
     ...: f(a, b)
     ...: 
     ...: 
81.1 µs ± 103 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

If you don't know about otypes, you haven't read the np.vectorize docs well enough.

Thanks @hpaulj for your answer. I though I was maybe missing a subtlety of numpy but apparently not. — Theo, Oct 19 '21 at 07:31

Is there a better way than using np.vectorize to use function on ndarray of ndarray of different shape?

1 Answers1