通过nump apply_along_axis计算多维高斯分布

Question

我想计算多波段图像的多维高斯。为此，我需要将每个像素及其在不同频带中的值视为向量来应用高斯，其中包含 V^T*inv(covariance)*V 的矩阵乘法。

import numpy as np
import numpy.linalg as lg
X,Y=np.meshgrid(np.arange(-10,11,1/10),np.arange(-10,11,1/10),indexing='xy')
Z=np.stack((X,Y))
cov=np.array([[10,-0.4],[-0.4,1]])
W=lg.inv(cov)
def ee(x):
    global W
    x=x.reshape((-1,1))
    
    return lg.matmul(lg.matmul(x.transpose(),W),x)[0,0]
tt=np.apply_along_axis(ee,0,Z)
p=np.exp(-0.5*tt)/(np.power(2*np.pi,cov.shape[0]/2)*np.power(lg.det(cov),0.5))
import matplotlib.pyplot as plt
plt.imshow(p)

在此输入图片描述

我找到了 apply_along_axis 作为解决方案，但我的图像大小是（3,1000,1000），并且 apply_along_axis 在像素上应用函数来查找每个像素的概率非常慢。我怎样才能在现实世界的应用程序中加速它。谢谢

Answer 1

您可以使用以下命令，而不是沿着轴循环：

a = Z.reshape(2,-1)
t = (a.T @ W * a.T).sum(1).reshape(*X.shape)

检查是否产生预期结果：

np.allclose(tt, t)
True