使向量化的 numpy 函数表现得像 ufunc

Question

假设我们有一个 Python 函数，它接收 Numpy 数组并返回另一个数组：

import numpy as np

def f(x, y, method='p'):
    """Parameters:  x (np.ndarray) , y (np.ndarray), method (str)
    Returns: np.ndarray"""
    z = x.copy()    
    if method == 'p':
        mask = x < 0
    else:
        mask = x > 0
    z[mask] = 0
    return z*y

虽然实际的实现并不重要。我们可以假设x 和y 总是形状相同的数组，并且输出的形状与x 相同。

问题是包装此类函数的最简单/最优雅的方法是什么，以便它可以与 ND 数组 (N>1) 和标量参数一起使用，其方式有点类似于 universal functions in Numpy。

例如，上述函数的预期输出应该是，

In [1]: f_ufunc(np.arange(-1,2), np.ones(3), method='p') 
Out[1]: array([ 0.,  0.,  1.]) # random array input -> output of the same shape

In [2]: f_ufunc(np.array([1]), np.array([1]), method='p') 
Out[2]: array([1])   # array input of len 1 -> output of len 1

In [3]: f_ufunc(1, 1, method='p')
Out[3]: 1  # scalar input -> scalar output

• 函数f 无法更改，如果为x 或y 提供标量参数，它将失败。

• 当x 和y 是标量时，我们将它们转换为一维数组，进行计算，然后在最后将它们转换回标量。

• f 已针对数组进行了优化，标量输入主要是为了方便。因此，编写一个使用标量的函数然后使用np.vectorize 或np.frompyfunc 是不可接受的。

实现的开始可能是，

def atleast_1d_inverse(res, x):
    # this function fails in some cases (see point 1 below).
    if res.shape[0] == 1:
        return res[0]
    else:
        return res

def ufunc_wrapper(func, args=[]):
    """ func:  the wrapped function
        args:  arguments of func to which we apply np.atleast_1d """

    # this needs to be generated dynamically depending on the definition of func
    def wrapper(x, y, method='p'):
        # we apply np.atleast_1d to the variables given in args
        x = np.atleast_1d(x)
        y = np.atleast_1d(x)

        res = func(x, y, method='p')

        return atleast_1d_inverse(res, x)

    return wrapper

f_ufunc = ufunc_wrapper(f, args=['x', 'y'])

这主要是有效的，但会在上面的测试 2 中失败，产生一个标量输出而不是一个向量输出。如果我们想解决这个问题，我们需要对输入类型添加更多测试（例如isinstance(x, np.ndarray)、x.ndim>0 等），但我害怕忘记那里的一些极端情况。此外，上述实现还不够通用，无法用不同数量的参数包装函数（请参见下面的第 2 点）。

在使用 Cython / f2py 函数时，这似乎是一个相当普遍的问题，我想知道在某个地方是否有通用的解决方案？

编辑：在@hpaulj 的 cmets 之后更精确一点。本质上，我正在寻找

1. 一个与np.atleast_1d相反的函数，例如 atleast_1d_inverse( np.atleast_1d(x), x) == x，其中第二个参数仅用于确定原始对象x 的类型或维数。返回 numpy 标量（即带有 ndim = 0 的数组）而不是 python 标量是可以的。

2. 一种检查函数 f 并生成与其定义一致的包装器的方法。例如，这样的包装器可以用作，

   f_ufunc = ufunc_wrapper(f, args=['x', 'y'])

   如果我们有不同的函数def f2(x, option=2): return x**2，我们也可以使用

   f2_ufunc = ufunc_wrapper(f2, args=['x']).

注意： 与 ufuncs 的类比可能有点有限，因为这对应于相反的问题。我没有一个我们转换为接受向量和标量输入的标量函数，而是有一个设计用于处理向量的函数（可以看作是以前向量化的东西），我想再次接受标量，而不改变原来的功能。

Answer 1

这并不能完全回答使矢量化函数真正表现得像ufunc 的问题，但我最近确实对numpy.vectorize 感到有点烦恼，这听起来与您的问题相似。即使传递了标量输入，该包装器仍坚持返回 array（带有 ndim=0 和 shape=()）。

但似乎下面的做法是正确的。在这种情况下，我正在对一个简单的函数进行矢量化，以将浮点值返回到一定数量的有效数字。

def signif(x, digits):
    return round(x, digits - int(np.floor(np.log10(abs(x)))) - 1)

def vectorize(f):
    vf = np.vectorize(f)

    def newfunc(*args, **kwargs):
        return vf(*args, **kwargs)[()]
    return newfunc

vsignif = vectorize(signif)

这给了

>>> vsignif(0.123123, 2)
0.12
>>> vsignif([[0.123123, 123.2]], 2)
array([[   0.12,  120.  ]])
>>> vsignif([[0.123123, 123.2]], [2, 1])
array([[   0.12,  100.  ]])