#!/usr/bin/env python
# coding: utf-8

# ## Inference: Differential Evolution (DE) MCMC
# 
# This example shows you how to perform Bayesian inference on a time series, using [Differential Evolution MCMC](http://pints.readthedocs.io/en/latest/mcmc_samplers/differential_evolution_mcmc.html).
# 
# It follows on from the [first sampling example](./sampling-first-example.ipynb).

# In[3]:


from __future__ import print_function
import pints
import pints.toy as toy
import pints.plot
import numpy as np
import matplotlib.pyplot as plt

# Load a forward model
model = toy.LogisticModel()

# Create some toy data
real_parameters = [0.015, 500]
times = np.linspace(0, 1000, 1000)
org_values = model.simulate(real_parameters, times)

# Add noise
noise = 10
values = org_values + np.random.normal(0, noise, org_values.shape)
real_parameters = np.array(real_parameters + [noise])

# Get properties of the noise sample
noise_sample_mean = np.mean(values - org_values)
noise_sample_std = np.std(values - org_values)

# Create an object with links to the model and time series
problem = pints.SingleOutputProblem(model, times, values)

# Create a log-likelihood function (adds an extra parameter!)
log_likelihood = pints.GaussianLogLikelihood(problem)

# Create a uniform prior over both the parameters and the new noise variable
log_prior = pints.UniformLogPrior(
    [0.01, 400, noise * 0.1],
    [0.02, 600, noise * 100],
)

# Create a posterior log-likelihood (log(likelihood * prior))
log_posterior = pints.LogPosterior(log_likelihood, log_prior)

# Choose starting points for 3 mcmc chains
xs = [
    real_parameters * 1.1,
    real_parameters * 0.9,
    real_parameters * 1.15,
    real_parameters * 1.15,
    real_parameters * 1.05,
    real_parameters * 0.85,
]

# Create mcmc routine
mcmc = pints.MCMCController(log_posterior, 6, xs, method=pints.DifferentialEvolutionMCMC)

# Add stopping criterion
mcmc.set_max_iterations(8000)

# Disable logging
mcmc.set_log_to_screen(False)

# Set number of iterations between which gamma=1 (for a single iteration)
mcmc.sampler().set_gamma_switch_rate(2000)

# Set to uniform error process rather than normal
mcmc.sampler().set_gaussian_error(False)

# Run!
print('Running...')
chains = mcmc.run()
print('Done!')

# Check convergence using rhat criterion
print('R-hat:')
print(pints.rhat_all_params(chains))

# Show traces and histograms
pints.plot.trace(chains)

# Discard warm up
chains = chains[:, 5000:, :]

# Apply thinning
chains = chains[:, ::10, :]

# Look at distribution in chain 0
pints.plot.pairwise(chains[0], kde=True)

# Show graphs
plt.show()