python - 如何正确地将对向量进行操作的函数推广到对多维数组进行操作的函数？

Question

我经常最终实现对一维数组进行操作的例程，然后对它们进行泛化，这样，如果在矩阵上调用它们，它们就会独立处理每一行。例如，假设我们想要一个函数，从它减去向量的平均值，然后在矩阵的每一行上执行此操作，如果它的输入有多个维度。这可以按如下方式实现：

import numpy as np

def _sub_mean(a):
    """ Subtract the mean from vector elements """

    assert isinstance(a, np.ndarray) and a.ndim == 1

    return a - np.mean(a)

def sub_mean(a):
    """ Subtract the mean

    If `a` is a 1D array it returns a 1D array (obtained by subtracting the
    mean of `a` from each of its elements).  If `a` has more than one dimension
    each row is treated separately

    """

    a = np.asanyarray(a)

    if a.ndim <= 1:
        a = np.atleast_1d(a)

        return np.squeeze(_sub_mean(a))

    retval = np.empty(a.shape)

    for idx in np.ndindex(*a.shape[:-1]):
        retval[idx] = _sub_mean(a[idx])

    return retval

以下是 的输出的几个示例sub_mean：

>>> sub_mean(5)
array(0.)
>>> sub_mean([1, 2])
array([-0.5,  0.5])
>>> sub_mean([[1, 2], [4, 6]])
array([[-0.5,  0.5],
       [-1. ,  1. ]])
>>> 

请注意，“核心”计算发生在私有函数_sub_mean中。实际上，相同的代码sub_mean可用于将任何对一维数组进行操作的函数泛化为对任意维数进行操作的函数，方法是替换_sub_mean. 还可以考虑进一步的概括，例如添加一个axis参数，指定函数在哪个轴上运行和/或在展平的输入数组上运行的可能性。

我想知道 NumPy 是否已经提供了一个装饰器来将对向量进行操作的函数泛化为对多维数组进行操作的函数？即，如果我可以通过以下方式替换上面的代码：

import numpy as np

@np.some_decorator
def sub_mean(a):        
    """ Subtract the mean

    If `a` is a 1D array it returns a 1D array (obtained by subtracting the
    mean of `a` from each of its elements).  If `a` has more than one dimension
    each row is treated separately  """


    assert isinstance(a, np.ndarray) and a.ndim == 1

    return a - np.mean(a)

（并且，显然，得到相同的输出。）

---

更新：我最终编写了以下装饰器：

class _Expand:
    def __init__(self, func1d):
        functools.update_wrapper(self, func1d)
        self._func1d = func1d

    def __call__(self, arr, *args, **kwargs):
        arr = np.asanyarray(arr)

        axis = kwargs.pop('axis', -1)

        return np.apply_along_axis(self._func1d, axis, arr, *args, **kwargs)

这将使我能够编写sub_mean（或任何其他在 1D 数组上运行的复杂函数）为：

@_Expand
def sub_mean(a):        
    """ Subtract the mean

    If `a` is a 1D array it returns a 1D array (obtained by subtracting the
    mean of `a` from each of its elements).  If `a` has more than one dimension
    each row is treated separately  """


    assert isinstance(a, np.ndarray) and a.ndim == 1

    return a - np.mean(a)

（请注意，_Expand允许选择执行操作的轴——这比我需要的稍微通用。）

然而，我仍然很想知道这样的装饰器是否已经在 NumPy 中实现了？

标签： pythonnumpynumpy-ndarray

Answer 1

这是一个相当普遍的问题，但没有完美的答案。真正的答案是您最终必须单独定制每个功能。

假设您想坚持使用纯 numpy（和 vanilla python）而不是使用 numba 或 cython，请记住以下注意事项：

• np.apply_along_axis主要作为python循环实现。除了更紧凑的符号之外，它没有提供任何真正的优势。
• 几乎所有 numpy 函数都接受一个axis参数。从 numpy v1.15 开始，您可以为ufuncs以及在它们之上构建的许多函数提供一个元组多个同时轴。
• Ufunc 还具有可让您将它们应用到特定位置和组合的方法。
• 归约函数，比如np.mean经常有一个keepdims参数。这允许您将归约操作与原始数组结合起来。不支持的功能keepdims可以通过手动插入轴来缓解，例如，np.expand_dims.
• 有些函数专门对索引执行操作；他们的大多数名字都以arg*: np.nonzero, np.argmin, np.argmax, np.argpartition,np.argsort等开头。

总之，您可以对几乎任何可以为单个维度编写的函数进行矢量化。有时结果需要额外的聪明才智，有时是完全不可能的，但通常相当简单。我能想到的一个不平凡的例子是涉及游程编码的任何事情。那里的问题是，在某些步骤中，您最终会得到一个参差不齐的数组，需要 numpy 之外的工具。

你的具体例子可以写成如下：

def sub_mean(a, axis):
    a = np.array(a, copy=False, subok=True)
    return a - np.mean(a, axis=axis, keepdims=True)