Auxiliary functions for reproducing Riess' 1998 paper results for dark energy

sn_riess.py

from random import uniform
from time import perf_counter
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad
import getdist
from getdist import plots

c = 299_792.458 # km/s

class SN:
    def __init__(self, mu, sigma_mu, z=None, logcz=None, name=""):
        self.name = name
        if logcz is not None:
            self.z = 10**logcz/c
        else:
            self.z = z
        self.mu = mu
        self.sigma_mu = sigma_mu
        self.DL = 10**((self.mu - 25)/5) # in Mpc
        self.sigma_DL = 10**((self.mu - 25)/5)*np.log(10)/5*sigma_mu # in Mpc

mlcs = [
    SN(logcz=3.734, mu=34.72, sigma_mu=0.16, name="1992bo"),
    SN(logcz=3.779, mu=34.87, sigma_mu=0.11, name="1992bc"),
    SN(logcz=4.481, mu=38.41, sigma_mu=0.15, name="1992aq"),
    SN(logcz=4.350, mu=37.80, sigma_mu=0.17, name="1992ae"),
    SN(logcz=3.896, mu=35.76, sigma_mu=0.13, name="1992P"),
    SN(logcz=4.178, mu=36.53, sigma_mu=0.15, name="1990af"),
    SN(logcz=3.859, mu=35.39, sigma_mu=0.18, name="1994M"),
    SN(logcz=3.685, mu=34.27, sigma_mu=0.12, name="1994S"),
    SN(logcz=4.030, mu=36.19, sigma_mu=0.21, name="1994T"),
    SN(logcz=3.398, mu=33.01, sigma_mu=0.13, name="1995D"),
    SN(logcz=3.547, mu=33.60, sigma_mu=0.17, name="1995E"),
    SN(logcz=4.166, mu=36.85, sigma_mu=0.13, name="1995ac"),
    SN(logcz=3.820, mu=35.15, sigma_mu=0.16, name="1995ak"),
    SN(logcz=3.679, mu=34.15, sigma_mu=0.19, name="1995bd"),
    SN(logcz=3.924, mu=35.98, sigma_mu=0.20, name="1996C"),
    SN(logcz=4.572, mu=39.01, sigma_mu=0.13, name="1996ab"),
    SN(logcz=3.891, mu=35.37, sigma_mu=0.23, name="1992ag"),
    SN(logcz=3.625, mu=33.92, sigma_mu=0.23, name="1992al"),
    SN(logcz=4.024, mu=36.26, sigma_mu=0.23, name="1992bg"),
    SN(logcz=4.130, mu=36.91, sigma_mu=0.23, name="1992bh"),
    SN(logcz=4.111, mu=36.26, sigma_mu=0.23, name="1992bl"),
    SN(logcz=4.379, mu=37.65, sigma_mu=0.23, name="1992bp"),
    SN(logcz=4.418, mu=38.21, sigma_mu=0.23, name="1992br"),
    SN(logcz=4.283, mu=37.61, sigma_mu=0.23, name="1992bs"),
    SN(logcz=3.871, mu=35.20, sigma_mu=0.23, name="1993H"),
    SN(logcz=4.189, mu=37.03, sigma_mu=0.23, name="1993O"),
    SN(logcz=4.177, mu=36.80, sigma_mu=0.23, name="1993ag"),
]

highz = [
    SN(z=0.43, mu=41.74, sigma_mu=0.28, name="1996E"),
    SN(z=0.62, mu=42.98, sigma_mu=0.17, name="1996H"),
    SN(z=0.57, mu=42.76, sigma_mu=0.19, name="1996I"),
    SN(z=0.30, mu=41.38, sigma_mu=0.24, name="1996J"),
    SN(z=0.38, mu=41.63, sigma_mu=0.20, name="1996K"),
    SN(z=0.43, mu=42.55, sigma_mu=0.25, name="1996U"),
    SN(z=0.44, mu=41.95, sigma_mu=0.17, name="1997ce"),
    SN(z=0.50, mu=42.40, sigma_mu=0.17, name="1997cj"),
    SN(z=0.97, mu=44.39, sigma_mu=0.30, name="1997ck"),
    SN(z=0.45, mu=42.45, sigma_mu=0.17, name="1995K"),
]

all = mlcs + highz
zs = np.array([sn.z for sn in all])
mus = np.array([sn.mu for sn in all])
sigmas = np.array([sn.sigma_mu for sn in all])

def S_k(x, omegak):
    sqrt_omegak = np.sqrt(np.abs(omegak))
    if omegak == 0: return x
    elif omegak > 0: return np.sinh(sqrt_omegak*x)/sqrt_omegak
    elif omegak < 0: return np.sin(sqrt_omegak*x)/sqrt_omegak

def lum_dist(z, omegam, omegal, H0):
    omegak = 1 - omegam - omegal
    integrand = lambda z: 1/np.sqrt((1+z)**2*(1+omegam*z) - z*(2+z)*omegal)
    integral = quad(integrand, 0, z)[0]
    return c*(1+z)/H0 * S_k(integral, omegak)

sigma_vel_disp_lowz = 5/np.log(10) * 200 / 299_792.458
sigma_vel_disp_highz = 5/np.log(10) * np.sqrt((200**2 + 2500**2)) / 299_792.458
sigma_vel_disp = np.array([sigma_vel_disp_lowz/z if z < 0.2 else sigma_vel_disp_highz/z for z in zs])

def chi2(model):
    global mus, sigmas, sigma_vel_disp
    om, ol, h0 = model
    if np.any(om*(1+zs)**3 + (1 - om - ol)*(1+zs)**2 + ol < 0): return np.nan
    lum_dists = np.array([lum_dist(z, *model) for z in zs])
    mus_th = 5*np.log10(lum_dists) + 25
    return np.sum((mus_th - mus)**2/(sigmas**2 + sigma_vel_disp**2))

class MCMCWalker:
    def __init__(self):
        # Hard-coding an initial point for now
        initial_om = 0.3
        initial_ol = 0.7
        initial_h0 = 70
        initial_params = [initial_om, initial_ol, initial_h0]
        initial_chi2 = chi2(initial_params)
        if initial_chi2 == np.nan:
            self.__init__()
            return
        initial_sample = {
            'params': initial_params,
            'chi2': initial_chi2,
            'weight': 1,
        }
        self.samples = [initial_sample]

    def accept_sample(self, params, chi2):
        sample = {
            'params': params,
            'chi2': chi2,
            'weight': 1
        }
        self.samples.append(sample)

    def step(self):
        while True:
            current_chi2 = self.samples[-1]['chi2']
            new_om = self.samples[-1]['params'][0] + uniform(-0.25, 0.25)
            new_ol = self.samples[-1]['params'][1] + uniform(-0.25, 0.25)
            new_h0 = self.samples[-1]['params'][2] + uniform(-2.0, 2.0)
            new_params = [new_om, new_ol, new_h0]
            if new_h0 < 60 or new_h0 > 80 or new_om < 0 or new_om > 1 or new_ol < -1 or new_ol > 3:
                # Reject point outside the prior
                self.samples[-1]['weight'] += 1
                continue
            new_chi2 = chi2(new_params)
            if new_chi2 == np.nan:
                # Reject points that have problematic chi2
                self.samples[-1]['weight'] += 1
                continue
            if new_chi2 < current_chi2:
                self.accept_sample(new_params, new_chi2)
                break
            else:
                r = uniform(0, 1)
                if r < np.exp(-(new_chi2 - current_chi2)/2):
                    self.accept_sample(new_params, new_chi2)
                    break
                else:
                    self.samples[-1]['weight'] += 1 # Increment weight
                    continue
    
    def gelman_rubin(self, n_split):
        all_params = np.array(
            [sample['params'] for sample in self.samples]
        )[:-(len(self.samples)%n_split)]
        np.random.shuffle(all_params)
        split_params = np.split(all_params, n_split)
        avg = np.mean(split_params, axis=1)
        std = np.std(split_params, axis=1)
        avg_of_std = np.mean(std, axis=0)
        std_of_avg = np.std(avg, axis=0)
        R_minus_one = std_of_avg/avg_of_std
        return np.max(R_minus_one)

def run_mcmc():
  w = MCMCWalker()
  print("Starting MCMC")
  start = perf_counter()
  while True:
      for _ in range(1000): w.step()
      R_minus_one = w.gelman_rubin(4)
      print(f"At {len(w.samples)} samples, R-1 = {R_minus_one}")
      if R_minus_one < 0.01: break 
  print(f"MCMC Converged! Took {perf_counter() - start:.2f} seconds")
  return w

def getdist_chain(w):
  mcmc = getdist.MCSamples(
      samples=np.array([sample['params'] + [sample['chi2']] for sample in w.samples]),
      weights=np.array([sample['weight'] for sample in w.samples]),
      names=["Omega_m", "Omega_Lambda", "H0", "chi2"],
      labels=["\\Omega_m", "\\Omega_\\Lambda", "H_0", "\\chi^2"],
      ranges={"Omega_m": (0, None)}
  )
  mcmc.removeBurn(0.3)
  mcmc.addDerived(mcmc["Omega_m"]/2 - mcmc["Omega_Lambda"], name="q0", label="q_0")
  mcmc.addDerived(1 - mcmc["Omega_Lambda"] - mcmc["Omega_m"], name="Omega_k", label="\\Omega_k")
  return mcmc

def plot_chain(chain):
    g = plots.get_single_plotter(width_inch=5, ratio=30/25)
    g.settings.num_plot_contours = 3
    ax = g.get_axes()
    chain.updateSettings({"contours": [0.68, 0.95, 0.997]})
    g.plot_2d(chain, "Omega_m", "Omega_Lambda", ls='-', lws=1)
    ax.set_xlim([0, 2.5])
    ax.set_ylim([-1, 3])

    # Deceleration parameter lines
    om = np.linspace(0, 2.5, 50)
    ol = om/2 # q_0 = 0
    ax.plot(om, ol, color="black", ls="--", lw=1, dashes=(4, 4))
    ax.text(2.0, ol[-7], r"$q_0 = 0$", rotation=22)
    ol = om/2 - 0.5 # q_0 = 0.5
    ax.plot(om, ol, color="black", ls="--", lw=1, dashes=(4, 4))
    ax.text(2.0, ol[-7], r"$q_0 = 0.5$", rotation=22)
    ol = om/2 + 0.5 # q_0 = -0.5
    ax.plot(om, ol, color="black", ls="--", lw=1, dashes=(4, 4))
    ax.text(2.0, ol[-7], r"$q_0 = -0.5$", rotation=22)

    # Closed/Open line
    ol = 1 - om
    ax.plot(om, ol, color="black", ls="-", lw=1)
    ax.text(1.2, ol[-23], "Closed", rotation=-35)
    ax.text(1.15, ol[-21], "Open", rotation=-35)

    # Omega_Lambda = 0
    ol = 0*om
    ax.plot(om, ol, color="black", ls=":", lw=1)
    ax.text(2.1, -0.15, r"$\Omega_\Lambda = 0$")

    ax.set_xticks(np.arange(0, 2.51, 0.1), minor=True)
    ax.set_yticks(np.arange(-1, 3, 0.1), minor=True)
    ax.tick_params(direction="in", length=10, top=True, right=True, pad=10)
    ax.tick_params(which="minor", direction="in", length=5, top=True, right=True)
    g.export("riess_reproduction.pdf")