沿一个轴应用 np.vectorize答案

【问题标题】：Apply np.vectorize along one axis沿一个轴应用 np.vectorize
【发布时间】：2021-09-05 09:53:44
【问题描述】：

假设我有两个数组arr1 和arr2：

arr1 = [0, 1, 2]

arr2 = [
    [0, 1, 2],
    [3, 4, 5],
    [6, 7, 8],
]

假设我有一个函数可以对这个数组的元素做一些事情：

def func(arr):
    new_arr = arr.copy()
    new_arr[0] = new_arr[0] * 2
    new_arr[1] = new_arr[1] * 10
    new_arr[2] = new_arr[2] * 100
    return new_arr

现在我想对其进行矢量化处理，以便它适用于 arr1 和 arr2：

func(arr1)
# returns [0, 10, 200]
func(arr2)
# returns
# [0, 10, 200],
# [6, 40, 500],
# [12, 70, 800],

np.vectorize 不起作用，因为它分解了我的数组参数中的每个元素。我希望它仅沿第一个轴应用该功能。

np.apply_along_axis 几乎可以工作，只是它不会将一维数组参数视为单个参数。

最好的方法是什么？

【问题讨论】：

矢量化和应用不是性能工具。如果它们不适合您，请执行其他操作。即使是嵌套循环也更好。
别怕测试维数，把arr1改成(1,3)形状。

标签： python numpy vectorization

【解决方案1】：

您可以直接将数组相乘。感谢numpy broadcasting：

factor = np.array([2, 10, 100])

arr1 * factor
array([  0,  10, 200])

arr2 * factor
array([[  0,  10, 200],
       [  6,  40, 500],
       [ 12,  70, 800]])

【讨论】：

【解决方案2】：

如果您花时间阅读np.vectorize 文档，您最终会遇到signature 选项：

In [27]: f= np.vectorize(func, signature='(n)->(n)')
In [28]: f(arr1)
Out[28]: array([  0,  10, 200])
In [29]: f(arr2)
Out[29]: 
array([[  0,  10, 200],
       [  6,  40, 500],
       [ 12,  70, 800]])

进一步阅读，您会发现有关性能的警告。

【讨论】：

【解决方案3】：

这样做：

import numpy as np

a = np.array([0, 1, 2])

b = np.array([
    [0, 1, 2],
    [3, 4, 5],
    [6, 7, 8],
])

c = np.array([2, 10, 100])

print(a*c)

print(b*c)

输出：

[  0  10 200]
[[  0  10 200]
 [  6  40 500]
 [ 12  70 800]]

【讨论】：