%matplotlib inline import numpy as np import matplotlib.pyplot as plt z, mag = np.loadtxt('sne.txt').T dmag = 1.0 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) 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) 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)) from time import time t0 = time() sampler.run_mcmc(pos, 200, rstate0=np.random.get_state()) print("finished in %.1f sec" % (time() - t0)) sampler.chain.shape 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) from sklearn.mixture import GMM gmm = GMM(1, 'full') gmm.fit(chain.reshape(-1, 3)) gmm.sample(100, random_state=0).shape 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()