NumPy中的‘composition’组合

Question

我的代码包含以下定义：

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

关键的一点是，它是ufunc，因为我进一步调用了skew_min.accumulate。但是这个定义真的很糟糕。它工作缓慢(因为从本机代码调用python)并产生一个错误类型的结果( pyobject数组而不是float数组)。我的梦想是，应该有一些FP风格的组合器来从其他的ufunc中构建ufunc，类似于这样：

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

(相应地，组合函数compose_1和subst_1执行ufunc-tional组合和替换)。但我在NumPy文档中没有找到任何类似于NumPy的东西。

那么，为了构建一个新的ufunc，是否有一种很好的方法来组合两个？

Answer 1

您可以使用numba的@numba.vectorize或@numba.guvectorize (numba博士)：

>>> 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.  ]])

另外，numba编译函数，通常比使用numpy.frompyfunc或numpy.vectorize的速度更快。