StochasticForceInference 2.0.3__tar.gz → 2.0.4__tar.gz

{stochasticforceinference-2.0.3 → stochasticforceinference-2.0.4}/AGENTS.md

@@ -45,6 +45,7 @@ SFI has mature, well-tested abstractions for every common operation.
 | Wrap a user function as a state function | `SFI.statefunc.make_basis` / `make_sf` / `make_psf`     |
 | Run force / diffusion inference          | `SFI.OverdampedLangevinInference` / `UnderdampedLangevinInference` |
 | Select a sparse model                    | `inf.sparsify_force(criterion="PASTIS", p=0.1)`         |
+| Select a sparse diffusion model          | `inf.sparsify_diffusion(criterion="PASTIS", p=0.05)`    |
 | Compare to ground truth                  | `inf.compare_to_exact(model_exact=...)`                 |
 | Plot in the SFI style                    | `SFI.utils.plotting.SFI_COLORS`, `dark_fig`, `phase2d`  |
 | Save / load a fitted model               | `SFI.inference.save_model` / `load_model`               |
@@ -78,6 +79,16 @@ from SFI import (
 The dynamics-order classifier (`classify_dynamics`, returning a
 `DynamicsOrderReport`) decides overdamped vs underdamped from raw
 positions; see `SFI.diagnostics` and `examples/gallery/dynamics_order_demo.py`.
+`SFI.diagnostics.finite_dt_scale(inf)` reports the per-step drift-gradient
+scale Δt‖∇F̂‖ of a fitted force (the size of the linear estimators' O(Δt)
+biases) with a `fine` / `coarse` / `unreliable` verdict; `assess(inf)` includes
+it and flags the coarse regimes. Both linear engines also record the force
+Gram condition number (`metadata["force_gram_cond"]`) and warn when the
+dictionary is ill-conditioned on the data.
+`SFI.diagnostics.psd_fraction(D_sf, coll)` reports the fraction of data
+points where a diffusion field is not positive semi-definite;
+`simulate_bootstrapped_trajectory` runs it on the fitted diffusion and warns
+(the simulator clips negative eigenvalues silently).
 
 ### 4.2 `SFI.statefunc` — composable state functions
 
@@ -89,7 +100,7 @@ positions; see `SFI.diagnostics` and `examples/gallery/dynamics_order_demo.py`.
 | `make_sf(func, dim, rank, ...)`  | Wrap a Python callable as an `SF`         |
 | `make_psf(func, dim, rank, ...)` | Wrap a callable with parameters as `PSF`  |
 | `make_basis(func, dim, ...)`     | Wrap a vector-valued callable as `Basis`  |
-| `Interactor`, `make_interactor`  | Multi-scale interaction graphs            |
+| `Interactor`, `make_interactor`  | Multi-scale interaction graphs (pairs, hyperedges; `per_slot=True` for many-body forces) |
 
 ### 4.3 `SFI.inference` — inference engines
 

{stochasticforceinference-2.0.3 → stochasticforceinference-2.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: StochasticForceInference
-Version: 2.0.3
+Version: 2.0.4
 Summary: Stochastic Force Inference for Langevin SDEs
 Author-email: Pierre Ronceray <pierre.ronceray@univ-amu.fr>
 License: MIT
@@ -42,6 +42,7 @@ Requires-Dist: sphinxcontrib-video; extra == "docs"
 Requires-Dist: sphinxcontrib-mermaid; extra == "docs"
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-xdist; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: pyright; extra == "dev"
 Requires-Dist: StochasticForceInference[docs]; extra == "dev"
@@ -57,6 +58,7 @@ Dynamic: license-file
 [![Documentation](https://readthedocs.org/projects/stochasticforceinference/badge/?version=latest)](https://stochasticforceinference.readthedocs.io/en/latest/)
 [![PyPI](https://img.shields.io/pypi/v/StochasticForceInference.svg)](https://pypi.org/project/StochasticForceInference/)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/ronceray/StochasticForceInference/blob/main/LICENSE)
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.21532842.svg)](https://doi.org/10.5281/zenodo.21532842)
 
 **Infer force and diffusion fields from stochastic trajectory data.**
 
@@ -197,8 +199,14 @@ guidance.
 ## Citation
 
 If you use SFI, please cite the software (see [`CITATION.cff`](CITATION.cff))
-together with the method paper(s) relevant to the features you use. See
-**[How to cite](https://stochasticforceinference.readthedocs.io/en/latest/#how-to-cite)** in the
+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](https://doi.org/10.5281/zenodo.21532842), which always
+resolves to the latest release; each version also has its own DOI (v2.0.3:
+[10.5281/zenodo.21532843](https://doi.org/10.5281/zenodo.21532843)).
+
+See **[How to cite](https://stochasticforceinference.readthedocs.io/en/latest/#how-to-cite)** in the
 documentation for the full, feature-specific list (overdamped, underdamped/ULI,
 trapeze, and PASTIS references).
 

{stochasticforceinference-2.0.3 → stochasticforceinference-2.0.4}/README.md

@@ -3,6 +3,7 @@
 [![Documentation](https://readthedocs.org/projects/stochasticforceinference/badge/?version=latest)](https://stochasticforceinference.readthedocs.io/en/latest/)
 [![PyPI](https://img.shields.io/pypi/v/StochasticForceInference.svg)](https://pypi.org/project/StochasticForceInference/)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/ronceray/StochasticForceInference/blob/main/LICENSE)
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.21532842.svg)](https://doi.org/10.5281/zenodo.21532842)
 
 **Infer force and diffusion fields from stochastic trajectory data.**
 
@@ -143,8 +144,14 @@ guidance.
 ## Citation
 
 If you use SFI, please cite the software (see [`CITATION.cff`](CITATION.cff))
-together with the method paper(s) relevant to the features you use. See
-**[How to cite](https://stochasticforceinference.readthedocs.io/en/latest/#how-to-cite)** in the
+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](https://doi.org/10.5281/zenodo.21532842), which always
+resolves to the latest release; each version also has its own DOI (v2.0.3:
+[10.5281/zenodo.21532843](https://doi.org/10.5281/zenodo.21532843)).
+
+See **[How to cite](https://stochasticforceinference.readthedocs.io/en/latest/#how-to-cite)** in the
 documentation for the full, feature-specific list (overdamped, underdamped/ULI,
 trapeze, and PASTIS references).
 

{stochasticforceinference-2.0.3 → stochasticforceinference-2.0.4}/SFI/bases/pairs.py

@@ -29,8 +29,8 @@ from ..statefunc import Basis, Interactor, make_basis, make_interactor
 # ═══════════════════════════════════════════════════════════════════════
 
 
-def pbc_displacement(xj, xi, box):
-    """Minimum-image displacement ``xj - xi`` under periodic boundaries.
+def pbc_displacement(xj, xi, box, ref=None):
+    r"""Minimum-image displacement ``xj - xi`` under periodic boundaries.
 
     Works in any dimension.  All inputs are plain JAX arrays of shape
     ``(d,)`` (or broadcastable).
@@ -41,14 +41,49 @@ def pbc_displacement(xj, xi, box):
         Positions (or sub-positions) of two particles.
     box : array, shape ``(d,)``
         Box lengths along each axis.
+    ref : tuple ``(ref_j, ref_i)`` of arrays, optional
+        Reference positions of the same two particles (same axes as
+        ``xj``, ``xi``).  The periodic image is the minimum image of
+        ``ref_j - ref_i`` instead of ``xj - xi``.  Re-picking the image at
+        every frame makes direction-carrying features such as
+        :math:`\phi(r)\,\hat r_{ij}` jump when a pair crosses the cut
+        ``|Δx| = L/2``, and a program combining a feature at several
+        frames (Stratonovich trapezoid, shifted Gram, ULI) integrates the
+        jumps into an :math:`O(1)` bias.  The integrator therefore hands
+        such programs one reference frame per row (a frame preceding every
+        increment of the row, see
+        :meth:`~SFI.integrate.Integrand.pbc_reference_in`) as the reserved
+        per-particle extra ``extras["pbc_ref"]``.  Hand-written pair
+        interactors declare ``particle_extras=("pbc_ref",)``, receive it
+        with shape ``(K, d)`` and pass ``ref=(pr[1], pr[0])`` (sliced to
+        the spatial axes) when present.
 
     Returns
     -------
     dx : array, shape ``(d,)``
-        ``xj - xi`` folded via minimum-image convention.
+        ``xj - xi`` folded via minimum-image convention (of ``ref``
+        when given).
     """
     dx = xj - xi
-    return dx - box * jnp.round(dx / box)
+    d_img = dx if ref is None else ref[0] - ref[1]
+    return dx - box * jnp.round(d_img / box)
+
+
+# Reserved per-particle extra: reference positions fixing the periodic image
+# of every pair (see :func:`pbc_displacement`); SFI.integrate.integrand.PBC_REF_KEY.
+_PBC_REF_KEY = "pbc_ref"
+
+
+def _pbc_ref(extras):
+    """The ``pbc_ref`` reference positions ``(2, d)`` of a pair, or None."""
+    if not extras:
+        return None
+    return extras.get(_PBC_REF_KEY)
+
+
+def _pbc_particle_extras(box):
+    """Periodic pair primitives gather the optional image reference per pair."""
+    return (_PBC_REF_KEY,) if box is not None else ()
 
 
 def wrap_angle(a):
@@ -56,7 +91,7 @@ def wrap_angle(a):
     return (a + jnp.pi) % (2.0 * jnp.pi) - jnp.pi
 
 
-def _pairwise_dr(XK, *, box=None, spatial_dims=None, eps=1e-12):
+def _pairwise_dr(XK, *, box=None, spatial_dims=None, eps=1e-12, ref=None):
     """Compute displacement, distance and unit vector for a K=2 pair.
 
     Parameters
@@ -70,6 +105,9 @@ def _pairwise_dr(XK, *, box=None, spatial_dims=None, eps=1e-12):
         Default: all.
     eps : float
         Regularisation to avoid division by zero.
+    ref : array, shape ``(2, dim)``, optional
+        Reference states of the pair fixing the periodic image (see
+        :func:`pbc_displacement`).
 
     Returns
     -------
@@ -83,7 +121,9 @@ def _pairwise_dr(XK, *, box=None, spatial_dims=None, eps=1e-12):
     else:
         xi_s, xj_s = xi, xj
     if box is not None:
-        dx = pbc_displacement(xj_s, xi_s, box)
+        if ref is not None:
+            ref = (ref[1], ref[0]) if spatial_dims is None else (ref[1][spatial_dims], ref[0][spatial_dims])
+        dx = pbc_displacement(xj_s, xi_s, box, ref=ref)
     else:
         dx = xj_s - xi_s
     r = jnp.sqrt(jnp.sum(dx**2) + eps)
@@ -303,7 +343,7 @@ def radial_pair_basis(
             b = extras["box"]
             if spatial_dims is not None:
                 b = b[spatial_dims] if b.ndim > 0 else b
-        dx, r, rhat = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims)
+        dx, r, rhat = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims, ref=_pbc_ref(extras))
 
         # Evaluate kernels → (n_feat,)
         vals = jnp.stack([fn(r) for fn in fns])  # (F,)
@@ -331,6 +371,7 @@ def radial_pair_basis(
         K=2,
         n_features=n_feat,
         extras_keys=extras_keys,
+        particle_extras=_pbc_particle_extras(box),
         labels=list(labels),
         descriptor="radial-pair-basis",
     )
@@ -377,7 +418,7 @@ def scalar_pair_basis(
             b = extras["box"]
             if spatial_dims is not None:
                 b = b[spatial_dims] if b.ndim > 0 else b
-        _, r, _ = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims)
+        _, r, _ = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims, ref=_pbc_ref(extras))
         return jnp.stack([fn(r) for fn in fns])  # (F,)
 
     extras_keys = ("box",) if _use_extras_box else ()
@@ -388,6 +429,7 @@ def scalar_pair_basis(
         K=2,
         n_features=n_feat,
         extras_keys=extras_keys,
+        particle_extras=_pbc_particle_extras(box),
         labels=list(labels),
         descriptor="scalar-pair-basis",
     )
@@ -460,7 +502,7 @@ def angular_pair_basis(
             b = extras["box"]
             if spatial_dims is not None:
                 b = b[spatial_dims] if b.ndim > 0 else b
-        _, r, _ = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims)
+        _, r, _ = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims, ref=_pbc_ref(extras))
         # angle coupling
         dth = wrap_angle(xj[angle_index] - xi[angle_index])
         g = coupling_fn(dth)
@@ -478,6 +520,7 @@ def angular_pair_basis(
         K=2,
         n_features=n_feat,
         extras_keys=extras_keys,
+        particle_extras=_pbc_particle_extras(box),
         labels=list(labels),
         descriptor="angular-pair-basis",
     )
@@ -575,7 +618,7 @@ def dyadic_pair_basis(
             b = extras["box"]
             if spatial_dims is not None:
                 b = b[spatial_dims] if b.ndim > 0 else b
-        _, r, rhat = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims)
+        _, r, rhat = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims, ref=_pbc_ref(extras))
         vals = jnp.stack([fn(r) for fn in fns])  # (F,)
         # rhat ⊗ rhat: (d, d) then weight by each kernel → (d, d, F)
         rr = jnp.outer(rhat, rhat)  # (d, d)
@@ -589,6 +632,7 @@ def dyadic_pair_basis(
         K=2,
         n_features=n_feat,
         extras_keys=extras_keys,
+        particle_extras=_pbc_particle_extras(box),
         labels=list(labels),
         descriptor="dyadic-pair-basis",
     )
@@ -656,7 +700,7 @@ def pair_direction(
             b = extras["box"]
             if spatial_dims is not None:
                 b = b[spatial_dims] if b.ndim > 0 else b
-        _, _, rhat = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims)
+        _, _, rhat = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims, ref=_pbc_ref(extras))
         if embed_dim is not None:
             full = jnp.zeros(out_d, dtype=XK.dtype)
             for k, ax in enumerate(embed_axes):
@@ -672,6 +716,7 @@ def pair_direction(
         K=2,
         n_features=1,
         extras_keys=extras_keys,
+        particle_extras=_pbc_particle_extras(box),
         labels=("r̂_ij",),
         descriptor="pair-direction",
     )
@@ -828,7 +873,7 @@ def vision_gate(
             b = extras["box"]
             if spatial_dims is not None:
                 b = b[spatial_dims] if b.ndim > 0 else b
-        dx, _, _ = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims)
+        dx, _, _ = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims, ref=_pbc_ref(extras))
         if dx.shape[0] != 2:
             raise ValueError(
                 f"vision_gate requires exactly 2 spatial dimensions, got {dx.shape[0]}. "
@@ -848,6 +893,7 @@ def vision_gate(
         K=2,
         n_features=1,
         extras_keys=extras_keys,
+        particle_extras=_pbc_particle_extras(box),
         labels=("vision",),
         descriptor="vision-gate",
     )
@@ -895,7 +941,7 @@ def parametric_radial_kernel(
             b = extras["box"]
             if spatial_dims is not None:
                 b = b[spatial_dims] if b.ndim > 0 else b
-        _, r, _ = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims)
+        _, r, _ = _pairwise_dr(XK, box=b, spatial_dims=spatial_dims, ref=_pbc_ref(extras))
         return kernel_fn(r, params)[None]  # (1,) — rank-0, 1 feature
 
     extras_keys = ("box",) if _use_extras_box else ()
@@ -907,6 +953,7 @@ def parametric_radial_kernel(
         n_features=1,
         params=params,
         extras_keys=extras_keys,
+        particle_extras=_pbc_particle_extras(box),
         labels=("k(r)",),
         descriptor="parametric-radial-kernel",
     )

{stochasticforceinference-2.0.3 → stochasticforceinference-2.0.4}/SFI/diagnostics/__init__.py

@@ -15,7 +15,10 @@ like an independent ``N(0, 1)`` sample:
 * Ljung--Box autocorrelation of the residuals and their squares;
 * Kolmogorov--Smirnov normality against ``N(0, 1)``;
 * predicted-vs-realised normalised mean squared error of the force
-  (a sampling-noise-aware chi-square check).
+  (a sampling-noise-aware chi-square check);
+* the finite sampling-interval scale ``dt*|grad F|`` of the fitted force
+  (:func:`finite_dt_scale`), the size of the O(dt) biases of the linear
+  estimators.
 
 See :mod:`SFI.diagnostics.residual_tests`. Two preset levels are exposed:
 
@@ -36,6 +39,8 @@ Ljung & Box (1978); Diebold, Gunther & Tay (1998).
 
 from .assess import assess
 from .dynamics_order import DynamicsOrderReport, classify_dynamics
+from .finite_dt import finite_dt_scale
+from .psd import psd_fraction
 from .plotting import (
     plot_dynamics_order,
     plot_qq,
@@ -51,10 +56,12 @@ __all__ = [
     "classify_dynamics",
     "DiagnosticsReport",
     "DynamicsOrderReport",
+    "finite_dt_scale",
     "parametric_four_point_diagnostic",
     "plot_dynamics_order",
     "plot_qq",
     "plot_residual_acf",
     "plot_residual_histogram",
     "plot_summary",
+    "psd_fraction",
 ]

{stochasticforceinference-2.0.3 → stochasticforceinference-2.0.4}/SFI/diagnostics/assess.py

@@ -8,6 +8,7 @@ from __future__ import annotations
 from typing import Optional
 
 from .report import DiagnosticsReport
+from .finite_dt import finite_dt_scale
 from .residual_tests import (
     autocorrelation_tests,
     mse_consistency,
@@ -74,6 +75,13 @@ def assess(
         residuals["normality"] = normality_test(bundle)
         residuals["mse_consistency"] = mse_consistency(inferer, bundle)
 
+    finite_dt: dict = {}
+    if level == "standard":
+        try:
+            finite_dt = finite_dt_scale(inferer)
+        except Exception as exc:  # never let the guard break the residual report
+            finite_dt = {"available": False, "reason": f"{type(exc).__name__}: {exc}", "verdict": "unavailable"}
+
     meta = {
         "backend": bundle.backend,
         "regime": bundle.regime,
@@ -84,4 +92,4 @@ def assess(
         "inferer_class": type(inferer).__name__,
     }
 
-    return DiagnosticsReport(residuals=residuals, meta=meta)
+    return DiagnosticsReport(residuals=residuals, meta=meta, finite_dt=finite_dt)

{stochasticforceinference-2.0.3 → stochasticforceinference-2.0.4}/SFI/diagnostics/dynamics_order.py

@@ -29,71 +29,70 @@ cross-check (:func:`_cross_check`) into an OD / UD / inconclusive verdict.
 from __future__ import annotations
 
 import json
+import warnings
 from dataclasses import dataclass, field
 from typing import Optional
 
-import jax.numpy as jnp
 import numpy as np
 from scipy.optimize import least_squares
 
 from SFI.diagnostics.report import _to_jsonable
-from SFI.integrate import Integrand, Term, TimeOperand, integrate, timeop
 
 # --------------------------------------------------------------------------- #
-# Lag-0/1/2 displacement-covariance time-operators.
+# Lag-0/1/2 displacement covariances.
 #
-# Each returns a per-particle (..., N, d, d) tensor, mirroring the diffusion /
-# measurement-noise estimators in SFI.inference.overdamped (``_Lambda`` etc.).
-# The dataset stream layer supplies the increments:
-#   dX        = X[t+1] - X[t]      (offset 0..+1)
-#   dX_minus  = X[t]   - X[t-1]    (offset -1..0)
-#   dX_plus   = X[t+2] - X[t+1]    (offset +1..+2)
-# so <dX . dX_minus> is the lag-1 covariance and <dX_minus . dX_plus> is the
-# lag-2 covariance (displacements one step apart).
+# Plain moment sums over the increments dX[t] = X[t+1] - X[t], with the same
+# row semantics as ``integrate(..., reduce="mean")`` on the dataset streams
+# (dX_minus = dX[t-1], dX_plus = dX[t+1]): a row (t, i) counts when the
+# dynamic mask holds at t and the static mask at every other frame its
+# increments touch; rows are weighted by dt and the dataset weight.  Computed
+# in NumPy: each coarse-graining stride has its own length, and a JAX path
+# recompiles for every one of them.
 # --------------------------------------------------------------------------- #
 
-
-@timeop(name="dyn_cov_lag0", batch_safe=True)
-def _cov_lag0(**streams):
-    """Lag-0 displacement covariance ``dX (x) dX``."""
-    dX = streams["dX"]
-    return _outer(dX, dX)
-
-
-@timeop(name="dyn_cov_lag1", batch_safe=True)
-def _cov_lag1(**streams):
-    """Symmetrised lag-1 covariance ``1/2 (dX (x) dX^- + dX^- (x) dX)``."""
-    dX = streams["dX"]
-    dXm = streams["dX_minus"]
-    return 0.5 * (_outer(dX, dXm) + _outer(dXm, dX))
-
-
-@timeop(name="dyn_cov_lag2", batch_safe=True)
-def _cov_lag2(**streams):
-    """Symmetrised lag-2 covariance ``1/2 (dX^- (x) dX^+ + dX^+ (x) dX^-)``."""
-    dXm = streams["dX_minus"]
-    dXp = streams["dX_plus"]
-    return 0.5 * (_outer(dXm, dXp) + _outer(dXp, dXm))
-
-
-_cov_lag0._requires = frozenset({"dX"})  # type: ignore[attr-defined]
-_cov_lag1._requires = frozenset({"dX", "dX_minus"})  # type: ignore[attr-defined]
-_cov_lag2._requires = frozenset({"dX_minus", "dX_plus"})  # type: ignore[attr-defined]
-
-
-def _outer(a, b):
-    """Outer product over the trailing spatial axis: ``a_m b_n``."""
-    return jnp.einsum("...m,...n->...mn", a, b)
-
-
-# --------------------------------------------------------------------------- #
-# Covariance backbone
-# --------------------------------------------------------------------------- #
-def _integrate_tensor(coll, op, alias):
-    """Time- and particle-averaged ``(d, d)`` value of a covariance time-op."""
-    operand = TimeOperand(op, alias=alias)
-    prog = Integrand(times=[operand], terms=[Term(eq="imn->imn", ops=(alias,))])
-    return np.asarray(integrate(coll, prog, reduce="mean"))
+# rows t in [first, T - tail), increments (dX[t+a], dX[t+b]), static-mask frames t+off
+_LAGS = {
+    "C0": (0, 1, (0, 0), (1,)),
+    "C1": (1, 1, (0, -1), (-1, 1)),
+    "C2": (1, 2, (-1, 1), (-1, 1, 2)),
+}
+
+
+def _dataset_steps(ds, T):
+    """Per-interval steps ``(T-1,)`` of a dataset (``t[k+1] - t[k]``)."""
+    if ds.t is not None:
+        return np.diff(np.asarray(ds.t, dtype=float))
+    if ds.dt is None:
+        raise ValueError("Both 't' and 'dt' are None; one must be provided.")
+    dt = np.asarray(ds.dt, dtype=float)
+    return np.broadcast_to(dt, (T - 1,)) if dt.ndim == 0 else dt[: T - 1]
+
+
+def _lag_covariances(coll):
+    """Time- and particle-averaged ``(d, d)`` lag-0/1/2 covariances of ``coll``."""
+    acc = dict.fromkeys(_LAGS, 0.0)
+    norm = dict.fromkeys(_LAGS, 0.0)
+    for ds, w in zip(coll.datasets, np.asarray(coll.weights, dtype=float)):
+        X = np.asarray(ds._X3d(), dtype=float)  # (T, N, d)
+        M = np.asarray(ds._M2d())  # static mask (T, N)
+        Mdyn = np.asarray(ds._dynamic_M2d())
+        T = X.shape[0]
+        steps = _dataset_steps(ds, T)
+        dX = np.diff(X, axis=0)
+        for key, (first, tail, (a, b), frames) in _LAGS.items():
+            t = np.arange(first, T - tail)
+            if t.size == 0:
+                continue
+            m = Mdyn[t].copy()
+            for off in frames:
+                m &= M[t + off]
+            u, v = dX[t + a], dX[t + b]
+            # NaN-safe masking: masked frames may hold NaN placeholders.
+            uv = np.einsum("kim,kin->kimn", u, v)
+            uv = np.where(m[..., None, None], 0.5 * (uv + np.swapaxes(uv, -1, -2)), 0.0)
+            acc[key] = acc[key] + w * np.einsum("k,kimn->mn", steps[t], uv)
+            norm[key] = norm[key] + w * float(np.sum(steps[t] * m.sum(axis=1)))
+    return {k: acc[k] / norm[k] if norm[k] > 0 else np.asarray(0.0) for k in _LAGS}
 
 
 def _mean_dt(coll):
@@ -112,13 +111,9 @@ def _mean_dt(coll):
     return float(np.mean(dts))
 
 
-def _n_eff(coll, require):
-    """Number of valid (masked) increment samples available for ``require``."""
-    n = 0
-    for ds in coll.datasets:
-        idx = ds.valid_indices(frozenset(require))
-        n += int(np.asarray(idx).size) * int(ds.N)
-    return n
+def _n_eff(coll):
+    """Number of lag-2 increment pairs (rows ``t = 1 .. T-3``, all particles)."""
+    return sum(max(int(ds.T) - 3, 0) * int(ds.N) for ds in coll.datasets)
 
 
 def _increment_covariances(coll):
@@ -135,15 +130,13 @@ def _increment_covariances(coll):
         ``C0``, ``C1``, ``C2`` : ``(d, d)`` displacement covariance tensors;
         ``dt`` : representative scalar step; ``n_eff`` : valid sample count.
     """
-    C0 = _integrate_tensor(coll, _cov_lag0, "c0")
-    C1 = _integrate_tensor(coll, _cov_lag1, "c1")
-    C2 = _integrate_tensor(coll, _cov_lag2, "c2")
+    cov = _lag_covariances(coll)
     return {
-        "C0": C0,
-        "C1": C1,
-        "C2": C2,
+        "C0": cov["C0"],
+        "C1": cov["C1"],
+        "C2": cov["C2"],
         "dt": _mean_dt(coll),
-        "n_eff": _n_eff(coll, {"dX_minus", "dX_plus"}),
+        "n_eff": _n_eff(coll),
     }
 
 
@@ -432,7 +425,6 @@ def _cross_check(data, *, force_order: int = 3):
     quick fit fails for any reason.
     """
     try:
-        from SFI.bases.constants import unit_vector_basis
         from SFI.bases.monomials import monomials_up_to
         from SFI.diagnostics.assess import assess
         from SFI.inference.overdamped import OverdampedLangevinInference
@@ -441,8 +433,12 @@ def _cross_check(data, *, force_order: int = 3):
         inf = OverdampedLangevinInference(data)
         inf.compute_diffusion_constant(method="auto")
         basis = monomials_up_to(order=force_order, dim=d, include_constant=True, include_x=True, include_v=False)
-        inf.infer_force_linear(basis * unit_vector_basis(d))
-        inf.compute_force_error()
+        with warnings.catch_warnings():
+            # A corroboration fit on a deliberately misfit model: its Gram
+            # conditioning is not the user's concern.
+            warnings.simplefilter("ignore", UserWarning)
+            inf.infer_force_linear(basis.vectorize(d))
+            inf.compute_force_error()
         report = assess(inf, level="standard")
         lb = report.residuals.get("autocorr", {}).get("ljung_box", {})
         return {"ljung_box_pvalue": lb.get("pvalue"), "ljung_box_stat": lb.get("statistic")}