This repository accompanies the JCAP-submitted paper Benchmarking field-level cosmological inference from galaxy surveys.
Field-level inference is a powerful approach that allows maximum information extraction from cosmological surveys (e.g. galaxy surveys) by modeling the entire observed field rather than just its summary statistics. This project provides a benchmark suite to evaluate different preconditioning strategies and high-dimensional sampling methods, and to compare their performance on an idealized but standardized galaxy clustering model.
For standard installation of the package and its dependencies, use:
pip install git+https://github.com/hsimonfroy/benchmark-field-level.gitFor development purposes, install in editable mode:
pip install -e git+https://github.com/hsimonfroy/benchmark-field-level.git#egg=flbenchThis package relies on JAX for GPU-accelerated computations. Please note that its installation is left to the user. Follow the official JAX installation guide to set it up for your specific hardware.
The benchmarked galaxy clustering model is fast (jit-compiled and GPU accelerated) and differentiable. It includes:
- Linear matter field generation
- Structure formation, selected among
- Linear growth
- Lagrangian Perturbation Theory (1LPT or 2LPT) displacement
- Particle Mesh (PM) N-body displacement, with BullFrog or FastPM solvers
- Redshift-Space Distortions (RSD)
- Second order Lagrangian galaxy bias
- Observational Noise
The model also includes different preconditioning strategies for the linear matter field:
- Prior-preconditioning in real space
- Prior-preconditioning in Fourier space
- Static posterior-preconditioning assuming a Kaiser model
- Dynamic posterior-preconditioning assuming a Kaiser model
Interfaced MCMC samplers include:
- Hamiltonian Monte Carlo (HMC)
- No-U-Turn Sampler (NUTS)
- NUTS within Gibbs (NUTSwG)
- Metropolis-Adjusted Microcanonical Sampler (MAMS, i.e. adjusted MCHMC)
- MicroCanonical Langevin Monte Carlo (MCLMC, unadjusted)
The cosmology
The examples/ directory contains Jupyter notebooks that demonstrate how to use the benchmark tools:
-
- Experimental setup:
- Instantiate a field-level cosmological model
- Generate observation and condition the model on it
- Perform field-level inference:
- Warmup phase only inferring the field with MCLMC sampler
- Warmup phase inferring all parameters jointly with any implemented sampler
- Experimental setup:
-
- Assess convergence:
- Chain diagnostics
- Inspection at the field-level
- Quantify performance
- Assess convergence: