EM(Expectation-Maimization) algorithm for find multi Gaussian means.

em.py

import numpy as np
import matplotlib.pyplot as plt

SIGMA = 1

def gauss(x, mu, sigma, eps=1e-8):
    return 1./np.sqrt(np.pi*2*sigma**2+eps) * np.exp(-(x-mu)**2/(2*sigma**2+eps))

def gauss_vec(x, mu, sigma, eps=1e-8):
    """ x:  Nx1
        mu: 1xm"""
    x = x.reshape(-1, 1)
    mu = mu.reshape(1, -1)
    x = x.repeat(mu.shape[1], axis=1)
    mu = mu.repeat(x.shape[0], axis=0)
    return gauss(x, mu, sigma, eps)

def init_h_of_mu(k, begin, stop):
    return (np.random.rand(k)*(stop-begin)+begin).reshape(1, -1)

def step1_calculate_Expectation(x, h_of_mu):
    # posibility
    p = gauss_vec(x, h, SIGMA)
    p = p/np.sum(p, axis=1, keepdims=True)
    return p    

def step2_calculate_h_of_mu(x, expetation):
    return np.sum(x * expetation, axis=0)/np.sum(expetation, axis=0)

def get_multi_gauss(mu=[1, 7, 15]):
    samples = 200
    
    mu = np.array([mu]).reshape(1, -1)
    mu = mu.repeat(samples, axis=0)
    ds = np.random.randn(samples, mu.shape[1])*SIGMA
    ds += np.random.randn(*ds.shape)*0.01
    return (ds + mu).reshape(-1)
    
t = get_multi_gauss([1, 7, 20])
plt.plot(t, np.zeros_like(t), '.')


k = 3
h = init_h_of_mu(k, -20, 20)
print('init state', h)
for i in range(8):
    p = step1_calculate_Expectation(t.reshape(-1, 1), h)
    h = step2_calculate_h_of_mu(t.reshape(-1, 1), p)
    print(h)