Metadata-Version: 2.4
Name: fiestaEM
Version: 0.2.0
Summary: Fast inference of electromagnetic signals with JAX
Author-email: Hauke Koehn <hauke.koehn@uni-potsdam.de>, Thibeau Wouters <thibeauwouters@gmail.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/nuclear-multimessenger-astronomy/fiestaEM
Keywords: sampling,inference,astrophysics,kilonovae,gamma-ray bursts,afterglow,surrogates
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: dill
Requires-Dist: jax>=0.4.14
Requires-Dist: jaxlib>=0.4.14
Requires-Dist: flax
Requires-Dist: numpy
Requires-Dist: pandas
Requires-Dist: jaxtyping
Requires-Dist: beartype
Requires-Dist: tqdm
Requires-Dist: scipy
Requires-Dist: ml_collections
Requires-Dist: matplotlib
Requires-Dist: astropy
Requires-Dist: sncosmo
Requires-Dist: flowMC
Requires-Dist: h5py
Requires-Dist: corner
Requires-Dist: huggingface_hub>=0.13.0
Provides-Extra: grb
Requires-Dist: afterglowpy; extra == "grb"
Provides-Extra: gpu
Requires-Dist: jax[cuda13]; extra == "gpu"
Provides-Extra: docs
Requires-Dist: sphinx>=7.0; extra == "docs"
Requires-Dist: sphinx-autobuild; extra == "docs"
Requires-Dist: sphinx-book-theme>=1.0; extra == "docs"
Requires-Dist: sphinx-design>=0.5; extra == "docs"
Requires-Dist: sphinx-autodoc-typehints>=2.0; extra == "docs"
Requires-Dist: sphinx-copybutton>=0.5; extra == "docs"
Requires-Dist: myst-parser>=2.0; extra == "docs"
Requires-Dist: dust-extinction; extra == "docs"
Dynamic: license-file

# fiesta 🎉

`fiesta`: **F**ast **I**nference of **E**lectromagnetic **S**ignals and **T**ransients with j**A**x

![fiesta logo](docs/fiesta_logo.jpeg)

**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 git@github.com: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](https://docs.jax.dev/en/latest/installation.html#installation) 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:

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

To run a single test file:

```bash
python -m pytest tests/test_models.py
```

### Building the docs locally

Install the documentation dependencies:

```bash
pip install -e ".[docs]"
```

Then build the HTML docs from the `docs/` directory:

```bash
sphinx-build docs docs/_build/html
```

Open the result in your browser:

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

For a live-reloading development server:

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

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

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

## Citing fiesta

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

```bibtex
@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](https://ideogram.ai/).
