Stochastic Process Learning via Operator Flow Matching

This repository contains the code for "Stochastic Process Learning via Operator Flow Matching". https://arxiv.org/abs/2501.04126

Our implementation relies on torchcfm , neuraloperator packages and benefits from repositories FFM and OpFlow.

Example of applying OFM for Gaussian Process (GP) and non-GP data regression

GP data example

non-GP black hole example

Environment

Our implementation uses Anaconda and Jupyter Notebook. To set up the environment, create a conda environment:

# clone project
git clone https://github.com/yzshi5/SPL_OFM.git
cd SPL_OFM

# create conda environment
conda env create -f environment.yml

# Activate the `ofm` environment
conda activate ofm

Install the ipykernel to run the code in a jupyter notebook

conda install -c anaconda ipykernel

pip install ipykernel

python -m ipykernel install --user --name=ofm

Description for folders and files

ofm_OT_likelihood.py, serves as the key file, see comments in the file for instructions

util folder contains the GP prior implementation and other helper functions

model folder includes FNO implementation, we also provide FNO with differential kernel

prior_learning folder contains all prior learning tasks

regression folder contains all regression tasks

sampling_FSGLD folder contains the code for SGLD sampling

Reference

@article{shi2025stochastic,
  title={Stochastic Process Learning via Operator Flow Matching},
  author={Shi, Yaozhong and Ross, Zachary E and Asimaki, Domniki and Azizzadenesheli, Kamyar},
  journal={arXiv preprint arXiv:2501.04126},
  year={2025}
}