#!/usr/bin/env python

"""
Multivariate gaussian mixture model.

Demonstrates bimodal sampling from a multidimensional space.

Change the relative weights of the modes to see how effectively DREAM samples
from lower probability regions.  In particular, with low *w2*, long chains
(thinning*generations*cycles) and small population, the second mode can
get lost.
"""
from pylab import *
from dream import *

n = 4
pop = 10
w1,w2 = 5,3
mu1 = -5 * ones(n)
mu2 = 5 * ones(n)
#mu2[0] = -3  # Watch marginal distribution for p1 overlap between modes
sigma = eye(n)
model = Mixture(MVNormal(mu1,sigma), w1, MVNormal(mu2,sigma), w2)
#model = MVNormal(zeros(n),sigma)

# TODO: with large number of samples, the 1/6 weight peak is lost
sampler = Dream(model=model, population=randn(pop,n,n),
                #use_delayed_rejection=False,
                #outlier_test='IQR',
                thinning=1, draws=20000)
state = sampler.sample()
save_state(filename='mixture',state=state)
state = load_state('mixture')
plot_all(state, portion=1)
show()