Numpy中每一行的`vectorize`答案

【问题标题】：`vectorize` for each row in NumpyNumpy中每一行的`vectorize`
【发布时间】：2016-09-01 15:37:31
【问题描述】：

假设我有一个 n x m 矩阵，并且想在它的每个元素上调用一个函数 fct。我可以这样做：

A = numpy.array(...)
vec_func = numpy.vectorize(fct)
A_out = vec_func(A)

这将严格地将函数应用于每个矩阵元素，fct 将是一个函数：

def fct(a_ij):
  # do something with matrix element a(i, j)

现在我想要相同的，但对于矩阵的每一行：

def fct(row_i):
  # do something with matrix row(i)

有没有办法用numpy.vectorize 或类似的方法来做到这一点？

【问题讨论】：

我刚刚读到：“提供 vectorize 函数主要是为了方便，而不是为了性能。实现本质上是一个 for 循环。”所以这个函数反正也没那么有用。
没有“功能无关”的矢量化方法可以显着提高性能（如您所见，np.vectorize 只是语法糖）。您可以重新编写 fct 以接受数组作为输入参数，或者切换到已编译的实现（例如 Cython、numba）。
要添加到@ali_m 的建议，您可以尝试 numba 的广义 ufunc numba.pydata.org/numba-doc/dev/user/…。除了速度提升之外，它还可以实现您想要的智能广播。

标签： python performance numpy scipy vectorization

【解决方案1】：

编辑：看起来np.apply_along_axis 做了你想做的事。例如：

import numpy as np

def f(x):
    return x * x.sum()

X = np.arange(12).reshape(2, 2, 3)
np.apply_along_axis(f, -1, X)
# array([[[  0,   3,   6],
#         [ 36,  48,  60]],
#
#        [[126, 147, 168],
#         [270, 300, 330]]])

我在下面的原始回复中关于性能的说明仍然适用。

原始回复：

没有内置的，但是 Python 可以很容易地自己定义这样的上下文管理器。例如：

import numpy as np
from contextlib import wraps

def row_vectorize(f):
    @wraps(f)
    def wrapped_f(X):
        X = np.asarray(X)
        rows = X.reshape(-1, X.shape[-1])
        return np.reshape([f(row) for row in rows],
                          X.shape[:-1] + (-1,))
    return wrapped_f


@row_vectorize
def func(row):
    return row * row.sum()

现在您可以在任何非零维度的数组上使用它：

>>> X_1D = np.arange(3)
>>> func(X_1D)
array([0, 3, 6])

>>> X_2D = np.arange(6).reshape(2, 3)
>>> func(X_2D)
array([[ 0,  3,  6],
       [36, 48, 60]])

>>> X_3D = np.arange(12).reshape((2, 2, 3))
>>> func(X_3D)
array([[[  0,   3,   6],
        [ 36,  48,  60]],

       [[126, 147, 168],
        [270, 300, 330]]])

在性能方面，np.vectorize 正在做一些非常相似的事情。

如果您需要对跨数组应用的自定义函数进行更快的循环，您通常可以根据 numpy 元素操作和聚合操作来构建您的方法；例如，这个函数完成了与上面的行向量化函数相同的事情，但在大输入时会更快：

def func2(X):
    return X * X.sum(-1, keepdims=True)

如果您想要跨数组行应用更复杂的操作并且循环的性能是一个瓶颈，那么最好的选择可能是使用numba 或cython。

【讨论】：