%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:
z, mag = np.loadtxt('sne.txt').T
dmag = 1.0
First we define our likelihood function:
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:
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:
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
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
sampler.chain.shape
(50, 200, 3)
Let's take a look at the marginalized histograms of each parameter:
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:
from sklearn.mixture import GMM
gmm = GMM(1, 'full')
gmm.fit(chain.reshape(-1, 3))
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:
gmm.sample(100, random_state=0).shape
(100, 3)
Now we need a single function which uses everything above to compute a sample given a set of data.
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)
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)
for Npts in [10, 100, 1000, 13000]:
plt.figure(figsize=(8, 6))
plot_results(Npts, 50)
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()