Skip to main content
 首页 » 编程设计

python之如何在 numpy 行上应用通用函数

2025年05月04日18lori

在您将此标记为重复之前,让我向您解释一下我读过 this page和许多其他人一样,我仍然没有找到解决问题的办法。

这就是我遇到的问题:给定两个二维数组,我想在这两个数组上应用函数 F。 F 将两个一维数组作为输入。

import numpy as np 
a = np.arange(15).reshape([3,5]) 
b = np.arange(30, step=2).reshape([3,5]) 
 
# what is the 'numpy' equivalent of the following? 
np.array([np.dot(x,y) for x,y in zip(a,b)]) 

请注意,np.dot 仅用于演示。这里真正的问题是任何适用于两组一维数组的通用函数 F。

  • 矢量化要么完全失败并出现错误,要么逐个元素而不是逐个数组(或逐行)应用函数
  • np.apply_along_axis 迭代地应用函数;例如,使用上面定义的变量,它执行 F(a[0], b[0]) 并将其与 F(a[0], b[1]) 和 F(a[0], b[2])。这不是我要找的。理想情况下,我希望它只停在 F(a[0], b[0])
  • 索引切片/高级切片也不符合我的要求。首先,如果我做类似 np.dot(a[np.arange(3)], b[np.arange(3)]) 的事情,这会抛出一个 ValueError 表示形状 (3,5 ) 和 (3,5) 没有对齐。我不知道如何解决这个问题。

我尝试尽我所能解决这个问题,但我想出的唯一可行的解​​决方案是使用列表理解。但是我担心由于使用列表理解而导致的性能成本。如果可能的话,我想使用 numpy 操作来达到同样的效果。我该怎么做?

请您参考如下方法:

此类问题已在 SO 上被打死,但我会尝试用您的框架来说明这些问题:

In [1]: a = np.arange(15).reshape([3,5]) 
   ...: b = np.arange(30, step=2).reshape([3,5]) 
   ...:  
In [2]: def f(x,y): 
   ...:     return np.dot(x,y) 

压缩理解

列表理解方法将 f 应用于 a 和 b 的 3 行。也就是说,它迭代 2 个数组,就像它们是列表一样。每次调用时,您的函数都会获得 2 个一维数组。 dot 可以接受其他形状,但目前我们假设它只适用于一对 1ds

In [3]: np.array([f(x,y) for x,y in zip(a,b)]) 
Out[3]: array([  60,  510, 1460]) 
In [4]: np.dot(a[0],b[0]) 
Out[4]: 60 

向量化/frompyfunc

np.vectorize 迭代输入(使用广播 - 这很方便),并给出函数标量值。我将用 frompyfunc 来说明返回一个对象 dtype 数组(并由 vectorize 使用):

In [5]: vf = np.frompyfunc(f, 2,1) 
In [6]: vf(a,b) 
Out[6]:  
array([[0, 2, 8, 18, 32], 
       [50, 72, 98, 128, 162], 
       [200, 242, 288, 338, 392]], dtype=object) 

所以结果是(3,5)数组;顺便求和跨列得到想要的结果

In [9]: vf(a,b).sum(axis=1) 
Out[9]: array([60, 510, 1460], dtype=object) 

np.vectorize 不做任何速度 promise 。

沿轴应用

我不知道您是如何尝试使用 apply_along_axis 的。它只需要一个数组。经过大量设置后,它最终会执行以下操作(对于像 a 这样的二维数组):

for i in range(3): 
    idx = (i, slice(None)) 
    outarr[idx] = asanyarray(func1d(arr[idx], *args, **kwargs)) 

对于 3d 和更大的尺寸,它使“其他”轴上的迭代更简单;对于 2d,这是矫枉过正。在任何情况下,它都不会加快计算速度。还是迭代。

(apply_along_axis 采用arr 和*args。它迭代arr,但使用* args 整个。)。

索引

np.dot(a[np.arange(3)], b[np.arange(3)]) 

与

相同
np.dot(a, b) 

dot 是矩阵乘积,(3,5) 与 (5,3) 一起产生 (3,3)。它将 1d 作为一种特殊情况处理(参见文档),(3,) 与 (3,) 产生 (3,)。

迭代

对于真正通用的f(x,y),压缩列表理解的唯一替代方法是像这样的索引循环:

In [18]: c = np.zeros((a.shape[0])) 
In [19]: for i in range(a.shape[0]): 
    ...:    c[i] = f(a[i,:], b[i,:]) 
In [20]: c 
Out[20]: array([   60.,   510.,  1460.]) 

速度会差不多。 (可以使用 cython 将该操作移动到编译代码,但我认为您还没有准备好深入研究。)

如注释中所述,如果数组为 (N,M),并且 N 与 M 相比较小,则此迭代并不昂贵。也就是说,在一个大任务上循环几次就可以了。如果简化大型阵列内存管理,它们甚至可能更快。

最佳

理想的解决方案是使用 numpy 编译函数重写通用函数,使其适用于二维数组。

在矩阵乘法的情况下,einsum 在编译代码中实现了“乘积和”的通用形式:

In [22]: np.einsum('ij,ij->i',a,b) 
Out[22]: array([  60,  510, 1460]) 

matmul 也概括了产品,但最适用于 3d 数组:

In [25]: a[:,None,:]@b[:,:,None]    # needs reshape 
Out[25]:  
array([[[  60]], 
 
       [[ 510]], 
 
       [[1460]]])