astroemu — Next Generation Emulators for Cosmology and Astrophysics

Under active development — interfaces may change between versions.

astroemu is a generalised framework for emulating spectral signals in cosmology and astrophysics, inspired by the globalemu package. Neural network emulators are implemented in JAX, with an optax-based training loop and PyTorch-style dataloaders.

Installation

pip install astroemu

How it works

The core idea (shared with globalemu) is to input independent variables alongside physical model parameters, predicting a single spectral value per call. Full spectra are recovered via a vectorised call over all $x$ points.

Given a signal $y = f(x, \theta)$ with $N$ parameter samples and $m$ independent variable points, the training data is tiled so that parameters and independent variables are concatenated as inputs:

Input	Output
$[\theta_0,\ x_0]$	$y_0$
$[\theta_0,\ x_1]$	$y_1$
$[\theta_0,\ \ldots]$	$\ldots$
$[\theta_0,\ x_m]$	$y_m$
$[\ldots,\ \ldots]$	$\ldots$
$[\theta_N,\ x_m]$	$y_m$

For more details see the globalemu paper. A paper demonstrating applications of astroemu to a broad range of astrophysical signals is in preparation.

Documentation

git clone git@github.com:htjb/astroemu.git
pip install ".[docs]"
mkdocs serve

or browse the hosted docs at astroemu.readthedocs.io.

Contributions

Contributions are welcome! Please open an issue to discuss and read the contribution guidelines before submitting a pull request.