AMPLFI

Accelerated Multi-messenger Parameter estimation with Likelihood Free Inference

Framework for performing rapid parameter estimation of gravitational wave events using likelihood free inference

Environment Setup

Note: this repository is a WIP. Please open up an issue if you encounter bugs, quirks, or any undesired behavior

Note: Running AMPLFI out-of-the-box requires access to an enterprise-grade GPU(s) (e.g. P100, V100, T4, A[30,40,100], etc.). There are several nodes on the LIGO Data Grid which meet these requirements.

Please see the ml4gw quickstart for help on setting up your environment on the LIGO Data Grid (LDG) and for configuring access to Weights and Biases, and the Nautilus hypercluster. This quickstart includes a Makefile and instructions for setting up all of the necessary software, environment variables, and credentials required to run AMPLFI.

Once setup, create a fork of this repository, and clone it.

Note: Ensure that you have added a github ssh key to your account

git clone git@github.com:albert-einstein/pe.git

AMPLFI utilizes git submodules. Make sure to initialize and update those

git submodule update --init

When pulling changes from this repository, it's recommended to use the --recurse-submodules flag to pull any updates from the submodules as well.

Next, install the amplfi.law submodule, which is used for launching amplfi workflows with law

cd amplfi/law
poetry install

Finally, build the train and data project apptainer images. Set the $AMPLFI_CONTAINER_ROOT environment variable to a location where you want the images to be stored, e.g. ~/amplfi/images.

Make sure you are in the respective projects home directory when building the container. For example, from the amplfi repositories home directory

cd projects/data
apptainer build $AMPLFI_CONTAINER_ROOT/data.sif apptainer.def

will build the data image. Make sure you do the same for the train image!

Generating Data

Training and testing background strain data can be generated with the amplfi.law.DataGeneration workflow.

cd amplfi/law/

export AMPLFI_DATADIR=~/amplfi/my-first-run/data/
export AMPLFI_CONDORDIR=~/amplfi/my-first-run/condor

LAW_CONFIG_FILE=config.cfg poetry run law run amplfi.law.DataGeneration --workers 2

An example configuration file for the strain data generation can be found at amplif/law/config.cfg

Training

With training data in hand, a flow training run can be launched via the train poetry environment:

cd amplfi/projects/train/
export AMPLFI_OUTDIR=~/amplfi/my-first-run/training/
poetry run python train/cli/flow.py --config configs/flow/cbc.yaml

or by running inside the container

APPTAINERENV_AMPLFI_OUTDIR=$AMPLFI_OUTDIR APPTAINERENV_AMPLFI_DATADIR=$AMPLFI_DATADIR \
    apptainer run --nv $AMPLFI_CONTAINER_ROOT/train.sif python /opt/amplfi/projects/train/train/cli/flow.py fit \
    --config /path/to/config.yaml

Testing

Sample corner plots and probability-probability plots (PP-plots) can be generated by running the test subcommand. Remember to pass your trained model weights, which can be found at the $AMPLFI_OUTDIR directory. In this case, we pass the weights corresponding to the best validation score, which are automatically saved at $AMPLFI_OUTDIR/train_logs/best.ckpt

APPTAINERENV_AMPLFI_OUTDIR=$AMPLFI_OUTDIR APPTAINERENV_AMPLFI_DATADIR=$AMPLFI_DATADIR \
    apptainer run --nv $AMPLFI_CONTAINER_ROOT/train.sif python /opt/amplfi/projects/train/train/cli/flow.py test \
    --config /path/to/config.yaml --model.checkpoint=$AMPLFI_OUTDIR/train_logs/best.ckpt