#!/usr/bin/env python

"""
Example model with strong correlations between the fitted parameters.

We use a*x = y + N(0,1) made complicated by defining a=p1+p2.

The expected distribution for p1 and p2 will be uniform, with p2 = a-p1 in
each sample.  Because this distribution is inherently unbounded, artificial
bounds are required on a least one of the parameters for finite duration
simulations.

The expected distribution for p1+p2 can be determined from the linear model
y = a*x.  This is reported along with the values estimated from MCMC.
"""

from dream import *  # sampler functions
from pylab import *  # Numeric functions and plotting

# Create the correlation function and generate some fake data
x = linspace(-1.0, 1, 40)
fn = lambda p: sum(p) * x
bounds = (-20, -inf), (40, inf)
sigma = 1
data = fn((1, 1)) + randn(*x.shape) * sigma  # Fake data


# Sample from the posterior density function
n = 2
model = Simulation(f=fn, data=data, sigma=sigma, bounds=bounds, labels=["x", "y"])
sampler = Dream(
    model=model,
    population=randn(5 * n, 4, n),
    thinning=1,
    draws=20000,
)
mc = sampler.sample()
mc.title = "Strong anti-correlation"

# Create a derived parameter without the correlation
mc.derive_vars(lambda p: (p[0] + p[1]), labels=["x+y"])

# Compare the MCMC estimate for the derived parameter to a least squares fit
from bumps.wsolve import wpolyfit

poly = wpolyfit(x, data, degree=1, origin=True)
print("x+y from linear fit", poly.coeff[0], poly.std[0])
points, logp = mc.sample(portion=0.5)
print("x+y from MCMC", mean(points[:, 2]), std(points[:, 2], ddof=1))

# Plot the samples
plot_all(mc, portion=0.5)
show()
