# Optimal transport unlocks end-to-end learning for single-molecule localization

This repository implements our approach to solve SMLM using optimal transport as described in our submitted paper for ICLR2026. 
It provides tools for training, evaluation, and localization prediction export.
We named our method **SHOT**, for SMLM at High-ensity with Optimal Transport.

## Installation

This project was developed using python 3.13. It should work with **python version >=3.11**.
You can create a virtual environment with the tool of your choice (below is an example with python's `venv` module),
and install the project with the following commands:

```bash
python -m venv .venv
source .venv/bin/activate  
pip install external/smlmsim -e .
```

## Usage

This project uses Hydra for configuration management and Lightning Fabric for multi-GPU training.  
Main scripts are located in the `apps` subdirectory. 
Hydra configurations are available in the `config` directory and can be customized to adjust simulation parameters. 
There is no default configuration, so user must provide a configuration name with the flag `--config-name` at runtime. 

**Example configuration**

We provide an example `shot_MT0N1HDAS` configuration that corresponds to the EPFL MT0N1HDAS dataset.
Evaluation images are available at: [https://srm.epfl.ch/srm/dataset/challenge-3D-simulation/MT0.N1.HD/index.html](https://srm.epfl.ch/srm/dataset/challenge-3D-simulation/MT0.N1.HD/index.html).
You can also use the following script to download the MT0N1HDAS dataset:
```bash
./data/EPFL/pull.sh
```

**PSF**

We recommend calibrating the PSF using *SMAP* on a set of bead files.
The output `.map` file can be directly refer in the PSF config file, at `config/psf`.

You can download a pre-calibrated astigmatic PSF from EPFL using the provided script:
```bash
./psf/pull_EPFL.sh
```

**Configuration:**

Configurations are managed with Hydra and located in the `config/` directory.
Feel free to explore and play with all parameters.

**1. Training:**

Start a training with the `apps/train.py` script. Don't forget to specify a configuration (e.g., `shot_MT0N1HDAS`):

```bash
python apps/train.py --config-name=shot_MT0N1HDAS
```

Output files and logs will be save in `logs/<name_of_the_model>/v<version_number>/` .
Pytorch compilation is enabled by default; if your machine does not support it, the flag `++compile=false` turns it off.

*Note:* Backpropagation through the simulator is memory intensive. Adjust `batch_size`, `n_accum_steps`, `ds_train.length`, and `n_epochs` if needed.
```bash
python apps/train.py --config-name=shot_MT0N1HDAS ++batch_size=16 ++n_epochs=50
```
If an OOM error occurs, the training script will automatically catch it and retry using half the batch size and twice the number of gradient accumulation steps.

We provide pre-trained weights fot the MT0N1HDAS dataset at `logs/shot_MT0N1HDAS/v0/last.tar`.

**2. Evaluation:**

The `apps/eval.py` script allows to evaluate a model against the `ds_test` defined in it's config. Example with `shot_MT0N1HDAS`:

```bash
python apps/eval.py --config-name=shot_MT0N1HDAS ++weights_path=logs/shot_MT0N1HDAS/v0/last.tar ++batch_size=1
# Output should be equal or very close to:
# jac: 0.888; prec: 0.991; rec: 0.896; rmse_lat: 25.7; rmse_axial: 28.4; rmse_vol: 39.4; Elat: 0.702; Eax: 0.8; E: 0.751; offset: [-2.0, 2.0, 0.0]; wooblecorr: [-53.0, 43.0, -0.0]
```

**3. Export:**

Use `apps/export.py` to export a `.csv` in ThunderSTORM format. 
It can be read with multiple softwares like SMAP, ThunderSTORM, GDSC, ShareLoc, etc ...

Example with `shot_MT0N1HDAS`:
```bash
python apps/export.py --config-name=shot_MT0N1HDAS ++weights_path=logs/shot_MT0N1HDAS/v0/last.tar +writer=csv ++writer.filepath=exports/shot_MT0N1HDAS.csv ++batch_size=1
```
