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.

Project 4: Building a Universal Inference Engine

San Diego State University

ASTR 596: Modeling the Universe Instructor: Dr. Anna Rosen Due: Monday, November 10, 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


AI Usage Policy Reminder

This project operates under Phase 2 AI Policy: You may use AI tools strategically after your baseline implementation works. You must document all AI usage:

  1. In code: Include a 3-line header comment in any file where AI assisted:

    # AI Usage: [Tool name] assisted with [specific task]
    # Date: [Date]
    # Extent: [Brief description of what was AI-generated vs. your work]
  2. In Growth Memo: Include an “AI Usage Reflection” section discussing how AI tools helped or hindered your learning

See the course AI policy document for complete guidelines.


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)H00zdzE(z)D_L(z) = \frac{c(1+z)}{H_0} \int_0^z \frac{dz'}{E(z')}

where:

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

and H0=100hH_0 = 100h km s1^{-1} Mpc1^{-1}, c=299,792.458 km s1c = 299{,}792.458\ \text{km s}^{-1}.

The distance modulus is:

μ=255log10(h)+5log10(DLMpc)\mu = 25 - 5\log_{10}(h) + 5\log_{10}\left(\frac{D_L^*}{\text{Mpc}}\right)

where DLDL(h=1)D_L^* \equiv D_L(h=1) is the luminosity distance computed with h=1h=1.

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Ωm10.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]

where:

η(a,Ωm)=2s3+1[1a40.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}

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 within 0.4% or better. 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.

Validation: Your implementations should agree with Ned Wright’s Cosmology Calculator to within ~0.1%. For the worked example in the Scientific Background (z=0.5z=0.5, Ωm=0.3\Omega_m=0.3, h=0.7h=0.7), you should get DL1,983D_L^* \approx 1{,}983 Mpc and μ42.26\mu \approx 42.26 mag (±0.02).


1.2: Likelihood Function

Module: likelihood.py

Required functions:

Mathematical Background:

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

lnL(θ)=12rC1r\ln \mathcal{L}(\theta) = -\frac{1}{2}\mathbf{r}^\top \mathbf{C}^{-1} \mathbf{r}

where the residual vector is:

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

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

Note on constants: The full Gaussian log-likelihood includes normalization terms 12lnCn2ln(2π)-\frac{1}{2}\ln|\mathbf{C}| - \frac{n}{2}\ln(2\pi). Since C\mathbf{C} is the fixed data covariance (independent of model parameters θ\theta), these are parameter-independent constants that cancel in Metropolis-Hastings and HMC acceptance ratios. We omit them for computational efficiency.

Implementation notes - CRITICAL:

Do NOT invert the covariance matrix. Instead, use Cholesky decomposition for numerical stability and efficiency:

from scipy.linalg import cho_factor, cho_solve

# One-time setup (do this once, outside your likelihood function):
C_factor = cho_factor(C)  # Returns (c, lower) tuple

# In your likelihood function (reuse C_factor):
# To compute C^{-1} @ r efficiently:
C_inv_r = cho_solve(C_factor, r)
chi_squared = np.dot(r, C_inv_r)
log_like = -0.5 * chi_squared

Why Cholesky?

Prior specification:

Use flat (uniform) priors with physically motivated bounds:

Return -np.inf for parameters outside these bounds, return 1 for parameters inside since the constant cancels out in the acceptance ratio.

Justification: These bounds are informed by independent observations (CMB, BAO, local distance ladder) but are broad enough to let the supernova data speak. You can justify different bounds in your memo if you have good physical reasons.


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. Compute the Cholesky factorization once during data loading, then pass the factored form to 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(θ(i1),Σprop)\theta^* \sim \mathcal{N}(\theta^{(i-1)}, \Sigma_{\text{prop}}) (multivariate normal proposal)

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

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

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

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

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

  3. Return chain and acceptance rate

Returns:

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

See the [Appendix] for more information on tuning your MCMC proposal distribution.

The Multivariate Proposal Distribution

For the proposal distribution, you can use either:

For MCMC where you’re drawing thousands of proposals, the NumPy options are typically faster. We recommend using np.random.default_rng() for better reproducibility and independent random streams across chains.


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)}

Effective Sample Size:

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

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

Gelman-Rubin Statistic (split-R^\hat{R}): Compares within-chain to between-chain variance. For mm chains of length nn, compute the potential scale reduction factor:

R^=V^W\hat{R} = \sqrt{\frac{\hat{V}}{W}}

where V^\hat{V} combines within-chain and between-chain variance, and WW is the average within-chain variance. Modern implementations use “split-R^\hat{R}” which splits each chain in half to increase the number of chains for more robust diagnostics.

Convergence is achieved when R^<1.1\hat{R} < 1.1 for all parameters.

See Module 5 Part 4 for implementation details and formulas.


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)

where:

Hamilton’s equations:

dθdt=Hp=M1p\frac{d\theta}{dt} = \frac{\partial H}{\partial p} = M^{-1}p

dpdt=Hθ=U(θ)=lnp(θD)\frac{dp}{dt} = -\frac{\partial H}{\partial \theta} = -\nabla U(\theta) = \nabla \ln p(\theta|D)

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 θlnp(θD)\nabla_\theta \ln p(\theta|D). Two approaches:

Finite differences:

fθif(θ+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}

Use h105h \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: pN(0,M)p \sim \mathcal{N}(0, M)

    • Set (θ,p)=(θ(i1),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: ppϵ2U(θ)p \leftarrow p - \frac{\epsilon}{2}\nabla U(\theta)

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

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

    • Compute energy change: ΔH=H(θnew,pnew)H(θ(i1),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 uUniform(0,1)u \sim \text{Uniform}(0, 1)

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

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

  3. Return chain and acceptance rate

CRITICAL NOTE ON SIGNS: Since U(θ)=lnp(θD)U(\theta) = -\ln p(\theta|D) and Hamilton’s equations give dp/dt=Udp/dt = -\nabla U, the leapfrog momentum updates use subtraction:

This is opposite from N-body simulation where you had gravitational acceleration entering with a negative sign. Here, U\nabla U points “uphill” in the negative log-posterior, so we subtract it to move “downhill” toward higher probability regions.

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ΩKSK(ΩK0zdzE(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)

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)={sinh(x)ΩK>0 (open)xΩK=0 (flat)sin(x)ΩK<0 (closed)S_K(x) = \begin{cases} \sinh(x) & \Omega_K > 0 \text{ (open)} \\ x & \Omega_K = 0 \text{ (flat)} \\ \sin(x) & \Omega_K < 0 \text{ (closed)} \end{cases}

Compare your ΩK\Omega_K constraint to Planck+BAO’s measurement: ΩK=0.001±0.002\Omega_K = 0.001 \pm 0.002 (consistent with flat).


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: hN(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 (split-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, typically 5-10 pages) 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.


7.3: Growth Memo

Submit a Growth Memo (1-2 pages) reflecting on your learning process. Address:

This memo is about metacognition—thinking about your thinking. Be honest and specific. See the course Growth Memo template for detailed guidance.


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 with different colors

    • Include axis labels with units, legend identifying chains

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

    • Shows how quickly samples become independent

    • Compare MCMC vs. HMC autocorrelation length explicitly

    • Mark where autocorrelation drops below 0.1 or 0.05

  3. Corner plot: Joint and marginal posterior distributions

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

    • Shows parameter correlations and credible intervals (16th, 50th, 84th percentiles)

    • Overlay MCMC and HMC results to compare (use transparency/different colors)

  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 (different colors/styles)

    • Add legend, axis labels with units

  5. Diagnostics:

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

    • Energy conservation for HMC: histogram of ΔH\Delta H (should be centered near 0)

    • Acceptance rate vs. proposal scale/step size (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 (e.g., overlaid trace plots showing faster mixing)

Caption requirements: Each figure needs a caption with:

Quality standards: All plots should be publication-ready with:


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:


Appendix: Tuning MCMC Proposal Distribution

Tuning Your MCMC Sampler: A Practical Guide

Getting good performance from Metropolis-Hastings requires tuning the proposal covariance Σprop\Sigma_{\text{prop}}. The goal is to find a proposal scale that balances exploration (proposing new regions) with efficiency (not wasting time on rejected proposals).

The Target: Acceptance Rate

Aim for acceptance rate r[0.2,0.5]r \in [0.2, 0.5] (20-50%)

Why this range?

Step 1: Choose a Reasonable Starting Point

Initialize your chain near physically plausible values: θ(0)=(Ωm,h)(0.30,0.70)\theta^{(0)} = (\Omega_m, h) \approx (0.30, 0.70)

Step 2: Set Initial Proposal Scale

Start with a diagonal covariance matrix where each diagonal element reflects a small percentage of the parameter’s typical scale:

Σprop(0)=(σΩm200σh2)\Sigma_{\text{prop}}^{(0)} = \begin{pmatrix} \sigma_{\Omega_m}^2 & 0 \\ 0 & \sigma_h^2 \end{pmatrix}

Reasonable starting values: σΩm0.05\sigma_{\Omega_m} \approx 0.05 (5% of Ωm\Omega_m) and σh0.02\sigma_h \approx 0.02 (2-3% of hh).

Physical intuition: You’re proposing to move by a few percent of each parameter’s meaningful range per step.

Step 3: Iterate to Find Good Scale

The tuning process is iterative:

  1. Run a short test chain (500-1000 steps)

  2. Compute acceptance rate r=Naccepted/Ntotalr = N_{\text{accepted}} / N_{\text{total}}

  3. Adjust proposal scale:

    • If r<0.20r < 0.20: Scale is too large → shrink Σprop\Sigma_{\text{prop}} by factor γ2\gamma^2 with γ[0.6,0.9]\gamma \in [0.6, 0.9]

    • If r>0.50r > 0.50: Scale is too small → grow Σprop\Sigma_{\text{prop}} by factor γ2\gamma^2 with γ[1.1,1.5]\gamma \in [1.1, 1.5]

    • If r[0.20,0.50]r \in [0.20, 0.50]: You’re in the target range → proceed to Step 4

Mathematically: Σpropnew=γ2Σpropold\Sigma_{\text{prop}}^{\text{new}} = \gamma^2 \Sigma_{\text{prop}}^{\text{old}}

Repeat this adjustment process until acceptance falls in the target range.

Step 4 (Optional): Capture Parameter Correlations

Once you have reasonable acceptance with a diagonal proposal, you can optionally improve efficiency by accounting for correlations between Ωm\Omega_m and hh.

The idea: If parameters are correlated in the posterior, proposing along the correlation structure is more efficient than axis-aligned steps.

Method:

  1. Run a longer warm-up chain (e.g., 5,000-10,000 steps) with your tuned diagonal Σprop\Sigma_{\text{prop}}

  2. Discard burn-in (first 20-30% of samples)

  3. Compute the empirical covariance Σ^\widehat{\Sigma} from the warm-up samples

  4. Use Gelman’s optimal scaling:

Σprop=2.382dΣ^\Sigma_{\text{prop}} = \frac{2.38^2}{d} \widehat{\Sigma}

where dd is the parameter dimension (here d=2d=2).

Why 2.382/d2.38^2/d? This comes from asymptotic theory showing this scaling gives optimal acceptance (~23%) for multivariate Gaussian target distributions (Roberts, Gelman & Gilks 1997).

Optional refinement (if your warm-up sample is small or noisy): Apply shrinkage to regularize the off-diagonal terms:

Σprop=2.382d[(1λ)Σ^+λdiag(Σ^)]\Sigma_{\text{prop}} = \frac{2.38^2}{d}\left[(1-\lambda)\widehat{\Sigma} + \lambda \cdot \text{diag}(\widehat{\Sigma})\right]

with λ[0.05,0.2]\lambda \in [0.05, 0.2]. This interpolates between the full empirical covariance and a diagonal version.

Step 5: Run Production Chains

With your tuned Σprop\Sigma_{\text{prop}} (either diagonal or with correlations), run multiple independent chains:

Monitor that acceptance rates remain in the 20-50% range across all chains.

Step 6: Verify Convergence

Use your diagnostic functions to ensure chains have converged:


Quick Tuning Workflow

Initialize:  θ⁽⁰⁾ = (0.30, 0.70)
             Σ_prop = diag(0.05², 0.02²)

Test loop:   Run 500-1000 steps
             Compute acceptance rate r
             If r < 20%:  Shrink Σ_prop (multiply by 0.7²)
             If r > 50%:  Grow Σ_prop (multiply by 1.3²)
             If r ∈ [20%, 50%]: Exit loop

Optional:    Run warm-up (5000+ steps)
             Compute empirical covariance
             Apply 2.38²/d scaling

Production:  Run multiple chains (fixed Σ_prop)
             Monitor diagnostics
             Verify convergence

Important Considerations

Diagonal vs. correlated proposals: For this 2-parameter problem, a well-tuned diagonal proposal is often sufficient. The correlation capture (Step 4) can improve efficiency but isn’t strictly necessary for convergence.

Hard boundaries: Your flat priors have hard bounds (Ωm[0,1.2]\Omega_m \in [0, 1.2], h[0.5,0.9]h \in [0.5, 0.9]). Proposals outside these bounds return -np.inf log-posterior and are automatically rejected. Near boundaries, your effective acceptance rate will be lower—this is expected and not a problem.

When to stop tuning: You’ve tuned sufficiently when you achieve:

Don’t over-tune: Perfect tuning isn’t necessary. MCMC is remarkably robust — even suboptimal proposals will eventually converge (just more slowly). Get “good enough” and move on.


Computational Implementation Notes

You’ll need to implement:

Think about how to structure your tuning workflow: Will you build an automated tuning function, or tune interactively? How will you store and update Σprop\Sigma_{\text{prop}} between tuning iterations?


“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