# Hierarchical Neural Simulation-Based Inference Over Event Ensembles: Supplemental Material

## Code

Files necessary to reproduce the experiments presented in the paper are in the `code` folder. Details for each experiment is given below.

### 1. Simple multi-variate normal likelihood

The notebook `notebooks/01_mvn_conjugate_prior.ipynb` can be used to run this experiment in a self-contained manner. The codes `models/deep_set.py` and `models/transformer.py` contain auxiliary code to define the models used.

### 2. Mixture models in particle physics: Bayesian inference for a narrow resonance

To reproduce Fig. 4 (left): run `notebooks/11_bump_hunt_plots.ipynb`

To reproduce Fig. 4 (center, right): The below instructions assume access to a `condor` batch system supporting `singularity`.

* Setup `singularity` images:
```
singularity build /path/to/pytorch.sif notebooks/pytorch_batch/pytorch.def
singularity build /path/to/pymc.sif notebooks/pymc_batch/pymc.def
```
* Setup environment:
```
source setup.sh
```
* Run MCMC campaign: run `notebooks/pymc_batch/RunMCMCCampaign.py --campaigndir path/to/MCMC/campaign --imagepath path/to/pymc.sif`
* Run deep set training: `notebooks/pytorch_batch/RunTrainingCampaign.py --campaigndir path/to/training/campaign --imagepath `path/to/pytorch.sif`
* Make comparison plots: Update paths to MCMC and training campaign directories and run `notebooks/10_MCMC_deepset_comparison.ipynb`

### 3. Astrophysics example: Strong gravitational lensing

The notebook `notebooks/05_lensing.ipynb` can be used to run this experiment in a self-contained manner. The notebook runs simulations of 
lensing images drawn from the forward model, implements the hierarchical deep set for learning global and local parameters simultaneously, and performs inference over local and global parameters. 

The notebook `notebooks/13_lensing_coverage.ipynb` can be used to reproduce the SBI calibration plots in Fig. 6.

### 4. Mixture models in particle physics: frequentist treatment

The notebook `notebooks/04_FrequentistTrain.ipynb` can be used to run train the deepset network used for learning a frequentist test statistic. The model is implemented in `models/deep_set_freq.py`.
