NumPy 中的 `ufunc` 组合

Question

My code contains the following definition:

skew_min = np.frompyfunc(lambda x, y: min(x + 1, y), nin=2, nout=1)

The essential point, is that it is ufunc , because I call skew_min.accumulate further. But this definition is really awful. It works slow (because of python-from-native-code calls) and produces an ill-typed result (array of pyobject instead of array of float). How I dream, there should be some FP-flavored combinators to construct ufunc from other ufunc s, something like that:

skew_min = compose_1(np.minimum, subst_1(np.add, 1))

(where made-up functions compose_1 and subst_1 perform ufunc -tional composition and substitution, correspondingly). But I have failed to find anything resembling compose in NumPy docs.

So, is there a nice way to compose two ufunc s in order to build a new one?

Answer 1

You could use numba's @numba.vectorize or @numba.guvectorize ( numba docs ):

>>> import numpy as np
>>> import numba
>>> @numba.vectorize(['f8(f8, f8)', 'i8(i8, i8)'])
... def skew_min(x: float, y: float) -> float:
...     return min(x + 1, y)
...
>>> skew_min(3, 3.5)
3.5
>>> skew_min([2, 3, 4], 3.5)
array([3. , 3.5, 3.5])
>>> skew_min.outer(np.arange(3), [1.62, 2.72, 3.14])
array([[1.  , 1.  , 1.  ],
       [1.62, 2.  , 2.  ],
       [1.62, 2.72, 3.  ]])

As a bonus, numba jit-compiles the function, making it generally faster than using eg numpy.frompyfunc or numpy.vectorize .