Search code examples
pythonnumpynumpy-einsum

numpy: get rid of for loop by broadcasting

I am trying to implement the Expectation Maximization Algorithm for Gaussian Mixture Model in python.

I have following line to compute the gaussian probability p of my data X given the mean mu and covariance sigma of the Gaussian distribution:

for i in range(len(X[0])):  
   p[i] = scipy.stats.multivariate_normal.pdf(X[:,i],mu,sigma)

I wanted to know if I somehow could get rid of the for loop to get something like

p[:] = scipy.stats.multivariate_normal.pdf(X[:,:]??)

I did some research on broadcasting and was thinking about using the numpy.einsum function but can't figure out how it would work in this case.

Solution

  • Flatten, use the pdf call and reshape back -

    from scipy import stats

out = stats.multivariate_normal.pdf(X.ravel(),mu,sigma).reshape(-1,len(X[0])).T