fiestaEM 0.2.0

Fast inference of electromagnetic signals with JAX

fiesta 🎉

fiesta: Fast Inference of Electromagnetic Signals and Transients with jAx

NOTE: fiesta is currently under development. We have some basic documentation available under ./docs. Feel free to contact us for any questions.

Installation

fiesta can be installed from pip via

pip install fiestaEM

Alternatively you can install it directly from source by cloning

git clone [e-mail]:nuclear-multimessenger-astronomy/fiestaEM.git

and then run

pip install -e .

in the cloning directory.

Note, that by default only the cpu version of jax is installed. If you want to use GPU acceleration, run

pip install fiestaEM[gpu]

or install jax[cuda12] as indicated on the jax webpage manually.

To obtain a set of recommended built-in surrogates (see below), additionally run

python -c "from fiesta.surrogates import download_recommended_surrogates; download_recommended_surrogates()"

Loading surrogate models

Surrogates can be loaded in the python shell as

from fiesta.inference.lightcurve_model import FluxModel

model = FluxModel(name=name, filters=filters, directory=directory)

This means you have to explicitely provide the directory argument that points to the .pkl files storing the neural network and surrogate metadata. However, if a surrogate is stored under src/fiesta/surrogates/KN or src/fiesta/surrogates/GRB, it can be loaded as a built-in surrogate without providing the directory argument explicitely. To see which surrogates are built-in and can simply be loaded without providing directory explicitely, you can run

python -c "from fiesta.surrogates import print_built_in_surrogates; print_built_in_surrogates()"

Initially, there are no built-in surrogates, unless you run the command mentioned above. You can also download additional models from the github main branch. For instance,

from fiesta.surrogates import download_surrogate
download_surrogate("pbag_tophat_CVAE")

will download the pbag_tophat_CVAE surrogate. To print a list of surrogates available for download, run

python -c "from fiesta.surrogates import print_downloadable_surrogates; print_downloadable_surrogates()"

Documentation

The documentation is available at https://nuclear-multimessenger-astronomy.github.io/fiestaEM/

Training surrogate models

To train your own surrogate models, have a look at some of the example scripts in the repository for inspiration. You can find them under ./surrogates/GRB/ and ./surrogates/KN/ in the respective model folders. The example section on training is currently work in progress.

Examples

The ./examples/training/ directory contains scripts that show how to use the fiesta API to train a flux density surrogate either with a CVAE architecture or with a simple feed-forward NN. Mock training data is provided as well, but note that the data set is reduced significantly and thus you will not be able to get a decent surrogate from it. If you want to use our training data, please contact us so we can figure out a way how to deliver the heavy files (> 10GB) to you. We also have example scripts for running an inference on AT2017gfo + GRB170817A. They can be found in ./examples/inference/. Note that for all of these example scripts it is highly recommended to use GPU-acceleration, since otherwise the runtime will rather long.

For Developers

Running tests

Install the package in editable mode and run the test suite with:

pip install -e .
python -m coverage run --source fiesta -m pytest tests/*.py

To run a single test file:

python -m pytest tests/test_models.py

Building the docs locally

Install the documentation dependencies:

pip install -e ".[docs]"

Then build the HTML docs from the docs/ directory:

sphinx-build docs docs/_build/html

Open the result in your browser:

open docs/_build/html/index.html   # macOS
xdg-open docs/_build/html/index.html  # Linux

For a live-reloading development server:

sphinx-autobuild docs docs/_build/html
# Visit http://127.0.0.1:8000

To check for broken references (mirrors the CI check):

sphinx-build -W --keep-going docs docs/_build/html

Citing fiesta

If you use fiesta in your research, please cite our paper:

@article{Koehn:2025zzb,
    author = "Koehn, Hauke and Wouters, Thibeau and Pang, Peter T. H. and Bulla, Mattia and Rose, Henrik and Wichern, Hannah and Dietrich, Tim",
    title = "{Efficient Bayesian analysis of kilonovae and gamma ray burst afterglows with FIESTA}",
    eprint = "2507.13807",
    archivePrefix = "arXiv",
    primaryClass = "astro-ph.HE",
    doi = "10.1051/0004-6361/202556626",
    journal = "Astron. Astrophys.",
    volume = "704",
    pages = "A55",
    year = "2025"
}

Acknowledgements

The logo was created by ideogram AI.