#!/usr/bin/env python
from netlab import *
import numpy as np
import matplotlib.pyplot as pp
from enthought.mayavi import mlab

#"""DEMGMM4 Demonstrate density modelling with a Gaussian mixture model.
#
#	Description
#	 The problem consists of modelling data generated by a mixture of
#	three Gaussians in 2 dimensions with a mixture model using full
#	covariance matrices.  The priors are 0.3, 0.5 and 0.2; the centres
#	are (2, 3.5), (0, 0) and (0,2); the variances are (0.16, 0.64) axis
#	aligned, (0.25, 1) rotated by 30 degrees and the identity matrix. The
#	first figure contains a scatter plot of the data.
#
#	A Gaussian mixture model with three components is trained using EM.
#	The parameter vector is printed before training and after training.
#	The user should press any key to continue at these points.  The
#	parameter vector consists of priors (the column), and centres (given
#	as (x, y) pairs as the next two columns).  The covariance matrices
#	are printed separately.
#
#	The second figure is a 3 dimensional view of the density function,
#	while the third shows the axes of the 1-standard deviation ellipses
#	for the three components of the mixture model.
#
#	See also
#	GMM, GMMINIT, GMMEM, GMMPROB, GMMUNPAK
#

#	Copyright (c) Ian T Nabney (1996-2001)
#       and Neil D. Lawrence (2009) (translation to python)"""


# Generate the data

ndata = 500

# Fix the seeds for reproducible results
np.random.seed(42)
data = np.random.randn(ndata, 2)
prior = np.array([0.3, 0.5, 0.2])
# Mixture model swaps clusters 1 and 3
datap = np.array([0.2, 0.5, 0.3])
datac = np.array([[0, 2],[0, 0], [2, 3.5]])
datacov = np.zeros((2, 2, 3))
for i in range(3):
    datacov[:, :, i] = np.eye(2)
data1 = data[0:round(prior[0]*ndata),:]
data2 = data[round(prior[0]*ndata):round((prior[1]+prior[0])*ndata), :]
data3 = data[round((prior[0]+prior[1])*ndata):ndata, :]

# First cluster has axis aligned variance and centre (2, 3.5)
data1[:, 0] = data1[:, 0]*0.4 + 2.0
data1[:, 1] = data1[:, 1]*0.8 + 3.5
datacov[:, :, 2] = np.array([[0.4*0.4, 0],[ 0, 0.8*0.8]])

# Second cluster has variance axes rotated by 30 degrees and centre (0, 0)
rotn = np.array([[np.cos(np.pi/6), -np.sin(np.pi/6)], [np.sin(np.pi/6), np.cos(np.pi/6)]])
data2[:, 1] = data2[:, 1]*0.5
data2 = np.dot(data2, rotn)
datacov[:, :, 1] = np.dot(np.dot(rotn.T, np.array([[0.25, 0],[0, 1]])), rotn)

# Third cluster is at (0,2) 
data3 = data3 + np.tile([0, 2], (round(prior[2]*ndata), 1))

# Put the dataset together again
data = np.r_[data1, data2, data3]

clc()
print "This demonstration illustrates the use of a Gaussian mixture model"
print "with full covariance matrices to approximate the unconditional "
print "probability density of data in a two-dimensional space."
print "We begin by generating the data from a mixture of three Gaussians and"
print "plotting it."
print " "
print "The first cluster has axis aligned variance and centre (0, 2)."
print "The second cluster has variance axes rotated by 30 degrees"
print "and centre (0, 0).  The third cluster has unit variance and centre"
print "(2, 3.5)."
print " "
raw_input('Press return to continue.')

fh1 = pp.figure()
pp.axes(axisbg='w', frameon=True)
pp.plot(data[:, 0], data[:, 1], 'o')

# Set up mixture model
ncentres = 3
input_dim = 2
mix = gmm(input_dim, ncentres, 'full')

# Initialise the model parameters from the data
options = foptions()
options[13] = 5	# Just use 5 iterations of k-means in initialisation
mix.init(data, options)

# Print out model
clc()
print "The mixture model has three components and full covariance"
print "matrices.  The model parameters after initialisation using the"
print "k-means algorithm are as follows"
print "    Priors        Centres"
print np.c_[mix.priors, mix.centres]
print "Covariance matrices are"
print mix.covars
pp.show()
raw_input('Press return to continue.')
pp.ioff()

# Set up vector of options for EM trainer
options = np.zeros(18)
options[0]  = 1		# Prints out error values.
options[13] = 50		# Number of iterations.

print "We now train the model using the EM algorithm for 50 iterations."
print " "
raw_input('Press return to continue.')

errlog = mix.em(data, options)

# Print out model
print " "
print "The trained model has priors and centres:"
print "    Priors        Centres"
print np.c_[mix.priors, mix.centres]
print "The data generator has priors and centres"
print "    Priors        Centres"
print np.c_[datap, datac]
print "Model covariance matrices are"
print mix.covars[:, :, 0]
print mix.covars[:, :, 1]
print mix.covars[:, :, 2]
print "Data generator covariance matrices are"
print datacov[:, :, 0]
print datacov[:, :, 1]
print datacov[:, :, 2]
print "Note the close correspondence between these parameters and those"
print "of the distribution used to generate the data.  The match for"
print "covariance matrices is not that close, but would be improved with"
print "more iterations of the training algorithm."
print " "
raw_input('Press return to continue.')

clc()
print "We now plot the density given by the mixture model as a surface plot."
print " "
raw_input('Press return to continue.')

# Plot the result
x0=-4.0
x1=5.0
step=0.2
y0=x0
y1=x1

x = np.r_[x0:x1:step]
y = np.r_[y0:y1:step]
# 					
# Generate the grid
# 
X, Y = np.mgrid[x0:x1:step,y0:y1:step]

grid = np.c_[X.flatten(), Y.flatten()]
Z = mix.prob(grid)
Z = np.reshape(Z, (len(x), len(y))).T
s = mlab.surf(x, y, Z)
#mlab.show()
# c = mesh(x, y, Z);
# hold on
mlab.title('Surface plot of probability density')
# hold off
# drawnow

clc()
print "The final plot shows the centres and widths, given by one standard"
print "deviation, of the three components of the mixture model.  The axes"
print "of the ellipses of constant density are shown."
print " "
pp.show()
raw_input('Press return to continue.')
pp.ioff()

# Try to calculate a sensible position for the second figure, below the first
# fig1_pos = get(fh1, 'Position')
# fig2_pos = fig1_pos
# fig2_pos(2) = fig2_pos(2) - fig1_pos(4) - 30
# fh2 = figure('Position', fig2_pos)
fh2 = pp.figure()
h3 = pp.plot(data[:, 0], data[:, 1], 'bo')
pp.axis('equal')
pp.hold(True)
pp.title('Plot of data and covariances')
for i in range(ncentres):
    d, v = la.eig(mix.covars[:,:,i])
    for j in range(2):
        # Ensure that eigenvector has unit length
        v[:,j] = v[:,j]/la.norm(v[:,j])
        start=mix.centres[i,:]-np.sqrt(d[j])*(v[:,j].T)
        endpt=mix.centres[i,:]+np.sqrt(d[j])*(v[:,j].T)
        linex = np.array([start[0], endpt[0]])
        liney = np.array([start[1], endpt[1]])
        pp.plot(linex, liney, color='k', linewidth=3)

    # Plot ellipses of one standard deviation
    theta = np.r_[0:2*np.pi:0.02]
    x = np.sqrt(d[0])*np.cos(theta)
    y = np.sqrt(d[1])*np.sin(theta)
    # Rotate ellipse axes
    ellipse = np.dot(v,(np.c_[x, y].T)).T
    # Adjust centre
    ellipse = ellipse + mix.centres[i,:]
    pp.plot(ellipse[:,0], ellipse[:,1], 'r-')

pp.hold(False)

print "Note how the data cluster positions and widths are captured by"
print "the mixture model."
print " "
pp.show()
raw_input('Press return to end.')
pp.ioff()

pp.close(fh1)
pp.close(fh2)

