- Albatross is a Simulation-based Inference (SBI) library designed to perform analysis on Milky Way stellar streams. It is built on top of the swyft code, which implements neural ratio estimation to efficiently access marginal posteriors for all parameters of interest.
- Related paper: The details regarding the implementation of the TMNRE algorithm and the specific demonstration for mock GD1-like stellar streams can be found in arxiv:2304.02032.
- Modelling Code: In paralell, we develop a
jax
-accelerated modelling codesstrax
which is available for download from this repo. - Key benefits: We showed in the above paper that albatross is extremely sample efficient when constraining e.g. the 16 parameters in our current model, requiring only 350,000 simulations to perform inference across the full parameter space. The method is also an 'implicit likelihood' technique, so it inherits all the associated advantages such as the fact that it does not require an explicit likelihood to be written down. This opens up the possibility of using albatross to analyse a wide range of interesting physical effects relevant to stellar streams, their environment and evolution history.
- Contacts: For questions and comments on the code, please contact either James Alvey, Mathis Gerdes or Christoph Weniger. Alternatively feel free to open an issue.
- Citation: If you use albatross in your analysis, or find it useful, we would ask that you please use the following citation.
@article{Alvey:2023pkx,
author = "Alvey, James and Gerdes, Mathis and Weniger, Christoph",
title = "{Albatross: A scalable simulation-based inference pipeline for analysing stellar streams in the Milky Way}",
eprint = "2304.02032",
archivePrefix = "arXiv",
primaryClass = "astro-ph.GA",
month = "4",
year = "2023"
}
The safest way to install the dependencies for albatross
is to create a virtual environment from python>=3.8
Option 1 (venv):
python3 -m venv /your/choice/of/env/path/
- Source the new environment
source /your/choice/of/env/path/bin/activate
Option 2 (conda):
conda create -n your_env_name python=3.x (python>=3.8 required)
conda activate your_env_name
- Clone the peregrine repo into location of choice
cd /path/to/your/code/store/
git clone git@github.com:undark-lab/albatross.git # for ssh install
(or git clone https://github.com/undark-lab/albatross.git # for https install)
pip install git+https://github.com/undark-lab/swyft.git@f036b15dab0664614b3e3891dd41e25f6f0f230f
pip install tensorboard psutil configparser pathlib
cd /path/to/your/code/store/
git clone git@github.com:undark-lab/sstrax.git # for ssh install
(or git clone https://github.com/undark-lab/sstrax.git # for https install)
cd sstrax
pip install .
Additional instructions, documentation and examples for the sstrax
code can be found at the corresponding repo
Key run files:
generate_observation.py
- Generates a test observation from a configuration file given a set of injection parameterstmnre.py
- Runs the TMNRE algorithm given the parameters in the specified configuration filecoverage.py
- Runs coverage tests on the logratio estimators that have been generated bytmnre.py
Example Run Scheme:
- Step 1: Generate a configuration file following the instructions in the examples directory. To just do a test run, you will only need to change the
store_path
andobs_path
options to point to the desired location in which you want to save your data. - Step 2: Change directory to
albatross/albatross
where the run scripts are stored - Step 3: Generate an observation using
python generate_observation.py /path/to/config/file.txt
or point to a desired observation in the configuration file - Step 4: Run the inference algorithm using
python tmnre.py /path/to/config/file.txt
, this will produce a results directory as described below - Step 5: (optional): Run the coverage tests using
python coverage.py /path/to/config/file.txt n_coverage_samples
(n_coverage_samples = 2000
is usually a good start)
Result output:
config_[run_id].txt
- copy of the config file used to generate the runbounds_[run_id]_R[k].txt
- bounds on the individual parameters from Roundk
of the algorithmcoverage_[run_id]/
- directory containing the coverage samples ifcoverage.py
has been runlogratios_[run_id]/
- directory containing the logratios and samples for each round of inference (stored in fileslogratios_R[k]
for each roundk
. These can be loaded using thepickle
python library)observation_[run_id]
-pickle
file containing the observation used for this run as aswyft.Sample
object. The same observation is used for both the TMNRE algorithm and any traditional sampling approach.param_idxs_[run_id].txt
- A list of parameter IDs that can be matched to the logratios results files and used for plotting purposes.simulations_[run_id]_R[k]/
-Zarrstore
directory containing the simulations for Roundk
of inferencetrainer_[run_id]_R[k]/
- directory containing the information and checkpoints for Roundk
of training the inference network. This directory can also be passed totensorboard
astensorboard --logdir trainer_[run_id]_R[k]
to investigate the training and validation performance.log_[run_id].log
- Log file containing timing details of run and any errors that were raised
- v0.0.1 | April 2023 | Initial release based on arxiv:2304.02032. Corresponding
sstrax
release:v0.0.1
.