from numpy import *
from scipy import *
from pylab import *

coords = array(meshgrid(linspace(-5,5,200),linspace(-5,5,200)))
coords = transpose(coords,[1,2,0])
grid = coords.reshape([prod(coords.shape[:-1]),coords.shape[-1]])
def G(a): return a.reshape(coords.shape[:2])

def normal(x,a,sigma):
    invsigma = inv(sigma)
    d = len(x[0])
    delta = x-a[newaxis,:]
    raw = exp(-0.5*sum(delta*dot(invsigma,delta.T).T,axis=1))
    result = raw * (linalg.det(invsigma) / (2*pi)**d)**0.5
    return result

from numpy.random import multivariate_normal as N

mu0 = ones(2); sigma0= diag([3,1])
mu1 = -ones(2); sigma1 = diag([1,2])

samples0 = N(mu0,sigma0,1000)
samples1 = N(mu1,sigma1,1000)
samples0[:5]

figsize(6,6)
plot(samples0[:,0],samples0[:,1],'+')
plot(samples1[:,0],samples1[:,1],'+')

mean0 = mean(samples0,axis=0)
sigma0 = var(samples0,axis=0)
mean0,sigma0

mean1 = mean(samples1,axis=0)
sigma1 = var(samples1,axis=0)
mean1,sigma1

figsize(12,6); subplot(121)
imshow(-G(normal(grid,mu0,diag(sigma0))),origin='lower',extent=(-5,5,-5,5),cmap=cm.RdBu,vmin=-0.1,vmax=0.1)
subplot(122)
imshow(G(normal(grid,mu1,diag(sigma1))),origin='lower',extent=(-5,5,-5,5),cmap=cm.RdBu,vmin=-0.1,vmax=0.1)

imshow(-G(normal(grid,mu0,diag(sigma0)))+G(normal(grid,mu1,diag(sigma1))),
       origin='lower',extent=(-5,5,-5,5),cmap=cm.RdBu,vmin=-0.1,vmax=0.1)
plot(samples0[:,0],samples0[:,1],'r+')
plot(samples1[:,0],samples1[:,1],'b+')

imshow((G(normal(grid,mu0,diag(sigma0)))<G(normal(grid,mu1,diag(sigma1))))-0.5,
       origin='lower',extent=(-5,5,-5,5),cmap=cm.RdBu,vmin=-2,vmax=2)
plot(samples0[:,0],samples0[:,1],'r+')
plot(samples1[:,0],samples1[:,1],'b+')

# covariances
sigma0 = cov(samples0.T)
sigma1 = cov(samples1.T)
print sigma0
print sigma1

# full covariance matrices
imshow((G(normal(grid,mu0,sigma0))<G(normal(grid,mu1,sigma1)))-0.5,
       origin='lower',extent=(-5,5,-5,5),cmap=cm.RdBu,vmin=-2,vmax=2)
plot(samples0[:,0],samples0[:,1],'r+')
plot(samples1[:,0],samples1[:,1],'b+')