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

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

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

"""
Explanation: Tutorial: The Cepheid Period-Luminosity Relation for Multiple Galaxies
So far (in the cepheids and cepheids_one_galaxy notebooks), we have fit a hierarchical model describing the period-luminosity relation and its intrinsic scatter to data from a single galaxy. Next, we're interested in how similar, or not, these scaling relation parameters ($a$, $b$ and $\sigma$) are among the galaxies in the data set.
A sensible place to start is to do an identical fit independently to each galaxy to see how compatible the parameter values are. Then we'll fit a model with another level of hierarchy that assumes these parameters come from a parent distribution. The width of that parent distribution will be a parameter that tells us how similar things are from galaxy to galaxy.
Start by restoring the previous notebook:
End of explanation
"""


ngc_numbers

"""
Explanation: 1. Data
We should now have all the data loaded, named as it was before. As a reminder, these are the NGC numbers of the galaxies in the data set:
End of explanation
"""


class singleFitter:
    def __init__(self, ngc):
        '''
        ngc: NGC identifier of the galaxy to fit
        '''
        self.ngc = ngc
        self.data = data[ngc] # from global scope
        # reproducing this for paranoia's sake
        self.param_names = ['a', 'b', 'sigma']
        self.param_labels = [r'$a$', r'$b$', r'$\sigma$']
    def _logpost_vecarg_B(self, pvec):
        params = {name:pvec[i] for i,name in enumerate(self.param_names)}
        return log_posterior(self.data, log_likelihood_B, **params)
    def fit(self, guess, nsteps=7500):
        npars = len(self.param_names)
        nwalkers = 2*npars
        sampler = emcee.EnsembleSampler(nwalkers, npars, self._logpost_vecarg_B)
        start = np.array([np.array(guess)*(1.0 + 0.01*np.random.randn(npars)) for j in range(nwalkers)])
        %time sampler.run_mcmc(start, nsteps)
        plt.rcParams['figure.figsize'] = (16.0, 3.0*npars)
        fig, ax = plt.subplots(npars, 1);
        cr.plot_traces(sampler.chain[:min(8,nwalkers),:,:], ax, labels=self.param_labels);
        self.sampler = sampler
        self.nwalkers = nwalkers
        self.npars = npars
        self.nsteps = nsteps
    def burnin(self, burn=1000, maxlag=1000):
        tmp_samples = [self.sampler.chain[i,burn:,:] for i in range(self.nwalkers)]
        print('R =', cr.GelmanRubinR(tmp_samples))
        print('neff =', cr.effective_samples(tmp_samples, maxlag=maxlag))
        print('NB: Since walkers are not independent, these will be optimistic!')
        self.samples = self.sampler.chain[:,burn:,:].reshape(self.nwalkers*(self.nsteps-burn), self.npars)
        del self.sampler
        # make it simpler/more readable to access the parameter samples
        # (could have been fancier and more robust by using self.param_names here)
        self.a = self.samples[:,0]
        self.b = self.samples[:,1]
        self.sigma = self.samples[:,2]
    def thin(self, thinto=1000):
        j = np.round(np.linspace(0, self.samples.shape[0]-1, thinto)).astype(int)
        self.a = self.samples[j,0]
        self.b = self.samples[j,1]
        self.sigma = self.samples[j,2]

"""
Explanation: 2. Independent fits for each galaxy
This class will package up the fitting using the "4b" method from the previous notebook (emcee plus analytic integration). In particular, it relies on the log_prior, log_posterior and log_likelihood_B functions (as well as the data, among other previous global-scope definitions). If you want to use a different approach instead, feel free.
There are various defaults here (e.g. nsteps, burn, maxlag) that you might want to tweak, but in principle they should work well enough for this problem.
End of explanation
"""


independent_fits = [singleFitter(ngc) for ngc in ngc_numbers]

independent_fits[0].fit(guessvec)

independent_fits[1].fit(guessvec)

independent_fits[2].fit(guessvec)

independent_fits[3].fit(guessvec)

independent_fits[4].fit(guessvec)

independent_fits[5].fit(guessvec)

independent_fits[6].fit(guessvec)

independent_fits[7].fit(guessvec)

independent_fits[8].fit(guessvec)

"""
Explanation: Let's set up and run each of these fits, which hopefully shouldn't take too long. As always, you are responsible for looking over the trace plots and making sure everything is ok.
End of explanation
"""


TBC(1) # burn = ...

for f in independent_fits:
    print('NGC', f.ngc)
    f.burnin(burn=burn) # optionally, set maxlag here also
    print('')

"""
Explanation: Based on the plots above, remove some burn-in. Check that the quantitative diagnostics are acceptable as they are printed out.
End of explanation
"""


plotGTC([f.samples for f in independent_fits], paramNames=param_labels, 
        chainLabels=['NGC'+str(f.ngc) for f in independent_fits],
        figureSize=8, customLabelFont={'size':12}, customTickFont={'size':12}, customLegendFont={'size':16});

"""
Explanation: Now we'll use pygtc to plot all the individual posteriors, and see how they compare.
End of explanation
"""


param_names_all  = ['mu_a', 'tau_a', 'mu_b', 'tau_b', 'mu_s', 'tau_s']
param_labels_all = [r'$\mu_a$', r'$\tau_a$', r'$\mu_b$', r'$\tau_b$', r'$\mu_\sigma$', r'$\tau_\sigma$']

"""
Explanation: Visually, would you say that it's likely that all the scaling parameters, or some subset, are universal?

TBC commentary

2. A hierarchical model for all galaxies
On the basis of the last section, it should be clear that at least one of the scaling parameters in question is not universal amongst galaxies in the data set, and at least one may well be. Further, it isn't obvious that there is any particular correlation or anticorrelation between the galaxy-to-galaxy differences in these parameters. If we were doing this as a research project, the latter would be an important thing to investigate, along with possible physical explanations for outliers. But we'll keep it relatively simple here.
Let's add a level of hierarchy to the model by assuming that the values of $a$ for each galaxy come from a normal distribution with mean $\mu_a$ and standard deviation $\tau_a$, and similarly $b$ and $\sigma$ come from their own normal distributions. We will not consider the possibility that, for example, all 3 come from a joint, multivariate normal distribution with possible correlations between them, although that could easily be justified. In practice, fitting for independent distributions for each parameter is a reasonable first step, much as fitting each galaxies data independently in Section 1 was a reasonable zeroth step.
Make the relatively simple modifications to your PGM and probabilistic expressions from Section 2 of the previous notebook to accomodate this model.

TBC probabilistic expressions and PGM

We will adopt wide, uniform priors on the new hyperparameters of the model, to make life easier.
3. Strategy
Even more than last time, the total number of free parameters in the model is, technically, staggering. We already know some ways of reducing the overhead associated with each galaxy. For example, using the analytic integration approach from the previous notebook, we have have only 3 parameters left to sample per galaxy, for a total of $3N_\mathrm{gal}+6=33$ parameters. Brute force sampling of these 33 parameters is not unthinkable, although in practice it may or may not be a headache.
Another option is to make use of the samples we obtained in Section 1. These are samples of the posterior (for each galaxy) when the priors on the scaling parameters are very wide and uniform, i.e. constant over the domain where the likelihood is significantly non-zero. They are, therefore, also samples from a PDF that is proportional to the likelihood function. To see why that might be helpful, consider the posterior for the hyperparameters of the new model, $\vec{\alpha} = (\mu_a,\tau_a,\mu_b,\tau_b,\mu_\sigma,\tau_\sigma)$, marginalized over all the pesky $a_i$, $b_i$ and $\sigma_i$ parameters:
$p(\vec{\alpha}|\mathrm{data}) \propto p(\vec{\alpha}) \prod_{i=1}^{N_\mathrm{gal}} \int da_i db_i d\sigma_i \, p(a_i,b_i,\sigma_i|\vec{\alpha}) \, p(\mathrm{data}|a_i,b_i,\sigma_i)$.
To restate what we said above, our individual fits (with uniform priors) give us samples from PDFs
$q(a_i,b_i,\sigma_i|\mathrm{data}) \propto p(\mathrm{data}|a_i,b_i,\sigma_i)$.
We can do this integral by simple monte carlo as
$p(\vec{\alpha}|\mathrm{data}) \propto p(\vec{\alpha}) \prod_{i=1}^{N_\mathrm{gal}} \frac{1}{n_i}\sum_{k=1}^{n_i} p(a_{ik},b_{ik},\sigma_{ik}|\vec{\alpha})$,
where the $n_i$ samples of $(a_{ik},b_{ik},\sigma_{ik}) \sim q(a_i,b_i,\sigma_i|\mathrm{data})$. Our samples from Section 1 happen to satisfy this. (Had we used a non-uniform prior before, we could do something similar, but would need to divide by that prior density in the sum above.) This approach has the advantage that we only need to sample the 6 parameters in $\vec{\alpha}$ to constrain our hierarchical model, since a lot of work is already done. On the other hand, carrying out the sums for each galaxy can become its own numerical challenge.
If we're really stuck in terms of computing power, we could consider a more compressed version of this, by approximating the posterior from each individual galaxy fit as a 3-dimensional Gaussian, or some other simple function. This approximation may or may not be a significant concession on our parts; here it's clearly a bit sketchy in the case of $\sigma$, which has a hard cut at $\sigma=0$ that at least one individual galaxy is consistent with. But, with this approximation, the integral in the first equation above could be done analytically, much as we simplified things for the single-galaxy analysis.
Finally, not that this is an exhaustive list, we could again consider whether conjugate Gibbs sampling is an option. Since the normal distribution has nice conjugacies, we could consider a scheme where we sample $\mu_a|\tau_a,{a_i}$, then $\tau_a|\mu_a,{a_i}$, then similarly for $\mu_b$, $\tau_b$, $\mu_\sigma$ and $\tau_\sigma$, and then all the individual $a_i$, $b_i$, $\sigma_i$ and $M_{ij}$ parameters as we did with LRGS in the previous notebook (accounting for the normal "prior" on $a_i$ given by $\mu_a$ and $\tau_a$, etc.). Or we could conjugate-Gibbs sample the $\mu$'s and $\tau$'s, while using some other method entirely for the galaxy-specific parameters. (We will not actually walk through this, since (a) LRGS (in python) doesn't implement Gaussian priors on the intercept/slope parameters, even though it's a simple addition; (b) I don't feel like dragging yet another code into the mix; and (c) the Gaussian parent distribution is not conjugate for the $\sigma$ parameters, so we'd have to use a different sampling method for those parameters anyway.)
4. Obtain the posterior
4a. Brute force
Let's again start by trying brute force, although in this case we'll still use the analytic integration method from the last notebook rather than the brutest force, which would have a free absolute magnitude for every cepheid in every galaxy. We can make use of our array of singleFitter objects, and specifically their _logpost_vecarg_B methods to do that part of the calculation.
The prototypes below assume the 33 parameters are ordered as: $(\mu_a,\tau_a,\mu_b,\tau_b,\mu_\sigma,\tau_\sigma,a_1,b_1,\sigma_1,a_2,b_2,\sigma_2,\ldots)$. Also, Let's... not include all the individual galaxy parameters in these lists of parameter names:
End of explanation
"""


def log_likelihood_all_A(mu_a, tau_a, mu_b, tau_b, mu_s, tau_s, galparams):
    lnp = np.sum([f._logpost_vecarg_B(galparams[(0+3*i):(3+3*i)]) for i,f in enumerate(independent_fits)])
    TBC() # lnp += ... more stuff ...
    return lnp

TBC_above()

"""
Explanation: Complete the log-likelihood function for this part. Similarly to the way we dealt with Mtrue before, the galparams argument will end up being an array containing $(a_1,b_1,\sigma_1,a_2,b_2,\sigma_2,\ldots)$, from which we can extract arrays of $a_i$, $b_i$ and $\sigma_i$ if we want. The line given to you accounts for the $\prod_{i=1}^{N_\mathrm{gal}} p(\mathrm{data}|a_i,b_i,\sigma_i)$ part, ultimately calling log_likelihood_B and log_prior from the last notebook (see comments below).
End of explanation
"""


def log_prior_all(mu_a, tau_a, mu_b, tau_b, mu_s, tau_s, galparams=None):
    TBC()

TBC_above()

"""
Explanation: As a consequence of the code above calling _logpost_vecarg_B (note post), the old priors for the $a_i$, $b_i$ and $\sigma_i$ will be included in the return value. This is ok only because we're using uniform priors, so in the log those priors are either a finite constant or $-\infty$. In general, we would need to divide the old priors out somewhere in the new posterior calculation. Even better, we would not write such dangerously lazy code.
But for our limited purposes, it should work. The bottom line is that we don't need to worry about the priors for the $a_i$, $b_i$ and $\sigma_i$ in the function below, just the hyperparameters of their parent distributions.
Again like the last notebook, we will make galparams an optional argument to the log-prior function, so we can re-use the function later, when the $a_i$, $b_i$ and $\sigma_i$ are not being sampled.
End of explanation
"""


def log_posterior_all(loglike, **params):
    lnp = log_prior_all(**params)
    if lnp != -np.inf:
        lnp += loglike(**params)
    return lnp

def logpost_vecarg_all_A(pvec):
    params = {name:pvec[i] for i,name in enumerate(param_names_all)}
    params['galparams'] = pvec[len(param_names_all):]
    return log_posterior_all(log_likelihood_all_A, **params)

"""
Explanation: You can have the log-posterior functions.
End of explanation
"""


TBC() # guess_all = [list of hyperparameter starting values]

guess_all_A = np.array(guess_all + guessvec*9)

"""
Explanation: Based on the triangle plot in the first section, guess rough starting values for $(\mu_a,\tau_a,\mu_b,\tau_b,\mu_\sigma,\tau_\sigma)$. (NB: make this a list rather than the usual dictionary.) We'll re-use the previous guess for the galaxy-specific parameters.
End of explanation
"""


logpost_vecarg_all_A(guess_all_A)

"""
Explanation: Quick check that the functions above work:
End of explanation
"""


%%time

nsteps = 100 # or whatever

npars = len(guess_all_A)
nwalkers = 2*npars
sampler = emcee.EnsembleSampler(nwalkers, npars, logpost_vecarg_all_A)
start = np.array([np.array(guess_all_A)*(1.0 + 0.01*np.random.randn(npars)) for j in range(nwalkers)])
sampler.run_mcmc(start, nsteps)
print('Yay!')

"""
Explanation: Below, we run emcee as before.
IMPORTANT
You should find this to be more tractable than the "brute force" solution in the previous notebook, but still very slow compared to what we normally see in class. Again, you do not need to run this version long enough to get what we would normally consider acceptable results, in terms of convergence and number of independent samples. Just convince yourself that it's functioning, and see how it performs. Again, please do not turn in a notebook where the sampling cell below takes longer than $\sim30$ seconds to evaluate.
End of explanation
"""


npars = len(guess_all)+3
plt.rcParams['figure.figsize'] = (16.0, 3.0*npars)
fig, ax = plt.subplots(npars, 1);
cr.plot_traces(sampler.chain[:min(8,nwalkers),:,:npars], ax, labels=param_labels_all+param_labels);
npars = len(guess_all_A)

"""
Explanation: Look at the traces (we'll only include one of the galaxy's scaling parameters).
End of explanation
"""


TBC()
# burn = ...
# maxlag = ...

tmp_samples = [sampler.chain[i,burn:,:9] for i in range(nwalkers)]
print('R =', cr.GelmanRubinR(tmp_samples))
print('neff =', cr.effective_samples(tmp_samples, maxlag=maxlag))
print('NB: Since walkers are not independent, these will be optimistic!')
print("Plus, there's a good chance that the results in this section are garbage...")

"""
Explanation: Go through the usual motions, making sure to set burn and maxlag to something appropriate for the length of the chain.
End of explanation
"""


samples_all_A = sampler.chain[:,burn:,:].reshape(nwalkers*(nsteps-burn), npars)
plotGTC([samples_all_A[:,:9]], paramNames=param_labels_all+param_labels, chainLabels=['emcee/brute'],
        figureSize=12, customLabelFont={'size':12}, customTickFont={'size':12}, customLegendFont={'size':16});

"""
Explanation: As before, we'll be comparing the posteriors from the methods we attempt:
End of explanation
"""


def log_likelihood_all_B(mu_a, tau_a, mu_b, tau_b, mu_s, tau_s):
    TBC()
    
TBC_above()

"""
Explanation: To be more thorough, we would also want to see how well the new hierarchical part of the model fits, meaning whether the posteriors of $a_i$, $b_i$ and $\sigma_i$ are collectively consistent with being drawn from their respective fitted Gaussians. Things might look slightly different than the plots we made above, since those fits used uniform priors rather than the hierarchical model. With only 9 galaxies, it seems unlikely that we could really rule out a Gaussian distribution, and it's tangential to the point of this tutorial. So this can be an exercise for the reader, if you want.
4b. Sampling with numerical marginalization
Let's see how we do trying to marginalize out the per-galaxy parameters by simple monte carlo, as described above,
$p(\mathrm{data}|\vec{\alpha}) = \prod_{i=1}^{N_\mathrm{gal}} \frac{1}{n_i}\sum_{k=1}^{n_i} p(a_{ik},b_{ik},\sigma_{ik}|\vec{\alpha})$.
Note that, because we are taking a sum of probabilities above, we do actually need to work with probabilities, as opposed to log-probabilities. You might reasonably worry about numerical stability here, but in this case a naive implementation seems to be ok. (In general, what we would need to check is whether the summands contributing most of the sum are easily floating-point representable, i.e. not so tiny that they underflow. We could always renormalize the summands to avoid this, since we will just end up taking the log afterwards.)
Implement the log-likelihood for this approach below.
End of explanation
"""


def logpost_vecarg_all_B(pvec):
    params = {name:pvec[i] for i,name in enumerate(param_names_all)}
    return log_posterior_all(log_likelihood_all_B, **params)

"""
Explanation: This is for free:
End of explanation
"""


logpost_vecarg_all_B(guess_all)

"""
Explanation: The usual sanity check:
End of explanation
"""


nsteps = 10

npars = len(guess_all)
nwalkers = 2*npars
sampler = emcee.EnsembleSampler(nwalkers, npars, logpost_vecarg_all_B)
start = np.array([np.array(guess_all)*(1.0 + 0.01*np.random.randn(npars)) for j in range(nwalkers)])
%time sampler.run_mcmc(start, nsteps)
print('Yay?')

"""
Explanation: Let's get an idea of how computationally expensive all these sums are by running a very short chain.
End of explanation
"""


for f in independent_fits:
    f.thin(500)

"""
Explanation: For me this comes out to about 7 seconds for 10 steps - slower than we'd ideally like, at least without more serious computing resources than my laptop. (If you run longer, though, you should see performance better than in part A.)
However, its worth asking if we can get away with using fewer samples. In principle, we are well justified in doing this, since the effective number of independent samples estimated for some of the individual fits are only $\sim500$ (when I ran them, anyway).
Note that the cell below is destructive, in that we can't easily get the original chains back after running it. Keep that in mind if you plan to play around, or improve on the code at the start of the notebook.
End of explanation
"""


%%time

TBC() # nsteps = 

npars = len(guess_all)
nwalkers = 2*npars
sampler = emcee.EnsembleSampler(nwalkers, npars, logpost_vecarg_all_B)
start = np.array([np.array(guess_all)*(1.0 + 0.01*np.random.randn(npars)) for j in range(nwalkers)])
sampler.run_mcmc(start, nsteps);
print('Yay!')

"""
Explanation: With only 500 samples left in the sum for each galaxy, it should be possible to get results that appear basically converged with a couple of minutes runtime (and you should do so). Nevertheless, before turning in the notebook, please reduce the number of steps such that the sampling cell below takes longer than $\sim30$ seconds to evaluate. (You can leave a comment saying what number of steps you actually used, if you like.)
End of explanation
"""


plt.rcParams['figure.figsize'] = (16.0, 3.0*npars)
fig, ax = plt.subplots(npars, 1);
cr.plot_traces(sampler.chain[:min(8,nwalkers),:,:npars], ax, labels=param_labels_all);

"""
Explanation: Let's see how it does:
End of explanation
"""


TBC()
# burn = ...
# maxlag = ...

tmp_samples = [sampler.chain[i,burn:,:] for i in range(nwalkers)]
print('R =', cr.GelmanRubinR(tmp_samples))
print('neff =', cr.effective_samples(tmp_samples, maxlag=maxlag))
print('NB: Since walkers are not independent, these will be optimistic!')

"""
Explanation: The sampler is probably struggling to move around efficiently, but you could imagine running patiently for a while and ending up with something useful. Let's call this approach viable, but not ideal. Still, make sure you have reasonable convergence before continuing.
End of explanation
"""


samples_all_B = sampler.chain[:,burn:,:].reshape(nwalkers*(nsteps-burn), npars)
plotGTC([samples_all_A[:,:len(param_names_all)], samples_all_B], paramNames=param_labels_all, chainLabels=['emcee/brute', 'emcee/SMC'],
        figureSize=10, customLabelFont={'size':12}, customTickFont={'size':12}, customLegendFont={'size':16});

"""
Explanation: And now the burning question: how does the posterior compare with the brute force version?
End of explanation
"""


sol = np.loadtxt('solutions/ceph2.dat.gz')
plotGTC([sol, samples_all_B], paramNames=param_labels_all, chainLabels=['solution', 'my emcee/SMC'],
        figureSize=8, customLabelFont={'size':12}, customTickFont={'size':12}, customLegendFont={'size':16});

"""
Explanation: Checkpoint: Your posterior is compared with our solution by the cell below. Keep in mind they may have very different numbers of samples - we let ours run for several minutes.
End of explanation
"""


