StructureFlow: Simulation-free Structure Learning for Stochastic Dynamics

Description

StructureFlow is a simulation-free framework for jointly learning structure and stochastic population dynamics of complex dynamical systems from noisy and partial measurements. Our method builds upon advances in score and flow matching and entropy regularized optimal transport, and frames the joint inference task as a multi-marginal schr"odinger Bridge problem. StructureFlow learns a probability flow ODE that models continuous population evolution while directly encoding the directional dependencies among system variables. There are four key components: (1) the Neural Graphical Model (NGM) which models the autonomous drift and captures (stationary) system structure, (2) a time-dependent score function that captures the time-dependent stochastic dynamics, (3) an interventional model to capture conditional dynamics, and (4) simulation-free training to avoid costly trajectory simulations and help scale to high-dimensional biological systems.

If you find this code useful in your research, please cite our paper:

N. El Rimawi-Fine, A. Stecklov, L. Nelson, M. Blanchette, A. Tong, S. Y. Zhang, L. Atanackovic. "Simulation-free Structure Learning for Stochastic Dynamics." arXiv preprint arXiv:2510.16656, 2025.

@article{rimawi2025simulation,
  title={Simulation-free Structure Learning for Stochastic Dynamics},
  author={Rimawi-Fine, Noah El and Stecklov, Adam and Nelson, Lucas and Blanchette, Mathieu and Tong, Alexander and Zhang, Stephen Y and Atanackovic, Lazar},
  journal={arXiv preprint arXiv:2510.16656},
  year={2025}
}

Installation

Pip

# clone project
git clone https://github.com/NoahElRimawiFine/StructureFlow.git
cd StructureFlow

# [OPTIONAL] create conda environment
conda create -n structflow python=3.12
conda activate structflow

# install pytorch according to instructions
# https://pytorch.org/get-started/

# install requirements
pip install -r requirements.txt

Conda

# clone project
git clone https://github.com/NoahElRimawiFine/StructureFlow.git
cd StructureFlow

# create conda environment and install dependencies
conda env create -f environment.yaml -n myenv

# activate conda environment
conda activate myenv

Project Structure

The directory structure of new project looks like this:

├── .github                   <- Github Actions workflows
│
├── configs                   <- Hydra configs
│   ├── callbacks                <- Callbacks configs
│   ├── data                     <- Data configs
│   ├── debug                    <- Debugging configs
│   ├── experiment               <- Experiment configs
│   ├── extras                   <- Extra utilities configs
│   ├── hparams_search           <- Hyperparameter search configs
│   ├── hydra                    <- Hydra configs
│   ├── local                    <- Local configs
│   ├── logger                   <- Logger configs
│   ├── model                    <- Model configs
│   ├── paths                    <- Project paths configs
│   ├── trainer                  <- Trainer configs
│   │
│   ├── eval.yaml             <- Main config for evaluation
│   └── train.yaml            <- Main config for training
│
├── data                   <- Project data
│
├── logs                   <- Logs generated by hydra and lightning loggers
│
├── notebooks              <- Jupyter notebooks
│                             
│
├── scripts                <- Shell scripts
│
├── src                    <- Source code
│   ├── data                     <- Data scripts
│   ├── models                   <- Model scripts
│   ├── utils                    <- Utility scripts
│   │
│   ├── eval.py                  <- Run evaluation
│   └── train.py                 <- Run training
│
├── tests                  <- Tests of any kind
|
├── tools                  <- Tooling to run other baselines
|   ├── SINCERITIES               <- Run SINCERITIES
|   ├── OTVelo                    <- Run OTVelo
|   ├── TIGON                     <- Run TIGON
│
├── .env.example              <- Example of file for storing private environment variables
├── .gitignore                <- List of files ignored by git
├── .pre-commit-config.yaml   <- Configuration of pre-commit hooks for code formatting
├── .project-root             <- File for inferring the position of project root directory
├── environment.yaml          <- File for installing conda environment
├── Makefile                  <- Makefile with commands like `make train` or `make test`
├── pyproject.toml            <- Configuration options for testing and linting
├── requirements.txt          <- File for installing python dependencies
├── setup.py                  <- File for installing project as a package
└── README.md

How to run

Train model with default configuration

# train on CPU
python src/train.py trainer=cpu

# train on GPU
python src/train.py trainer=gpu

Train model with chosen experiment configuration from configs/experiment/

python src/train.py experiment=experiment_name.yaml

You can override any parameter from command line like this

python src/train.py trainer.max_epochs=20 data.batch_size=64

Running StructureFlow

python src/train.py experiment=synthetic_dyn-TF.yaml

This will run StructureFlow on a BoolODE dataset (dyn-TF)

Running Trajectory/Knockout Inference Experiments

Trajectory Inference Experiments (runs Structure, SF2M, RF)

python trajectory_inference_experiments.py

Knockout Inference Experiments (runs Structure, SF2M, RF)

python ko_inference_experiments.py

Trajectory Inference Experiments (OTVelo + TIGON)

cd tools/OTVelo
bash otvelo_runner.sh

cd tools/TIGON
bash tigon_runner.sh

Scaling Experiments

cd scaling_experiment
python scaling_experiment.py

License

Lightning-Hydra-Template is licensed under the MIT License.

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

Contributions

Have a question? Found a bug? Missing a specific feature? Feel free to file a new issue, discussion or PR with respective title and description.

Before making an issue, please verify that:

• The problem still exists on the current main branch.
• Your python dependencies are updated to recent versions.

Suggestions for improvements are always welcome!