Neural Ratio Estimation (NRE) for cosmological parameter inference from strong gravitational lensing images.
This repository provides code for the paper:
Cosmology Inference from Strong Gravitational Lensing using Neural Ratio Estimation
We train a ResNet-based binary classifier on ~2 million simulated DES-noise strong lensing images to learn the likelihood ratio r(x, θ) = p(x|θ)/p(x) for the dark energy equation of state w and matter density Ω_m. Individual per-lens log-ratios are summed across a population of lenses for joint cosmological inference. We demonstrate rank-histogram-based post-hoc calibration and show the method is ~100× more data-efficient than analytical likelihood approaches.
Clone the repository and create the conda environment:
git clone https://github.com/deepskies/CosmoLensNRE.git
cd CosmoLensNRE
conda env create -f environment.yml
conda activate <env-name>Or with Poetry (Python ≥ 3.10):
pip install poetry
poetry installCosmoLensNRE/
├── notebooks/
│ ├── train_2M.ipynb # Train NRE classifier on 2M lensing images
│ ├── compare_analytical_nre_likelihood.ipynb # Compare NRE vs. analytical Einstein-radius likelihood
│ ├── plot_dataset.ipynb # Visualize training images
│ ├── plot_image_posterior.ipynb # Per-lens posteriors: uncalibrated vs. calibrated
│ ├── plot_parity_residuals_regions.ipynb # Parity plots and residuals across parameter regions
│ ├── plot_population_posteriors.ipynb # Population-level (w, Ω_m) posteriors
│ ├── visualize_grid.ipynb # Test data regions in (Ω_m, w) parameter space
│ └── w0_om0_degeneracy.ipynb # w–Ω_m degeneracy in Einstein radius
├── src/scripts/
│ ├── evaluate_with_calibration.py # MCMC sampling with calibrated log-ratio
│ └── evaluate_without_calibration.py # MCMC sampling with uncalibrated NRE log-ratio
├── environment.yml
├── pyproject.toml
└── LICENSE.txt
Lensing images are simulated using deeplenstronomy with DES-like noise on a 32×32 pixel grid. The training set contains ~2 million images spanning a grid of (w, Ω_m) values.
The train, test, and calibrated dataset is available on Zenodo.
- Sreevani Jarugula (Fermilab)
- Brian Nord (Fermilab / University of Chicago)
- Aleksandra Ćiprijanović (Fermilab / University of Chicago / SkAI Institute)
- Shubhendu Trivedi (University of Chicago)
If you use this code, please cite the paper:
This project is licensed under the MIT License — see LICENSE.txt for details.
