cosmulator

Accelerating cosmological inference with JAX machine-learning emulators on GPU

I am the principal investigator of the UKRI-funded DiRAC high-performance computing project behind cosmulator, awarded 10,000 GPU hours on the DiRAC Tursa supercomputer to train machine-learning emulators.

cosmulator trains JAX normalising-flow emulators of nested-sampling posteriors on GPU, so a distribution can be resampled instantly for fast marginalisation and Bayesian evidence, and manages the resulting emulators (training, validation, storage and serving) across a grid of cosmological models and surveys.

An emulator at work. A posterior emulator (orange), trained on the DiRAC Tursa supercomputer, reproduces the Dark Energy Survey Year 1 constraints on the cosmological parameters of the ΩkΛCDM model from its full analysis (blue). The diagonal panels show the one-dimensional marginalised posteriors, and the inner and outer contours in the lower triangular panels correspond to the 68% (1σ) and 95% (2σ) credible regions.

Installation

pip install cosmulator            # inspect and use emulators (numpy only)
pip install "cosmulator[train]"   # adds the JAX training stack

Requires Python 3.11 or later.

Resources