Stochastic Force Inference (SFI)

Learn the stochastic equation behind your time series.

SFI is a JAX-based Python package that infers the drift (force) and diffusion of Langevin stochastic differential equations from trajectories and other continuous time series: tracked particles, cells and animals, population abundances, gene-expression and other molecular levels, concentration fields. It returns an interpretable equation with error bars, selects the terms the data supports, and simulates the fitted model. It handles overdamped and underdamped dynamics, interacting particles and spatially extended (SPDE) systems.

Is SFI for my data?

You want to…	SFI does it with
estimate drift and diffusion (first and second Kramers–Moyal coefficients, a Fokker–Planck model)	OverdampedLangevinInference: infer_force_linear, infer_diffusion_linear
discover the equation of a noisy stochastic system (sparse regression, "stochastic SINDy")	a candidate basis + sparsify_force(criterion="PASTIS")
map forces and diffusivity from single-particle tracking, with localization error	the overdamped engine; the defaults correct white localization noise
model cell migration, animal movement or active particles (persistent velocities)	UnderdampedLangevinInference, from positions alone
measure interaction forces between particles, collective motion	pair bases (SFI.bases.pairs)
model population dynamics: microbial or ecological abundances, growth rates, species interactions	the overdamped engine on log-abundances + sparsify_force (Lotka–Volterra example)
model gene-expression, protein or metabolite levels, whose noise scales with the level	the overdamped engine with a diffusion field D(x) (infer_diffusion_linear)
fit noisy oscillators and chaotic signals (genetic or circadian clocks, hair bundles)	a nonlinear basis; UnderdampedLangevinInference when one coordinate of an inertial oscillator is observed
quantify irreversibility: entropy production, broken detailed balance	compute_entropy_production()
decide whether the dynamics are overdamped or underdamped	SFI.classify_dynamics

SFI is not the right tool for categorical data, small counts or events, for forecasting without an equation, for snapshots without time correlation, or for processes with jumps. More routes: Start here.

Designed for experimental data. SFI is built for real experimental trajectories (tracked particles, cells, organisms, …) where no dynamical model pre-exists. Two first-class estimator families share one API: fast closed-form linear estimators that require no initial guess, and parametric estimators that model measurement noise and finite sampling explicitly — robust where real data is hard. The PASTIS information criterion rigorously identifies which terms the data actually supports. Synthetic examples in the gallery serve as runnable demonstrations.

Key features

• Two estimator families — closed-form linear estimators (infer_force_linear, seconds even on large datasets; the defaults correct localization noise and the leading finite-time-step bias) and parametric estimators (infer_force / infer_diffusion: noise-aware likelihood fit, robust to localization error and coarse sampling, accepts nonlinear models such as neural-network drifts).
• Overdamped & underdamped inference — works with position-only data; velocities are reconstructed automatically for underdamped systems.
• Composable state functions — build force/diffusion models from monomials, custom basis functions, pair interactions, or arbitrary parametric families, all with automatic differentiation via JAX.
• Sparse model selection — Pareto-front beam search with AIC / BIC / PASTIS information criteria.
• Entropy production — dissipation rate with error bar, for overdamped and underdamped dynamics.
• Simulation — simulate Langevin SDEs from the same model objects used for inference, including bootstrapped trajectories of a fitted model.
• Trajectory toolkit — mask-aware increments, synthetic degradation (noise, downsampling, data loss, motion blur), streaming for large datasets, I/O (CSV / Parquet / HDF5, tracking tables).
• Diagnostics — residual tests with action hints, an overdamped-vs-underdamped classifier, predicted errors, comparison to exact models.

Installation

pip install StochasticForceInference      # import name: SFI

For development (editable install with test/doc dependencies):

git clone https://github.com/ronceray/StochasticForceInference.git
cd StochasticForceInference
pip install -e ".[dev,io]"

Note: SFI requires Python ≥ 3.11 and JAX ≥ 0.10.

Quick start

For experimental data, load your trajectories and infer directly:

import numpy as np
from SFI import OverdampedLangevinInference, TrajectoryCollection
from SFI.bases import monomials_up_to

# Load your tracked data (positions: T×d array, dt: time step)
positions = np.load("my_experiment.npz")["positions"]  # shape (T, d)
coll = TrajectoryCollection.from_arrays(X=positions, dt=0.01)

# Infer using a polynomial basis — no model needed
B = monomials_up_to(order=3, dim=2, rank="vector")
inf = OverdampedLangevinInference(coll)
inf.infer_force_linear(B)               # also estimates D and the measurement noise
inf.sparsify_force(criterion="PASTIS")  # optional: find the minimal model
inf.compute_force_error()
inf.print_report()                      # coefficients with error bars
inf.diagnose().print_summary()          # residual tests, with hints if one fails

# Coarsely sampled or very noisy recordings? Use the parametric estimator
# instead: inf.infer_force(B) — it models the measurement noise natively.

For a synthetic example (useful for validation):

import jax.numpy as jnp
from jax import random

from SFI import OverdampedLangevinInference
from SFI.bases import X, monomials_up_to
from SFI.langevin import OverdampedProcess

# 1. Build a force model (Ornstein–Uhlenbeck: F = -k x) and simulate a trajectory.
#    X(dim) is the identity basis x ↦ x; theta_F holds its coefficient, so F(x) = -x.
proc = OverdampedProcess(F=X(dim=2), D=jnp.eye(2) * 0.5, theta_F=jnp.array([-1.0]))
proc.initialize(jnp.zeros(2))
coll = proc.simulate(dt=0.01, Nsteps=10_000, key=random.PRNGKey(0))

# 2. Infer from the trajectory using a generic polynomial basis
B = monomials_up_to(order=2, dim=2, rank="vector")
inf = OverdampedLangevinInference(coll)
inf.infer_force_linear(B)
inf.compute_force_error()

# 3. Inspect results and validate against the known model
inf.print_report()
inf.compare_to_exact(model_exact=proc)

Using SFI with an AI agent

SFI ships a workflow guide for AI coding agents (Claude Code, Codex, Cursor, …) inside the package, matched to the installed version:

python -m SFI guide            # Markdown: workflow, decisions, pitfalls, API names

# Install it as an Agent Skill (path shown for Claude Code; other tools that support skills use their own folder)
mkdir -p ~/.claude/skills/stochastic-force-inference
python -m SFI guide --skill > ~/.claude/skills/stochastic-force-inference/SKILL.md

The documentation has an llms.txt index. Agents that modify SFI itself should read AGENTS.md.

Coming from the original SFI scripts (StochasticForceInference.py, StochasticTrajectoryData, or the separate UnderdampedLangevinInference package)? SFI 2 is a rewrite with a new API; the guide maps the old calls to the new ones.

Documentation

Full documentation is available at stochasticforceinference.readthedocs.io, including:

• Start here — is SFI right for my data, and which route to take
• Tutorial — end-to-end walkthrough
• Examples gallery — worked examples covering sparsity, multi-particle, and SPDE inference
• Model building — composable state functions
• Trajectory handling — data ingestion, masking, degradation
• Running inference — estimator choice, overdamped/underdamped engines
• API reference — full autodoc

Package structure

Subpackage	Purpose
SFI.statefunc	Composable state functions: Basis, PSF, SF
SFI.inference	Overdamped & underdamped inference engines
SFI.trajectory	TrajectoryCollection / TrajectoryDataset
SFI.langevin	Langevin simulators (OverdampedProcess, …)
SFI.bases	Ready-made basis builders (monomials, pairs, SPDE, …)
SFI.diagnostics	Residual tests, dynamics classifier, error checks
SFI.integrate	Time-averaging integration engine
SFI.utils	Math helpers, formatting, plotting

Citation

If you use SFI, please cite the software (see CITATION.cff) together with the method paper(s) relevant to the features you use.

The software is archived on Zenodo. Cite the concept DOI 10.5281/zenodo.21532842, which always resolves to the latest release; each version also has its own DOI (v2.1.0: 10.5281/zenodo.24144932).

See How to cite in the documentation for the full, feature-specific list (overdamped, underdamped/ULI, trapeze, and PASTIS references).

Funding

Funded by the European Union (ERC, SuperStoc, 101117322). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency. Neither the European Union nor the granting authority can be held responsible for them. P. R. also acknowledges funding from the French National Research Agency (ANR-16-CONV-0001) and from the Excellence Initiative of Aix-Marseille University A*MIDEX.

License

MIT — see LICENSE for details.