scikit-learn/examples/mixture/plot_gmm.py

45 lines
1.1 KiB
Python
Raw Normal View History

"""
2010-06-25 22:57:40 +08:00
=================================
Gaussian Mixture Model Ellipsoids
=================================
2010-06-25 22:57:40 +08:00
Plot the confidence ellipsoids of a mixture of two gaussians.
"""
import numpy as np
from scikits.learn import mixture
import itertools
import pylab as pl
import matplotlib as mpl
2010-06-25 22:51:00 +08:00
n, m = 300, 2
2010-06-25 22:51:00 +08:00
# generate random sample, two components
np.random.seed(0)
2010-06-25 22:51:00 +08:00
C = np.array([[0., -0.7], [3.5, .7]])
X = np.r_[np.dot(np.random.randn(n, 2), C),
2010-06-22 22:13:08 +08:00
np.random.randn(n, 2) + np.array([3, 3])]
clf = mixture.GMM(n_states=2, cvtype='full')
clf.fit(X)
splot = pl.subplot(111, aspect='equal')
color_iter = itertools.cycle (['r', 'g', 'b', 'c'])
Y_ = clf.predict(X)
2010-06-25 22:51:00 +08:00
for i, (mean, covar, color) in enumerate(zip(clf.means, clf.covars, color_iter)):
v, w = np.linalg.eigh(covar)
u = w[0] / np.linalg.norm(w[0])
pl.scatter(X[Y_==i, 0], X[Y_==i, 1], .8, color=color)
angle = np.arctan(u[1]/u[0])
angle = 180 * angle / np.pi # convert to degrees
2010-06-25 22:57:40 +08:00
ell = mpl.patches.Ellipse (mean, v[0], v[1], 180 + angle, color=color)
ell.set_clip_box(splot.bbox)
ell.set_alpha(0.5)
splot.add_artist(ell)
pl.show()