如何numpy矢量化要应用于numpy矩阵每一行的函数答案

【问题标题】：How to numpy vectorize a function to be applied to every row of numpy matrix如何numpy矢量化要应用于numpy矩阵每一行的函数
【发布时间】：2021-06-30 13:17:15
【问题描述】：

我写了这个函数，用于计算给定多元高斯参数的 x 的概率。其中 x 是具有 2 个特征的特征，mu 是 2 个特征的向量，sigma 是 2x2。

def prob(x, mu, sigma):
    n = len(x)
    var = x - mu
    sigma_inv = np.linalg.inv(sigma)
    rhs = np.exp(-0.5*np.matmul(np.matmul(var.T, sigma_inv),var))
    lhs = 1/(((2*np.pi)**(n/2))*np.linalg.det(sigma)**.5)
    return rhs*lhs

但这仅在 x 是一维数组时才有效。我希望能够针对目前我拥有的多维 x（例如 x 为 100x2）进行矢量化和优化。

for i in range(len(x)):
    curr_prob = prob(x[i], mu, sigma)
        
    if i == 0:
        prob = curr_prob
    else:
        prob = np.append(prob, curr_prob)

但这很慢。我听说有一种方法可以为此使用 np.vectorize 或 np.pyfunc，但我不确定如何应用它们。

【问题讨论】：

首先，在循环中使用np.append 很慢。 np.array([prob(i, mu, sigma) for i in x]) 应该更快。但是加速矢量化需要重写prob 以直接使用二维x。请注意，inv, matmul 和 det 可以处理 3d 数组，有效地“批量”矩阵。
@hpaulj 你介意发布如何在 3D 数组上应用 inv、matmul 和 det 吗？
@hpaulj 抱歉，我的意思是二维数组。基本上，如果您可以发布如何进行批处理操作，我将非常感激。

标签： python arrays numpy gaussian

【解决方案1】：

看看形状在当前函数中是如何流动的

x (2,), mu (2,), sigma (2,2)

def prob(x, mu, sigma):
    n = len(x)             # 2
    var = x - mu           # (2,)-(2,)
    sigma_inv = np.linalg.inv(sigma)    # (2,2) no change w/ x
    rhs = np.exp(-0.5*np.matmul(np.matmul(var.T, sigma_inv),var))
    # var is (2,), var.T is the same (2,)
    # (2,)@(2,2)=>(2,); (2,)@(2,)=> scalar
    # np.exp(-0.5 * var @ sigma_inv @ var

    lhs = 1/(((2*np.pi)**(n/2))*np.linalg.det(sigma)**.5)
    # np.linalg.def(sigma)**.5 - no dependence on x
    # the whole lhs doesn't vary with x
    # lhs is scalar

    return rhs*lhs    # scalar

现在考虑如果x 是 (100,2) 会发生什么变化

def prob(x, mu, sigma):
    n = x.shape[-1]    

    sigma_inv = np.linalg.inv(sigma)    # (2,2) or (n,n)?
    lhs = 1/(((2*np.pi)**(n/2))*np.linalg.det(sigma)**.5) # scalar

    var = x - mu           # (100,2)-(2,)
    # by broadcasting this is (100,2)-(1,2)=>(100,2) 
    # no change needed

    rhs = var@sigma_inv@var
    # (100,2) @ (2,2) => (100,2)
    # (100,2) @ (100,2)  oops
    # var@sigma_inv@var.T   
    # (100,2) with (2,100)=>(100,100)   no! 
    # np.einsum('ij,jk,ik->i',var,sigma_inv,var) 
    # 
    rhs = np.exp(-0.5*rhs)    # (100,2)
    return rhs*lhs

在我第一次尝试时，var@sigma_inv@var 与 (2,) 一起工作，但在 (n,2) 上给了我一个错误。但是einsum 表达式得到了i，批处理维度是正确的。我可能也可以纠正双 @ 以使其正确。

def prob(x, mu, sigma):
    n = x.shape[-1]    
    sigma_inv = np.linalg.inv(sigma)
    lhs = 1/(((2*np.pi)**(n/2))*np.linalg.det(sigma)**.5)

    var = x - mu
    rhs = np.einsum('ij,jk,ik->i',var,sigma_inv,var)
    rhs = np.exp(-0.5*rhs)
    return rhs*lhs

测试：

In [495]: x = np.array([1,2.]); mu=np.array([.5,.5]);
In [496]: sigma = np.array([[1,.5],[.5,3]])
In [497]: X = np.array([x,x,x+1])
In [498]: prob(X,mu,sigma)
Out[498]: array([0.06375113, 0.06375113, 0.01785457])

您的函数还会为 x 生成 0.06375。

einsum 可以写成：

np.squeeze(var[:,None,:]@sigma_inv@var[:,:,None])

但这并没有更清楚，而且速度上可能没有太大差异。

【讨论】：

【解决方案2】：

用numpy.apply_along_axis怎么样：

np.apply_along_axis(prob, 1, x, mu, sigma)

【讨论】：

是的，我调查过，我相信它也会将轴参数传递给 probs，而它不需要
apply_along_axis 不会加快速度。