EM求解GMM代码

import numpy as np
from numpy import linalg
import math
from sklearn.mixture import GaussianMixture


class GMM(object):
    def __init__(self, X, K):
        self.N, self.D = X.shape
        self.K = K
        self.Mu, self.Sigma, self.Pi, self.Gamma = self.init_params()
    
    def init_params(self):
        Mu = np.random.rand(self.K, self.D)  # K*D
        Sigma = np.array([np.eye(self.D)] * self.K) * 0.1  # K*D*D
        Pi = np.array([1.0 / self.K] * self.K)  # K
        Gamma = np.random.rand(self.N, self.K)
        return Mu, Sigma, Pi, Gamma
    
    def e_step(self, X):
        prob = []
        for i in range(self.K):
            prob.append(get_Gaussian_prob(X, self.Mu[i], self.Sigma[i]))
        prob = np.array(prob).T
        Pi_Gussian = prob * self.Pi
        sum_K = np.sum(Pi_Gussian, axis=1)
        return (Pi_Gussian.T / sum_K).T  # Gamma
    
    def m_step(self, X):
        NK = np.sum(self.Gamma, axis=0)
        _1_NK = 1 / NK
        # get new Mu
        m = np.dot(self.Gamma.T, X)
        newMu = (_1_NK * m.T).T
        # get new Sigma
        newSigma = []
        for i in range(self.K):
            newSigma.append(_1_NK[i] * np.dot((self.Gamma[:, i] * (X - newMu[i]).T), (X - newMu[i])))
        newSigma = np.array(newSigma)
        # get Pi
        newPi = NK / self.N
        return newMu, newSigma, newPi
    
    def train(self, X, times):
        for i in range(times):
            self.Gamma = self.e_step(X)
            self.Mu, self.Sigma, self.Pi = self.m_step(X)
        return self.Mu, self.Sigma, self.Pi, self.Gamma


def get_Gaussian_prob(X, mu, sigma):
    """已知一个高斯分布的参数，求其概率"""
    D = len(sigma)
    if X.size == D:
        n = 1
    else:
        n, _ = X.shape
    det = linalg.det(sigma)
    
    norm_const = 1.0 / (math.pow(2 * math.pi, float(D) / 2) * math.pow(det, 1.0 / 2))
    x_mu = np.array(X - mu)
    inv = linalg.inv(sigma)
    result = np.power(math.e, -0.5 * np.dot(np.dot(x_mu, inv), x_mu.T))
    if n != 1:
        result = np.diag(result)
    return norm_const * result