Parametric estimators — concepts¶

This page presents the mathematical foundations of the parametric estimators used by infer_force() and infer_force().

Throughout we use the notation established in the Physics Reference.

Motivation¶

The linear estimators (infer_force_linear()) solve a linear regression:

\[M_a = \sum_b G_{ab}\, c_b\]

where the force moments \(M_a\) and Gram matrix \(G_{ab}\) are computed at each time step from finite-difference velocities and user-chosen basis functions. This estimator is fast, non-iterative, and unbiased at leading order in \(\Delta t\). However:

Finite-:math:`Delta t` bias: the linear regression is exact only in the \(\Delta t \to 0\) limit. At finite sampling interval the force moments carry systematic \(O(\Delta t)\) errors, especially in regions of high curvature.
Measurement noise: when positions are corrupted by localization error \(\Lambda\), the finite-difference velocity \(v = \Delta x / \Delta t\) is dominated by \(\Lambda / \Delta t\) rather than by the true velocity.
Nonlinear models: the regression is restricted to forces that are linear in their parameters. Parametric models such as neural networks or rational functions cannot be fitted directly.

The parametric estimator addresses all three limitations. It replaces the finite-difference velocity by an exact numerical flow, accounts for measurement noise through a joint Gaussian likelihood evaluated exactly, and iterates via Gauss–Newton to handle nonlinear parameter dependence.

The observation model¶

We assume the observed trajectory \(y_0, y_1, \dots, y_{T-1}\) is related to the true latent state \(x_t\) by

(1)¶\[y_t = x_t + \eta_t, \qquad \eta_t \sim \mathcal{N}(0, \Lambda)\]

where \(\Lambda\) is the measurement-noise covariance (localization error). The latent state evolves according to the Langevin SDE.

Overdamped (position only):

\[\mathrm{d}x = F(x;\,\theta)\,\mathrm{d}t + \sqrt{2\,D(x)}\;\mathrm{d}W_t .\]

Underdamped (position + velocity):

\[\dot{x} = v, \qquad \mathrm{d}v = F(x, v;\,\theta)\,\mathrm{d}t + \sqrt{2\,D}\;\mathrm{d}W_t .\]

Deterministic-flow linearisation¶

Instead of the Euler–Maruyama secant, we linearise each observation interval around the deterministic flow of the drift. Given the current state \(x_t\), we integrate the noise-free ODE \(\dot{z} = F(z;\theta)\), \(z(0) = x_t\), over a full step \(\Delta t_k\) — solved with \(n_{\text{sub}}\) RK4 micro-steps \(h_k = \Delta t_k / n_{\text{sub}}\) (default \(n_{\text{sub}} = 1\), a single 4th-order step per observation interval) — and model the stochastic part as a single Gaussian increment accumulated over that interval. There is no operator splitting: one deterministic flow step and one noise increment per observation. The step \(\Delta t_k\) is per interval: datasets sampled non-uniformly (a time vector t or a per-step dt array) are handled by exactly the same construction, with no interpolation or resampling.

The resulting flow map \(\Phi(x;\,\theta)\) is the RK4-integrated displacement:

(2)¶\[\Phi(x;\,\theta) = z(\Delta t_k) - x, \qquad \dot{z} = F(z;\,\theta), \quad z(0) = x.\]

Under this linearisation the transition distribution over one step is

(3)¶\[x_{t+1} \;\big|\; x_t \;\sim\; \mathcal{N}\!\Big( x_t + \Phi(x_t;\,\theta), \; V_t \Big)\]

where the process-noise covariance \(V_t\) is obtained by integrating the Lyapunov equation along the same RK4 stages as the flow:

(4)¶\[\dot{Q} = J\,Q + Q\,J^{\!\top} + 2\,D(z), \qquad Q(0) = 0, \qquad V_t = Q(\Delta t_k),\]

with \(J\) the running flow Jacobian. This is exact for linear drift and \(O((k\Delta t)^2)\)-accurate in general — replacing the endpoint-trapezoid approximation \(V_t \approx \Delta t\,(J_t D J_t^{\!\top} + D)\) of earlier versions, which is recovered with SFI_EXACT_UPGRADES=0. The improvement matters chiefly for diffusion estimates at coarse sampling (\(k\,\Delta t \gtrsim 0.1\), with \(k\) a typical drift gradient); force estimates are insensitive to it.

In the leading-order approximation used by the Gauss–Newton solver, the log-determinant contribution of \(V_t\) is treated as \(\theta\)-independent. This is justified because the \(\theta\)-dependent part of \(\log |V_t|\) enters only at \(O(\Delta t^3)\).

Residuals and the exact joint Gaussian¶

Combining (1) and (3) yields residuals

(5)¶\[r_t = y_{t+1} - y_t - \Phi(x_t;\,\theta)\]

that follow a block-tridiagonal multivariate Gaussian. The diagonal blocks of the joint covariance are

\[A_t = V_t + J_t \, \Lambda \, J_t^{\!\top} + \Lambda\]

and the off-diagonal blocks (from measurement noise coupling consecutive residuals) are

\[C_t = -\Lambda \, J_t^{\!\top} .\]

The NLL (negative log-likelihood up to constants) is:

(6)¶\[-\log p(y \mid \theta) = \tfrac{1}{2} \bigl[ \log\det \Sigma + r^{\!\top} \Sigma^{-1} r \bigr]\]

where \(\Sigma\) is the block-banded matrix with blocks \((A_t, C_t)\). Both the log-determinant and the quadratic form are evaluated exactly by a reverse-time block-LDLᵀ innovations recursion: sweeping the band once produces whitened residuals \(\tilde r_t = L_t^{-1}(r_t - \text{carry}_t)\) whose squared norm is the quadratic form and whose factor log-determinants sum to \(\log\det\Sigma\) — \(O(T d^3)\) time, one jax.lax.scan pass, no truncation and no tuning knobs. The same sweep whitens the parameter sensitivities, so the normal equations below are assembled against the exact precision. Masked frames and tracking gaps simply restart the recursion: only bandwidth residuals are lost per gap.

Parameter sensitivities without autodiff through the flow¶

The force \(F(x;\,\theta) = \sum_\alpha \theta_\alpha \, \phi_\alpha(x)\) (or any differentiable parametric model) enters nonlinearly through the flow map \(\Phi\). The Gauss–Newton solver needs the test functions (parameter sensitivities)

\[\psi_\alpha(x_t) = \frac{\partial \Phi(x_t;\,\theta)}{\partial \theta_\alpha} .\]

These are computed by an exact per-stage θ-recursion: the RK4 combination rules are applied to the stage values of \(F\), \(\partial F/\partial x\) and \(\partial F/\partial \theta\), propagating \(\psi\) alongside the flow. This gives the same \(\psi\) as forward-mode autodiff at the cost of one force / Jacobian / parameter-gradient evaluation per stage — instead of \(n_{\rm params}\) forward tangents through the integrator — with a zero-autodiff fast path for linear-in-θ bases. It also removes the second-order-autodiff compile wall for bases containing derivatives of fitted fields.

Errors-in-variables: the skip-trick instrument. Under measurement noise the regressor \(\psi(y_t)\) is evaluated at a noisy point, which correlates it with the residual and biases the symmetric normal equations (the classical errors-in-variables effect). The estimator therefore replaces the left factor of the normal equations by an η-clean instrument: the sensitivity evaluated at a lagged point whose noise is independent of the residual’s,

\[\psi_{\rm inst} = \frac{\partial \Phi}{\partial \theta} \Big|_{\text{stop-grad clean point}} ,\]

giving an asymmetric estimating equation whose root remains consistent under noise (eiv=True, the default; eiv=False recovers the plain MLE).

Note

Coarse sampling, (near-)clean data. A Δt-independent measurement noise \(\Lambda\) is identifiable against process noise \(\propto \Delta t\) only when the sampling resolves the dynamics. At coarse effective sampling on nonlinear systems the (D, Λ) profile can fit a spurious \(\Lambda\) on clean data, degrading \(\hat D\) and the force fit. eiv="auto" does not inspect the noise level — when you know the data is clean and the sampling is coarse, pass eiv=False (which also holds \(\Lambda \equiv 0\) in the profile) and use n_substeps=2.

The whitened Gram matrix and whitened right-hand side are

(7)¶\[\begin{split}G_{\alpha\beta} &= \sum_t \tilde\psi_\alpha^{\!\top} \, \tilde\psi_\beta, \\ f_\alpha &= \sum_t \tilde\psi_\alpha^{\!\top} \, \tilde\rho_t,\end{split}\]

where \(\rho_t = r_t + \psi \cdot \theta\) is the augmented residual and the tildes denote the exact innovations whitening of (6) (see Parametric estimators — algorithm and parameters for the assembly).

One GN step updates

\[\theta^{(k+1)} = G^{-1} f\]

(not an increment — \(f\) already contains the shifted residual). The iteration converges in 3–10 steps for typical problems.

Regularisation. When \(\operatorname{cond}(G) > \texttt{gram\_cond\_max}\) the system is Tikhonov-damped: \(G_{\rm reg} = G + \lambda I\) with \(\lambda = \sigma_1(G) / \texttt{gram\_cond\_max}\).

Comparison with the linear estimators¶

Property	`infer_force_linear()`	`infer_force()` (parametric)
Discretisation	Euler (\(O(\Delta t)\))	RK4 parametric (\(O(\Delta t^2)\))
Measurement noise	Heuristic subtraction	Joint likelihood (6)
Nonlinear models	No (linear in \(\theta\))	Yes (Gauss–Newton)
Non-uniform sampling	Yes (per-step increments)	Yes (per-interval flow + covariance)
(D, Λ) estimation	Pre-computed constant	Moment init + profiled exact NLL
Cost per iteration	\(O(T n^2 d)\)	\(O(T n d^3)\) (banded)
Iterations needed	1 (exact solve)	2–5 GN steps

The parametric approach is more accurate at larger \(\Delta t\), handles measurement noise natively, and can fit nonlinear parametric models. The cost is a modest iteration loop. For small \(\Delta t\) and clean data, the linear regression is often sufficient.

Overdamped vs. underdamped specifics¶

Overdamped parametric¶

The overdamped estimator uses a tridiagonal covariance structure (bandwidth 1): each residual couples only to its immediate neighbours through the off-diagonal block \(C_t = -\Lambda J_t^\top\).

The residuals are:

\[r_t = y_{t+1} - y_t - \Phi(y_t;\,\theta) .\]

The solve proceeds in three steps:

Moment init: \((\mathbf{D}, \Lambda)\) from closed-form noise-robust moment estimators (Vestergaard diffusion + lag-1 increment anticorrelation) — one trajectory pass, no optimiser. Both stencils are exactly unbiased under non-uniform sampling (per-sample two-interval normalisation).
Gauss–Newton loop: iterated whitened least-squares at the fixed noise levels.
Reprofile: one exact banded-NLL refinement of \((\mathbf{D}, \Lambda)\) at the fitted \(\theta\), run on cached fixed-θ tensors (the flow residuals and stage Jacobians are computed once; each L-BFGS iteration of the profile costs linear algebra only — zero basis evaluations), followed by a warm re-solve (the iteratively-reweighted fixed point for linear-in-\(\theta\) models).

Underdamped parametric (ULI)¶

The underdamped estimator works in \((x, v)\) phase space — only positions are observed, so the velocity at each residual is resolved by shooting: one Newton step finds the initial velocity whose flow hits the next observed position. The 3-point residuals have a pentadiagonal covariance structure (bandwidth 2), and the process noise enters the position at \(O(\Delta t^3)\) — which makes the Lyapunov-exact covariance (4) mandatory for non-uniform sampling (the leading-order blocks assume equal neighbouring intervals).

The skip-trick instrument is built from the clean lagged position pair and is particularly effective here: the shooting velocity divides the position noise by \(\Delta t\), so the velocity errors-in-variables bias of the plain MLE is strong.

Note

A single Euler step cannot support the underdamped skip-trick: the Euler position update \(x + v\,\Delta t\) does not depend on the force, so the instrument is identically zero. The solver detects integrator="euler", n_substeps=1 with eiv enabled, warns, and falls back to the plain MLE. The RK4 default does not have this degeneracy.

State-dependent diffusion inference¶

After force inference, infer_diffusion() fits a state-dependent diffusion tensor \(D(x;\,\theta_D)\) by optimising the exact banded NLL with the force held fixed, on the same cached fixed-θ tensors as the profile: per-point process covariances for any candidate \(D\) are rebuilt from the cached RK4 stage Jacobians ((4)), so each L-BFGS iteration costs linear algebra only — zero force or basis evaluations.

Two parametrisations are available:

a rank-2 Basis: \(D(x) = \sum_j \theta_j \, d_j(x)\) (linear in \(\theta_D\); the default symmetric_matrix_basis(d) spans all constant symmetric matrices);
a PSF: \(D(x) = \text{PSF}(x;\,\theta_D)\) directly — in the underdamped case the model may be velocity-dependent, \(D(x, v)\), evaluated at the shooting velocity.

The optimisation is initialised by projecting the moment-estimated constant \(\hat{\mathbf{D}}\) onto the model (a zero start is a singular covariance and cannot be escaped).

Key references¶

Frishman, A. & Ronceray, P., Learning force fields from stochastic trajectories, Phys. Rev. X 10, 021009 (2020). — The original SFI linear regression estimator.
Vestergaard, C. L., Blainey, P. C. & Flyvbjerg, H., Optimal estimation of diffusion coefficients from single-particle trajectories, Phys. Rev. E 89, 022726 (2014). — Measurement-noise-aware diffusion estimation.
Brückner, D. B., Ronceray, P. & Broedersz, C. P., Inferring the dynamics of underdamped stochastic systems, Phys. Rev. Lett. 125, 058103 (2020). — Underdamped Langevin Inference (ULI), extended here to the parametric estimator.
Amri, S. et al., Inferring geometrical dynamics of cell nucleus translocation, Phys. Rev. Res. 6, 043030 (2024). — Trapezoidal integration scheme.