exec(open('tbc.py').read()) # define TBC and TBC_above
import dill

# may need to change the load path
TBC() # dill.load_session('../ignore/model_evaluation.db')

exec(open('tbc.py').read()) # (re-)define TBC and TBC_above

# additional imports go here
from scipy.special import logsumexp
import dynesty
from dynesty import plotting as dyplot

"""
Explanation: Tutorial: Model Selection
In the previous tutorial, we evaluated the "goodness of fit" of two different models on a data set representing the distribution of centering offsets in a galaxy cluster sample. Here we will continue by doing two styles of quantitative test to compare the two models, with the goal of deciding whether the improved goodness of fit actually justifies the additional complexity of the second model. Specifically, we will compute:
1. The Deviance Information Criterion, one of several possible information criteria, which has the advantages of being relatively simple and having a straightforward Bayesian interpretation.
2. The Bayesian Evidence, a more principled but more complex approach fully in the Bayesian framework.
First, we'll restore everything from the last notebook, including the model class definitions and instances, and the Markov chains they hold.
End of explanation
"""


def DIC(Model):
    """
    Compute the Deviance Information Criterion for the given model.
    (In a less pedagogical world, this would logically be a method of the base Model class.)
    """
    # Compute the deviance D for each sample
    D = -2.0*np.array([ Model.log_likelihood(*params) for params in Model.samples ])
    #pD = 
    #DIC = 
    #return DIC, pD

TBC_above(1)

"""
Explanation: 1. Calculate the DIC for each model
Recall that the Deviance Information Criterion is given by:
$\mathrm{DIC} = \langle D(\theta) \rangle + p_D; \quad p_D = \langle D(\theta) \rangle - D(\langle\theta\rangle)$
where $\theta$ are the parameters of a model, the deviance $D(\theta)=-2\log P(\mathrm{data}|\theta)$, and averages $\langle\rangle$ are over the posterior distribution of $\theta$.
Write this function, and execute it for the simple model.
End of explanation
"""


DIC1, pD1 = DIC(Model1)
print('Model 1:')
print("Effective number of fitted parameters =", pD1)
print("DIC =", DIC1)

"""
Explanation: Compute the DIC for each model.
End of explanation
"""


DIC2, pD2 = DIC(Model2)
print('Model 2:')
print("Effective number of fitted parameters =", pD2)
print("DIC =", DIC2)

"""
Explanation: Self-check: For Model 1 (the exponential), I get $p_D \approx 0.9$ and DIC $\approx 1668$. As with anything else computed from chains, there will be some stochasticity to the values you compute.
End of explanation
"""


DIC1 - DIC2

"""
Explanation: Do your values of $p_D$ make intuitive sense?

TBC

Now, to interpret this, we compare the reduction (hopefully) in the DIC of Model 2 compared with Model 1 to the Jeffreys scale (see the notes). By this metric, is your second model better at explaining the data than the exponential model?
End of explanation
"""


def log_evidence(Model, N=1000):
    """
    Compute the log evidence for the model using N samples from the prior
    """
    TBC()
    
TBC_above()

"""
Explanation: TBC: your answer

2. Compute the evidence by Monte Carlo integration
To do this, note that
$p(\mathrm{data}|H)=\int d\theta \, p(\mathrm{data}|\theta,H) \, p(\theta|H)$
can be approximated by an average over samples from the prior:
$p(\mathrm{data}|H) \approx \frac{1}{m}\sum_{k=1}^m p(\mathrm{data}|\theta_k,H)$, with $\theta_k\sim p(\theta|H)$.
This estimate is better than trying to use samples from the posterior to calculate the evidence, if only because it's unbiased. But in general, and especially for large-dimensional parameter spaces, it is very inefficient (because the likelihood typically is large in only a small fraction of the prior volume). Still, let's give it a try.
Write a function to draw a large number of samples from the prior and use them to calculate the evidence. To avoid numerical over/underflows, use the special scipy function logsumexp (which we imported directly, way at the top of the notebook) to do the sum. As the name implies, this function is equivalent to log(sum(exp(...))), but is more numerically stable.
End of explanation
"""


log_evidence(Model1, N=2), log_evidence(Model2, N=2)

"""
Explanation: Do a quick test to check for NaNs:
End of explanation
"""


print('Model 1:')
for Nevidence in [1, 10, 100, 1000, 10000]: # Do you need even more than this?
    %time logE1 = log_evidence(Model1, N=Nevidence)
    print("From", Nevidence, "samples, the log-evidence is", logE1)

"""
Explanation: Roughly how precisely do we need to know the log Evidence, to be able to compare models? Run log_evidence with different values of N (the number of prior samples in the average) to until you're satisfied that you're getting a usefully accurate result for each model.
End of explanation
"""


print('Model 2:')
for Nevidence in [1, 10, 100, 1000, 10000]: # Do you need even more than this?
    %time logE2 = log_evidence(Model2, N=Nevidence)
    print("From", Nevidence, "samples, the log-evidence is", logE2)

"""
Explanation: Self-check: For the exponential model, I converge to a log-evidence of $\approx -837$.
End of explanation
"""


logE2 - logE1

"""
Explanation: So, we have log evidences computed for each model. Now what? We just compare their difference to the Jeffreys scale again:
End of explanation
"""


prior_H1 = 0.9 # or your choice
prior_H2 = 1.0 - prior_H1 # assuming only these two options

log_post_H1 = np.log(prior_H1) + logE1
log_post_H2 = np.log(prior_H2) + logE2

print('Difference of log posteriors (H2-H1):', log_post_H2-log_post_H1)
print('Ratio of posteriors (H2/H1):', np.exp(log_post_H2-log_post_H1)) # NB this one might over/underflow

"""
Explanation: Note that we might end up with a different conclusion as to the strength of any preference of Model 2, compared with the DIC! The reason is that the evidence explicitly accounts for the information in the prior (which, recall, counts as part of the model definition), while the DIC does this much less directly.
We could also be good Bayesians and admit that there should be a prior distribution in model space. For example, maybe I have a very compelling theoretical reason why the offset distribution should be exponential (I don't, but just for example). Then, I might need some extra convincing that an alternative model is required.
We would then compute the ratio of the posterior probabilities of the models as follows:
End of explanation
"""


def ptform(u):
    '''
    Input: a vector in the unit cube, 0 <= u[i] <= 1, that conforms to the prior volume.
    Output: a vector in our parameter space.
    '''
    TBC()
    
TBC_above()

"""
Explanation: Depending on the fitness of the alternative model you chose, you may find that only an extremely lopsided prior in model space would influence your conclusion.
Comment on what you find from the evidence, compared with the DIC.

TBC: your answer

3. Compute the evidence with dynesty
To get some experience with a package that uses nested sampling to compute the evidence, let's repeat Section 2 using dynesty.
Looking at the docs, we first need a function that maps the unit cube onto our prior. For a uniform prior, this is a simple translation and rescaling, but for other priors it involves going through the prior distribution's quantile function. Fortunately, we assumed uniform priors for Model 1! However, if you followed the advice in the previous notebook to keep a dictionary of univariate priors in the Model object, and those priors are scipy.stats distributions, it's relatively easy to handle this slightly more general case. However you choose to do it, implement a function that performs this transformation for the 1-parameter Exponential model (and priors) below.
End of explanation
"""


print(ptform([0.25]), ptform([0.5]), ptform([0.75]))

"""
Explanation: Do a sanity check - the cell below should return the quartiles of the uniform prior.
End of explanation
"""


def dyn_log_like(params):
    return Model1.log_likelihood(*params)

"""
Explanation: We'll also need a log-likelihood function that takes a vector of parameters as input, in contrast to how it's written in the Model class. This could be solved similarly to what's done in the Model.log_posterior method, or we could just define:
End of explanation
"""


%%time
sampler = dynesty.NestedSampler(dyn_log_like, ptform, len(Model1.param_names))
sampler.run_nested()
results = sampler.results

"""
Explanation: Then we can just go ahead and run.
End of explanation
"""


# Plot a summary of the run.
rfig, raxes = dynesty.plotting.runplot(results)

"""
Explanation: And look at some stuff (see the Dynesty documentation for an explanation of what all this is):
End of explanation
"""


plt.rcParams['figure.figsize'] = (12.0, 4.0)
Ndiscard = 100
plt.plot(results['logz'][Ndiscard:]);
plt.xlabel("Iteration - "+str(Ndiscard), fontsize=14);
plt.ylabel("Log Evidence", fontsize=14);

"""
Explanation: The Evidence panel appears to be utterly unhelpful in this case (at least when I ran it), so below is a plot of the log-evidence as a function of iteration, discarding the beginning when the accumulated evidence is truly tiny. It should show us converging to a value similar to what we found earlier.
End of explanation
"""


results['logz'][-1]

"""
Explanation: This extracts the log-evidence at the last iteration, which may or not be the approved method.
End of explanation
"""


logE1 - results['logz'][-1]

"""
Explanation: Let's compare with what we got from simple monte carlo:
End of explanation
"""


def ptform2(u):
    '''
    Input: a vector in the unit cube, 0 <= u[i] <= 1, that conforms to the prior volume.
    Output: a vector in our parameter space.
    '''
    TBC()
    
TBC_above()

"""
Explanation: That was so easy, let's go ahead and do the same for Model 2. Implement the transformation from the unit cube to your model's parameter space:
End of explanation
"""


def dyn_log_like2(params):
    return Model2.log_likelihood(*params)

"""
Explanation: And run this elegant definition:
End of explanation
"""


%%time
sampler = dynesty.NestedSampler(dyn_log_like2, ptform2, len(Model2.param_names))
sampler.run_nested()
results = sampler.results

"""
Explanation: Now we can go ahead and run the sample.
End of explanation
"""


rfig, raxes = dynesty.plotting.runplot(results)

plt.rcParams['figure.figsize'] = (12.0, 4.0)
Ndiscard = 100
plt.plot(results['logz'][Ndiscard:]);
plt.xlabel("Iteration - "+str(Ndiscard), fontsize=14);
plt.ylabel("Log Evidence", fontsize=14);

"""
Explanation: ... and produce the same set of plots.
End of explanation
"""


results['logz'][-1]

"""
Explanation: Finally, extract the final evidence,
End of explanation
"""


logE2 - results['logz'][-1]

"""
Explanation: ... and compare it to the SMC result.
End of explanation
"""


