falsesync 0.1.2__py3-none-any.whl

falsesync/features.py

@@ -0,0 +1,164 @@
+"""Feature extraction for regime classification (Phase 3).
+
+Every feature is a rescaled, roughly O(1) quantity so the multinomial
+logistic classifier is numerically stable. Features are documented in
+diagnostic_specification.md §4 (s1..s4) plus Phase-3 additions.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .aggregation import weighted_aggregate
+from .changepoints import aggregate_breakpoint, ls_breakpoint, multibreak
+from .inference import deconvolve_timing, unit_break_uncertainties
+
+FEATURE_NAMES = [
+    "agg_strength", "excess_sd_ratio", "n_modes", "mode_mass2",
+    "s1_width", "s3_alignment", "s4_weighting", "noise_floor_ratio",
+    "fwhm_ratio", "agg_skew", "log_n", "log_T", "frac_valid_units",
+    "binseg_n_breaks", "unit_strength_med",
+]
+
+
+def _kde(samples, grid, bw):
+    u = (grid[:, None] - samples[None, :]) / bw
+    return np.exp(-0.5 * u * u).mean(axis=1) / (bw * np.sqrt(2 * np.pi))
+
+
+def _mode_count(f, thr=0.3, min_sep_frac=0.08):
+    """Count modes with prominence threshold and minimum separation.
+
+    A mode must exceed `thr` * max(f) locally and be at least
+    `min_sep_frac` * len(f) grid points away from a higher mode
+    (greedy, highest peak first).
+    """
+    if f.size < 3 or f.max() <= 0:
+        return 0, []
+    cand = [i for i in range(1, len(f) - 1)
+            if f[i] >= f[i - 1] and f[i] > f[i + 1] and f[i] > thr * f.max()]
+    cand.sort(key=lambda i: -f[i])
+    sep = max(1, int(min_sep_frac * len(f)))
+    chosen = []
+    for i in cand:
+        if all(abs(i - j) >= sep for j in chosen):
+            chosen.append(i)
+    return len(chosen), sorted(chosen)
+
+
+def extract_features(
+    t: np.ndarray,
+    values: np.ndarray,
+    weights: np.ndarray | None = None,
+    windows: np.ndarray | None = None,
+    aggregate: np.ndarray | None = None,
+    rng: np.random.Generator | None = None,
+    unit_boot: int = 30,
+) -> tuple[np.ndarray, dict]:
+    """Feature vector + aux dict used by the classifier and the result object."""
+    rng = rng or np.random.default_rng(0)
+    t = np.asarray(t, dtype=float)
+    n = values.shape[0]
+    if aggregate is None:
+        aggregate = weighted_aggregate(values, weights)
+    valid_u = np.isfinite(values).sum(axis=1) >= 8
+    agg_f = np.isfinite(aggregate)
+    agg_y = np.nan_to_num(aggregate, nan=np.nanmean(aggregate[agg_f])) if agg_f.any() else np.zeros(t.size)
+
+    # aggregate break
+    res_agg = ls_breakpoint(t, agg_y)
+    dm = np.gradient(agg_y, t)
+    half = np.max(dm) / 2
+    idx = np.where(dm >= half)[0]
+    w_agg = float(t[idx[-1]] - t[idx[0]]) if idx.size > 1 else np.inf
+
+    # unit breaks
+    taus = np.full(n, np.nan)
+    strengths = np.full(n, np.nan)
+    for i in np.where(valid_u)[0]:
+        y = values[i]
+        fill = np.nanmean(y[np.isfinite(y)])
+        try:
+            r = ls_breakpoint(t, np.nan_to_num(y, nan=fill))
+            taus[i] = r.location
+            strengths[i] = r.strength
+        except Exception:
+            pass
+    tv = taus[np.isfinite(taus)]
+    sv = strengths[np.isfinite(taus)]
+
+    s_i = unit_break_uncertainties(t, values[np.isfinite(taus)], tv, rng=rng, n_boot=unit_boot) if tv.size >= 5 else np.array([])
+    noise_floor = float(np.nanmedian(s_i)) if s_i.size and np.isfinite(s_i).any() else np.nan
+
+    tau_sd = float(np.std(tv)) if tv.size else np.nan
+    excess = np.sqrt(max(tau_sd**2 - (noise_floor if np.isfinite(noise_floor) else 0) ** 2, 0)) if np.isfinite(tau_sd) else np.nan
+    scale = np.ptp(t)
+    excess_ratio = float(excess / scale) if np.isfinite(excess) else np.nan
+    floor_ratio = float(noise_floor / max(tau_sd, 1e-9)) if np.isfinite(noise_floor) and np.isfinite(tau_sd) else np.nan
+
+    grid = np.linspace(t.min(), t.max(), 300)
+    bw = max(1.06 * (tau_sd if np.isfinite(tau_sd) else 1.0) * max(len(tv), 2) ** -0.2, 3 * np.median(np.diff(t)))
+    f_hat = _kde(tv, grid, bw) if tv.size >= 5 else np.zeros(grid.size)
+    f_dec_raw = deconvolve_timing(tv, s_i, grid) if tv.size >= 5 else np.zeros(grid.size)
+    # smooth deconvolved density before mode counting (deconvolution is noisy)
+    k = max(3, int(0.04 * grid.size) | 1)
+    kern = np.exp(-0.5 * (np.arange(k) - k // 2) ** 2 / (k / 4) ** 2)
+    kern /= kern.sum()
+    f_dec = np.convolve(np.pad(f_dec_raw, k // 2, mode="edge"), kern, mode="valid")
+    f_dec = np.clip(f_dec, 0, None)
+    # modes counted on the (smoother) KDE; deconvolved density is used for
+    # dispersion correction, where its noise matters less
+    n_modes, modes = _mode_count(f_hat, thr=0.3, min_sep_frac=0.08)
+    mode_mass2 = float(np.sort(f_dec)[::-1][:2].sum() / max(f_dec.sum(), 1e-12)) if f_dec.sum() > 0 else 0.0
+
+    s1 = float(tau_sd / w_agg) if np.isfinite(tau_sd) and np.isfinite(w_agg) and w_agg > 0 else np.nan
+    s3 = float(np.mean(np.abs(sv) > 2.0)) if sv.size else np.nan
+
+    s4 = 0.0
+    if weights is not None and tv.size >= 5:
+        wv = np.asarray(weights)
+        if wv.ndim == 2:
+            wv = wv.mean(axis=1)
+        try:
+            s4 = max(s4, abs(float(np.corrcoef(wv[np.isfinite(taus)], tv)[0, 1])))
+        except Exception:
+            pass
+    if windows is not None and tv.size >= 5:
+        span = np.asarray(windows)[np.isfinite(taus)]
+        span = span[:, 1] - span[:, 0]
+        try:
+            s4 = max(s4, abs(float(np.corrcoef(span, tv)[0, 1])))
+        except Exception:
+            pass
+
+    agg_skew = float(((agg_y - agg_y.mean()) ** 3).mean() / max(agg_y.std(), 1e-9) ** 3)
+    try:
+        nb = len(multibreak(t, agg_y, n_bkps=5).extras.get("all_break_indices", []))
+    except Exception:
+        nb = 1
+
+    feats = np.array([
+        res_agg.strength,
+        excess_ratio if np.isfinite(excess_ratio) else 0.0,
+        float(n_modes),
+        mode_mass2,
+        s1 if np.isfinite(s1) else 0.0,
+        s3 if np.isfinite(s3) else 0.0,
+        float(s4),
+        floor_ratio if np.isfinite(floor_ratio) else 1.0,
+        float(w_agg / scale) if np.isfinite(w_agg) else 1.0,
+        agg_skew,
+        float(np.log(max(n, 1))),
+        float(np.log(t.size)),
+        float(valid_u.mean()),
+        float(nb),
+        float(np.nanmedian(np.abs(sv))) if sv.size else 0.0,
+    ])
+    aux = {
+        "agg_break": res_agg.location, "agg_strength": res_agg.strength,
+        "taus": taus, "unit_strengths": strengths, "s_i": s_i,
+        "noise_floor": noise_floor, "tau_sd": tau_sd, "f_dec": (grid, f_dec),
+        "f_hat": (grid, f_hat), "w_agg": w_agg, "n_modes": n_modes,
+        "aggregate": agg_y,
+    }
+    return feats, aux

falsesync/inference.py

@@ -0,0 +1,90 @@
+"""Uncertainty: unit-level bootstrap and error-law deconvolution.
+
+See resampling_plan.md and deconvolution_assessment.md. No asymptotic
+coverage claims — intervals are bootstrap intervals.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def unit_break_uncertainties(
+    t: np.ndarray,
+    values: np.ndarray,
+    taus: np.ndarray,
+    rng: np.random.Generator | None = None,
+    n_boot: int = 60,
+) -> np.ndarray:
+    """Residual bootstrap around each unit's two-segment fit -> s_i (the F_e law)."""
+    from .changepoints import ls_breakpoint
+
+    rng = rng or np.random.default_rng(0)
+    t = np.asarray(t, dtype=float)
+    n = values.shape[0]
+    out = np.full(n, np.nan)
+    for i in range(n):
+        y = np.asarray(values[i], dtype=float)
+        ok = np.isfinite(y)
+        if ok.sum() < 8:
+            continue
+        ti, yi = t[ok], y[ok]
+        res = ls_breakpoint(ti, yi)
+        c = res.index
+        fit = np.concatenate([np.full(c, yi[:c].mean()), np.full(yi.size - c, yi[c:].mean())])
+        resid = yi - fit
+        resid = resid - resid.mean()
+        boots = []
+        for _ in range(n_boot):
+            yb = fit + rng.choice(resid, resid.size, replace=True)
+            try:
+                boots.append(ls_breakpoint(ti, yb).location)
+            except Exception:
+                continue
+        if len(boots) >= 10:
+            out[i] = float(np.std(boots))
+    return out
+
+
+def deconvolve_timing(
+    taus: np.ndarray, s_i: np.ndarray, grid: np.ndarray, damp: float = 0.02
+) -> np.ndarray:
+    """Recover f_tau from noisy unit break estimates (P8c machinery).
+
+    Gaussian errors with per-unit sd s_i: f_hat_tau = f_tau * phi_bar, where
+    phi_bar is the average error kernel. Damped Fourier deconvolution with
+    fixed Tikhonov factor `damp` (a priori, not tuned).
+    """
+    taus = np.asarray(taus, dtype=float)
+    s_i = np.asarray(s_i, dtype=float)
+    grid = np.asarray(grid, dtype=float)
+    dt = grid[1] - grid[0]
+    n_g = grid.size
+    hist, _ = np.histogram(taus, bins=np.append(grid, grid[-1] + dt), density=True)
+    sbar2 = float(np.nanmean(s_i**2)) if np.isfinite(s_i).any() else 0.0
+    F_hist = np.fft.rfft(hist)
+    freqs = np.fft.rfftfreq(n_g, dt)
+    phi_bar = np.exp(-2 * np.pi**2 * freqs**2 * sbar2)  # E over units of error char. fn (approx: mean sd)
+    denom = phi_bar**2 + damp
+    F_f = F_hist * phi_bar / denom
+    f = np.fft.irfft(F_f, n_g).real
+    return np.clip(f, 0, None)
+
+
+def cluster_bootstrap_band(
+    taus: np.ndarray,
+    grid: np.ndarray,
+    bw: float,
+    rng: np.random.Generator | None = None,
+    n_boot: int = 200,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Cluster (unit-resample) bootstrap band for the KDE of f_tau."""
+    rng = rng or np.random.default_rng(0)
+    taus = np.asarray(taus, dtype=float)
+    dens = []
+    for _ in range(n_boot):
+        s = rng.choice(taus, taus.size, replace=True)
+        u = (grid[:, None] - s[None, :]) / bw
+        dens.append(np.exp(-0.5 * u * u).mean(axis=1) / (bw * np.sqrt(2 * np.pi)))
+    d = np.stack(dens)
+    return np.quantile(d, 0.025, axis=0), np.quantile(d, 0.975, axis=0)

falsesync/model.py

@@ -0,0 +1,221 @@
+"""FalseSynchronyModel — Phase-3 research-grade workflow (spec §17).
+
+    from falsesync.model import FalseSynchronyModel
+    model = FalseSynchronyModel(classifier=clf)
+    result = model.fit(t, values, weights=None, windows=None)
+
+`classifier` is a fitted falsesync.regimes.RegimeClassifier. If None, the
+result carries uncalibrated heuristic probabilities and a warning.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from .aggregation import weighted_aggregate
+from .alignment import event_time_realign, realign_summary
+from .changepoints import aggregate_breakpoint, ls_breakpoint, multibreak
+from .diagnostics import FalseSynchronyResult
+from .features import extract_features
+from .regimes import REGIME_CLASSES, RegimeClassifier
+
+
+@dataclass
+class FalseSynchronyModel:
+    classifier: RegimeClassifier | None = None
+    method: str = "ls"
+    operators: tuple = ("ls", "maxslope", "cusum", "binseg")
+    unit_boot: int = 30
+    agg_strength_crit: float = 3.0  # aggregate break evidence criterion
+    sync_prob_crit: float = 0.5     # calibrated synchrony-support threshold
+    seed: int = 0
+
+    def _trend_check(self, t, values, taus, rng) -> tuple[bool, dict]:
+        """Detrend each unit (per-unit linear fit), refit breaks, compare.
+
+        If unit break dispersion collapses after detrending, part of the
+        apparent timing heterogeneity is trend confounding (§7).
+        """
+        t = np.asarray(t, dtype=float)
+        tc = (t - t.mean()) / max(np.ptp(t), 1e-9)
+        taus_d = []
+        valid = np.isfinite(taus)
+        for i in np.where(valid)[0]:
+            y = values[i]
+            ok = np.isfinite(y)
+            if ok.sum() < 8:
+                continue
+            slope = np.polyfit(tc[ok], y[ok], 1)[0]
+            yd = y.copy()
+            yd[ok] = y[ok] - slope * tc[ok]
+            fill = np.nanmean(yd[ok])
+            try:
+                r = ls_breakpoint(t, np.nan_to_num(yd, nan=fill))
+                taus_d.append((i, r.location, abs(slope)))
+            except Exception:
+                pass
+        if not taus_d:
+            return False, {"detrended_sd": np.nan, "raw_sd": np.nan}
+        idx = np.array([d[0] for d in taus_d], dtype=int)
+        raw_sd = float(np.std(taus[idx]))
+        det_sd = float(np.std([d[1] for d in taus_d]))
+        med_slope = float(np.median([d[2] for d in taus_d]))
+        info = {"raw_sd": raw_sd, "detrended_sd": det_sd, "median_slope": med_slope}
+        # confounding: detrending materially shrinks spread AND slopes nonzero
+        confounded = (raw_sd > 0 and det_sd < 0.6 * raw_sd and med_slope > 0.05)
+        return confounded, info
+
+    def _operator_spread(self, t, agg_y) -> dict:
+        locs, strengths = {}, {}
+        for m in self.operators:
+            try:
+                r = aggregate_breakpoint(t, agg_y, m) if m != "binseg" else multibreak(t, agg_y, 3)
+                locs[m] = r.location
+                strengths[m] = r.strength
+            except Exception:
+                pass
+        vals = np.array(list(locs.values()))
+        spread = float(vals.max() - vals.min()) if vals.size > 1 else 0.0
+        consensus = float(1 - spread / max(np.ptp(t), 1e-9))
+        return {
+            "locations": locs, "strengths": strengths,
+            "operator_spread": spread,
+            "operator_consensus": consensus,
+            "operator_sensitivity_warning": spread > 0.2 * np.ptp(t),
+        }
+
+    def fit(
+        self,
+        t: np.ndarray,
+        values: np.ndarray,
+        weights: np.ndarray | None = None,
+        windows: np.ndarray | None = None,
+        aggregate: np.ndarray | None = None,
+    ) -> FalseSynchronyResult:
+        rng = np.random.default_rng(self.seed)
+        t = np.asarray(t, dtype=float)
+        if aggregate is None:
+            aggregate = weighted_aggregate(values, weights)
+        if self.method == "ls":
+            prim_loc, prim_str = None, None  # taken from aux below
+        else:
+            prim = (multibreak(t, aggregate, 3) if self.method == "binseg"
+                    else aggregate_breakpoint(t, aggregate, self.method))
+            prim_loc, prim_str = prim.location, prim.strength
+        feats, aux = extract_features(
+            t, values, weights=weights, windows=windows,
+            aggregate=aggregate, rng=rng, unit_boot=self.unit_boot,
+        )
+        warnings: list[str] = []
+        if self.classifier is not None and self.classifier.fitted:
+            probs = self.classifier.predict_proba(feats)[0]
+            regime_probabilities = dict(zip(REGIME_CLASSES, probs.tolist()))
+        else:
+            regime_probabilities = {k: np.nan for k in REGIME_CLASSES}
+            warnings.append("no calibrated classifier: regime probabilities not estimable")
+
+        ops = self._operator_spread(t, aux["aggregate"])
+        if ops["operator_sensitivity_warning"]:
+            warnings.append(
+                f"operator spread {ops['operator_spread']:.3g} > 20% of window: "
+                "no method-independent aggregate breakpoint (P7); report operator."
+            )
+
+        indeterminate = False
+        obs_corr = np.nan
+        if windows is not None:
+            win = np.asarray(windows, dtype=float)
+            if win.ndim == 2 and win.shape[1] == 2:
+                lo = np.nanmin(t)
+                trunc_start = win[:, 0] > lo + 0.02 * max(np.ptp(t), 1e-9)
+                m_ok = trunc_start & np.isfinite(aux["taus"])
+                if m_ok.sum() >= 5 and trunc_start.mean() >= 0.2:
+                    obs_corr = float(np.corrcoef(win[m_ok, 0], aux["taus"][m_ok])[0, 1])
+                    indeterminate = obs_corr >= 0.5
+
+        confounded, trend_info = self._trend_check(t, values, aux["taus"], rng)
+        if confounded:
+            warnings.append(
+                "trend confounding: detrending shrinks tau_hat dispersion "
+                f"{trend_info['raw_sd']:.3g} -> {trend_info['detrended_sd']:.3g}; "
+                "timing inference unreliable until trends modeled (§7)."
+            )
+
+        if indeterminate:
+            warnings.append(
+                "indeterminate observation process: unit entry times correlate "
+                f"with estimated break times (corr={obs_corr:.2f}); synchrony "
+                "inference is not supported under a possibly informative "
+                "observation process (Phase-4 F1 mitigation)."
+            )
+        agg_loc = aux["agg_break"] if self.method == "ls" else prim_loc
+        agg_str = aux["agg_strength"] if self.method == "ls" else prim_str
+        p_sync = regime_probabilities.get("SYNCHRONOUS", np.nan)
+        warn_fire = (
+            np.isfinite(agg_str) and agg_str >= self.agg_strength_crit
+            and np.isfinite(p_sync) and p_sync < self.sync_prob_crit
+        ) or indeterminate
+        if warn_fire:
+            warnings.append(
+                f"false common-event warning: aggregate strength "
+                f"{agg_str:.2f} >= {self.agg_strength_crit} but "
+                f"P(SYNCHRONOUS)={p_sync:.2f} < {self.sync_prob_crit}."
+            )
+
+        tv = aux["taus"][np.isfinite(aux["taus"])]
+        grid, f_dec = aux["f_dec"]
+        comps = {
+            "s1_width": feats[4], "s2_shape": feats[2],
+            "s3_alignment": feats[5], "s4_weighting": feats[6],
+            "noise_floor_ratio": feats[7],
+        }
+        scale = np.ptp(t)
+        a_sync = p_sync if np.isfinite(p_sync) else np.nan
+        synchrony_score = a_sync  # calibrated: score IS P(synchrony)
+
+        aligned = event_time_realign(t, values[np.isfinite(aux["taus"])], tv) if tv.size else np.zeros((0, 0))
+        s_dt = float(np.median(np.diff(t)))
+        ea = realign_summary(aligned, s_dt) if aligned.size else {}
+        ea.update({"n_aligned": int(tv.size),
+                   "tau_spread_before": float(np.std(tv)) if tv.size else np.nan,
+                   "est_noise_floor": aux["noise_floor"],
+                   "trend_confounded": confounded, "trend_info": trend_info})
+
+        timing_dispersion = {
+            "sd": float(np.std(tv)) if tv.size else np.nan,
+            "iqr": float(np.quantile(tv, 0.75) - np.quantile(tv, 0.25)) if tv.size else np.nan,
+            "sd_corrected": float(np.sqrt(max(np.std(tv) ** 2 - (aux["noise_floor"] or 0) ** 2, 0)))
+            if tv.size else np.nan,
+        }
+        cluster_structure = {
+            "n_components": aux["n_modes"],
+            "centers": [float(grid[i]) for i in np.argsort(f_dec)[::-1][: aux["n_modes"]]] if f_dec.size else [],
+        }
+        return FalseSynchronyResult(
+            aggregate_breakpoint=agg_loc,
+            aggregate_break_strength=agg_str,
+            unit_breakpoints=aux["taus"],
+            unit_breakpoint_uncertainty=aux["s_i"],
+            timing_distribution={"grid": grid, "density": aux["f_hat"][1], "deconvolved": f_dec},
+            timing_dispersion=timing_dispersion,
+            cluster_structure=cluster_structure,
+            synchrony_score=synchrony_score,
+            regime_probabilities=regime_probabilities,
+            event_aligned_summary=ea,
+            interpretation_warning=warnings,
+            method=self.method,
+            components={
+                **comps,
+                "operator_locations": ops["locations"],
+                "operator_strengths": ops["strengths"],
+                "operator_spread": ops["operator_spread"],
+                "operator_consensus": ops["operator_consensus"],
+                "operator_sensitivity_warning": ops["operator_sensitivity_warning"],
+                "trend_confounding_warning": confounded,
+                "indeterminate_observation_process": indeterminate,
+                "obs_process_corr": obs_corr,
+                "false_common_event_warning": warn_fire,
+            },
+        )

falsesync/plotting.py

@@ -0,0 +1,59 @@
+"""Figures for Phase 2 outputs (matplotlib, non-interactive)."""
+
+from __future__ import annotations
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def aggregate_curve_figure(t, m_pop, agg_obs, path, title="Aggregate curve"):
+    fig, ax = plt.subplots(figsize=(7, 4))
+    ax.plot(t, m_pop, label="population m(t)", lw=2)
+    if agg_obs is not None:
+        ax.plot(t, agg_obs, label="observed aggregate", alpha=0.7)
+    ax.set_xlabel("t")
+    ax.set_title(title)
+    ax.legend()
+    fig.tight_layout()
+    fig.savefig(path, dpi=140)
+    plt.close(fig)
+
+
+def timing_distribution_figure(grid, densities: dict, path, title="Timing distribution"):
+    fig, ax = plt.subplots(figsize=(7, 4))
+    for label, d in densities.items():
+        ax.plot(grid, d, label=label)
+    ax.set_xlabel("tau")
+    ax.set_title(title)
+    ax.legend()
+    fig.tight_layout()
+    fig.savefig(path, dpi=140)
+    plt.close(fig)
+
+
+def operator_comparison_figure(t, m, results: dict, path):
+    fig, ax = plt.subplots(figsize=(7, 4))
+    ax.plot(t, m, lw=1.5, label="m(t)")
+    for name, loc in results.items():
+        ax.axvline(loc, ls="--", label=f"{name}: {loc:.2f}")
+    ax.set_xlabel("t")
+    ax.legend(fontsize=8)
+    fig.tight_layout()
+    fig.savefig(path, dpi=140)
+    plt.close(fig)
+
+
+def aligned_panel_figure(s_grid, aligned, path, title="Event-time aligned panel"):
+    fig, ax = plt.subplots(figsize=(7, 4))
+    for row in aligned:
+        ax.plot(s_grid, row, color="0.7", lw=0.5)
+    ax.plot(s_grid, np.nanmean(aligned, axis=0), color="C0", lw=2, label="aligned mean")
+    ax.set_xlabel("event time s = t - tau_i")
+    ax.set_title(title)
+    ax.legend()
+    fig.tight_layout()
+    fig.savefig(path, dpi=140)
+    plt.close(fig)

falsesync/regimes.py

@@ -0,0 +1,124 @@
+"""Calibrated regime classifier (Phase 3).
+
+Multinomial logistic regression on simulation-grid features, probability-
+calibrated by temperature scaling fitted on the calibration grid itself.
+The locked evaluation grid is never touched during fitting.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+REGIME_CLASSES = [
+    "SYNCHRONOUS", "NEAR_SYNCHRONOUS", "DIFFUSE_ASYNCHRONOUS",
+    "CLUSTERED", "NO_TRANSITION",
+]
+
+
+def regime_label(truth: str, taus: np.ndarray, cluster_gap: float = 1.5) -> str:
+    """Map engine truth + sampled taus to classifier classes."""
+    if truth == "no_transition":
+        return "NO_TRANSITION"
+    if truth == "synchronous":
+        return "SYNCHRONOUS"
+    if truth == "near_synchronous":
+        return "NEAR_SYNCHRONOUS"
+    if truth in ("two_cluster", "three_cluster"):
+        return "CLUSTERED"
+    if truth == "diffuse":
+        return "DIFFUSE_ASYNCHRONOUS"
+    return truth
+
+
+@dataclass
+class RegimeClassifier:
+    """Logistic + temperature-calibrated probabilities over REGIME_CLASSES."""
+
+    C: float = 1.0
+    max_iter: int = 2000
+    model: object | None = None
+    scaler_mean: np.ndarray | None = None
+    scaler_sd: np.ndarray | None = None
+    temperature: float = 1.0
+    fitted: bool = False
+    meta: dict = field(default_factory=dict)
+
+    def _fit_logreg(self, X, y):
+        from sklearn.linear_model import LogisticRegression
+        from sklearn.preprocessing import LabelEncoder
+
+        self._le = LabelEncoder().fit(y)
+        yv = self._le.transform(y)
+        self.scaler_mean = X.mean(axis=0)
+        self.scaler_sd = X.std(axis=0)
+        self.scaler_sd[self.scaler_sd == 0] = 1.0
+        Xs = (X - self.scaler_mean) / self.scaler_sd
+        self.model = LogisticRegression(
+            C=self.C, max_iter=self.max_iter, multi_class="multinomial",
+            class_weight="balanced",
+        )
+        self.model.fit(Xs, yv)
+        self.fitted = True
+        self._classes_in_model = list(self._le.classes_)
+
+    def fit(self, X: np.ndarray, y: np.ndarray) -> "RegimeClassifier":
+        self._fit_logreg(np.asarray(X), np.asarray(y))
+        # temperature scaling on the same (held-out-to-temp data should be
+        # supplied separately via calibrate() when available)
+        return self
+
+    def calibrate(self, X: np.ndarray, y: np.ndarray) -> float:
+        """Fit temperature on a held-out calibration split; returns T."""
+        logits = self._logits(X)
+        best_t, best_loss = 1.0, np.inf
+        for T_c in np.geomspace(0.2, 5.0, 60):
+            p = self._softmax(logits / T_c)
+            yi = self._le.transform(np.asarray(y))
+            loss = -np.log(np.clip(p[np.arange(len(yi)), yi], 1e-12, 1)).mean()
+            if loss < best_loss:
+                best_loss, best_t = loss, T_c
+        self.temperature = float(best_t)
+        return self.temperature
+
+    def _logits(self, X: np.ndarray) -> np.ndarray:
+        Xs = (np.asarray(X) - self.scaler_mean) / self.scaler_sd
+        return self.model.decision_function(Xs)
+
+    @staticmethod
+    def _softmax(z: np.ndarray) -> np.ndarray:
+        z = z - z.max(axis=1, keepdims=True)
+        e = np.exp(z)
+        return e / e.sum(axis=1, keepdims=True)
+
+    def predict_proba(self, X: np.ndarray) -> np.ndarray:
+        """(n, len(REGIME_CLASSES)) probability matrix, columns in class order."""
+        p = self._softmax(self._logits(np.atleast_2d(X)) / self.temperature)
+        out = np.zeros((p.shape[0], len(REGIME_CLASSES)))
+        for j, cls in enumerate(self._classes_in_model):
+            if cls in REGIME_CLASSES:
+                out[:, REGIME_CLASSES.index(cls)] = p[:, j]
+        return out
+
+    def predict(self, X: np.ndarray) -> np.ndarray:
+        p = self.predict_proba(X)
+        return np.array(REGIME_CLASSES)[p.argmax(axis=1)]
+
+
+def multiclass_brier(P: np.ndarray, y_idx: np.ndarray, n_classes: int) -> float:
+    Y = np.zeros((len(y_idx), n_classes))
+    Y[np.arange(len(y_idx)), y_idx] = 1.0
+    return float(np.mean(np.sum((P - Y) ** 2, axis=1)))
+
+
+def expected_calibration_error(P: np.ndarray, y_idx: np.ndarray, n_bins: int = 10) -> float:
+    conf = P.max(axis=1)
+    pred = P.argmax(axis=1)
+    ece = 0.0
+    for b in range(n_bins):
+        lo, hi = b / n_bins, (b + 1) / n_bins
+        m = (conf > lo) & (conf <= hi)
+        if m.any():
+            ece += m.mean() * abs((pred[m] == y_idx[m]).mean() - conf[m].mean())
+    return float(ece)