pythonnumpygaussian

calculating multi dimensional gaussian distribution by nump apply_along_axis


I want to calculate multi dimensional Gaussian for multi-band image. For this purpose I need to consider each pixel with its value in different bands as vector to apply Gaussian which contains matrix multiplication of 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)

enter image description here

I have found apply_along_axis as solution but my image size is (3,1000,1000) and apply_along_axis is so slow to apply function on pixels to find each pixel's probability. How can i speed it up fore real world applications.thank you


Solution

  • Instead of looping along the axis, You could use the following:

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

    Check whether it produces the desired results:

    np.allclose(tt, t)
    True