# Iterative Smoothing Bridge (ISB) Model

A minimal implementation of the Iterative Smoothing Bridge method, from the submission *"Transport with Support: Data-Conditional Diffusion Bridges"*. See the instructions below for how to run the reference implementation.

NB: To ensure double-blindness, the files have been stripped off of names and acknowledgements for those codes that are not publicly available.

## Set up the environment

### Create a conda environment using the file isb-model-env.yaml
In the project root folder, run
```
conda create env -f isb-model-env.yaml
```

### Install the project as a package
Activate the environment
```
conda activate isb-model-env
```

and install the project package
```
pip install -e ./src
```

## Running the training script

### Data sets
We have provided the S-shape observations in a csv file in the folder `./data/s_shape/observations`.

### Configs
We have a provided a config file for running the S-shape experiment (Figure 2 in manuscript). Modify the base_folder argument to match the root folder of this project, and keep the other config parameters fixed.

### Running the task
To run the the training script, run the following two lines in the project root folder
```
conda activate isb-model-env
```
and
```
python run_iterative_smoother.py 
```

### Output
Running the script creates three kinds of outputs: plots, model pickles, and config logs.

#### Plots
The script will write plots to folder `./plots/isb_model/s_shape/iterative_training`, each model run will be placed in a separate folder, named based on the current time. Within each experiment folder, there are two subfolders, `trajectories` and `videos`. After completing the model run, `trajectories` will be empty, as it is used as a staging folder for the videos. The videos with a name starting with `init_uncontrolled_{number}` show the marginal densities from the current SDE model. The final model video has a name starting with `trained_model`.

#### Model pickles
Model pickles are saved in the folder "model", with the subfolders named as the config arguments `config.model.fwd_drift_file_name` and `config.model.bwd_drift_file_name`.

#### Config logs
The config manager hydra stores the config file used for each run in the folder `outputs`.

## License
Code to be released under the MIT license upon acceptance.

