68 lines
1.9 KiB
Python
68 lines
1.9 KiB
Python
"""
|
|
Simple Gaussian-mixture model example
|
|
======================================
|
|
"""
|
|
|
|
import numpy as np
|
|
from numpy.random import seed
|
|
|
|
from scikits.learn.em import GM, GMM, EM
|
|
|
|
seed(2)
|
|
|
|
k = 4
|
|
d = 2
|
|
mode = 'diag'
|
|
nframes = 1e3
|
|
|
|
################################################################################
|
|
# Create an artificial GMM model, samples it
|
|
################################################################################
|
|
w, mu, va = GM.gen_param(d, k, mode, spread = 1.0)
|
|
gm = GM.fromvalues(w, mu, va)
|
|
|
|
# Sample nframes frames from the model
|
|
data = gm.sample(nframes)
|
|
|
|
################################################################################
|
|
# Learn the model with EM
|
|
################################################################################
|
|
|
|
# List of learned mixtures lgm[i] is a mixture with i+1 components
|
|
lgm = []
|
|
kmax = 6
|
|
bics = np.zeros(kmax)
|
|
em = EM()
|
|
for i in range(kmax):
|
|
lgm.append(GM(d, i+1, mode))
|
|
|
|
gmm = GMM(lgm[i], 'kmean')
|
|
em.train(data, gmm, maxiter = 30, thresh = 1e-10)
|
|
bics[i] = gmm.bic(data)
|
|
|
|
print "Original model has %d clusters, bics says %d" % (k, np.argmax(bics)+1)
|
|
|
|
################################################################################
|
|
# Draw the model
|
|
################################################################################
|
|
import pylab as pl
|
|
pl.subplot(3, 2, 1)
|
|
|
|
for k in range(kmax):
|
|
pl.subplot(3, 2, k+1)
|
|
level = 0.9
|
|
pl.plot(data[:, 0], data[:, 1], '.', label = '_nolegend_')
|
|
|
|
# h keeps the handles of the plot, so that you can modify
|
|
# its parameters like label or color
|
|
h = lgm[k].plot(level = level)
|
|
[i.set_color('r') for i in h]
|
|
h[0].set_label('EM confidence ellipsoides')
|
|
|
|
h = gm.plot(level = level)
|
|
[i.set_color('g') for i in h]
|
|
h[0].set_label('Real confidence ellipsoides')
|
|
|
|
pl.legend(loc = 0)
|
|
pl.show()
|