Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

Modeling the Universe
ASTR 596 - Fall 2025

Project 4: Building a Universal Inference Engine

San Diego State University

ASTR 596: Modeling the Universe Instructor: Dr. Anna Rosen Due: Friday, November 7, 2025, 11:59 PM
AI Policy: Phase 2 (strategic use after baseline implementation works)

“The most important questions of life are indeed, for the most part, really only problems of probability.”
— Pierre-Simon Laplace, Théorie Analytique des Probabilités (1812)

“It is a capital mistake to theorize before one has data. Insensibly one begins to twist facts to suit theories, instead of theories to suit facts.”
— Arthur Conan Doyle, A Scandal in Bohemia

The Big Picture: What You’re Really Building

This project is about building general-purpose inference machinery, not just analyzing one dataset. You’re creating tools that work for any scientific model, any data, any dimensionality.

The Core Insight: Most scientific problems share the same structure:

  1. You have a forward model that predicts observations from parameters

  2. You have data (noisy measurements of those observations)

  3. You need to infer the parameters that best explain the data

  4. You need to quantify uncertainty (not just a best fit, but a full probability distribution)

Whether you’re:

The inference machinery is the same. Only the forward model and likelihood function change.

What Makes This Powerful: Modularity

You’ll build a two-layer architecture:

Layer 1: Universal Samplers (general, reusable)

These modules know nothing about cosmology. They just need:

Layer 2: Problem-Specific Components (this project)

The beauty: Tomorrow you can swap in a different likelihood (exoplanets, stellar dynamics, whatever) and your samplers still work. This is the glass-box philosophy in action — you’re not learning “how to analyze supernovae,” you’re learning “how to build inference engines.”

Connection to Your Learning Journey

This project synthesizes everything you’ve learned:

Module 1 - How the Universe Computes:

Module 5 - Statistical Inference:

Project 2: N-body Simulation

Everything connects. This isn’t isolated knowledge—it’s a unified computational framework for science.

The Scientific Application: Measuring Dark Energy

To make this concrete, you’ll apply your inference engine to Type Ia supernova cosmology—the 2011 Nobel Prize-winning discovery of cosmic acceleration.

Read first: Scientific Background Document

That document explains:

Your task: Use Bayesian inference (MCMC and HMC) to infer (Ωm,h)(\Omega_m, h) from real supernova data, assuming a flat universe (Ωm+ΩΛ=1\Omega_m + \Omega_\Lambda = 1).

This is the same analysis that revealed dark energy dominates the cosmos. You’re not just practicing methods—you’re doing real science with Nobel Prize data.

Project Structure and Workflow

Important: The parts below are not strictly sequential. As you implement your MCMC sampler (Part 2), you’ll need to develop diagnostic tools (Part 3) in parallel to assess convergence and tune your algorithm. Think of this project as iterative: build, test, diagnose, refine. By the time you move to HMC (Part 4), you’ll have a robust testing framework that you can immediately reuse—same diagnostic functions, different sampler.

Part 1: The Forward Model and Likelihood

1.1: Forward Model Implementation

Module: cosmology.py

You need to compute the theoretical distance modulus μ(z;Ωm,h)\mu(z; \Omega_m, h) for a flat universe.

Required functions:

Mathematical Background:

For a flat universe (ΩΛ=1−Ωm\Omega_\Lambda = 1 - \Omega_m), the luminosity distance is:

DL(z)=c(1+z)H0∫0zdz′E(z′)D_L(z) = \frac{c(1+z)}{H_0} \int_0^z \frac{dz'}{E(z')}
(1)

where:

E(z)=Ωm(1+z)3+(1−Ωm)E(z) = \sqrt{\Omega_m(1+z)^3 + (1-\Omega_m)}
(2)

and H0=100hH_0 = 100h km/s/Mpc, c=3×105c = 3 \times 10^5 km/s.

The distance modulus is:

μ=25−5log⁡10(h)+5log⁡10(DL∗Mpc)\mu = 25 - 5\log_{10}(h) + 5\log_{10}\left(\frac{D_L^*}{\text{Mpc}}\right)
(3)

where DL∗≡DL(h=1)D_L^* \equiv D_L(h=1) is the luminosity distance with hh factored out.

Implementation approaches:

  1. Numerical integration: Use scipy.integrate.quad to evaluate the integral directly

  2. Fitting formula: Use the Pen (1999) approximation (accurate to 0.4% for 0.2≤Ωm≤10.2 \leq \Omega_m \leq 1):

DL(z)=cH0(1+z)[η(1,Ωm)−η(11+z,Ωm)]D_L(z) = \frac{c}{H_0}(1+z)\left[\eta(1, \Omega_m) - \eta\left(\frac{1}{1+z}, \Omega_m\right)\right]
(4)

where:

η(a,Ωm)=2s3+1[1a4−0.1540sa3+0.4304s2a2+0.19097s3a+0.066941s4]−1/8\eta(a, \Omega_m) = 2\sqrt{s^3 + 1}\left[\frac{1}{a^4} - 0.1540\frac{s}{a^3} + 0.4304\frac{s^2}{a^2} + 0.19097\frac{s^3}{a} + 0.066941s^4\right]^{-1/8}
(5)

and s3≡(1−Ωm)/Ωms^3 \equiv (1-\Omega_m)/\Omega_m.

What to do: Implement both methods. The fitting formula is much faster (no numerical integration), but you should validate both approaches against each other to ensure they agree. Then experiment with timing: which method is fast enough for MCMC where you’ll call this function thousands of times? You’ll likely choose the approximation for production runs, but understanding both builds intuition about the trade-offs between accuracy and speed.

1.2: Likelihood Function

Module: likelihood.py

Required functions:

Mathematical Background:

The log-likelihood for Gaussian errors with covariance C\mathbf{C} is:

ln⁡L(θ)=−12∑i,j=1nri [C−1]ij rj\ln \mathcal{L}(\theta) = -\frac{1}{2}\sum_{i,j=1}^{n} r_i\, [\mathbf{C}^{-1}]_{ij}\, r_j
(6)

where the residual vector is:

ri=μiobs−μitheory(zi;θ)r_i = \mu_i^{\text{obs}} - \mu_i^{\text{theory}}(z_i; \theta)
(7)

and θ=(Ωm,h)\theta = (\Omega_m, h).

Implementation notes:

1.3: Data Loading

Data files (in same directory as this document):

Implement data loading functions to read these files and prepare them for your likelihood function.

Part 2: Metropolis-Hastings MCMC

Module: mcmc.py

This is completely general—it knows nothing about cosmology and should work for any inference problem.

Required function:

Algorithm:

  1. Initialize: θ(0)=θinit\theta^{(0)} = \theta_{\text{init}}

  2. For i=1,2,…,Ni = 1, 2, \ldots, N:

    • Propose: θ∗∼N(θ(i−1),Σprop)\theta^* \sim \mathcal{N}(\theta^{(i-1)}, \Sigma_{\text{prop}}) (example: multivariate normal proposal =target distribution)

    • Compute log-probabilities: ln⁡p(θ∗)\ln p(\theta^*) and ln⁡p(θ(i−1))\ln p(\theta^{(i-1)})

    • Compute acceptance ratio: α=min⁡(1,exp⁡[ln⁡p(θ∗)−ln⁡p(θ(i−1))])\alpha = \min\left(1, \exp[\ln p(\theta^*) - \ln p(\theta^{(i-1)})]\right)

    • Draw u∼Uniform(0,1)u \sim \text{Uniform}(0, 1)

    • If u<αu < \alpha: accept θ(i)=θ∗\theta^{(i)} = \theta^*

    • Else: reject θ(i)=θ(i−1)\theta^{(i)} = \theta^{(i-1)}

  3. Return chain and acceptance rate

Returns:

Tuning: Target acceptance rate 20-50%. Adjust proposal covariance to achieve this.

The Multivariate Proposal Distribution

scipy.stats.multivariate_normal

Part 3: Convergence Diagnostics

Module: diagnostics.py

Why build these now: As you develop and tune your MCMC sampler, you need tools to assess whether it’s working. These diagnostic functions are universal—they’ll work for MCMC, HMC, and any other sampler you build in the future. Develop these tools alongside your MCMC implementation in Part 2, then reuse them for HMC in Part 4.

MCMC samples are correlated. You need to assess convergence and effective sample size.

Required functions:

Mathematical Background:

Autocorrelation: Measures correlation between samples separated by lag kk:

ρk=E[(θt−θˉ)(θt+k−θˉ)]Var(θ)\rho_k = \frac{\mathbb{E}[(\theta_t - \bar{\theta})(\theta_{t+k} - \bar{\theta})]}{\text{Var}(\theta)}
(8)

Effective Sample Size:

ESS=N1+2∑k=1Kρk\text{ESS} = \frac{N}{1 + 2\sum_{k=1}^{K}\rho_k}
(9)

where KK is chosen such that ρk\rho_k becomes negligible.

Gelman-Rubin Statistic: Compares within-chain to between-chain variance. For mm chains of length nn:

R^=V^W\hat{R} = \sqrt{\frac{\hat{V}}{W}}
(10)

Convergence achieved when R^<1.1\hat{R} < 1.1.

See Module 5 Part 4 for implementation details.

Part 4: Hamiltonian Monte Carlo

Module: hmc.py

Before you start: Make sure your MCMC (Part 2) is working and you have diagnostic tools (Part 3) functioning. You’ll use the same diagnostics.py functions to assess your HMC chains—this demonstrates the power of modular design.

4.1: The Connection to N-body Simulation

HMC solves Hamilton’s equations in parameter space. This is the same physics you implemented in Project 2!

The Hamiltonian:

H(θ,p)=U(θ)+K(p)H(\theta, p) = U(\theta) + K(p)
(11)

where:

Hamilton’s equations:

dθdt=∂H∂p=M−1p\frac{d\theta}{dt} = \frac{\partial H}{\partial p} = M^{-1}p
(12)

dpdt=−∂H∂θ=−∇U(θ)=∇ln⁡p(θ∣D)\frac{dp}{dt} = -\frac{\partial H}{\partial \theta} = -\nabla U(\theta) = \nabla \ln p(\theta|D)
(13)

These are exactly analogous to position and momentum evolution in N-body dynamics. You can repurpose your leapfrog integrator from Project 2.

4.2: Computing Gradients

HMC requires ∇θln⁡p(θ∣D)\nabla_\theta \ln p(\theta|D). Two approaches:

Finite differences:

∂f∂θi≈f(θ+hei)−f(θ−hei)2h\frac{\partial f}{\partial \theta_i} \approx \frac{f(\theta + h\mathbf{e}_i) - f(\theta - h\mathbf{e}_i)}{2h}
(14)

Use h∼10−5h \sim 10^{-5} to 10-7.

Automatic differentiation (see Extensions for JAX): More accurate and efficient, but requires learning JAX basics.

Implement gradient computation using at least one of these methods.

4.3: The HMC Algorithm

Required function:

Algorithm:

  1. Initialize: θ(0)=θinit\theta^{(0)} = \theta_{\text{init}}

  2. For i=1,2,…,Ni = 1, 2, \ldots, N:

    • Sample momentum: p∼N(0,M)p \sim \mathcal{N}(0, M)

    • Set (θ,p)=(θ(i−1),p)(\theta, p) = (\theta^{(i-1)}, p)

    • Simulate Hamiltonian dynamics for LL leapfrog steps:

      • For j=1,2,…,Lj = 1, 2, \ldots, L:

        • Half-step momentum: p←p+ϵ2∇U(θ)p \leftarrow p + \frac{\epsilon}{2}\nabla U(\theta)

        • Full-step position: θ←θ+ϵM−1p\theta \leftarrow \theta + \epsilon M^{-1}p

        • Half-step momentum: p←p+ϵ2∇U(θ)p \leftarrow p + \frac{\epsilon}{2}\nabla U(\theta)

    • Compute energy change: ΔH=H(θnew,pnew)−H(θ(i−1),pinit)\Delta H = H(\theta_{\text{new}}, p_{\text{new}}) - H(\theta^{(i-1)}, p_{\text{init}})

    • Compute acceptance probability: α=min⁡(1,exp⁡(−ΔH))\alpha = \min(1, \exp(-\Delta H))

    • Draw u∼Uniform(0,1)u \sim \text{Uniform}(0, 1)

    • If u<αu < \alpha: accept θ(i)=θnew\theta^{(i)} = \theta_{\text{new}}

    • Else: reject θ(i)=θ(i−1)\theta^{(i)} = \theta^{(i-1)}

  3. Return chain and acceptance rate

Returns:

Tuning:

Resources: Betancourt (2017), “A Conceptual Introduction to Hamiltonian Monte Carlo” (arXiv:1701.02434)

Part 5: Extensions (Choose 1, or Propose Your Own)

You must complete one extension beyond the baseline MCMC/HMC implementation. Discuss with the instructor if you want to propose something not listed here.

Extension Option 1: JAX for Automatic Differentiation

Replace finite-difference gradients with automatic differentiation using JAX.

Why: JAX provides exact gradients (no approximation error) and is typically much faster than finite differences for gradient computation.

Key JAX function: jax.grad() computes gradients automatically. Example usage:

import jax
grad_fn = jax.grad(your_function)
gradient = grad_fn(theta)

See the JAX documentation for details on making your functions JAX-compatible.

Resources:

What to implement: JAX-compatible versions of your likelihood and gradient functions. Compare computational performance (wall-clock time) and accuracy to finite differences. Quantify the speedup and discuss when autodiff is worth the additional dependency.

Extension Option 2: Non-Flat Universe

Relax the flatness assumption: let Ωm\Omega_m, ΩΛ\Omega_\Lambda, hh all vary independently (3 parameters instead of 2).

Why: Tests whether data alone can constrain spatial curvature ΩK=1−Ωm−ΩΛ\Omega_K = 1 - \Omega_m - \Omega_\Lambda.

What changes:

Physics:

DL(z)=c(1+z)H0∣ΩK∣Sk(∣ΩK∣∫0zdz′E(z′))D_L(z) = \frac{c(1+z)}{H_0\sqrt{|\Omega_K|}} S_k\left(\sqrt{|\Omega_K|} \int_0^z \frac{dz'}{E(z')}\right)
(15)

where E(z)=Ωm(1+z)3+ΩΛ+ΩK(1+z)2E(z) = \sqrt{\Omega_m(1+z)^3 + \Omega_\Lambda + \Omega_K(1+z)^2} and:

Sk(x)={sin⁡(x)ΩK>0 (closed)xΩK=0 (flat)sinh⁡(x)ΩK<0 (open)S_k(x) = \begin{cases} \sin(x) & \Omega_K > 0 \text{ (closed)} \\ x & \Omega_K = 0 \text{ (flat)} \\ \sinh(x) & \Omega_K < 0 \text{ (open)} \end{cases}
(16)

Compare your ΩK\Omega_K constraint to Planck’s ∣ΩK∣<0.001|\Omega_K| < 0.001.

Extension Option 3: No-U-Turn Sampler (NUTS)

Standard HMC requires manual tuning of trajectory length LL. NUTS adapts LL dynamically by detecting when the trajectory starts doubling back.

Why: State-of-the-art method used in Stan and PyMC.

What to implement: Research the NUTS algorithm, implement it, compare performance to fixed-LL HMC.

Resources: Hoffman & Gelman (2014), “The No-U-Turn Sampler” (arXiv:1111.4246)

Extension Option 4: Importance Sampling with Planck Prior

Weight MCMC samples by an informative Gaussian prior on hh from Planck CMB measurements: h∼N(0.674,0.0052)h \sim \mathcal{N}(0.674, 0.005^2).

Why: Demonstrates combining information from multiple datasets without re-running chains.

What to implement: Reweight posterior samples, compare results with/without prior. How much do constraints on Ωm\Omega_m tighten?

Extension Option 5: Propose Your Own

Have an idea? Want to explore something specific? Talk to the instructor. Possibilities include:

Part 6: Analysis and Results

6.1: Running Your Chains

For both MCMC and HMC:

  1. Run multiple independent chains (at least 3-4) with different random seeds

  2. Discard burn-in (first 20% is typical)

  3. Check convergence using Gelman-Rubin (R^<1.1\hat{R} < 1.1)

  4. Compute effective sample size

  5. Pool chains for final posterior analysis

Tuning guidance:

6.2: Interpreting Your Results

Parameter estimates:

Comparison to literature:

Are your results consistent? How do uncertainties compare?

Physical interpretation:

6.3: Efficiency Comparison (MCMC vs. HMC)

Metrics to compare:

  1. Effective samples per second: ESS / wall-clock time

  2. Autocorrelation length: How many steps between independent samples?

  3. Convergence speed: How quickly does R^→1\hat{R} \to 1?

Expected: HMC should be 5-20× more efficient than random-walk Metropolis-Hastings.

Why? HMC uses gradient information to propose distant points in high-probability regions, while MCMC random walks diffuse slowly.

Part 7: Deliverables

7.1: Code Repository

Submit a well-organized code repository with the following structure:

project4/
├── src/
│   ├── cosmology.py
│   ├── likelihood.py
│   ├── mcmc.py
│   ├── hmc.py
│   └── diagnostics.py
├── scripts/
│   ├── run_mcmc.py
│   ├── run_hmc.py
│   └── analyze_results.py
├── tests/
│   ├── test_cosmology.py
│   ├── test_mcmc.py
│   └── test_hmc.py
├── outputs/
│   ├── chains/
│   └── figures/
├── README.md
└── requirements.txt

Validation tests: Include tests that verify:

README: Document installation, usage, and how to reproduce your results.

7.2: Research Memo

Write a research memo (format and length at your discretion) that presents your analysis. Your memo should:

Style: This is a technical research memo — more formal than a lab report, but more concise than a journal article. Write for a peer who understands MCMC fundamentals but wants to learn your specific implementation choices (forward model approach, tuning methods, gradient computation) and scientific findings. Include technical details needed for reproducibility (chain lengths, acceptance rates, convergence metrics) and maintain formal scientific tone.

Figures: Include publication-quality plots (see Part 8 below). Each figure must have a detailed caption that enables standalone understanding: state what is plotted, include key technical details (e.g., “4 chains of 20,000 samples, 20% burn-in discarded”), and summarize the main takeaway. Captions are where you practice technical writing precision.

Part 8: Visualization Requirements

Generate the following plots for both MCMC and HMC (where applicable):

Required Plots:

  1. Trace plots: Parameter values vs. iteration for each chain

    • Shows mixing, burn-in, stationarity

    • One plot per parameter, all chains overlaid

  2. Autocorrelation plots: ρk\rho_k vs. lag kk for each parameter

    • Shows how quickly samples become independent

    • Compare MCMC vs. HMC autocorrelation length

  3. Corner plot: Joint and marginal posterior distributions

    • Use corner.py package if desired: pip install corner (documentation)

    • Shows parameter correlations and credible intervals

    • Overlay MCMC and HMC results to compare

  4. Data vs. model: μ(z)\mu(z) with observations and posterior predictions

    • Plot JLA data points with error bars

    • Overplot 50-100 random posterior samples as thin lines (shows uncertainty)

    • Include both MCMC and HMC posteriors

  5. Diagnostics:

    • Gelman-Rubin R^\hat{R} vs. iteration (convergence monitoring)

    • Energy conservation for HMC: histogram of ΔH\Delta H

    • Acceptance rate vs. proposal scale (for tuning demonstration)

  6. Efficiency comparison:

    • Bar chart or table comparing ESS/second for MCMC vs. HMC

    • Illustrative figure showing why HMC explores parameter space more efficiently

Tip: All plots should have clear labels, legends, and be publication-ready. Consider using a consistent color scheme and style.

Part 9: Tips and Debugging

Common Issues:

MCMC stuck or not mixing:

HMC rejecting everything:

Parameters hitting boundaries:

Chains not converging (R^>1.1\hat{R} > 1.1):

Computational Efficiency:

“Science is not about building a body of known ‘facts’. It is a method for asking awkward questions and subjecting them to a reality-check, thus avoiding the human tendency to fallaciously reason from assumptions to conclusions.”
— Terry Pratchett

Project 3
Project 3 MCRT - Starter Code Structure
Project 4
Scientific Background: Measuring the Universe with Supernovae