# FNOPE: Simulation-based inference on function spaces with Fourier Neural Operators

This repository contains the code used to produce the results presented in the submission:
"FNOPE: Simulation-based inference on function spaces with Fourier Neural Opeartors". Some of
the attributions have been removed to preserve anonymity.

To install the python environment, create a virtual environment (e.g., using conda) with 
Python >= 3.11, and install the package and all its dependencies from `pyproject.toml` 
using:

```
pip install -e .
```

NOTE: to run the simulator for the Darcy Flow, additional dependencies should be
installed. Please see `simulators/Readme_darcy.md` for full details.

We use [Hydra](https://hydra.cc/) to submit multiple jobs in parallel and for
configuration management. To submit training and evaluation runs, use the scripts under
`src/experiments`. For example, to run submit the training and evaluation scripts for 
the baseline methods, run:

```
python src/experiments/run_baseline_sbi_${task_name}.py
```

where ${task_name}$ is one of `linear_gaussian`,`sir`,`darcy`, or `ice` for each of the 
experiments respectively.Similarly, to train and evaluate FNOPE and FNOPE (fix) methods,
run:

```
python src/experiments/run_fno_fmpe_${task_name}.py
```

Note that these files use the configurations defined by
`src/config/baseline_sbi_config.yaml` and `src/config/fno_fmpe_config.yaml`
respectively. To run the different experiments, make sure that the respective config
groups (data_config, model_config, evaluation_config, sweep) defined in these files
correspond to the experiment you wish to run. The `sweep` config group submits multiple
jobs, overwrriting the method parameters and the simulation budgets. If you wish to run 
only one job locally, simply comment out the `sweep` config group in these config files.
The results of these runs will be saved into the `out` folder.

To compile the results of individual runs into a summary file for an experiment, use the 
notebooks `compile_results_${task_name}.ipynb`. To generate the plots in the manuscript,
use the notebooks in the `plots` folder.

Note that not all the data is included in this repository due to the file size
constraint of OpenReview. The SIR and Darcy Flow experiment plots require also files
including posterior samples as well as posterior predictive samples, which are too large
to include.