exec(open('tbc.py').read()) # define TBC and TBC_above
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
%matplotlib inline
from glob import glob
import incredible as cr

"""
Explanation: Tutorial: MCMC Diagnostics
You should already have run two different MCMC algorithms to generate chains for the AGN photometry problem (or an approximation thereof). Let's work through the process of diagnosing whether these Markov chains are usefully sampling the posterior distribution.
The diagnostics discussed below include both qualitative and quantitative checks. We don't particularly think it's all that instructive to write the code that does the quatitative calculations - though there is surely room for improvement or expansion if you're interested - so intead we will demonstrate how to use functions provided by the incredible and pandas packages.
End of explanation
"""


TBC() # change path if need be
# chains = [np.loadtxt(f) for f in glob('../ignore/agn_gibbs_chain_*.txt')]

param_labels = [r'$\mu$', r'$\sigma$', r'$x_0$', r'$y_0$']

"""
Explanation: Gibbs samples
It's nice to start with the Gibbs sampled chains, since they almost certainly look nicer. First, read them in.
Read in your conjugate Gibbs chains.
End of explanation
"""


plt.rcParams['figure.figsize'] = (16.0, 12.0)
fig, ax = plt.subplots(len(param_labels), 1);
cr.plot_traces(chains, ax, labels=param_labels, Line2D_kwargs={'markersize':1.0})

"""
Explanation: Visual inspection
You've already used the most important method of vetting chains: visual inspection. The key questions are:
1. Do multiple, independent chains appear to be sampling the same distribution (have they converged to the same distribution)?
2. Is there a clear "burn-in" period before convergence that should be eliminated?
3. Are the chains highly autocorrelated (taking small steps compared with the width of the posterior)? This is not an issue per se, if the chains are long enough, although it means the sampler is not moving as efficiently as one might hope.
Plot the parameter traces below, and answer these questions (qualitatively) for the conjugate Gibbs sampling chains.
End of explanation
"""


burn = 0
for i in range(len(chains)):
    chains[i] = chains[i][range(burn, chains[i].shape[0]),:]

"""
Explanation: TBC
TBC
TBC

If you thought some burn-in should be removed, do so here by changing the lower limit of burn.
End of explanation
"""


cr.GelmanRubinR(chains)

"""
Explanation: Gelman-Rubin statistic
Recall from the notes that the Gelman-Rubin convergence statistic, $R$, quantitatively tests the similarlity of independent chains intended to sample the same PDF. To be meaningful, they should start from different locations and burn-in should be removed.
For a given parameter, $\theta$, the $R$ statistic compares the variance across chains with the variance within a chain. Intuitively, if the chains are random-walking in very different places, i.e. not sampling the same distribution, $R$ will be large.
We'd like to see $R\approx 1$; for example, $R<1.1$ is often used.
End of explanation
"""


plt.rcParams['figure.figsize'] = (16.0, 12.0)
fig, ax = plt.subplots(len(param_labels), 1);
for j,lab in enumerate(param_labels):
    pd.plotting.autocorrelation_plot(chains[0][:,j], ax=ax[j]);
    ax[j].set_ylabel(lab+' autocorrelation')

"""
Explanation: Checkpoint: If your Gibbs sampler works properly, $R$ for the chains we ran should be very close to 1 (we have differences of order 0.00001).
Autocorrelation
Similarly, the autocorrelation of a sequence, as a function of lag, $k$, can be quantified:
$\rho_k = \frac{\mathrm{Cov}i\left(\theta_i,\theta{i+k}\right)}{\mathrm{Var}(\theta)}$
The larger lag one needs to get a small autocorrelation, the less informative individual samples are.
The pandas function plotting.autocorrelation_plot() is useful for this. Note that seemingly random oscillations basically tell us the level of noise due to the fininte chain length. A coherent drop as a function of lag indicates a genuine autocorrelation, and the lag at which it drops to within the noise is an approximate autocorrelation length. If we needed to thin the chains to conserve disk space, this would be a reasonable factor to thin by.
End of explanation
"""


cr.effective_samples(chains, maxlag=500) # `maxlag' might be something you need to play with, in practice

"""
Explanation: Checkpoint: Again, for this problem, the Gibbs chains should be very well behaved. Our autocorrelation plots basically look like noise (almost all points within the horizontal, dashed lines that pandas provides as an estimate of the noise).
Effective number of independent samples
From $m$ chains of length $n$, we can also estimate the "effective number of independent samples" as
$n_\mathrm{eff} = \frac{mn}{1+2\sum_{t=1}^\infty \hat{\rho}_t}$, with
$\hat{\rho}_t = 1 - \frac{V_t}{2V}$ ($V$ as in the Gelman-Rubin calculation), and
$V_t = \frac{1}{m(n-t)} \sum_{j=0}^m \sum_{i=t+1}^n (\theta_{i,j} - \theta_{i-t,j})^2$.
In practice, the sum in $n_\mathrm{eff}$ is cut off when the estimates $\hat{\rho}t$ become "too noisy", e.g. when the sum of two successive values $\hat{\rho}_t$ and $\hat{\rho}{t+1}$ is negative. Roughly speaking, this should occur when the lag is of the order of the autocorrelation length.
The effective_samples function allows you to pass a guess at this maximum lag, since doing the calculation to arbitrarily long lags becomes very expensive. It will issue a warning if it thinks this maximum lag is too small, according to the criterion above.
End of explanation
"""


chain = np.concatenate(chains, axis=0)

"""
Explanation: As with the Gelman-Rubin statistic, this is a case where one might be interested in seeing the effective number of samples for the most degenerate linear combinations of parameters, rather than the parameters themselves.
Something to do
By now you are probably bored. Don't worry. Here is some work for you to do.
Let's get a sense of how many samples are really needed to, e.g., determine 1D credible intervals (as opposed to making the whole posterior look nice). Remember that the effective number of samples is less than the total, obviously.
At this point, we're done comparing the individual chains, so we can lump them all together into one massive list of MCMC samples.
End of explanation
"""


print(chain.shape[0], 'samples')
plt.rcParams['figure.figsize'] = (14.0, 5.0)
fig, ax = plt.subplots(1, 2);
h40k = cr.whist(chain[:,0], plot=ax[0])
ci40k = cr.whist_ci(h40k, plot=ax[1]);
ax[0].set_xlabel(r'$\mu$');
ax[1].set_xlabel(r'$\mu$');
ci40k

"""
Explanation: Let's have a look at the credible interval calculation for the first parameter. If you followed the notebooks as given, and didn't remove any burn-in, the full chain should be of length 40,000.
End of explanation
"""


TBC()
# No clues here, but it's pretty much cut and paste.
# Analogous to the cell above, save the output of `whist` in h10k, h1k, h100, and the output of
# whist_ci in ci10k, ci1k and ci100. This is so we can plot them all togther later.

"""
Explanation: The PDF estimate should look pretty reliable with so many samples. The question is, if we're going to reduce this to a statement like $\mu=X^{+Y}_{-Z}$, keeping only up to the leading significant figure of $Y$ and $Z$, how many did we actually need to keep?
Thin the chain by factors of 4, 40, and 400 (to produce chains of length about 10000, 1000 and 100), and see how the endpoints of the 68.3% credible intervals compare. We're looking at the endpoints rather than the values of $Y$ and $Z$ above because the latter are more volatile (depending also on the estimate of $X$).
Remember that thinning by a factor of 4 means that we keep only every 4th entry in the chain, not that we simply select the first 25% of samples. So we're not answering how long we needed to bother running the chain to begin with - that's a slightly different question. We're finding out how redundant our samples are, not just in the "effective independence" sense, but for the specific purpose of quantifying this credible interval.
End of explanation
"""


plt.rcParams['figure.figsize'] = (14.0, 5.0)
fig, ax = plt.subplots(1, 2);
ax[0].plot(h40k['x'], h40k['density'], '-', label='40k');
ax[0].plot(h10k['x'], h10k['density'], '-', label='10k');
ax[0].plot(h1k['x'], h1k['density'], '-', label='1k');
ax[0].plot(h100['x'], h100['density'], '-', label='100');
ax[0].legend();
ax[0].set_xlabel(r'$\mu$');
ax[1].plot(0.0, ci40k['mode'], 'o', color='C0', label='40k');
ax[1].plot([0.0]*2, [ci40k['min'][0],ci40k['max'][0]], '-', color='C0', linewidth=3);
ax[1].plot([0.0]*2, [ci40k['min'][1],ci40k['max'][1]], '--', color='C0');
ax[1].plot(1.0, ci10k['mode'], 'o', color='C1', label='10k');
ax[1].plot([1.0]*2, [ci10k['min'][0],ci10k['max'][0]], '-', color='C1', linewidth=3);
ax[1].plot([1.0]*2, [ci10k['min'][1],ci10k['max'][1]], '--', color='C1');
ax[1].plot(2.0, ci1k['mode'], 'o', color='C2', label='1k');
ax[1].plot([2.0]*2, [ci1k['min'][0],ci1k['max'][0]], '-', color='C2', linewidth=3);
ax[1].plot([2.0]*2, [ci1k['min'][1],ci1k['max'][1]], '--', color='C2');
ax[1].plot(3.0, ci100['mode'], 'o', color='C3', label='100');
ax[1].plot([3.0]*2, [ci100['min'][0],ci100['max'][0]], '-', color='C3', linewidth=3);
ax[1].plot([3.0]*2, [ci100['min'][1],ci100['max'][1]], '--', color='C3');
ax[1].legend();
ax[1].set_ylabel(r'$\mu$');

"""
Explanation: Checkpoint: Your mileage may vary, of course. But we got a difference of unity in one endpoint in the 100-sample case, and otherwise everything was identical.
... which is a little surprising, honestly, even though we knew the autocorrelation was quite low in this case. But here's a slightly different question: which of the possible results would you be confident enough to put in a paper? The cell below compares them visually.
End of explanation
"""


TBC() # change path if necessary
# chains = [np.loadtxt(f) for f in glob('../ignore/agn_metro_chain_*.txt')]

param_labels = [r'$x_0$', r'$y_0$', r'$\ln F_0$', r'$b$', r'$\sigma$']

"""
Explanation: Metropolis samples
Now, read in the Metropolis chains and perform the same checks.
End of explanation
"""


plt.rcParams['figure.figsize'] = (16.0, 12.0)
fig, ax = plt.subplots(len(param_labels), 1);
cr.plot_traces(chains, ax, labels=param_labels, Line2D_kwargs={'markersize':1.0})

"""
Explanation: Below we plot the traces. Address the same 3 questions posed for the Gibbs samples.
End of explanation
"""


TBC()
# burn = 
for i in range(len(chains)):
    chains[i] = chains[i][range(burn, chains[i].shape[0]),:]

"""
Explanation: TBC
TBC
TBC

Compare the two methods in these terms. (Though keep in mind that we solved slightly different problems in the two notebooks, making this comparison less than entirely fair. Or, go back and run Metropolis sampling on a background-free simulation if you really want to. We'll wait.)

Commentary TBC

On the basis of the traces above, choose a burn-in length to remove from the beginning of each chain.
End of explanation
"""


cr.GelmanRubinR(chains)

"""
Explanation: Here we compute the G-R criterion. Do the values make sense in light of your visual inspection?
End of explanation
"""


plt.rcParams['figure.figsize'] = (16.0, 12.0)
fig, ax = plt.subplots(len(param_labels), 1);
for j,lab in enumerate(param_labels):
    pd.plotting.autocorrelation_plot(chains[0][:,j], ax=ax[j]);
    ax[j].set_ylabel(lab+' autocorrelation')

"""
Explanation: Commentary TBC

Next, we'll look at the autocorrelation plot. If you had to guess an autocorrelation length, what would it be?
End of explanation
"""


cr.effective_samples(chains, maxlag=1000)

"""
Explanation: Commentary TBC

Next, the effective number of samples. How does it compare to the Gibbs case?
End of explanation
"""


