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为标量时，我们将它们转换为1D数组，然后进行计算，然后将它们转换回标量。
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失败，产生标量输出而不是矢量1。如果我们想要解决这个问题，我们需要在输入类型上添加更多测试（例如isinstance(x, np.ndarray)，x.ndim>0等），但我不敢忘记那里的一些极端情况。此外，上面的实现不够通用，不能包含具有不同数量参数的函数（参见下面的第2点）。

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

在 编辑：更精确一些。基本上，我正在寻找

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

使矢量化numpy函数的行为类似于ufunc

1 个答案: