Physics-based models: weighing the universe with supernovae

In E01-E06 the regression function was something we chose: a line, a spline, a GP. Here we do not get to choose. General relativity plus a list of what the universe contains predicts, with no free functional form, how bright a standard candle looks at redshift $z$. The theory is the regression function. That changes the character of every step:

• parameters have units and meaning (the fraction of the universe that is matter), so priors come from physics rather than from "weakly informative" folklore;
• extrapolation is principled - the curve beyond the data is a prediction of the theory, not the whim of a basis function;
• model comparison answers a scientific question. "Does the model with a cosmological constant predict better than the one without?" is the question that won the 2011 Nobel prize in physics: is the expansion of the universe accelerating?

The price is computational: the theory contains an integral with no closed form, and it has to live inside the model, differentiable, for 580 supernovae at once.

#1 · Question and data

A type Ia supernova is a white dwarf that explodes at (very nearly) a fixed mass, so its peak luminosity is (very nearly) the same everywhere: a standard candle. Compare how bright it is (absolute magnitude $M$) with how bright it looks (apparent magnitude $m$) and you have its distance, reported as the distance modulus

$$\mu = m - M = 5\log_{10}\!\frac{d_L}{10\ \text{pc}}$$

where $d_L$ is the luminosity distance. Five magnitudes are a factor 100 in flux, so $\Delta\mu = 0.1$ is about 5% in distance. The redshift $z$ of the host galaxy says how much the universe has expanded since the light left. Distance against redshift - the Hubble diagram - is therefore a record of the expansion history.

Nearly twelve magnitudes of dynamic range and a curve that looks like a boring logarithm. Almost all of that is just "further away is fainter". Everything interesting - the entire evidence for dark energy - is a deviation of a few tenths of a magnitude at $z > 0.3$, comparable to the error bar of a single supernova. We will need a better way to look at it (section 5) and a model that uses all 580 objects at once.

#2 · The physics, compactly

In a homogeneous, isotropic universe the expansion rate is $H(z) = H_0\,E(z)$, where the Friedmann equation gives the dimensionless function

$$E(z)^2 = \Omega_m (1+z)^3 + \Omega_k (1+z)^2 + \Omega_{\Lambda}\,(1+z)^{3(1+w)}, \qquad \Omega_k = 1 - \Omega_m - \Omega_\Lambda .$$

$\Omega_m$ is today's matter density and $\Omega_\Lambda$ the dark-energy density, both as fractions of the critical density; $\Omega_k$ is whatever is left over and shows up as spatial curvature. $w$ is the dark-energy equation of state (pressure / energy density): $w=-1$ is Einstein's cosmological constant, whose density does not dilute as space expands.

Light travelling towards us covers the dimensionless comoving distance

$$\chi(z) = \int_0^z \frac{dz'}{E(z')}$$

and the luminosity distance is

$$d_L(z) = (1+z)\,\frac{c}{H_0}\;\text{sinn}(\chi;\Omega_k), \qquad \text{sinn}(\chi;\Omega_k) = \begin{cases} \sinh(\sqrt{\Omega_k}\,\chi)/\sqrt{\Omega_k} & \Omega_k > 0 \ \text{(open)}\\ \chi & \Omega_k = 0 \ \text{(flat)}\\ \sin(\sqrt{-\Omega_k}\,\chi)/\sqrt{-\Omega_k} & \Omega_k < 0 \ \text{(closed)} \end{cases}$$

With $c$ in km/s and $H_0$ in km/s/Mpc, $c/H_0$ is in Mpc, and 1 Mpc = $10^5 \times$ 10 pc, so

$$\mu(z) = \underbrace{5\log_{10}\big[(1+z)\,\text{sinn}(\chi)\big]}_{\text{shape: depends on }\Omega_m,\,\Omega_\Lambda,\,w} \;+\; \underbrace{25 + 5\log_{10}\frac{c}{H_0\,[\text{Mpc}]}}_{\text{a constant}} .$$

Every model in this notebook is a constraint on the same function:

The quantity that answers the headline question is the deceleration parameter today, $q_0 = -\ddot a a/\dot a^2 = \tfrac12\Omega_m + \tfrac12(1+3w)\,\Omega_\Lambda$, which for $w=-1$ is $\Omega_m/2 - \Omega_\Lambda$. Matter pulls the brake; $q_0 < 0$ means the expansion is speeding up.

#3 · The computational core: an integral inside the model

$\chi(z)$ has no closed form for general $(\Omega_m, \Omega_\Lambda, w)$. NUTS needs the log-density and its gradient with respect to the parameters, so scipy.integrate.quad is out: PyTensor cannot differentiate through it. Options:

1. Fixed-order Gauss-Legendre quadrature per supernova. Map the $K$ nodes $x_k$ and weights $w_k$ from $[-1, 1]$ onto $[0, z_i]$: $\chi(z_i) \approx \frac{z_i}{2}\sum_k w_k \,/\, E\!\big(\tfrac{z_i}{2}(x_k+1)\big)$. That is a $(580 \times K)$ array of plain arithmetic - trivially vectorised, and differentiable because it is just arithmetic. $1/E$ is very smooth, so a small $K$ gives machine precision.
2. A cumulative trapezoid/Simpson rule on a fine shared $z$-grid, then interpolate to each $z_i$. Fine too, but the accuracy depends on grid and interpolation.
3. Write the integral in JAX and bring it in with pytensor.wrap_jax (the subject of E09). The right tool when the numerics need adaptive control flow; overkill here.

We take option 1. The functions below only use arithmetic, sqrt, log10 and maximum, so the same code runs on NumPy arrays (for checking and for posterior curves) and on PyTensor variables (inside the model) - pass xp=np or xp=pt, as C01 did.

One subtlety: $\text{sinn}$. The obvious pt.switch between sinh, sin and $\chi$ returns the right value everywhere, but its gradient at $\Omega_k = 0$ - exactly where flat models live - is NaN: the derivative of $\sqrt{|\Omega_k|}$ in the unused branches is infinite there and leaks through the switch. But $\text{sinn}$ is an entire function of $\Omega_k$: both branches are the same power series $\chi\,\sum_n (\Omega_k\chi^2)^n/(2n+1)!$. For $|\Omega_k\chi^2| \lesssim 2$, which covers anything remotely relevant, nine terms are exact to double precision. No branches, no NaN.

#Verify before you trust

A numerical approximation inside a likelihood is a silent source of bias, so we test it three ways before it goes anywhere near a sampler.

Units. A standard benchmark: for $\Omega_m = 0.3$, flat, $h = 0.7$, the distance modulus at $z = 0.5$ is 42.26. And the empty ("Milne") universe has the closed form $(H_0/c)\,d_L = z(1 + z/2)$, which exercises the $\sinh$ branch of the series.

Accuracy against adaptive quadrature, at every observed redshift, for cosmologies spread over (and beyond) the region the sampler will explore - including strongly open, closed and phantom ($w < -1$) cases - as a function of the number of nodes:

The smallest measurement error in the data is 0.08 mag. Four nodes would already be more than good enough; with 16 the quadrature error is at the level of floating-point rounding ($10^{-14}$ mag) for every test case, exotic ones included. We keep 16: a $580 \times 16$ array costs nothing.

Gradient. Finally, check that PyTensor differentiates through the whole thing, by comparing $\partial\mu/\partial\Omega_m$ with a finite difference.

Note the size of the signal: at $z = 1$, changing $\Omega_m$ by 0.1 (at fixed $\Omega_\Lambda$) moves $\mu$ by less than a tenth of a magnitude, while a single supernova is measured to about 0.2 mag. This is a many-objects, small-effect problem - exactly where a joint probabilistic model earns its keep.

#4 · Something goes wrong: $H_0$ and $M$

Let us write the "obvious" flat $\Lambda$CDM model. The theory has two parameters, $\Omega_m$ and $H_0$, so give both a prior: $\Omega_m \sim \text{Uniform}(0,1)$ and $H_0 \sim \text{Uniform}(40, 100)$ km/s/Mpc, which brackets every measurement of the last fifty years. And a standard candle is only useful if you know its wattage: the published $\mu$ were computed with an assumed absolute magnitude $M$, which could be off, so add $\delta M \sim \text{Normal}(0, 1)$ for the calibration error. Finally, supernovae are not perfectly standard: add an intrinsic scatter $\sigma_\text{int}$ in quadrature to the quoted errors, with a prior that knows type Ia scatter is of order 0.1 mag.

$$\mu_i^\text{obs} \sim \text{Normal}\Big(\mu_\text{shape}(z_i;\Omega_m) + 25 + 5\log_{10}\tfrac{c}{H_0} + \delta M,\ \ \sqrt{\sigma_i^2 + \sigma_\text{int}^2}\Big)$$

(A PyMC trap met on the way: pm.Uniform("H0", 40, 100) with integer bounds gives a NaN initial point and nutpie dies with "All initialization points failed" - small integer constants become int8, and $40 + 100$ overflows where the interval transform computes its starting point. Uniform(0, 1) is unaffected. Write bounds as floats.)

H0 and dM are in trouble: r_hat of 1.09 and an effective sample size below 50 out of 4000 draws, plus a few divergences. Even Om, which the data constrain well, mixes worse than it should. The pair plot shows why.

Every draw sits on the curve $\delta M - 5\log_{10}H_0 = \text{const}$. Look back at the likelihood: $H_0$ and $\delta M$ enter only through that sum. A universe that expands 10% faster with supernovae that are 0.2 mag fainter produces exactly the same data. No amount of supernovae can separate them; the data pin the combination to about a hundredth of a magnitude, and along the ridge the posterior is simply the prior (the posterior sd of $H_0$ is nearly that of Uniform(40, 100); the $\delta M$ prior trims the ends a little). A ridge fifty times longer than it is wide, curved, with hard walls at both ends, is miserable for a sampler with a diagonal mass matrix - hence the diagnostics.

The tempting reaction is to turn sampler knobs:

The diagnostics now look almost respectable, and nothing has been learned: the "posterior" of $H_0$ is still the prior. This is the dangerous version of the problem - more careful sampling hides the symptom, and a reader of the summary table would happily report $H_0 \approx 70 \pm 16$ "from supernovae". It is not a sampler problem. The model asks the data a question they cannot answer.

(This is why measuring $H_0$ needs a distance ladder: Cepheids or some other rung to calibrate $M$ independently. It is also why the tension between different $H_0$ measurements does not touch the evidence for acceleration, which lives in the shape.)

The fix is to model what is identified. Collapse the two into one nuisance parameter:

$$\text{offset} = \delta M - 5\log_{10}(h/0.7)$$

The Union2.1 distance moduli are published for a fiducial calibration equivalent to $h = H_0/100 = 0.7$, so we write $\mu = \mu_\text{shape} + \mu_0(h{=}0.7) + \text{offset}$ and expect an offset near zero. Its prior can come from physics: $h$ between 0.5 and 1 is $\pm 0.75$ mag, so Normal(0, 0.5) is honest. Because the offset enters linearly with Gaussian noise it could even be integrated out analytically, as many published analyses do; we sample it instead, so that the pointwise log-likelihood needed for LOO stays simple.

Every model below shares the same nuisance parameters and likelihood, so wrap them once:

#5 · Prior predictive check, in the Hubble diagram

$\Omega_m \sim \text{Uniform}(0, 1)$ is the physical range for a flat universe with non-negative matter and dark-energy densities. What does the prior claim about data?

The raw Hubble diagram hides everything (section 1), so cosmologists plot residuals against a reference: the empty universe ($\Omega_m = \Omega_\Lambda = 0$), which neither accelerates nor decelerates. Points above zero are fainter - further away - than coasting would predict: the expansion has sped up since the light left. Below zero: slowed down.

Left: the prior is a band about $\pm 1$ mag wide around the data - that is the offset prior, i.e. our ignorance of $h$ and $M$ - and the shape differences are invisible. Right: the same 100 curves with the offset removed. Now the prior's actual content is clear: it spans everything from a matter-only universe (yellow, always decelerating, $-0.5$ mag at $z = 1$) to pure dark energy (purple). All of them agree at $z \to 0$, where every cosmology reduces to Hubble's law, and fan out with redshift. The binned data (taking the published calibration at face value) already hint where this is going, and the dotted line marks the end of the data: everything right of it is extrapolation by the theory.

The helper binned_mean and the 11 redshift bins will come back for the posterior predictive checks.

#6 · The model ladder

#(a) Einstein-de Sitter: matter only

Until the mid-1990s the theoretically favoured universe was flat and made of matter: $\Omega_m = 1$. No cosmological parameter is free - only the two nuisance parameters.

From here on we sample with target_accept=0.95, for a stated reason: in the good models $\sigma_\text{int}$ turns out to be squeezed against zero, and at the default 0.8 nutpie typically reports a few divergences down there. Each model samples in a second or two, so the smaller step size is free.

The sampler is perfectly happy - two well-identified parameters. Clean diagnostics say nothing about whether the model is right. Two numbers are suspicious, though: an offset of a quarter of a magnitude (taken at face value, $h \approx 0.62$), and an intrinsic scatter as large as a typical measurement error. Hold that thought until we have something to compare with.

#(b) Flat $\Lambda \text{CDM}$

No divergences, r_hat at 1.00, four chains on top of each other, and the offset is now pinned to about $\pm 0.01$ mag. The effective sample size is "only" a third of the draws because $\Omega_m$ and the offset are correlated at about 0.7 (a lower $\Omega_m$ makes distant supernovae fainter; so does a higher offset) - plenty for our purposes. One consequence worth knowing: az.plot_rank, which E01 recommended, draws its envelope for independent draws, so with this much autocorrelation it tends to flag perfectly healthy chains. Thin the draws before using it, or rely on r_hat and ESS, which account for autocorrelation.

The physics:

• $\Omega_m \approx 0.28$: matter is only about 28% of the critical density, and if the universe is flat the other 72% is something that behaves like a cosmological constant.
• The offset is consistent with zero: the published $\mu$ indeed correspond to $h = 0.7$.
• $\sigma_\text{int}$ is consistent with zero, while Einstein-de Sitter needed 0.17 mag. Two lessons in one. First, the published errors evidently already contain an intrinsic-dispersion term (compilations like this one tune it so that the best fit has $\chi^2$ per degree of freedom near one - computed below), leaving nothing for ours to do. Second, a free scatter term is a misfit sponge: in model (a) it quietly soaked up the failure of the physics. A large "noise" estimate is a symptom worth chasing.

#Posterior predictive checks: residuals against redshift

A histogram-style PPC is useless here (both models reproduce "a pile of numbers between 34 and 45"). The informative check is structure in the residuals along $z$. First the raw version: standardised residuals with a running mean.

Einstein-de Sitter has a trend the offset cannot hide: nearby supernovae are brighter than predicted, distant ones fainter, because the offset can only shift the curve up and down to split the difference. Flat $\Lambda$CDM's running mean wobbles around zero with no trend (its excursions are what a mean of 40 unit-variance numbers does: $\pm 0.16$).

Now the proper posterior predictive version. For each model, simulate replicated datasets, reduce each to the 11 binned Hubble residuals relative to the empty universe, and ask whether the observed binned residuals look like a draw from that distribution.

This is the 1998 discovery in one figure. The observed points rise above the empty universe out to $z \approx 0.5$ - supernovae are fainter than even a coasting universe allows - and then turn over. A matter-only universe can only ever go down; even with its inflated scatter and its 0.25 mag offset, most observed bins fall outside its predictive band, in a systematic pattern (the model sits above the data at low $z$ and far below at high $z$). Flat $\Lambda$CDM tracks the rise and the turn-over.

#Heavy tails?

Supernova samples contain the odd misclassified or dust-reddened object, and a Gaussian likelihood lets one bad point drag the fit. The standard robustness check is a Student-t likelihood with the degrees of freedom $\nu$ free.

$\nu$ lands in the tens - essentially Gaussian - and $\Omega_m$ does not move. The heavy tails are not needed; the compilation has evidently already been through outlier rejection. We carry the Student-t model into the LOO table as a check and otherwise keep the Normal likelihood.

#(c) Curved $\Lambda\text{CDM}$: the banana

Drop the flatness assumption and let $\Omega_m$ and $\Omega_\Lambda$ vary independently. Physics gives $\Omega_m \ge 0$ but does not bound $\Omega_\Lambda$ (a negative cosmological constant is perfectly legal), so we use a generous box.

One corner of that box is unphysical in an interesting way: for large $\Omega_\Lambda$ and small $\Omega_m$, $E(z)^2$ goes negative at some redshift - a universe that bounced and never had a Big Bang. comoving_distance clips $E^2$ at a tiny positive number there, which makes the predicted distances absurdly large and the likelihood correspondingly terrible, so the sampler stays away. We check afterwards that it did.

We take 2000 draws per chain here and in (d): two-dimensional contours of a curved posterior need more draws than a marginal mean does.

The famous picture. Supernovae do not measure $\Omega_m$ and $\Omega_\Lambda$ separately: they measure, roughly, the difference between brake and accelerator, so the posterior is a long diagonal ridge - the two are correlated at about 0.9 - and the marginal uncertainties are 3.5 times ($\Omega_m$) and 6 times ($\Omega_\Lambda$) what they were in the flat model. Unlike the $H_0$-$M$ ridge, this one is a soft degeneracy: it has finite length because the redshift lever arm distinguishes the two terms a little, and NUTS handles it without complaint.

What the banana does and does not say:

• It lies entirely above the dashed $q_0 = 0$ line and above $\Omega_\Lambda = 0$. Acceleration does not depend on assuming flatness.
• It is consistent with flat: the solid line runs through the middle, and $\Omega_k = 0$ is well inside its posterior. But supernovae alone constrain curvature poorly ($\pm 0.2$); it is the CMB that nails flatness, which is why combining probes whose degeneracy directions cross is so powerful.
• Einstein-de Sitter is nowhere near.

#(d) Flat $w\text{CDM}$: is dark energy a cosmological constant?

Keep flatness, free the equation of state. The prior $w \sim \text{Uniform}(-3, 0)$ is deliberately agnostic: it includes "phantom" dark energy ($w < -1$) and components that would not accelerate anything ($w > -1/3$).

Another banana, this time genuinely bent: more matter can be compensated by more negative pressure. $w = -1$ passes straight through the middle, so the data are perfectly happy with a cosmological constant - but on their own they allow anything from about $-1.4$ to $-0.7$, and the 95% region stretches from $\Omega_m < 0.1$ to above 0.4. The whole posterior stays below the $q_0 = 0$ curve. An external constraint on $\Omega_m$ (BAO, CMB) would slice the banana vertically and tighten $w$ dramatically - see Try it yourself.

#LOO: which universe predicts held-out supernovae best?

Two very different verdicts, and both matter:

• Einstein-de Sitter is out, by more than 200 units of elpd - over ten standard errors of the difference. Note this is after it was allowed to inflate $\sigma_\text{int}$ to protect itself; without that freedom it would look far worse. Here LOO is answering a physics question: a universe with only matter in it does not predict supernova brightnesses, and one with a cosmological constant does.
• Among the dark-energy models LOO cannot choose, and says so: the differences are below one unit, far inside the $|\text{elpd\_diff}| < 4$ zone where ArviZ itself warns against reading anything into the ranking. Curvature and $w$ each cost about one effective parameter (p) and buy no predictive accuracy. That is not evidence that the universe is flat or that $w=-1$; it says these data cannot tell, which is exactly what the two bananas showed. The simplest adequate model is (b).

#7 · From posterior to scientific answers

Parameters with physical meaning can be pushed through more physics, draw by draw.

Is the expansion accelerating? $P(q_0 < 0)$ under each dark-energy model:

Not a single posterior draw in any of the three models decelerates, so all a Monte Carlo estimate can honestly say is that $P(\text{accelerating})$ exceeds $1 - 1/\text{draws}$. Relaxing flatness or $w = -1$ widens $q_0$ by a factor of three to four but leaves it centred near $-0.6$, many standard deviations from zero.

That probability is conditional on everything we assumed: a homogeneous universe obeying general relativity, supernovae that are the same kind of candle at $z = 1$ as nearby, and

• importantly - the error model discussed in section 8.

How much dark energy, when did acceleration begin, and how old is the universe? For $w = -1$, $\ddot a = 0$ when $\Omega_m (1+z)^3 = 2\,\Omega_\Lambda$, so the transition redshift is $z_t = (2\Omega_\Lambda/\Omega_m)^{1/3} - 1$. The age is one more integral,

$$t_0 = \frac{1}{H_0}\int_0^1 \frac{da}{a\,E(a)} = \frac{1}{H_0}\int_0^1 \frac{2u^2\,du}{\sqrt{\Omega_m + \Omega_k u^2 + \Omega_\Lambda u^6}} \qquad (a = \tfrac{1}{1+z} = u^2),$$

where the substitution removes the square-root singularity at the Big Bang so that Gauss-Legendre converges fast. No gradients are needed now, so this is plain NumPy on the posterior draws. It needs an $H_0$, which - section 4 - supernovae cannot supply: we condition on $h = 0.7$, and $t_0$ scales as $1/h$.

Under flat $\Lambda\text{CDM}$: dark energy is about 72% of the universe; the expansion switched from slowing down to speeding up at $z_t \approx 0.73$ - inside the redshift range of the data, which is why the turn-over in the residual plot is visible at all - and the universe is about 13.8 billion years old if $h = 0.7$. The last row is the assumption-free version of that statement: the age is $0.985 \pm 0.02$ Hubble times. (For Einstein-de Sitter it would be exactly $2/3$ of a Hubble time, 9.3 Gyr for $h = 0.7$ - younger than the oldest stars, one of the pre-1998 hints that something was missing.)

#The four universes, and the extrapolation

Posterior mean and 94% band of each model's curve (with its own offset) against the binned data, continued beyond the last supernova.

Inside the data the three dark-energy models are indistinguishable - as LOO said. Beyond $z = 1.4$ they are still all physically sensible curves: each bends back down, because in any of these universes matter dominated early on and the expansion was decelerating. The bands widen in the grey region, the more flexible models more so, and that widening is itself informative: it says precisely which future observations ($z > 1.5$ supernovae) would discriminate between them. A spline or polynomial fitted to the same points has no such anchor: nothing in it knows that matter must win at early times.

#8 · What we did not model

• Systematics. We used only the diagonal mu_err. The real analysis has a full $580 \times 580$ covariance matrix: calibration, dust and light-curve-model uncertainties shift groups of supernovae together, which is far more damaging to a measurement of shape than independent noise. In the Union2.1 paper, including it roughly doubles the uncertainty on $\Omega_m$. Our intervals are too narrow, and statements like "many standard deviations from zero" should be read with that in mind (the qualitative conclusions survive). With the matrix in hand the change to the model is small: pm.MvNormal with the full covariance in place of pm.Normal.
• The standardisation itself. The published $\mu$ are the output of a light-curve fit (stretch and colour corrections) whose coefficients were estimated together with a cosmology. A fully Bayesian treatment models light-curve parameters and cosmology jointly, hierarchically.
• Astrophysics. Supernovae in low-mass host galaxies are slightly fainter after standardisation; the column p_low_mass_host is there to model it (below).

#Recap

#Try it yourself

1. Host-mass step. Add a term $\Delta_\text{host}\cdot$p_low_mass_host to mu in hubble_likelihood (the column is the probability that the host galaxy has low stellar mass). Is $\Delta_\text{host}$ distinguishable from zero, which sign does it have, and does $\Omega_m$ move? Think about why a population drift with redshift would matter.
2. Break the $w$ degeneracy with an external probe. Replace the uniform prior on $\Omega_m$ in model (d) by an informative Gaussian standing in for BAO + CMB, say pm.Normal("Om", 0.30, 0.02) truncated to $(0, 1)$. How much does the interval on $w$ shrink? Redraw the banana with the old one behind it.
3. Time-varying dark energy. Implement the CPL parametrisation $w(a) = w_0 + w_a(1-a)$, for which the dark-energy term in $E(z)^2$ becomes $\Omega_\Lambda (1+z)^{3(1+w_0+w_a)}\exp\!\big(-3 w_a z/(1+z)\big)$. Only E2 needs to change. Expect the sampler to work hard: look at the $(w_0, w_a)$ pair plot, check r_hat, and decide what prior information on $w_a$ you can defend.

Next: E09, where the numerics are too involved for plain PyTensor and come in through pytensor.wrap_jax instead.