SFI – Stochastic Force Inference

Efficient, low‑bias force & diffusion estimation from Langevin time series

Version 1.9.x – public beta (API stabilised but not yet frozen; v2 features already under active testing)

Installation

pip install stochasticforceinference

---

What is SFI?

SFI is a Python toolkit for data‑driven inference of force (drift) and diffusion fields that underlie stochastic differential equations (SDEs) such as in physics, biology and active‑matter research. Starting from discrete, noisy time series (trajectories) it combines:

• Multi‑point finite‑difference estimators that remain low bias even at large sampling intervals and high localisation noise.
• A quasi-maximum‑likelihood approximation that separates diffusion and force inference – boosting robustness by decoupling ill‑conditioned parameters.
• A basis‑decomposition regression allowing flexible decomposition of force and diffusion fields into physically meaningful terms, for instance through they symmetries.
• A sparse model selection algorithm to reduce basis expansions into statistically significant terms via the PASTIS information criterion.
• JAX acceleration (vectorization with vmap, just-in-time compilation with jit, gradient evaluation with autodiff) with optional GPU/TPU offload, delivering ≳ 100× CPU speed‑ups over v1.

---

Mathematical setting

We consider trajectories ${\mathbf x(t_k)}$ generated by Langevin dynamics in the Itô convention with additive white Gaussian noise,

$$ \begin{aligned} \text{(OLI)}\qquad & \dot{\mathbf x} = \mathbf F(\mathbf x) + \sqrt{2\mathbf D(\mathbf x)}\boldsymbol{\xi}(t),\\ \text{(ULI)}\qquad & \dot{\mathbf x} = \mathbf v, \qquad \dot{\mathbf v} = \mathbf F(\mathbf x,\mathbf v) + \sqrt{2\mathbf D(\mathbf x,\mathbf v)}\boldsymbol{\xi}(t), \end{aligned} $$

with $\langle\xi_\mu(t),\xi_\nu(t')\rangle=\delta_{\mu\nu},\delta(t-t')$ and time step $\Delta t$ that need not be vanishingly small.

Basis‑decomposition inference

The core of the code relies on a decomposition of force and diffusion fields as linear combinations of basis functions:

$$ \mathbf D(\mathbf z) \approx \sum_a d_a ,\boldsymbol b_a(\mathbf z),\qquad \mathbf F(\mathbf z) \approx \sum_a c_a,\boldsymbol b_a(\mathbf z), $$

with $\mathbf z!=!\mathbf x$ (OLI) or $(\mathbf x,\mathbf v)$ (ULI). This is done by quasi-likelihood maximization corresponding to a linear regression of instantaneous diffusion and force estimators with smooth state-dependent function b(z).

Symmetries (radial, spherical, alignment…) are encoded in the vector/tensor structure of $\boldsymbol b_a$ or in pair‑interaction kernels for many‑body systems.

A forthcoming non‑linear variant will fit parametric forms $\mathbf F(\mathbf z;\boldsymbol\theta)$ and optimise $\boldsymbol\theta$ directly – see Roadmap.

PASTIS sparse selection

The PASTIS criterion maximises a quasi‑likelihood-based information criterion $I_{\text{PASTIS}}$ with a complexity penalty that scales with the logarithm of the total basis size. This extreme‑value‑theory‑motivated correction controls false positives when comparing many nested models, enabling principled sparse discovery of relevant force terms. The current implementation relies on a skyline identification (best model for each complexity) through a beam search, then maximization on that skyline.

---

Why SFI?

Key point	Impact
Diffusion - Force decoupling	Robust fits, avoids joint ill‑conditioning
Multi‑point (2, 3, 4‑pt) estimators	Low bias at large $\Delta t$ or high noise
Automatic estimator selection	Picks the right estimator for the data (ULI support incoming)
Single‑ & multi‑particle modes	Generic vector fields or pair interactions; automatic masking tolerates missing tracks & particle entry/exit
Symmetry‑aware bases	Enforce radial/spherical/tensor constraints
JAX backend (vmap·jit·grad)	≳ 100× CPU speed‑up; effortless GPU/TPU acceleration
CSV trajectory import	One‑liner load_trajectory_csv()

---

Quick start

# Load a 'trajectory.csv' file
meta, particle_idx, time_idx, xvals = SFI.SFI_utils.load_trajectory_csv("path/to/trajectory.csv")
data = SFI.StochasticTrajectoryData(xvals, time_idx, meta["dt"],
                                    particle_indices=particle_idx)

# Choose inference engine (OLI here)
S = SFI.OverdampedLangevinInference(data)

# Diffusion then force
S.compute_diffusion_constant(method="MSD")

# Get basis functions and descriptors
(basis_F, grad_F), descriptor = SFI.OLI_bases.basis_selector(
        {"type": "polynomial", "order": 3}, data.d, output="vector")

S.infer_force_linear(basis_linear=basis_F,
                      basis_linear_gradient=grad_F)
		    
# Model selection:		      
S.sparsify_force()

# Computing error and printing:
S.compute_force_error() 
S.print_report()

---

Single vs. multi‑particle systems

• Single‑particle: fits $\mathbf F(\mathbf x)$ from one or many trajectory segments; missing frames are masked automatically.
• Multi‑particle: infers pair‑interaction kernels $\mathbf F(\mathbf r)$ and single-particle forces; masking handles particles leaving/entering the field of view (some bias if attrition is extreme).

---

Library structure

Module	Purpose
SFI_data.py	Trajectory container & finite‑difference helpers
OLI_inference.py	Overdamped drift & diffusion inference
ULI_inference.py	Underdamped inference
OLI_bases.py, ULI_bases.py	Basis factories (polynomial, Fourier, pair interactions, custom with symmetry options)
SFI_sparsity.py	Sparse model selection (skyline construction, PASTIS/AIC/BIC maximization)
SFI_Langevin.py	Simple Langevin dynamics simulators for method benchmarking & predictivity assesment
SFI_utils.py	CSV I/O, convenience wrappers, pretty-printing
SFI_base_inference.py	Common backend for inference methods

---

Examples

Each folder * contains a *_model.py file with model definition and parameters; a *_simulate.py to generate a synthetic dataset (several presets: tiny for debugging, small/medium/large for benchmarking); and a OLI/ULI_demo_*.py file for a commented step-by-step analysis of the data set including loading, formatting, basis definition, inference, error estimation, sparsification; optional comparison to ground truth model and plotting; and finally bootstrapped simulation of new trajectories using the inferred force and diffusion fields.

Folder	Regime	System	Highlights
OLI_examples/Lorenz/	OLI	Lorenz attractor (3 D)	Rebuilds nonlinear polynomial forces - start here for OLI
OLI_examples/Lotka-Volterra/	OLI	Lotka–Volterra predator–prey	Recovers a sparse interaction matrix
OLI_examples/ActiveBrownianParticles	OLI	Active Brownian particles	Pair interactions: central forces + Vicsek-like torques
ULI_examples/VanDerPol	ULI	Van‑der‑Pol oscillator	Identify nonlinear oscillator - start here for ULI
ULI_examples/DampedHarmonicOscillator	ULI	Damped harmonic oscillator	Sparse coupling matrix & friction inference in d dimensions
ULI_examples/Flocking	ULI	Bird‑flock model	Infers alignment, central forces & propulsion

---

Changelog – v1.9 vs v1.0

• Unified OLI & ULI codebases; common API
• JAX backend with vmap/jit/autodiff ⇒ ≳ 100× CPU speed‑up
• Multi‑point low‑bias estimators; automatic selection (OLI)
• Sparse model selection via PASTIS skyline
• Symmetry‑aware, flexible basis specification
• Direct CSV import; richer metadata handling

---

Roadmap to v2.0  (features implemented & under active testing)

• Non‑linear parametric force and diffusion models $\mathbf F(\mathbf z;\boldsymbol\theta)$
• Sparse selection for diffusion tensors
• Analytical error estimates for diffusion
• Bootstrap/jackknife fallback error bars
• Extended documentation & real‑world examples
• Tutorial on custom basis design & selection

---

References

• SFI (OLI): A. Frishman & P. Ronceray, Learning force fields from stochastic trajectories, Phys. Rev. X 10, 021009 (2020).
• ULI: D. B. Brückner, P. Ronceray & C. P. Broedersz, Inferring the dynamics of underdamped stochastic systems, Phys. Rev. Lett. 125, 058103 (2020).
• PASTIS: A. Gerardos & P. Ronceray, Parsimonious model selection for stochastic dynamics, arXiv:2501.10339 (2025).

If you use SFI in academic work, please cite the relevant paper(s).

---

Released under the MIT License • © 2025 Pierre Ronceray