StochasticForceInference 2.0.0__py3-none-any.whl

SFI/inference/parametric_core/__init__.py

@@ -0,0 +1,34 @@
+# SFI/inference/parametric_core/ — stripped-to-the-bone parametric estimator.
+"""
+Parallel, minimal reimplementation of the parametric force/diffusion
+estimator (RK4-flow residuals + banded residual covariance + windowed
+precision + AD-minimised windowed NLL).
+
+Goal: a single, paper-ready, one-size-fits-all objective with a handful
+of parameters, integrating like ``infer_*_linear`` (Basis/PSF in →
+``InferenceResultSF`` out, multi-particle aware) and reusing the
+``SFI.integrate`` engine for chunking/masking/JIT.
+
+This package is intentionally independent of ``SFI.inference.parametric``
+(the existing implementation, kept as the benchmark baseline).
+
+Design note: ``docs/source/inference/parametric_core.rst`` (to be written);
+working draft in the session plan.
+"""
+
+from __future__ import annotations
+
+from .solve import (
+    DiffusionSolveResult,
+    ForceSolveResult,
+    solve_diffusion_od,
+    solve_diffusion_ud,
+    solve_force_od,
+    solve_force_ud,
+)
+
+__all__ = [
+    "solve_force_od", "solve_force_ud",
+    "solve_diffusion_od", "solve_diffusion_ud",
+    "ForceSolveResult", "DiffusionSolveResult",
+]

SFI/inference/parametric_core/covariance.py

@@ -0,0 +1,232 @@
+# SFI/inference/parametric_core/covariance.py
+"""
+Banded residual covariance for the parametric core.
+
+The flow residuals are correlated only over a finite lag set by the SDE
+order, so their covariance is **block-banded** with bandwidth ``B``:
+
+    overdamped (order 1) → B = 1  (tridiagonal)
+    underdamped (order 2) → B = 2  (pentadiagonal)
+    order-3 SDE          → B = 3  (heptadiagonal)   [future]
+
+A single container :class:`CovarianceBlocks` (diagonal blocks + a tuple
+of off-diagonal lag blocks) represents any bandwidth, so adding a
+higher-order builder is a *localized* addition — the precision engine and
+assembler below are bandwidth-generic and need no change.
+
+``D`` and ``Λ`` are full ``(d, d)`` matrices throughout (scalars / per-
+axis vectors are accepted and promoted).
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import jax.numpy as jnp
+
+__all__ = ["CovarianceBlocks", "build_od_blocks", "build_ud_blocks", "assemble_dense"]
+
+
+class CovarianceBlocks(NamedTuple):
+    """Block-banded covariance.
+
+    Parameters
+    ----------
+    A : array ``(n_res, d, d)``
+        Diagonal blocks ``Cov(r_n, r_n)``.
+    offdiag : tuple of arrays
+        ``offdiag[lag-1]`` has shape ``(n_res-lag, d, d)`` and holds
+        ``Cov(r_n, r_{n+lag})``.
+    bandwidth : int
+        Number of off-diagonal lags (``len(offdiag)``).
+    """
+
+    A: jnp.ndarray
+    offdiag: tuple
+    bandwidth: int
+
+
+def _rel_jitter(A_raw, jitter, I_d):
+    r"""Diagonal regulariser **scaled to the block magnitude**.
+
+    ``jitter`` is a *relative* coefficient: the added term is
+    ``jitter · ⟨diag(A_raw)⟩ · I`` where the scale is the mean diagonal
+    entry of the (pre-jitter) blocks.  An absolute floor would mask the
+    ``Δt^k``-small process variance at small ``Δt`` (e.g. the underdamped
+    ``(4/3)Δt³D ~ 1e-9``), making ``D`` unidentifiable; a relative floor
+    is always a tiny fraction of the residual variance, at any ``Δt``.
+    """
+    d = A_raw.shape[-1]
+    scale = jnp.mean(jnp.einsum("nii->n", A_raw)) / d
+    scale = jnp.maximum(scale, 1e-30)
+    return jitter * scale * I_d[None]
+
+
+def _as_dd(x, d, dtype):
+    """Promote scalar / ``(d,)`` / ``(d, d)`` to a ``(d, d)`` matrix."""
+    x = jnp.asarray(x, dtype=dtype)
+    if x.ndim == 0:
+        return x * jnp.eye(d, dtype=dtype)
+    if x.ndim == 1:
+        return jnp.diag(x)
+    return x
+
+
+def build_od_blocks(J, D, Lambda, dt, *, jitter=0.0, valid_mask=None):
+    r"""Overdamped (bandwidth-1) residual covariance blocks.
+
+    For the 2-point residual ``r_k = Y_{k+1} − Y_k − Φ_k(Y_k)`` the
+    leading-order covariance combines process noise (diffusion propagated
+    through the linearised flow) and measurement noise (shared between
+    neighbouring residuals):
+
+        A_k = Δt(J_k D J_kᵀ + D) + J_k Λ J_kᵀ + Λ + jitter·⟨diag A⟩·I
+        C_k = Cov(r_k, r_{k+1}) = −Λ J_{k+1}ᵀ
+
+    Parameters
+    ----------
+    J : array ``(n_res, d, d)``
+        Flow Jacobians at the residual base points.
+    D, Lambda : ``(d, d)`` (or scalar / ``(d,)``)
+        Diffusion and measurement-noise covariance matrices.
+    dt : float
+    jitter : float
+        *Relative* diagonal regulariser (a fraction of the mean diagonal
+        block magnitude — see :func:`_rel_jitter`), so it never masks the
+        ``Δt``-small process variance.
+    valid_mask : ``(n_res,)`` bool, optional
+        Where ``False``, the diagonal block is replaced by a neutral
+        ``(1+jitter)·I`` and adjacent off-diagonal couplings are zeroed,
+        decoupling masked residuals.
+
+    Returns
+    -------
+    CovarianceBlocks with ``bandwidth=1``.
+    """
+    J = jnp.asarray(J)
+    d = J.shape[-1]
+    dtype = J.dtype
+    I_d = jnp.eye(d, dtype=dtype)
+    D = jnp.asarray(D, dtype=dtype)
+    Lambda = _as_dd(Lambda, d, dtype)
+
+    if D.ndim == 3:  # per-step (state-dependent) D, shape (n_res, d, d)
+        JDJ = jnp.einsum("kij,kjl,kml->kim", J, D, J)
+        D_diag = D
+    else:
+        D = _as_dd(D, d, dtype)
+        JDJ = jnp.einsum("kij,jl,kml->kim", J, D, J)
+        D_diag = D[None]
+    JSJ = jnp.einsum("kij,jl,kml->kim", J, Lambda, J)
+    A_raw = dt * (JDJ + D_diag) + JSJ + Lambda[None]
+    A = A_raw + _rel_jitter(A_raw, jitter, I_d)
+
+    # lag-1: C_k = −Λ J_{k+1}ᵀ
+    C = -jnp.einsum("ij,klj->kil", Lambda, J)[1:]  # (n_res-1, d, d)
+
+    if valid_mask is not None:
+        m = valid_mask[:, None, None]
+        A = jnp.where(m, A, (1.0 + jitter) * I_d[None])
+        pair = (valid_mask[:-1] & valid_mask[1:])[:, None, None]
+        C = jnp.where(pair, C, jnp.zeros_like(C))
+
+    return CovarianceBlocks(A=A, offdiag=(C,), bandwidth=1)
+
+
+def build_ud_blocks(alpha_plus, alpha_zero, alpha_minus, D, Lambda, dt,
+                    *, jitter=0.0, valid_mask=None):
+    r"""Underdamped (bandwidth-2) residual covariance blocks.
+
+    The 3-point shooting residual ``r_n = α₋ ε_{n-1} + α₀ ε_n + α₊ ε_{n+1}
+    + ξ`` (measurement + process) has a pentadiagonal covariance:
+
+        A_n = (4/3)Δt³ D_sym + Σ_k α_k Λ α_kᵀ + jitter·⟨diag A⟩·I
+        C_n = (1/3)Δt³ D_sym + α₊ Λ α₀ᵀ|_{n,n+1} + α₀ Λ α₋ᵀ|_{n,n+1}
+        E_n = α₊ Λ α₋ᵀ|_{n,n+2}
+
+    The ``jitter`` regulariser is *relative* to the block magnitude
+    (:func:`_rel_jitter`): an absolute floor would dominate the ``Δt³``-small
+    process variance at small ``Δt`` and make ``D`` unidentifiable.
+
+    where ``D_sym = (D + Dᵀ)/2`` and the leading-order process noise scales
+    as ``Δt³`` (the velocity is integrated once more than the overdamped
+    increment).
+
+    Parameters
+    ----------
+    alpha_plus, alpha_zero, alpha_minus : ``(n_res, d, d)``
+        Measurement-noise propagators (``α₊ = I`` by construction).
+    D, Lambda : ``(d, d)`` (or scalar / ``(d,)``).
+    dt, jitter : float
+    valid_mask : ``(n_res,)`` bool, optional.
+
+    Returns
+    -------
+    CovarianceBlocks with ``bandwidth=2``.
+    """
+    ap = jnp.asarray(alpha_plus)
+    a0 = jnp.asarray(alpha_zero)
+    am = jnp.asarray(alpha_minus)
+    d = ap.shape[-1]
+    dtype = ap.dtype
+    I_d = jnp.eye(d, dtype=dtype)
+    D = jnp.asarray(D, dtype=dtype)
+    Lambda = _as_dd(Lambda, d, dtype)
+
+    if D.ndim == 3:  # per-step (state/velocity-dependent) D, shape (n_res, d, d)
+        D_sym = 0.5 * (D + jnp.swapaxes(D, -1, -2))
+        V_xi = (4.0 / 3.0) * dt**3 * D_sym          # (n_res, d, d)
+        C_proc = (1.0 / 3.0) * dt**3 * D_sym
+        V_xi_diag = V_xi
+        C_proc_lag1 = 0.5 * (C_proc[:-1] + C_proc[1:])
+    else:
+        D = _as_dd(D, d, dtype)
+        D_sym = 0.5 * (D + D.T)
+        V_xi_diag = ((4.0 / 3.0) * dt**3 * D_sym)[None]
+        C_proc_lag1 = ((1.0 / 3.0) * dt**3 * D_sym)[None]
+
+    se_pp = jnp.einsum("nij,jk,nlk->nil", ap, Lambda, ap)
+    se_00 = jnp.einsum("nij,jk,nlk->nil", a0, Lambda, a0)
+    se_mm = jnp.einsum("nij,jk,nlk->nil", am, Lambda, am)
+    A_raw = V_xi_diag + se_pp + se_00 + se_mm
+    A = A_raw + _rel_jitter(A_raw, jitter, I_d)
+
+    cross1a = jnp.einsum("nij,jk,nlk->nil", ap[:-1], Lambda, a0[1:])
+    cross1b = jnp.einsum("nij,jk,nlk->nil", a0[:-1], Lambda, am[1:])
+    C = C_proc_lag1 + cross1a + cross1b
+    E = jnp.einsum("nij,jk,nlk->nil", ap[:-2], Lambda, am[2:])
+
+    if valid_mask is not None:
+        m = valid_mask[:, None, None]
+        A = jnp.where(m, A, (1.0 + jitter) * I_d[None])
+        pair = (valid_mask[:-1] & valid_mask[1:])[:, None, None]
+        C = jnp.where(pair, C, jnp.zeros_like(C))
+        triple = (valid_mask[:-2] & valid_mask[1:-1] & valid_mask[2:])[:, None, None]
+        E = jnp.where(triple, E, jnp.zeros_like(E))
+
+    return CovarianceBlocks(A=A, offdiag=(C, E), bandwidth=2)
+
+
+def assemble_dense(blocks: CovarianceBlocks):
+    """Assemble a :class:`CovarianceBlocks` into a dense symmetric matrix.
+
+    Returns
+    -------
+    Sigma : array ``(n_res·d, n_res·d)``
+    """
+    A = blocks.A
+    n_res, d, _ = A.shape
+    Nd = n_res * d
+    S = jnp.zeros((Nd, Nd), dtype=A.dtype)
+
+    for i in range(n_res):
+        S = S.at[i * d:(i + 1) * d, i * d:(i + 1) * d].set(A[i])
+
+    for lag_idx, C_lag in enumerate(blocks.offdiag):
+        lag = lag_idx + 1
+        for i in range(n_res - lag):
+            r0, c0 = i * d, (i + lag) * d
+            S = S.at[r0:r0 + d, c0:c0 + d].set(C_lag[i])
+            S = S.at[c0:c0 + d, r0:r0 + d].set(C_lag[i].T)
+
+    return S

SFI/inference/parametric_core/driver.py

@@ -0,0 +1,272 @@
+# SFI/inference/parametric_core/driver.py
+"""
+IRLS driver for the parametric core.
+
+Minimises the windowed objective by alternating:
+
+1. **freeze** the precision ``P`` at the current ``θ`` (and the current
+   ``D, Λ``) and minimise the resulting quadratic
+   ``Q(θ) = ½ Σ rₙ(θ)ᵀ P rₙ(θ)`` over ``θ`` with L-BFGS — the inner solve;
+2. **reprofile** ``D, Λ`` from the residuals at the updated ``θ``.
+
+For linear-in-θ models this converges in ~1 outer iteration (it is the
+GLS fixed point); for nonlinear-in-θ it is ordinary IRLS.  The inner
+objective and its gradient are supplied by the caller (built from the
+integrate-engine loss program), so the driver is agnostic to dynamics and
+bandwidth.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from scipy.optimize import minimize as _sp_minimize
+
+from SFI.utils.maths import default_float_dtype
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["irls_minimize", "gn_minimize", "condition_cap_ridge"]
+
+
+def condition_cap_ridge(G, gram_cond_max=1e10):
+    r"""Tikhonov ridge that caps the condition number of the Gram ``G``.
+
+    Mirrors the existing Gauss–Newton driver: if ``cond(G) > gram_cond_max``
+    return ``λ = σ_max / gram_cond_max`` (so the regularised matrix
+    ``G + λI`` has condition number ``≤ gram_cond_max``), else ``0``.  Adding
+    ``½λ‖θ‖²`` to the inner objective then bounds ``θ`` along the
+    ill-conditioned / unconstrained directions instead of letting L-BFGS
+    drive them to ``±∞`` — graceful degradation under high measurement noise.
+
+    Returns a finite, non-negative scalar (``0`` for a non-finite / zero
+    Gram, where no meaningful curvature scale exists).
+    """
+    G = jnp.asarray(G)
+    svals = jnp.linalg.svdvals(G)
+    s_max = svals[0]
+    s_min = svals[-1]
+    finite = jnp.isfinite(s_max) & (s_max > 0)
+    over = s_max > gram_cond_max * s_min
+    lam = jnp.where(finite & over, s_max / gram_cond_max, 0.0)
+    return lam
+
+
+def irls_minimize(
+    theta0,
+    build_inner,
+    D0,
+    Lambda0,
+    *,
+    profile_fn=None,
+    max_outer=5,
+    reprofile_iters=None,
+    inner_maxiter=80,
+    inner_ftol=1e-12,
+    tol=1e-7,
+    label="",
+):
+    r"""Run the IRLS outer loop.
+
+    Parameters
+    ----------
+    theta0 : ``(n_params,)``
+    build_inner : callable
+        ``build_inner(theta_frozen, D, Lambda) -> inner(theta_live) ->
+        (value, grad)``.  Built once per outer iteration so the frozen
+        precision is hoisted out of the inner L-BFGS sweep.
+    D0, Lambda0 : ``(d, d)`` initial noise matrices.
+    profile_fn : callable or None
+        ``profile_fn(theta, D, Lambda) -> (D, Lambda)``; called after
+        the inner solve at the first ``reprofile_iters`` outer iterations
+        (warm-started from the current noise).  ``None`` holds ``(D, Λ)``
+        fixed.
+    max_outer, reprofile_iters, inner_maxiter, inner_ftol, tol : see module docstring.
+
+    Returns
+    -------
+    theta_best, D_best, Lambda_best, info
+    """
+    dtype = default_float_dtype()
+    theta = jnp.asarray(theta0, dtype=dtype)
+    D = jnp.asarray(D0, dtype=dtype)
+    Lambda = jnp.asarray(Lambda0, dtype=dtype)
+    if reprofile_iters is None:
+        reprofile_iters = max_outer
+
+    best_loss = float("inf")
+    best = (theta, D, Lambda)
+    converged = False
+    outer = 0
+
+    for outer in range(max_outer):
+        theta_prev = theta
+        inner = build_inner(jax.lax.stop_gradient(theta), D, Lambda)
+
+        def _obj(x):
+            v, g = inner(jnp.asarray(x, dtype=dtype))
+            return float(v), np.asarray(g, dtype=np.float64)
+
+        res = _sp_minimize(
+            _obj, np.asarray(theta, dtype=np.float64), jac=True, method="L-BFGS-B",
+            options={"maxiter": inner_maxiter, "ftol": inner_ftol, "gtol": inner_ftol},
+        )
+        theta = jnp.asarray(res.x, dtype=dtype)
+        loss = float(res.fun)
+
+        if profile_fn is not None and outer < reprofile_iters:
+            D, Lambda = profile_fn(theta, D, Lambda)
+
+        if loss < best_loss and bool(jnp.all(jnp.isfinite(theta))):
+            best_loss = loss
+            best = (theta, D, Lambda)
+
+        rel = float(jnp.linalg.norm(theta - theta_prev) / (jnp.linalg.norm(theta_prev) + 1e-30))
+        logger.info("[core-irls] %s outer %d: loss=%.6g  rel=%.3g", label, outer, loss, rel)
+        if rel < tol:
+            converged = True
+            break
+
+    info = {"outer_iterations": outer + 1, "converged": converged, "loss": best_loss}
+    return best[0], best[1], best[2], info
+
+
+def gn_minimize(
+    theta0,
+    gram_fn,
+    D0,
+    Lambda0,
+    *,
+    profile_fn=None,
+    nll_fn=None,
+    max_iter=8,
+    reprofile_iters=None,
+    gram_cond_max=1e10,
+    line_search_alphas=(1.0, 0.5, 0.25, 0.1),
+    tol=1e-7,
+    label="",
+    merit="nll",
+):
+    r"""Direct Gauss–Newton outer loop on the windowed Gram.
+
+    The linear-in-θ fast path: the windowed estimator already returns the GN
+    normal-equation pieces ``(G, f, nll)`` — with ``G = ψᵀPψ`` the GN Hessian,
+    ``f = ψᵀPr`` the score (``ψ = ∂r/∂θ``, so ``f`` is the **+gradient** of
+    ``½ rᵀP r``) — so the step is the closed-form ``δθ = −(G+λI+μ·diagG)⁻¹ f``.
+    No nested ``value_and_grad`` through the flow, so for a linear-in-θ force it
+    reaches the GLS fixed point in ~1–2 iterations instead of ~80 L-BFGS
+    objective evaluations.
+
+    Mirrors the legacy ``parametric_gn_iterate`` driver: Tikhonov condition cap
+    (:func:`condition_cap_ridge`), Levenberg–Marquardt damping with line-search
+    retry, ``(D, Λ)`` reprofiling at early iterations, and best-iterate
+    tracking.
+
+    Parameters
+    ----------
+    theta0 : ``(n_params,)``
+    gram_fn : callable ``(theta, D, Lambda) -> (G, f, nll)``.
+    D0, Lambda0 : ``(d, d)`` initial noise matrices.
+    profile_fn : callable or None  ``(theta, D, Λ) -> (D, Λ)``; called after
+        iter 0 for the first ``reprofile_iters`` iterations.
+    nll_fn : callable or None  scalar objective for the line search; falls back
+        to the Gram's own ``nll`` when ``None``.
+    merit : {"nll", "phi"}  line-search / best-iterate merit.  ``"nll"`` (the
+        symmetric MLE) descends the windowed NLL.  ``"phi"`` is the
+        estimating-equation residual ``‖f‖ = ‖⟨ψ_left P r⟩‖`` — the correct
+        merit when ``G`` is asymmetric (the EIV instrument path), where the IV
+        root does **not** minimise the NLL.  The condition guard is taken on the
+        symmetric part ``½(G+Gᵀ)`` so it stays valid for asymmetric ``G``.
+    max_iter, gram_cond_max, line_search_alphas, tol : GN budget / guards.
+
+    Returns
+    -------
+    theta_best, D_best, Lambda_best, info  (same contract as
+    :func:`irls_minimize`).
+    """
+    dtype = default_float_dtype()
+    theta = jnp.asarray(theta0, dtype=dtype)
+    D = jnp.asarray(D0, dtype=dtype)
+    Lambda = jnp.asarray(Lambda0, dtype=dtype)
+    n_params = int(theta.shape[0])
+    I_n = jnp.eye(n_params, dtype=dtype)
+    if reprofile_iters is None:
+        reprofile_iters = max_iter
+
+    def _merit(th, D_, Se_):
+        if merit == "phi":
+            return float(jnp.linalg.norm(gram_fn(th, D_, Se_)[1]))
+        if nll_fn is not None:
+            return float(nll_fn(th, D_, Se_))
+        return float(gram_fn(th, D_, Se_)[2])
+
+    best_nll = float("inf")
+    best = (theta, D, Lambda)
+    converged = False
+    lm = 0.0
+    outer = 0
+
+    for outer in range(max_iter):
+        if profile_fn is not None and 0 < outer < reprofile_iters + 1:
+            D, Lambda = profile_fn(theta, D, Lambda)
+
+        G, f, nll = gram_fn(theta, D, Lambda)
+        # The acceptance bar must be the SAME objective as the line-search trial
+        # (_merit): when nll_fn is given, the Gram program's own nll carries an
+        # arbitrary offset against it (different window/conditioning
+        # normalisation), and comparing trial-merit against gram-nll can reject
+        # every step — freezing θ at the init and reporting rel=0 "convergence".
+        if merit == "phi":
+            nll_val = float(jnp.linalg.norm(f))
+        elif nll_fn is not None:
+            nll_val = float(nll_fn(theta, D, Lambda))
+        else:
+            nll_val = float(nll)
+        if nll_val < best_nll and bool(jnp.all(jnp.isfinite(theta))):
+            best_nll = nll_val
+            best = (theta, D, Lambda)
+
+        # Tikhonov condition cap on the symmetric part (valid for asymmetric G).
+        lam = float(condition_cap_ridge(0.5 * (G + G.T), gram_cond_max))
+        G_reg = G + lam * I_n
+
+        theta_prev = theta
+        accepted = False
+        retries = 0
+        while not accepted and retries <= 8:
+            G_damped = G_reg + lm * jnp.diag(jnp.diag(G_reg)) if lm > 0 else G_reg
+            delta = -jnp.linalg.solve(G_damped, f)
+            if not bool(jnp.all(jnp.isfinite(delta))):
+                retries += 1
+                lm = max(lm * 4.0, 1e-3)
+                continue
+            for alpha in line_search_alphas:
+                theta_trial = theta + alpha * delta
+                if not bool(jnp.all(jnp.isfinite(theta_trial))):
+                    continue
+                nll_trial = _merit(theta_trial, D, Lambda)
+                if nll_trial < nll_val + 1e-8:
+                    theta = theta_trial
+                    accepted = True
+                    if nll_trial < best_nll:
+                        best_nll = nll_trial
+                        best = (theta, D, Lambda)
+                    break
+            if not accepted:
+                retries += 1
+                lm = max(lm * 4.0, 1e-3)
+
+        if accepted:
+            lm *= 0.25
+        rel = float(jnp.linalg.norm(theta - theta_prev) / (jnp.linalg.norm(theta_prev) + 1e-30))
+        logger.info("[core-gn] %s iter %d: nll=%.6g rel=%.3g lam=%.2e accepted=%s",
+                    label, outer, nll_val, rel, lam, accepted)
+        if not accepted or rel < tol:
+            converged = rel < tol
+            break
+
+    info = {"outer_iterations": outer + 1, "converged": converged, "loss": best_nll}
+    return best[0], best[1], best[2], info

SFI/inference/parametric_core/flow.py

@@ -0,0 +1,149 @@
+# SFI/inference/parametric_core/flow.py
+"""
+Deterministic flow, Jacobian, and residuals for the parametric core.
+
+The estimator linearises the SDE around the *deterministic* flow of the
+drift field.  For a single observation interval Δt the displacement and
+its sensitivity to the start point are
+
+    Φ(y;θ) = ψ_Δt(y;θ) − y,        J(y;θ) = I + ∂Φ/∂y = ∂ψ_Δt/∂y,
+
+where ``ψ_Δt`` is the RK4 (or single-step Euler) flow of ``ẋ = F(x;θ)``.
+Choosing ``integrator="euler", n_substeps=1`` recovers the bare Euler
+predictor ``Φ = F(y;θ)·Δt`` — useful for *independently* assessing the
+effect of flow propagation versus the covariance/precision model.
+
+These are pure, JAX-traceable functions (parameters are closed over in
+``drift_fn``); batching over windows/particles is done by the caller via
+``vmap`` or the ``SFI.integrate`` engine.
+"""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+
+from SFI.integrate.rk4 import select_flow
+
+__all__ = ["flow_step", "flow_displacement", "od_window_residuals", "od_instrument"]
+
+
+def flow_displacement(drift_fn, y, dt, n_substeps, integrator="rk4"):
+    """Deterministic displacement ``Φ(y) = flow(y) − y`` over one interval Δt.
+
+    Parameters
+    ----------
+    drift_fn : callable ``(d,) → (d,)``
+        Autonomous vector field with parameters already closed over.
+    y : array ``(d,)``
+    dt : float
+    n_substeps : int  (static)
+    integrator : {"rk4", "euler"}
+
+    Returns
+    -------
+    Phi : array ``(d,)``
+    """
+    flow = select_flow(integrator)
+    return flow(drift_fn, y, dt, n_substeps) - y
+
+
+def flow_step(drift_fn, y, dt, n_substeps, integrator="rk4"):
+    """Displacement ``Φ`` and flow Jacobian ``J = I + ∂Φ/∂y`` at a point.
+
+    Parameters
+    ----------
+    drift_fn : callable ``(d,) → (d,)``
+    y : array ``(d,)``
+    dt : float
+    n_substeps : int  (static)
+    integrator : {"rk4", "euler"}
+
+    Returns
+    -------
+    Phi : array ``(d,)``
+    J : array ``(d, d)``
+        ``J = ∂(flow)/∂y`` — the linear tangent map of the flow.
+    """
+    d = y.shape[-1]
+
+    def disp(z):
+        return flow_displacement(drift_fn, z, dt, n_substeps, integrator)
+
+    Phi = disp(y)
+    J = jax.jacfwd(disp)(y) + jnp.eye(d, dtype=y.dtype)
+    return Phi, J
+
+
+def od_window_residuals(drift_fn, X_w, dt, n_substeps, integrator="rk4"):
+    """Overdamped 2-point residuals and flow Jacobians over a position window.
+
+    For a window ``X_w`` of ``W`` consecutive positions, returns the
+    ``W-1`` residuals ``r_k = X_{k+1} − X_k − Φ_k(X_k)`` together with the
+    flow Jacobians ``J_k`` at each window base point.
+
+    Parameters
+    ----------
+    drift_fn : callable ``(d,) → (d,)``
+    X_w : array ``(W, d)``
+    dt : float
+    n_substeps : int  (static)
+    integrator : {"rk4", "euler"}
+
+    Returns
+    -------
+    r : array ``(W-1, d)``
+    J : array ``(W-1, d, d)``
+    """
+    def _step(y):
+        return flow_step(drift_fn, y, dt, n_substeps, integrator)
+
+    Phi, J = jax.vmap(_step)(X_w[:-1])
+    r = X_w[1:] - X_w[:-1] - Phi
+    return r, J
+
+
+def od_instrument(drift_of_theta, theta, x_base, dt, n_substeps, integrator="rk4"):
+    r"""η-clean instrument (left test function) for one overdamped residual.
+
+    The residual ``r_c = X_{c+1} − X_c − Φ(X_c;θ)`` has regressor
+    ``ψ_right = ∂r_c/∂θ = −∂Φ/∂θ|_{X_c}``, which is correlated with the
+    measurement noise ``η_c`` sitting in ``X_c`` (and in ``r_c``) — the
+    errors-in-variables bias.  The instrument replaces that left factor with
+    the same sensitivity built from a base point ``x_base`` whose noise is
+    *independent* of the residual's (the core uses the lagged ``X_{c-1}``),
+    flow-propagated forward one interval so it stays geometrically aligned
+    with the center:
+
+    .. math::
+
+        ψ_{\mathrm{inst}} = -\tfrac12\!\left(
+            \left.\frac{∂Φ}{∂θ}\right|_{x_{\mathrm{base}}}
+          + \left.\frac{∂Φ}{∂θ}\right|_{x_{\mathrm{base}}+Φ(x_{\mathrm{base}})}
+        \right).
+
+    The trapezoidal pair (base point and its deterministic forward image) is
+    the parametric analogue of the legacy lag-1 *shift* instrument that SFI
+    auto-selects under measurement noise.  The forward point is evaluated at
+    the current ``θ`` (stop-gradient), so the θ-derivative is taken of the
+    sensitivity *at fixed points* — a valid test function for ``∂r/∂θ``.
+
+    Parameters
+    ----------
+    drift_of_theta : callable  ``θ_flat → (drift_fn : (d,) → (d,))``
+    theta : array ``(n_params,)``
+    x_base : array ``(d,)``
+        η-clean base point (independent of the center residual's noise).
+    dt, n_substeps, integrator : flow controls (as elsewhere).
+
+    Returns
+    -------
+    psi_inst : array ``(d, n_params)``
+    """
+    def disp_theta_jac(y):
+        return jax.jacfwd(
+            lambda th: flow_displacement(drift_of_theta(th), y, dt, n_substeps, integrator))(theta)
+
+    Phi0 = flow_displacement(drift_of_theta(theta), x_base, dt, n_substeps, integrator)
+    x_fwd = jax.lax.stop_gradient(x_base + Phi0)
+    return -0.5 * (disp_theta_jac(x_base) + disp_theta_jac(x_fwd))