# 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
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)