"""An example of plot2Ddist.

Originally based on an example by Anand Patil:
http://groups.google.com/group/pymc/browse_thread/thread/afe91c4d9440d6da

"""

import pymc as pm
import numpy as np
import matplotlib.pyplot as pl

# Set up parameters. Our model is y = a + b x + c x^2
a = pm.Uninformative('a',0)
b = pm.Uninformative('b',0)
c = pm.Uninformative('c',0)
v = pm.OneOverX('v',1)

# Locations of observations.
nx = 50
span = 5.0
x = np.linspace(-span/2,span/2,nx)

# Generate mock set of data.
slope = 2.0
intercept = 1.0
curve = 3.0
sigma = 3.0
data = (x**2)*curve + x*slope + intercept + np.random.normal(size=len(x))*sigma

# Observed variable.
y = pm.Normal('y', a+b*x+c*x**2, 1./v,value=data,observed=True)
ypred = pm.Normal('ypred',a+b*x+c*x**2,1./v)

# Set up MCMC chain.
M = pm.MCMC([a,b,c,v,y,ypred])
M.use_step_method(pm.AdaptiveMetropolis, [a,b,c,v], tally=True)

# Sample.
M.isample(20000,5000,10)

# Plot the function and envelope.
pl.figure()
env = pm.Matplot.func_envelopes(ypred)
for e in env:
    e.display(x,.5,new=False)
pl.plot(x,data,'r.',markersize=4) 
pl.plot(x,data,'r.',markersize=4) 

# Plot the parameter distribution from the samples.
import plot2Ddist
scatterstyle={'color':'r', 'alpha':0.1}
styleargs = {'color':'k', 'scatterstyle':scatterstyle}
truevalues = dict(b = slope,
                  a = intercept,
                  c = curve,
                  v = sigma**2.)
#plot2Ddist.plot2DdistsAllPairs([a,b,c,v], truevalues=truevalues, 
#                               **styleargs)

plot2Ddist.plot2DdistsPairs(a, [b,c,v], truevalues=truevalues, 
                            **styleargs)
#plot2Ddist.plot_AM_correlation(M)
plot2Ddist.plot_AM_correlation(M, variables=[a,v])

pl.show()