Visualizing Cosmological Constraints with Supernovae

In [1]:
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

We need a model of the magnitude as a function of redshift.

For a cosmological model, we could use a lambda-CDM distance modulus $\mu$, which effectively gives

$$ M(z; M_0, w, \Omega_M) = M_0 + \mu(z; w, \Omega_M) $$

This means that the likelihood is

$$ \log L(M_0, \Omega_M, w) \propto - \sum_i \frac{[M_i - M_0 - \mu(z_i;w, \Omega_M)]^2}{2\sigma_i^2} $$

where $(z_i, M_i)$ are the data points used for the fit, and $(M_0, \Omega_M, w)$ are the model parameters.

We can then maximize this likelihood with respect to $w$, while marginalizing over $M_0$ and $\Omega_M$ subject to flat priors and find the posterior distribution in $w$.

Because the distance modulus takes so long to compute, let's start out with a simple polynomial fit to $\log(z)$:

$$ M(z; \theta_0, \theta_1, \theta_2) = \theta_0 + \theta_1\log(z) + \theta_2 [\log(z)]^2 $$

Let's first import the data:

In [2]:
z, mag = np.loadtxt('sne.txt').T
dmag = 1.0

First we define our likelihood function:

In [3]:
z, mag = np.loadtxt('sne.txt').T

if False:
    # real cosmology fit. Note that this takes ~200us per evaluation,
    # which for 10000 points becomes ~2s per likelihood evaluation
    from astropy import cosmology
    theta_guess = [20, 0.3, -1]

    def lnprior(theta):
        M0, Om0, w = theta
        if (0.0 <= Om0 <= 1.0) and (-2 <= w <= 0):
            return 0.0
        else:
            return -np.inf
        
    def compute_mag_single(theta, z):
        M0, Om0, w = theta
        cosmo = cosmology.FlatwCDM(H0=70, Om0=Om0, w0=w)
        return M0 - cosmo.distmod(z)

    def lnlike(theta, z, mag, dmag):
        M0, Om0, w = theta
        cosmo = cosmology.FlatwCDM(H0=70, Om0=Om0, w0=w)
        return -0.5 * np.sum(((mag - M0 - cosmo.distmod(z)) * 1.0 / dmag) ** 2)

else:
    # simpler model: a degree-2 polynomial
    theta_guess = np.polyfit(np.log(z), mag, 2)
    
    def lnprior(theta):
        return 0.0
    
    def compute_mag_single(theta, z):
        return np.polyval(theta, np.log(z))
    
def compute_mag(theta, z):
    thetas = np.atleast_2d(theta)
    mag = np.array([compute_mag_single(t, z) for t in thetas])
    return mag.reshape(theta.shape[:-1] + z.shape)

def lnlike(theta, z, mag, dmag):
    return -0.5 * np.sum(((mag - compute_mag_single(theta, z)) * 1.0 / dmag) ** 2)

def lnprob(theta, z, mag, dmag):
    return lnprior(theta) + lnlike(theta, z, mag, dmag)

def chi2(theta, z, mag, dmag):
    return -2 * lnprob(theta, z, mag, dmag)

Next we do a simple $\chi^2$ minimization to find the optimum:

In [4]:
from scipy import optimize
z, mag = np.loadtxt('sne.txt').T
dmag = 0.5

theta_best = optimize.fmin(chi2, theta_guess, args=(z, mag, dmag))
print(theta_best)
Optimization terminated successfully.
         Current function value: 951.749205
         Iterations: 43
         Function evaluations: 79
[  0.0955759    2.76563849  24.93284035]

Now let's set up the emcee package to solve this problem:

In [5]:
import emcee
ndim = len(theta_best)
nwalkers = 50
pos = np.array(theta_best) + 1E-4 * np.random.randn(nwalkers, ndim)
sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(z, mag, dmag))

Now we'll run the MCMC sampler

In [6]:
from time import time
t0 = time()
sampler.run_mcmc(pos, 200, rstate0=np.random.get_state())
print("finished in %.1f sec" % (time() - t0))
finished in 3.5 sec
In [7]:
sampler.chain.shape
Out[7]:
(50, 200, 3)

Let's take a look at the marginalized histograms of each parameter:

In [8]:
burnin = 50
chain = sampler.chain[:, burnin:, :]
fig, ax = plt.subplots(ndim, ndim, figsize=(10, 10), sharex='col', sharey='row')
for i in range(ndim):
    for j in range(ndim):
        ax[i, ndim - 1 - j].plot(chain[:, :, j], chain[:, :, i], 'ok', ms=2, alpha=0.05)

Looks like a nice multivariate normal. Let's cheat and use a Gaussian mixture model with a single component to find a quick fit for this distribution:

In [9]:
from sklearn.mixture import GMM
gmm = GMM(1, 'full')
gmm.fit(chain.reshape(-1, 3))
Out[9]:
GMM(covariance_type='full', init_params='wmc', min_covar=0.001,
  n_components=1, n_init=1, n_iter=100, params='wmc', random_state=None,
  thresh=0.01)

Now we have a nice posterior model that we can sample from:

In [10]:
gmm.sample(100, random_state=0).shape
Out[10]:
(100, 3)

Plotting Draws from the Model

Now we need a single function which uses everything above to compute a sample given a set of data.

In [11]:
def draw_sample_from_model(z, mag, dmag, N, nwalkers=50, nsamples=200, burnin=50, quiet=False):
    """Draw a sample of N points from the model trained on z and mag"""
    theta_best = optimize.fmin(chi2, theta_guess, args=(z, mag, dmag), disp=not quiet)
    ndim = len(theta_best)
    pos = np.array(theta_best) + 1E-4 * np.random.randn(nwalkers, ndim)
    sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=(z, mag, dmag))
    
    if not quiet:
        print("Running MCMC:")
        t0 = time()
    sampler.run_mcmc(pos, nsamples, rstate0=np.random.get_state())
    if not quiet:
        print("finished in %.1f sec" % (time() - t0))
    
    chain = sampler.chain[:, burnin:, :]
    gmm = GMM(1, 'full')
    gmm.fit(chain.reshape(-1, 3))
    
    return gmm.sample(N, random_state=0)
In [12]:
z, mag = np.loadtxt('sne.txt').T
dmag = 2.0

def plot_results(Ndata, Nlines=50, dmag=0.5, ax=None):
    sample = draw_sample_from_model(z[:Ndata], mag[:Ndata], dmag, Nlines, quiet=True)
    zfit = np.logspace(-1, 0, 10)
    
    # come up with some transparency gradient based on number of points
    alpha = min(0.3, 0.01 * 10000. / Ndata)
    size = 7
    
    if ax is None:
        ax = plt.gca()
        
    ax.semilogx(zfit, compute_mag(sample, zfit).T, '-b', alpha=0.1, zorder=3)
    ax.semilogx(z[:Ndata], mag[:Ndata], 'ok', ms=size, alpha=alpha, zorder=2)

    ax.set_xlim(0.1, 1)
    ax.set_ylim(19, 26)

    ax.set_xlabel('redshift')
    ax.set_ylabel('magnitude')
    
    ax.text(0.01, 0.99, "N = {0}".format(Ndata),
            size=18, ha='left', va='top',
            transform=ax.transAxes)
In [13]:
for Npts in [10, 100, 1000, 13000]:
    plt.figure(figsize=(8, 6))
    plot_results(Npts, 50)

Animating the Results

In [15]:
from matplotlib import animation

# enables automatic display of matplotlib animations
from JSAnimation import IPython_display


class AnimateResults(object):
    def __init__(self, z, mag, dmag, Nframes=30, interval=200):
        self.z = z
        self.mag = mag
        self.dmag = dmag
        self.N = np.logspace(1, 4.11, Nframes).astype(int)
        self.zfit = np.logspace(-1, 0, 10)
        self.interval = interval
        
        self.points_alpha = np.linspace(0.5, 0.05, len(self.N))
        self.lines_alpha = np.linspace(0.1, 0.5, len(self.N))
        
        self.fig, self.ax = plt.subplots(figsize=(8, 6))
        self.ax.set_xscale('log')
        
    def anim_init(self):
        # Here we set up the plot elements that we'll work with
        self.points, = self.ax.plot([1], [1], 'ok', ms=7, alpha=0.1)
        self.lines = self.ax.plot([1], np.ones((1, 50)), '-b', alpha=0.1)

        self.ax.set_xlim(0.1, 1)
        self.ax.set_ylim(19, 26)

        self.ax.set_xlabel('redshift')
        self.ax.set_ylabel('magnitude')
    
        self.text = self.ax.text(0.01, 0.99, "",
                                 size=18, ha='left', va='top',
                                 transform=self.ax.transAxes)
    
        for line in self.lines:
            line.set_data([], [])
        self.points.set_data([], [])
        return self.lines + [self.points, self.text]
    
    def anim_frame(self, i):
        Ni = self.N[i]
        zi = self.z[:Ni]
        magi = self.mag[:Ni]
        sample = draw_sample_from_model(zi, magi, self.dmag, 50, quiet=True)
        mag_fit = compute_mag(sample, self.zfit)
        for line, mfit in zip(self.lines, mag_fit):
            line.set_data(self.zfit, mfit)
            line.set_alpha(self.lines_alpha[i])
        self.points.set_data(zi, magi)
        self.points.set_alpha(self.points_alpha[i])
        self.text.set_text("N = {0}".format(Ni))
        return self.lines + [self.points, self.text]
        
    def show(self):
        return animation.FuncAnimation(self.fig, self.anim_frame, init_func=self.anim_init,
                                       frames=len(self.N), interval=self.interval)

        
z, mag = np.loadtxt('sne.txt').T
dmag = 0.5
anim = AnimateResults(z, mag, dmag)
anim.show()
Out[15]:


Once Loop Reflect