Finsler Flow Matching

This is the code for the paper Finsler Flow Matching: Dynamics-Aware Geodesic Interpolation for Single-Snapshot Trajectory Inference.

1. Geodesic interpolant — fit φ by minimising E F(x_t, ẋ_t)².
2. Finsler-cost OT — score each (x₀, x₁) pair by the mean Finsler energy along the frozen interpolant, then entropic Sinkhorn for the coupling π*.
3. Distillation — regress v_θ onto ẋ_t under π*, then RK4 from p₀.

Install and run

pip install -r requirements.txt    
jupyter lab notebooks/

Variable	Effect
NB_SMOKE=1	one seed, cut budgets, separate output trees
NB_CPU=1	force CPU

Five benchmarks

Notebook	Cloud	Task	Train cost
sheet_ffm_paper	4 000 points on a saddle in ℝ³	withheld strip between p₀ and p₁	~20 GPU-min
forksheet_ffm_paper	3 000 points, bifurcating sheet	withheld stretch where the arms resolve	~50 GPU-min
pancreas_ffm_paper	Endocrinogenesis, 3 696 cells	middle latent_time marginal withheld, d = 2/10/20/50	~90 GPU-min
erythroid_ffm_paper	Mouse gastrulation erythroid, 9 815 cells	3 spaces × 7 arms × 5 seeds	~4.5 GPU-h
itracer_r2_lineage_kernel	iTracer R2 hindbrain, 2 901 cells	P from expression × lineage × velocity kernels	~3.5 GPU-h

Every notebook trains its baselines, OT-CFM, MFM/LAND, Curly-FM, through the same trainer, budget and network as FFM, so a column difference is the geometry, not the harness.

Layout

FFM/
├── notebooks/        five benchmarks and the extension study bottom
│   ├── out/          figures (.png + .pdf)
│   └── tables/       the numbers they print
├── scripts/
│   ├── core/         geometry, arms, bridge, metrics, data loaders
│   │   └── vendor/   MIT-licensed upstream code for the MFM and Curly-FM baselines
│   ├── method/       the same maths in plain PyTorch, as the notebooks read it
│   └── experiments/  one package per benchmark
├── data/             derived caches (39 MB)
└── outputs/          

scripts/core/ produced the paper's tables; scripts/method/ is a second, independent implementation of the same mathematics. Both ship, deliberately.

Data and outputs

Path	What	Rebuild
pancreas_cache/	5 MB — the PCA/velocity cloud	python -m scripts.core.scvelo_data --rebuild --path <h5ad>
erythroid_cache/	9 MB — three PCA prefixes, d = 2/20/50	python -m scripts.experiments.erythroid.preprocess (needs the 1.4 GB Pijuan-Sala h5ad and a UniTVelo/TensorFlow env)
erythroid_sweep/	2 MB — the selection grid the sensitivity table reads	scripts/experiments/erythroid/tune.py
itracer/	23 MB — the R2 cloud and its velocity-alignment cache	needs the Mendeley supplement and a velocity h5ad
band/	95 KB — band-access records for the pancreas figure	scripts/experiments/pancreas/band.py

outputs/ is rebuilt rather than versioned, with eight exceptions (34 MB) that would otherwise cost two GPU-days before a single number appeared:

Path	Buys you
matched_bench/, pancreas_band/, pancreas_noise/	the pancreas ablation sections, ~6.5 GPU-h
erythroid_bench/	105 run records, ~4.5 GPU-h
itracer_bench/	108 run records (seeds 3–11 at d = 2 for itr_seeds), ~6 GPU-h
pancreas_review/, pancreas_review_hw/	the pancreas robustness grid, 2 123 runs, ~38 GPU-h
erythroid_review/	the erythroid seed pool and perturbations, 127 runs, ~1 GPU-h

Credits

The baselines are the authors' own code, vendored under their MIT licences in scripts/core/vendor/, with the upstream file and line recorded per function:

• Kapusniak et al. (2024), Metric Flow Matching for Smooth Interpolations on the Data Manifold
• Petrović et al. (2025), Curly Flow Matching the Curly-FM arm and the erythroid benchmark's shape.
• Tong et al. (2024), Simulation-free Schrödinger bridges via score and flow matching

Data: Bastidas-Ponce et al. (2019) pancreatic endocrinogenesis Pijuan-Sala et al. (2019) mouse gastrulation atlas; He et al. (2022), Lineage recording in human cerebral organoids.