StochasticForceInference 2.0.0__py3-none-any.whl

SFI/diagnostics/plotting.py

@@ -0,0 +1,226 @@
+"""Diagnostic plotting helpers.
+
+Three panels covering the common visual checks of a fitted inference
+object:
+
+* :func:`plot_qq` — normal Q--Q plot of pooled whitened residuals;
+* :func:`plot_residual_histogram` — histogram of pooled residuals
+  overlaid with the standard normal density;
+* :func:`plot_residual_acf` — autocorrelation of the residuals (and of
+  their squares) with the ``±1.96/√n`` band;
+* :func:`plot_summary` — a 1×3 figure combining the three panels.
+
+All helpers accept either a fitted inference object (calling
+:func:`~SFI.diagnostics.assess.assess` internally) or a pre-computed
+:class:`~SFI.diagnostics.report.DiagnosticsReport`.
+
+Plots use the SFI palette (:data:`SFI.utils.plotting.SFI_COLORS`) but do
+*not* call ``apply_style()`` — applying the gallery style is the caller's
+responsibility (see ``GALLERY_STYLE_GUIDE.md``).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from SFI.utils.plotting import SFI_COLORS
+
+from .assess import assess
+from .report import DiagnosticsReport
+
+
+def _coerce_report(obj: Any, *, level: str = "standard") -> DiagnosticsReport:
+    """Return a :class:`DiagnosticsReport` from either a report or inferer."""
+    if isinstance(obj, DiagnosticsReport):
+        return obj
+    return assess(obj, level=level)
+
+
+def _get_residuals(report: DiagnosticsReport) -> np.ndarray:
+    """Return pooled standardised residuals as a 1-D array."""
+    qq = report.residuals.get("normality", {}).get("qq", {})
+    sample = qq.get("sample_quantiles")
+    if sample is not None:
+        return np.asarray(sample).ravel()
+    raise ValueError(
+        "Diagnostics report does not contain raw residual samples. "
+        "Run assess() with level='standard' (the default) before plotting."
+    )
+
+
+def plot_qq(report_or_inferer, *, ax=None, level: str = "standard"):
+    """Normal Q--Q plot of pooled whitened residuals."""
+    import matplotlib.pyplot as plt
+    from scipy.stats import norm
+
+    rep = _coerce_report(report_or_inferer, level=level)
+    sample = np.sort(_get_residuals(rep))
+    n = sample.size
+    if n == 0:
+        raise ValueError("No residuals to plot.")
+    probs = (np.arange(1, n + 1) - 0.5) / n
+    theo = norm.ppf(probs)
+
+    if ax is None:
+        _, ax = plt.subplots(figsize=(4.5, 4.5))
+
+    lo = float(min(theo.min(), sample.min()))
+    hi = float(max(theo.max(), sample.max()))
+    ax.plot([lo, hi], [lo, hi], "--", color=SFI_COLORS["exact"], lw=1.2, label="$y = x$")
+
+    if n > 5000:  # subsample to keep the figure light
+        idx = np.linspace(0, n - 1, 5000).astype(int)
+        ax.scatter(theo[idx], sample[idx], s=6, alpha=0.6, color=SFI_COLORS["data"], label=f"residuals (n={n})")
+    else:
+        ax.scatter(theo, sample, s=8, alpha=0.7, color=SFI_COLORS["data"], label=f"residuals (n={n})")
+
+    ax.set_xlabel("theoretical quantile (N(0,1))")
+    ax.set_ylabel("sample quantile")
+    ax.set_title("Q--Q plot")
+    ax.legend(frameon=False, loc="best", fontsize=8)
+    return ax
+
+
+def plot_residual_histogram(report_or_inferer, *, ax=None, bins: int = 60, level: str = "standard"):
+    """Histogram of pooled residuals overlaid with N(0,1) density."""
+    import matplotlib.pyplot as plt
+    from scipy.stats import norm
+
+    rep = _coerce_report(report_or_inferer, level=level)
+    sample = _get_residuals(rep)
+    if ax is None:
+        _, ax = plt.subplots(figsize=(4.5, 3.5))
+
+    ax.hist(sample, bins=bins, density=True, alpha=0.65, color=SFI_COLORS["data"], label=f"residuals (n={sample.size})")
+    grid = np.linspace(sample.min(), sample.max(), 256)
+    ax.plot(grid, norm.pdf(grid), "--", color=SFI_COLORS["exact"], lw=1.4, label="N(0,1)")
+    ax.set_xlabel("standardised residual $z$")
+    ax.set_ylabel("density")
+    ax.set_title("Residual histogram")
+    ax.legend(frameon=False, loc="best", fontsize=8)
+    return ax
+
+
+def plot_residual_acf(report_or_inferer, *, ax=None, level: str = "standard"):
+    """Autocorrelation of the residuals (and of $z^2$).
+
+    Reads the ACF computed by
+    :func:`~SFI.diagnostics.residual_tests.autocorrelation_tests`
+    (stored on the report) rather than recomputing it.
+    """
+    import matplotlib.pyplot as plt
+
+    rep = _coerce_report(report_or_inferer, level=level)
+    ac = rep.residuals.get("autocorr", {})
+    acf = np.asarray(ac.get("acf", []), dtype=np.float64)
+    acf2 = np.asarray(ac.get("acf_squared", []), dtype=np.float64)
+    n_eff = int(ac.get("n_eff", 0))
+    if acf.size == 0:
+        raise ValueError("Report has no autocorrelation data; run assess(level='standard').")
+
+    if ax is None:
+        _, ax = plt.subplots(figsize=(4.5, 3.5))
+
+    lags = np.arange(1, acf.size + 1)
+    band = 1.96 / np.sqrt(max(n_eff, 1))
+    ax.axhline(0.0, color="0.5", lw=0.8)
+    ax.fill_between(
+        np.arange(0, acf.size + 1), -band, band,
+        color=SFI_COLORS["exact"], alpha=0.15, label="95% band",
+    )
+    ax.vlines(lags, 0.0, acf, color=SFI_COLORS["data"], lw=1.5, label="ACF($z$)")
+    if acf2.size == acf.size:
+        ax.plot(lags, acf2, "o-", color=SFI_COLORS["highlight"], ms=3, lw=1.0, label="ACF($z^2$)")
+    ax.set_xlabel("lag")
+    ax.set_ylabel("autocorrelation")
+    ax.set_title("Residual ACF")
+    ax.legend(frameon=False, loc="best", fontsize=8)
+    return ax
+
+
+def plot_summary(report_or_inferer, *, level: str = "standard", figsize=(13.0, 4.0)):
+    """1×3 summary figure: Q--Q plot, residual histogram, and ACF.
+
+    Returns the matplotlib :class:`Figure`.
+    """
+    import matplotlib.pyplot as plt
+
+    rep = _coerce_report(report_or_inferer, level=level)
+    fig, axes = plt.subplots(1, 3, figsize=figsize)
+    plot_qq(rep, ax=axes[0])
+    plot_residual_histogram(rep, ax=axes[1])
+    plot_residual_acf(rep, ax=axes[2])
+    fig.set_layout_engine("tight")
+    return fig
+
+
+def plot_dynamics_order(report, *, axes=None):
+    """Visualise an OD-vs-UD classification (:func:`classify_dynamics`).
+
+    Two panels versus the sampling step ``dt``:
+
+    * **lag-2 persistence** ``rho2 = C2/(C0+2C1)`` (noise-immune): tends to 0
+      for overdamped data, to a positive plateau for inertia;
+    * **apparent kinetic energy** ``K = (C0+2C1)/dt^2`` on log-log axes, with
+      reference slopes ``-1`` (overdamped, ``K ~ 2D/dt``) and ``0``
+      (underdamped); the fitted log-log slope ``beta`` is annotated.
+
+    The parametric-fit prediction is overlaid on both panels.  Accepts a
+    :class:`~SFI.diagnostics.dynamics_order.DynamicsOrderReport`.
+    """
+    import matplotlib.pyplot as plt
+
+    from .dynamics_order import _dynamics_model
+
+    dt = np.asarray(report.scan["dt"], dtype=float)
+    rho2 = np.asarray(report.scaling["rho2"], dtype=float)
+    K = np.asarray(report.scaling["K"], dtype=float)
+    beta = float(report.scaling.get("beta", float("nan")))
+    p = report.fit.get("params", {})
+    verdict = report.verdict
+
+    hg = np.geomspace(dt.min(), dt.max(), 100)
+    m0, m1, m2 = _dynamics_model(hg, p.get("sigma2", 0.0), p.get("D", 0.0), p.get("V", 0.0), p.get("gamma", 1.0))
+    nf = m0 + 2.0 * m1
+    with np.errstate(divide="ignore", invalid="ignore"):
+        rho2_fit = np.where(np.abs(nf) > 0, m2 / nf, np.nan)
+        K_fit = nf / hg**2
+
+    if axes is None:
+        fig, axes = plt.subplots(1, 2, figsize=(10.0, 4.0))
+    else:
+        fig = np.ravel(axes)[0].figure
+    ax_rho, ax_k = np.ravel(axes)[:2]
+
+    ax_rho.semilogx(dt, rho2, "o", color=SFI_COLORS["data"], label="data")
+    ax_rho.semilogx(hg, rho2_fit, "-", color=SFI_COLORS["inferred"], label="fit")
+    ax_rho.axhline(0.0, color="0.5", lw=0.8)
+    ax_rho.set_xlabel(r"sampling step $\Delta t$")
+    ax_rho.set_ylabel(r"$\rho_2 = C_2 / (C_0 + 2 C_1)$")
+    ax_rho.set_title(f"lag-2 persistence — verdict: {verdict}")
+    ax_rho.legend(frameon=False, loc="best", fontsize=8)
+
+    good = K > 0
+    anchor = float(K[good][0]) if good.any() else 1.0
+    ax_k.loglog(dt[good], K[good], "o", color=SFI_COLORS["data"], label="data")
+    ax_k.loglog(hg, K_fit, "-", color=SFI_COLORS["inferred"], label="fit")
+    ax_k.loglog(hg, anchor * (dt[0] / hg), ":", color=SFI_COLORS["exact"], lw=1.0, label="OD ref (slope $-1$)")
+    ax_k.loglog(hg, np.full_like(hg, anchor), "--", color=SFI_COLORS["highlight"], lw=1.0, label="UD ref (slope $0$)")
+    ax_k.set_xlabel(r"sampling step $\Delta t$")
+    ax_k.set_ylabel(r"apparent KE $\tilde K = (C_0 + 2 C_1)/\Delta t^2$")
+    ax_k.set_title(rf"scaling slope $\beta = {beta:.2f}$  (OD$\to-1$, UD$\to0$)")
+    ax_k.legend(frameon=False, loc="best", fontsize=8)
+
+    fig.set_layout_engine("tight")
+    return fig
+
+
+__all__ = [
+    "plot_qq",
+    "plot_residual_histogram",
+    "plot_residual_acf",
+    "plot_summary",
+    "plot_dynamics_order",
+]

SFI/diagnostics/report.py

@@ -0,0 +1,238 @@
+"""DiagnosticsReport dataclass — structured output of `assess()`."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass, field
+from typing import Any, Mapping
+
+import numpy as np
+
+#: One-line action hint per flag, appended to `flag_issues` messages.
+#: Keys are ``(section, test)``; ``(section, "")`` is the section fallback.
+#: Wording mirrors docs/source/diagnostics.rst ("Interpreting flags" and
+#: "When a flag points beyond the linear estimators") — keep in lockstep.
+_FLAG_HINTS: dict[tuple[str, str], str] = {
+    ("autocorr", "ljung_box"): (
+        "missing time-correlated feature — widen the basis; if it persists, "
+        "suspect coarse sampling: the parametric estimator (infer_force) "
+        "extends the usable Δt"
+    ),
+    ("autocorr", "ljung_box_squared"): (
+        "diffusion mis-estimated or state-dependent — try the other "
+        "compute_diffusion_constant method or a state-dependent diffusion "
+        "basis; the parametric estimators profile (D, Λ)"
+    ),
+    ("normality", ""): (
+        "non-Gaussian residuals — rare events not captured by the basis, "
+        "or a non-Gaussian noise structure"
+    ),
+    ("moments", "mean"): (
+        "non-zero residual mean — systematic drift bias; widen the basis"
+    ),
+    ("moments", "std"): (
+        "whitened std far from 1 — D̄ likely wrong: try both "
+        "compute_diffusion_constant methods, then the parametric estimator"
+    ),
+    ("mse_consistency", ""): (
+        "realised error above predicted — model bias; on experimental data "
+        "usually measurement noise: consider the parametric estimator "
+        "(infer_force)"
+    ),
+}
+
+
+def _hint(section: str, test: str = "") -> str:
+    return _FLAG_HINTS.get((section, test)) or _FLAG_HINTS.get((section, ""), "")
+
+
+def _to_jsonable(obj: Any) -> Any:
+    """Recursively convert numpy / jax arrays to plain Python objects."""
+    if isinstance(obj, Mapping):
+        return {k: _to_jsonable(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_to_jsonable(v) for v in obj]
+    if hasattr(obj, "tolist"):
+        try:
+            return obj.tolist()
+        except Exception:
+            pass
+    if isinstance(obj, (np.floating, np.integer)):
+        return obj.item()
+    return obj
+
+
+@dataclass
+class DiagnosticsReport:
+    """Container for residual-consistency test results.
+
+    Attributes
+    ----------
+    residuals : dict
+        Test results. Always holds ``"moments"``; at ``level="standard"``
+        also ``"autocorr"``, ``"normality"`` and ``"mse_consistency"``.
+    meta : dict
+        Backend tag, regime, ``n_obs``, ``n_particles``, ``d``, level.
+    """
+
+    residuals: dict = field(default_factory=dict)
+    meta: dict = field(default_factory=dict)
+
+    # ------------------------------------------------------------------ #
+    # Serialisation
+    # ------------------------------------------------------------------ #
+    def to_dict(self) -> dict:
+        """Return a JSON-serialisable representation of the report."""
+        return _to_jsonable(asdict(self))
+
+    def to_json(self, indent: int = 2) -> str:
+        """Serialise the report to a JSON string."""
+        return json.dumps(self.to_dict(), indent=indent, default=str)
+
+    # ------------------------------------------------------------------ #
+    # Issue flagging
+    # ------------------------------------------------------------------ #
+    def flag_issues(self, alpha: float = 0.01, *, hints: bool = True) -> list[str]:
+        """List human-readable warnings.
+
+        Returns one line per test whose p-value is below ``alpha`` or
+        whose statistic crosses a sane threshold (residual mean off zero,
+        std far from one, MSE-consistency ``|z| > 5``).
+
+        Parameters
+        ----------
+        alpha : float
+            Significance level for the p-value tests.
+        hints : bool
+            When True (default), each message carries a one-line action
+            hint (" — <what to do>"); set False for bare statistics
+            (machine parsing).
+        """
+        msgs: list[str] = []
+
+        def _emit(base: str, section: str, test: str = "") -> None:
+            h = _hint(section, test) if hints else ""
+            msgs.append(f"{base} — {h}" if h else base)
+
+        # Normality / autocorrelation p-values
+        for section_name, section in (
+            ("normality", self.residuals.get("normality", {})),
+            ("autocorr", self.residuals.get("autocorr", {})),
+        ):
+            for test_name, payload in section.items():
+                if not isinstance(payload, Mapping):
+                    continue
+                p = payload.get("pvalue")
+                if p is not None and p < alpha:
+                    _emit(
+                        f"[{section_name}/{test_name}] p={p:.2e} < {alpha}",
+                        section_name,
+                        test_name,
+                    )
+
+        # Residual moments
+        moments = self.residuals.get("moments", {})
+        m = moments.get("mean")
+        s = moments.get("std")
+        n = float(moments.get("n", 1))
+        band = 4.0 / max(n**0.5, 1.0)
+        if m is not None and s is not None and abs(m) > band:
+            _emit(
+                f"[moments/mean] |mean|={abs(m):.3g} (expected ~0, 4σ band={band:.3g})",
+                "moments",
+                "mean",
+            )
+        if s is not None and (s < 0.5 or s > 2.0):
+            _emit(f"[moments/std] std={s:.3g}", "moments", "std")
+
+        # Predicted vs realised MSE — flag on the chi^2 z-score of the
+        # residual excess, which is sampling-noise-aware (the raw ratio is
+        # too noisy at modest n_obs).  A decisive z (|z|>5) flags on its
+        # own; a moderately significant z (|z|>2) flags when the realised/
+        # predicted ratio is also an order of magnitude off — a large,
+        # consistently-signed excess at modest n_obs (e.g. a structurally
+        # misspecified force family) would otherwise stay silent.
+        consistency = self.residuals.get("mse_consistency", {})
+        excess_z = consistency.get("excess_z")
+        ratio = consistency.get("ratio")
+        if excess_z is not None and (
+            abs(excess_z) > 5.0
+            or (abs(excess_z) > 2.0 and ratio is not None and ratio > 10.0)
+        ):
+            _emit(
+                f"[mse_consistency] residual chi^2 z-score = {excess_z:+.2f}, "
+                f"realised/predicted NMSE = {ratio:.2g}",
+                "mse_consistency",
+            )
+
+        return msgs
+
+    # ------------------------------------------------------------------ #
+    # Pretty-printing
+    # ------------------------------------------------------------------ #
+    def print_summary(self, alpha: float = 0.01, *, hints: bool = True) -> None:
+        """Print a human-readable summary of the diagnostic report.
+
+        Each flagged issue carries a one-line action hint unless
+        ``hints=False``.
+        """
+        print("\n=== SFI diagnostics report ===")
+        meta = self.meta or {}
+        if meta:
+            print(f"backend  : {meta.get('backend', '?')}")
+            print(f"regime   : {meta.get('regime', '?')}")
+            print(
+                f"n_obs    : {meta.get('n_obs', '?')}  "
+                f"n_particles: {meta.get('n_particles', '?')}  "
+                f"d: {meta.get('d', '?')}"
+            )
+            print(f"level    : {meta.get('level', '?')}")
+
+        res = self.residuals or {}
+        if res:
+            print("\n-- Residuals --")
+            mom = res.get("moments", {})
+            if mom:
+                print(
+                    f"  mean = {mom.get('mean', float('nan')):+.4f}  "
+                    f"std = {mom.get('std', float('nan')):.4f}  "
+                    f"skew = {mom.get('skew', float('nan')):+.3f}  "
+                    f"kurt-3 = {mom.get('excess_kurt', float('nan')):+.3f}  "
+                    f"(n={mom.get('n', 0)})"
+                )
+            for sect_name, sect in (
+                ("normality", res.get("normality", {})),
+                ("autocorr", res.get("autocorr", {})),
+            ):
+                if not sect:
+                    continue
+                for test_name, payload in sect.items():
+                    if not isinstance(payload, Mapping):
+                        continue
+                    stat = payload.get("statistic")
+                    p = payload.get("pvalue")
+                    if stat is None and p is None:
+                        continue
+                    flag = "✗" if (p is not None and p < alpha) else "✓"
+                    s_str = f"stat={stat:.3g}" if stat is not None else ""
+                    p_str = f"p={p:.3g}" if p is not None else ""
+                    print(f"  {flag} {sect_name:9s} {test_name:18s} {s_str:14s} {p_str}")
+
+            mc = res.get("mse_consistency", {})
+            if mc:
+                pred = mc.get("predicted_NMSE")
+                real = mc.get("realised_NMSE")
+                excess_z = mc.get("excess_z")
+                pred_str = f"{pred:.3g}" if isinstance(pred, (int, float)) else str(pred)
+                real_str = f"{real:.3g}" if isinstance(real, (int, float)) else str(real)
+                z_str = f"{excess_z:+.2f}" if isinstance(excess_z, (int, float)) else str(excess_z)
+                print(f"  predicted NMSE = {pred_str}  realised NMSE = {real_str}  χ² z = {z_str} (|z|>5 ⇒ bias)")
+
+        issues = self.flag_issues(alpha=alpha, hints=hints)
+        print("\n-- Flags --")
+        if not issues:
+            print("  (no issues at α = {:.2g})".format(alpha))
+        else:
+            for msg in issues:
+                print(f"  ! {msg}")
+        print()