sctm2-r5-code-20260611 / package_code_imports /v5_scripts /evaluate_v5_matched_baselines.py
agoniii97's picture
Normalize datetime precision for HF P1 tensor build
ae6d94c verified
Raw
History Blame Contribute Delete
50 kB
from __future__ import annotations
import argparse
import json
import math
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any
import numpy as np
import pandas as pd
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.isotonic import IsotonicRegression
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import average_precision_score, roc_auc_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parents[1]
V4P4_SCRIPT_DIR = ROOT_DIR / "v4p4_world_model" / "scripts"
V3P5_SCRIPT_DIR = ROOT_DIR / "v3p5_static" / "scripts"
for path in (SCRIPT_DIR, V4P4_SCRIPT_DIR, V3P5_SCRIPT_DIR):
if str(path) not in sys.path:
sys.path.insert(0, str(path))
from action_ontology_v5 import build_action_ontology, load_json # noqa: E402
from build_v5_action_tensors import ( # noqa: E402
build_action_arrays,
build_target_trial_labels,
medication_flags_from_stage0,
resolve_split_path,
validate_medication_flags,
)
from run_service_process_target_trial_v5 import pre_action_features # noqa: E402
from scan_target_trial_support_v5 import era_tokens_np, future_event_within, load_or_build_action_arrays, unique_landmark_filter # noqa: E402
from target_trial_estimators_v5 import effective_sample_size # noqa: E402
try:
from lightgbm import LGBMClassifier
except Exception: # pragma: no cover
LGBMClassifier = None
CAUSE_NAMES = {0: "next_contact", 1: "death", 2: "disengagement"}
PWE_BIN_EDGES_DAYS = (7.0, 14.0, 30.0, 60.0, 90.0, 365.0, math.inf)
RISK_ENDPOINTS = {
"death": {"kind": "pwe", "cause_id": 1},
"disengagement": {"kind": "pwe", "cause_id": 2},
"primary_referral_relapse": {"kind": "event", "event_index": 0},
"service_escalation": {"kind": "event", "event_index": 3},
"clinical_deterioration": {"kind": "event", "event_index": 4},
"high_acuity_state": {"kind": "event", "event_index": 5},
}
RISK_THRESHOLD_GRID = {
"death": [0.005, 0.01, 0.02, 0.03, 0.05],
"primary_referral_relapse": [0.01, 0.03, 0.05, 0.10, 0.15],
"disengagement": [0.01, 0.02, 0.05, 0.08, 0.10],
"service_escalation": [0.03, 0.05, 0.10, 0.15, 0.20],
"clinical_deterioration": [0.30, 0.40, 0.50, 0.55, 0.60],
"high_acuity_state": [0.10, 0.20, 0.30, 0.40, 0.50],
"default": [0.01, 0.03, 0.05, 0.10, 0.15],
}
def load_npz_to_memory(path: Path) -> dict[str, np.ndarray]:
with np.load(path, allow_pickle=True) as z:
return {key: z[key] for key in z.files}
def json_ready(value: Any) -> Any:
if isinstance(value, Path):
return str(value)
if isinstance(value, dict):
return {str(k): json_ready(v) for k, v in value.items()}
if isinstance(value, (list, tuple)):
return [json_ready(v) for v in value]
if isinstance(value, np.generic):
return value.item()
if isinstance(value, float) and not math.isfinite(value):
return None
return value
def service_era_key(service: np.ndarray, years: np.ndarray) -> np.ndarray:
return service.astype(np.int64) * 10 + era_tokens_np(years).astype(np.int64)
def clip_prob(p: np.ndarray, eps: float = 1.0e-7) -> np.ndarray:
return np.clip(np.asarray(p, dtype=np.float64).reshape(-1), eps, 1.0 - eps)
def weighted_mean(x: np.ndarray, sample_weight: np.ndarray | None = None) -> float:
x = np.asarray(x, dtype=np.float64).reshape(-1)
if not x.size:
return math.nan
if sample_weight is None:
return float(np.mean(x))
w = np.asarray(sample_weight, dtype=np.float64).reshape(-1)
ok = np.isfinite(x) & np.isfinite(w) & (w > 0)
if not ok.any():
return math.nan
return float(np.sum(w[ok] * x[ok]) / np.sum(w[ok]))
def calibration_intercept_slope(y: np.ndarray, p: np.ndarray, sample_weight: np.ndarray | None = None) -> tuple[float, float]:
y = np.asarray(y, dtype=np.int8).reshape(-1)
p = clip_prob(p)
if y.size == 0 or np.unique(y).size < 2:
return math.nan, math.nan
try:
x = np.log(p / (1.0 - p)).reshape(-1, 1)
clf = LogisticRegression(C=1.0e6, solver="lbfgs", max_iter=500)
clf.fit(x, y, sample_weight=sample_weight)
return float(clf.intercept_[0]), float(clf.coef_[0, 0])
except Exception:
return math.nan, math.nan
def integrated_calibration_index(y: np.ndarray, p: np.ndarray, sample_weight: np.ndarray | None = None) -> float:
y = np.asarray(y, dtype=np.int8).reshape(-1)
p = clip_prob(p)
w = None if sample_weight is None else np.asarray(sample_weight, dtype=np.float64).reshape(-1)
if w is not None:
ok = np.isfinite(w) & (w > 0)
y = y[ok]
p = p[ok]
w = w[ok]
if y.size == 0 or np.unique(y).size < 2:
return math.nan
try:
model = IsotonicRegression(y_min=0.0, y_max=1.0, out_of_bounds="clip")
calibrated = model.fit_transform(p, y, sample_weight=w)
return weighted_mean(np.abs(calibrated - p), w)
except Exception:
return math.nan
def weighted_binary_metrics(
y: np.ndarray,
p: np.ndarray,
sample_weight: np.ndarray | None = None,
include_calibration_model: bool = True,
) -> dict[str, float]:
y = np.asarray(y, dtype=np.int8).reshape(-1)
p = clip_prob(p)
w = None if sample_weight is None else np.asarray(sample_weight, dtype=np.float64).reshape(-1)
if w is not None:
ok = np.isfinite(w) & (w > 0)
y = y[ok]
p = p[ok]
w = w[ok]
if y.size == 0:
return {
"auc": math.nan,
"average_precision": math.nan,
"brier": math.nan,
"ece": math.nan,
"ici": math.nan,
"mean_predicted": math.nan,
"observed_rate": math.nan,
"calibration_intercept": math.nan,
"calibration_slope": math.nan,
}
auc = float(roc_auc_score(y, p, sample_weight=w)) if np.unique(y).size >= 2 else math.nan
ap = float(average_precision_score(y, p, sample_weight=w)) if np.unique(y).size >= 2 else math.nan
order = np.argsort(p)
ece = 0.0
ece_den = float(y.size) if w is None else float(w.sum())
for idx in np.array_split(order, min(10, max(1, y.size))):
if idx.size:
bin_w = None if w is None else w[idx]
bin_weight = float(idx.size) if w is None else float(bin_w.sum())
ece += abs(weighted_mean(p[idx], bin_w) - weighted_mean(y[idx], bin_w)) * bin_weight / ece_den
intercept, slope = calibration_intercept_slope(y, p, w) if include_calibration_model else (math.nan, math.nan)
return {
"auc": auc,
"average_precision": ap,
"brier": weighted_mean((p - y) ** 2, w),
"ece": float(ece),
"ici": integrated_calibration_index(y, p, w) if include_calibration_model else math.nan,
"mean_predicted": weighted_mean(p, w),
"observed_rate": weighted_mean(y, w),
"calibration_intercept": intercept,
"calibration_slope": slope,
}
def net_benefit(y: np.ndarray, p: np.ndarray, threshold: float, sample_weight: np.ndarray | None = None) -> float:
y = np.asarray(y, dtype=np.int8).reshape(-1)
p = clip_prob(p)
if y.size == 0 or threshold <= 0.0 or threshold >= 1.0:
return math.nan
w = np.ones(y.size, dtype=np.float64) if sample_weight is None else np.asarray(sample_weight, dtype=np.float64).reshape(-1)
ok = np.isfinite(w) & (w > 0)
y = y[ok]
p = p[ok]
w = w[ok]
if y.size == 0 or float(w.sum()) <= 0.0:
return math.nan
pred_pos = p >= float(threshold)
den = float(w.sum())
tp = float(np.sum(w * pred_pos * (y == 1)) / den)
fp = float(np.sum(w * pred_pos * (y == 0)) / den)
return float(tp - fp * threshold / (1.0 - threshold))
def calibration_curve_rows(
*,
split: str,
endpoint: str,
horizon_days: float,
model: str,
y: np.ndarray,
p: np.ndarray,
n_bins: int = 10,
) -> list[dict[str, Any]]:
y = np.asarray(y, dtype=np.int8).reshape(-1)
p = clip_prob(p)
if y.size == 0:
return []
rows = []
for bin_idx, idx in enumerate(np.array_split(np.argsort(p), min(n_bins, max(1, y.size))), start=1):
if idx.size == 0:
continue
rows.append(
{
"split": split,
"endpoint": endpoint,
"horizon_days": float(horizon_days),
"model": model,
"metric_family": "calibration_curve",
"calibration_bin": int(bin_idx),
"n": int(idx.size),
"events": int(y[idx].sum()),
"mean_predicted": float(p[idx].mean()),
"observed_rate": float(y[idx].mean()),
"predicted_min": float(p[idx].min()),
"predicted_max": float(p[idx].max()),
}
)
return rows
def bootstrap_metric_ci(
y: np.ndarray,
p: np.ndarray,
patient_ids: np.ndarray,
n_bootstrap: int,
seed: int,
) -> dict[str, float]:
if n_bootstrap <= 0 or y.size == 0:
return {}
rng = np.random.default_rng(seed)
clusters, inv = np.unique(patient_ids.astype(str), return_inverse=True)
tracked: dict[str, list[float]] = defaultdict(list)
for _ in range(n_bootstrap):
counts = rng.multinomial(clusters.size, np.full(clusters.size, 1.0 / clusters.size))
w = counts[inv].astype(np.float64)
metrics = weighted_binary_metrics(y, p, sample_weight=w, include_calibration_model=False)
for key, value in metrics.items():
if math.isfinite(float(value)):
tracked[key].append(float(value))
out: dict[str, float] = {"bootstrap_n": int(n_bootstrap), "bootstrap_clusters": int(clusters.size)}
for key, vals in tracked.items():
arr = np.asarray(vals, dtype=np.float64)
out[f"{key}_ci_low"] = float(np.nanquantile(arr, 0.025)) if arr.size else math.nan
out[f"{key}_ci_high"] = float(np.nanquantile(arr, 0.975)) if arr.size else math.nan
return out
def bootstrap_net_benefit_ci(
y: np.ndarray,
p: np.ndarray,
patient_ids: np.ndarray,
threshold: float,
n_bootstrap: int,
seed: int,
) -> dict[str, float]:
if n_bootstrap <= 0 or y.size == 0:
return {}
rng = np.random.default_rng(seed)
clusters, inv = np.unique(patient_ids.astype(str), return_inverse=True)
vals = []
for _ in range(n_bootstrap):
counts = rng.multinomial(clusters.size, np.full(clusters.size, 1.0 / clusters.size))
vals.append(net_benefit(y, p, threshold, sample_weight=counts[inv].astype(np.float64)))
arr = np.asarray(vals, dtype=np.float64)
return {
"net_benefit_ci_low": float(np.nanquantile(arr, 0.025)) if arr.size else math.nan,
"net_benefit_ci_high": float(np.nanquantile(arr, 0.975)) if arr.size else math.nan,
}
def active_next_mask(arrays: dict[str, np.ndarray]) -> np.ndarray:
return arrays["valid_mask"][:, :-1].astype(bool) & arrays["valid_mask"][:, 1:].astype(bool)
def pwe_targets(arrays: dict[str, np.ndarray]) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
terminal = arrays["terminal_label"][:, 1:].astype(np.int64)
service = arrays["service_state"][:, 1:].astype(np.int64)
valid = active_next_mask(arrays)
cause = np.zeros_like(terminal, dtype=np.int64)
cause[(terminal == 1) | (service == 7)] = 1
cause[(terminal == 2) | (terminal == 3) | (service == 6)] = 2
censored = (terminal == 4) | (service == 5)
days = np.expm1(arrays["delta_t_next_log"][:, :-1].astype(np.float64)).clip(min=1.0e-6)
return cause, censored.astype(bool), valid, days
def build_conditional_specs(meta: dict[str, Any], vocab: dict[str, Any]) -> tuple[dict[str, Any], ...]:
cat_cols = list(meta.get("cat_cols", []))
def field_idx(name: str) -> int:
return cat_cols.index(name) if name in cat_cols else -1
specs = []
tables = [
("是否转诊", ["是"], ["转诊类型", "转诊原因", "转诊至机构"]),
("治疗方式", ["住院"], ["本次入院形式", "本次住院名称", "本次入院日期", "末次出院日期"]),
]
for parent, active_values, children in tables:
parent_idx = field_idx(parent)
child_indices = [field_idx(child) for child in children if field_idx(child) >= 0]
active_ids = [int(vocab[f"{parent}={value}"]) for value in active_values if f"{parent}={value}" in vocab]
if parent_idx >= 0 and child_indices and active_ids:
specs.append({"parent_idx": parent_idx, "active_value_ids": active_ids, "child_indices": child_indices})
return tuple(specs)
def conditional_status_np(cat_value_ids: np.ndarray, missing_ids: np.ndarray, specs: tuple[dict[str, Any], ...]) -> np.ndarray:
status = np.zeros_like(cat_value_ids, dtype=np.int64)
if not specs:
return status
observed = missing_ids.astype(np.int64) == 0
unknown = 1
for spec in specs:
parent_idx = int(spec["parent_idx"])
child_indices = [int(x) for x in spec["child_indices"]]
active_ids = np.asarray([int(x) for x in spec["active_value_ids"]], dtype=cat_value_ids.dtype)
if parent_idx < 0 or parent_idx >= cat_value_ids.shape[-1] or not child_indices or active_ids.size == 0:
continue
parent_val = cat_value_ids[:, :, parent_idx]
parent_observed = observed[:, :, parent_idx] & (parent_val != unknown)
parent_active = np.isin(parent_val, active_ids) & parent_observed
parent_unknown = ~parent_observed
for child_idx in child_indices:
if child_idx < 0 or child_idx >= cat_value_ids.shape[-1]:
continue
child_val = cat_value_ids[:, :, child_idx]
child_observed = observed[:, :, child_idx] & (child_val != unknown)
code = np.zeros_like(parent_val, dtype=np.int64)
code[parent_unknown] = 2
code[(~parent_unknown) & (~parent_active) & (~child_observed)] = 1
code[(~parent_unknown) & (~parent_active) & child_observed] = 3
status[:, :, child_idx] = code
return status
def mapped_missing_targets(
arrays: dict[str, np.ndarray],
meta: dict[str, Any],
vocab: dict[str, Any],
*,
start: int,
stop: int | None,
) -> np.ndarray:
"""Match the v4/v5 loss missingness target: 5 classes incl. structural/no-clinical."""
raw = arrays["missing_ids"][:, start:stop, :].astype(np.int64)
mapped = np.full_like(raw, 3, dtype=np.int64)
mapped[raw == 0] = 0
mapped[raw == 1] = 1
mapped[raw == 2] = 3
specs = build_conditional_specs(meta, vocab)
if specs:
cat = arrays["cat_value_ids"][:, start:stop, :].astype(np.int64)
conditional = conditional_status_np(cat, raw, specs)
mapped[conditional == 1] = 2
service = arrays["service_state"][:, start:stop].astype(np.int64)
terminal = arrays["terminal_label"][:, start:stop].astype(np.int64)
no_clinical = (terminal > 0) | (service >= 5)
mapped = np.where(no_clinical[:, :, None], 4, mapped)
return mapped.clip(0, 4)
def exposure_by_bin(days: np.ndarray) -> np.ndarray:
upper = np.asarray(PWE_BIN_EDGES_DAYS[:-1], dtype=np.float64)
lower = np.asarray((0.0, 7.0, 14.0, 30.0, 60.0, 90.0), dtype=np.float64)
finite = (np.minimum(days[..., None], upper) - lower).clip(min=0.0)
tail = (days - 365.0).clip(min=0.0)[..., None]
return np.concatenate([finite, tail], axis=-1)
def pwe_bin_index(days: np.ndarray) -> np.ndarray:
upper = np.asarray(PWE_BIN_EDGES_DAYS[:-1], dtype=np.float64)
return np.searchsorted(upper, days, side="left").clip(max=len(PWE_BIN_EDGES_DAYS) - 1)
def fit_semi_markov(train: dict[str, np.ndarray], smoothing: float) -> dict[str, Any]:
cause, censored, valid, days = pwe_targets(train)
cur_state = train["service_state"][:, :-1].astype(np.int64)
next_state = train["service_state"][:, 1:].astype(np.int64)
key = service_era_key(cur_state, train["visit_year"][:, :-1])
n_states = int(max(train["service_state"].max(), 7) + 1)
n_bins = len(PWE_BIN_EDGES_DAYS)
trans_counts: dict[int, np.ndarray] = {}
exposure: dict[int, np.ndarray] = {}
events: dict[int, np.ndarray] = {}
for k in np.unique(key[valid]):
m = valid & (key == k)
counts = np.bincount(next_state[m].clip(0, n_states - 1), minlength=n_states).astype(np.float64) + smoothing
trans_counts[int(k)] = counts / counts.sum()
ex = exposure_by_bin(days[m]).sum(axis=0)
ev = np.zeros((3, n_bins), dtype=np.float64)
bin_idx = pwe_bin_index(days[m])
c = cause[m]
z = censored[m]
for cause_id in range(3):
np.add.at(ev[cause_id], bin_idx[(~z) & (c == cause_id)], 1.0)
exposure[int(k)] = ex + smoothing
events[int(k)] = ev + smoothing
global_counts = np.bincount(next_state[valid].clip(0, n_states - 1), minlength=n_states).astype(np.float64) + smoothing
global_trans = global_counts / global_counts.sum()
global_ex = exposure_by_bin(days[valid]).sum(axis=0) + smoothing
global_ev = np.zeros((3, n_bins), dtype=np.float64) + smoothing
bin_idx = pwe_bin_index(days[valid])
for cause_id in range(3):
np.add.at(global_ev[cause_id], bin_idx[(~censored[valid]) & (cause[valid] == cause_id)], 1.0)
return {
"trans": trans_counts,
"exposure": exposure,
"events": events,
"global_trans": global_trans,
"global_exposure": global_ex,
"global_events": global_ev,
"n_states": n_states,
"n_bins": n_bins,
"smoothing": smoothing,
}
def semi_markov_metrics(model: dict[str, Any], arrays: dict[str, np.ndarray], split: str) -> list[dict[str, Any]]:
cause, censored, valid, days = pwe_targets(arrays)
cur_state = arrays["service_state"][:, :-1].astype(np.int64)
next_state = arrays["service_state"][:, 1:].astype(np.int64).clip(0, model["n_states"] - 1)
key = service_era_key(cur_state, arrays["visit_year"][:, :-1])
rows = []
key_flat = key[valid].astype(np.int64)
y = next_state[valid]
pred = np.empty(y.shape, dtype=np.int64)
ce_sum = 0.0
for k in np.unique(key_flat):
idx = np.flatnonzero(key_flat == int(k))
probs_k = model["trans"].get(int(k), model["global_trans"])
pred[idx] = int(np.argmax(probs_k))
ce_sum += float(-np.log(np.clip(probs_k[y[idx]], 1.0e-12, 1.0)).sum())
ce = ce_sum / y.size if y.size else math.nan
conf = np.zeros((model["n_states"], model["n_states"]), dtype=np.int64)
if y.size:
np.add.at(conf, (y, pred), 1)
f1 = []
for cls in range(model["n_states"]):
if conf[cls].sum() <= 0:
continue
tp = conf[cls, cls]
fp = conf[:, cls].sum() - tp
fn = conf[cls].sum() - tp
f1.append(2 * tp / max(1, 2 * tp + fp + fn))
rows.append(
{
"split": split,
"model": "Semi-Markov Care-Process",
"metric_family": "service_transition",
"n": int(y.size),
"cross_entropy": float(ce),
"accuracy": float((pred == y).mean()) if y.size else math.nan,
"macro_f1": float(np.mean(f1)) if f1 else math.nan,
}
)
cause_flat = cause[valid].astype(np.int64)
cens_flat = censored[valid].astype(bool)
days_flat = days[valid].astype(np.float64)
nll_sum = 0.0
nll_n = int(days_flat.size)
bin_flat = pwe_bin_index(days_flat)
for k in np.unique(key_flat):
idx = np.flatnonzero(key_flat == int(k))
rates = model["events"].get(int(k), model["global_events"]) / model["exposure"].get(int(k), model["global_exposure"])[None, :]
total_rate = rates.sum(axis=0)
ex = exposure_by_bin(days_flat[idx])
ll = -np.sum(ex * total_rate[None, :], axis=1)
event_idx = idx[~cens_flat[idx]]
if event_idx.size:
rel = np.flatnonzero(~cens_flat[idx])
ll[rel] += np.log(np.clip(rates[cause_flat[event_idx], bin_flat[event_idx]], 1.0e-12, None))
nll_sum += float((-ll).sum())
rows.append(
{
"split": split,
"model": "Semi-Markov Care-Process",
"metric_family": "pwe_time_to_next_process",
"n": nll_n,
"pwe_nll": nll_sum / nll_n if nll_n else math.nan,
}
)
return rows
def fit_empirical_grammar(train: dict[str, np.ndarray], meta: dict[str, Any], vocab_map: dict[str, Any], smoothing: float) -> dict[str, Any]:
valid = active_next_mask(train)
next_contact = valid & (train["service_state"][:, 1:] < 5)
key = service_era_key(train["service_state"][:, :-1], train["visit_year"][:, :-1])
n_missing = 5
missing_target = mapped_missing_targets(train, meta, vocab_map, start=1, stop=None)
vocab = int(meta.get("cat_vocab_size", int(train["cat_value_ids"].max()) + 1))
cbe_dim = int(meta.get("cbe_dim", 5))
out: dict[str, Any] = {"missing": {}, "cat": {}, "numeric": {}, "ordinal": {}, "global": {}, "target_contract": "v4_v5_mapped_5class_missingness"}
for k in np.unique(key[valid]):
m = valid & (key == k)
miss = missing_target[m].reshape(-1).astype(int)
out["missing"][int(k)] = (np.bincount(miss.clip(0, n_missing - 1), minlength=n_missing) + smoothing)
out["missing"][int(k)] /= out["missing"][int(k)].sum()
miss_global = missing_target[valid].reshape(-1).astype(int)
out["global"]["missing"] = (np.bincount(miss_global.clip(0, n_missing - 1), minlength=n_missing) + smoothing)
out["global"]["missing"] /= out["global"]["missing"].sum()
for field_idx in range(train["cat_value_ids"].shape[2]):
out["cat"][field_idx] = {}
observed = next_contact & (missing_target[:, :, field_idx] == 0)
vals_global = train["cat_value_ids"][:, 1:, field_idx][observed].astype(int)
g = np.bincount(vals_global.clip(0, vocab - 1), minlength=vocab).astype(np.float64) + smoothing
out["global"][f"cat_{field_idx}"] = g / g.sum()
for k in np.unique(key[observed]):
vals = train["cat_value_ids"][:, 1:, field_idx][observed & (key == k)].astype(int)
counts = np.bincount(vals.clip(0, vocab - 1), minlength=vocab).astype(np.float64) + smoothing
out["cat"][field_idx][int(k)] = counts / counts.sum()
for field_idx in range(train["numeric_values"].shape[2]):
out["numeric"][field_idx] = {}
observed = next_contact & train["numeric_mask"][:, 1:, field_idx]
vals = train["numeric_values"][:, 1:, field_idx][observed].astype(float)
out["global"][f"num_{field_idx}"] = float(np.nanmean(vals)) if vals.size else 0.0
for k in np.unique(key[observed]):
sub = train["numeric_values"][:, 1:, field_idx][observed & (key == k)].astype(float)
out["numeric"][field_idx][int(k)] = float(np.nanmean(sub)) if sub.size else out["global"][f"num_{field_idx}"]
for field_idx in range(train["ordinal_cbe"].shape[2]):
out["ordinal"][field_idx] = {}
observed = next_contact & train["ordinal_mask"][:, 1:, field_idx]
levels = train["ordinal_cbe"][:, 1:, field_idx, :][observed].sum(axis=1).astype(int)
counts = np.bincount(levels.clip(0, cbe_dim), minlength=cbe_dim + 1).astype(np.float64) + smoothing
out["global"][f"ord_{field_idx}"] = int(np.argmax(counts))
for k in np.unique(key[observed]):
sub = train["ordinal_cbe"][:, 1:, field_idx, :][observed & (key == k)].sum(axis=1).astype(int)
c = np.bincount(sub.clip(0, cbe_dim), minlength=cbe_dim + 1).astype(np.float64) + smoothing
out["ordinal"][field_idx][int(k)] = int(np.argmax(c))
return out
def empirical_grammar_metrics(model: dict[str, Any], arrays: dict[str, np.ndarray], split: str, meta: dict[str, Any], vocab: dict[str, Any]) -> list[dict[str, Any]]:
valid = active_next_mask(arrays)
next_contact = valid & (arrays["service_state"][:, 1:] < 5)
key = service_era_key(arrays["service_state"][:, :-1], arrays["visit_year"][:, :-1])
rows = []
missing_target = mapped_missing_targets(arrays, meta, vocab, start=1, stop=None)
miss_y = missing_target[valid].reshape(-1).astype(int)
miss_keys = np.repeat(key[valid], arrays["missing_ids"].shape[2])
miss_correct = 0
miss_brier_sum = 0.0
for k in np.unique(miss_keys):
idx = np.flatnonzero(miss_keys == int(k))
probs = model["missing"].get(int(k), model["global"]["missing"])
yk = miss_y[idx].clip(0, probs.shape[0] - 1)
miss_correct += int((int(np.argmax(probs)) == yk).sum())
miss_brier_sum += float((np.sum(probs * probs) + 1.0) * yk.size - 2.0 * np.sum(probs[yk]))
rows.append(
{
"split": split,
"model": "Empirical state-year grammar",
"family": "missingness",
"n": int(miss_y.size),
"accuracy": float(miss_correct / miss_y.size) if miss_y.size else math.nan,
"multiclass_brier": float(miss_brier_sum / miss_y.size) if miss_y.size else math.nan,
"target_contract": model.get("target_contract", "v4_v5_mapped_5class_missingness"),
}
)
cat_ce = []
cat_top1 = []
cat_top3 = []
cat_n = 0
for field_idx in range(arrays["cat_value_ids"].shape[2]):
observed = next_contact & (missing_target[:, :, field_idx] == 0)
y = arrays["cat_value_ids"][:, 1:, field_idx][observed].astype(int)
if y.size:
cat_n += int(y.size)
k_obs = key[observed].astype(np.int64)
ce_sum = 0.0
top1_sum = 0
top3_sum = 0
for k in np.unique(k_obs):
idx = np.flatnonzero(k_obs == int(k))
probs = model["cat"][field_idx].get(int(k), model["global"][f"cat_{field_idx}"])
yk = y[idx].clip(0, probs.shape[0] - 1)
ce_sum += float(-np.log(np.clip(probs[yk], 1.0e-12, 1.0)).sum())
top1 = int(np.argmax(probs))
top1_sum += int((yk == top1).sum())
k_top = min(3, probs.shape[0])
top3 = np.argpartition(-probs, kth=k_top - 1)[:k_top]
top3_sum += int(np.isin(yk, top3).sum())
cat_ce.append(ce_sum / y.size)
cat_top1.append(top1_sum / y.size)
cat_top3.append(top3_sum / y.size)
rows.append(
{
"split": split,
"model": "Empirical state-year grammar",
"family": "categorical",
"n": int(cat_n),
"mean_cross_entropy": float(np.nanmean(cat_ce)) if cat_ce else math.nan,
"mean_top1_accuracy": float(np.nanmean(cat_top1)) if cat_top1 else math.nan,
"mean_top3_accuracy": float(np.nanmean(cat_top3)) if cat_top3 else math.nan,
}
)
abs_err = []
sq_err = []
for field_idx in range(arrays["numeric_values"].shape[2]):
observed = next_contact & arrays["numeric_mask"][:, 1:, field_idx]
y = arrays["numeric_values"][:, 1:, field_idx][observed].astype(float)
pred = np.asarray([model["numeric"][field_idx].get(int(k), model["global"][f"num_{field_idx}"]) for k in key[observed]], dtype=float) if y.size else np.asarray([])
abs_err.append(np.abs(pred - y))
sq_err.append((pred - y) ** 2)
ae = np.concatenate(abs_err) if abs_err else np.asarray([])
se = np.concatenate(sq_err) if sq_err else np.asarray([])
rows.append(
{
"split": split,
"model": "Empirical state-year grammar",
"family": "numeric",
"n": int(ae.size),
"mae": float(ae.mean()) if ae.size else math.nan,
"rmse": float(math.sqrt(float(se.mean()))) if se.size else math.nan,
}
)
ord_abs = []
for field_idx in range(arrays["ordinal_cbe"].shape[2]):
observed = next_contact & arrays["ordinal_mask"][:, 1:, field_idx]
y = arrays["ordinal_cbe"][:, 1:, field_idx, :][observed].sum(axis=1).astype(int)
pred = np.asarray([model["ordinal"][field_idx].get(int(k), model["global"][f"ord_{field_idx}"]) for k in key[observed]], dtype=int) if y.size else np.asarray([])
ord_abs.append(np.abs(pred - y))
oa = np.concatenate(ord_abs) if ord_abs else np.asarray([])
rows.append({"split": split, "model": "Empirical state-year grammar", "family": "ordinal", "n": int(oa.size), "mae": float(oa.mean()) if oa.size else math.nan})
return rows
def horizon_label(arrays: dict[str, np.ndarray], endpoint: str, horizon: float) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
spec = RISK_ENDPOINTS[endpoint]
valid = active_next_mask(arrays)
row_idx, pos_idx = np.nonzero(valid)
if endpoint in ("death", "disengagement"):
cause, censored, _, days = pwe_targets(arrays)
cause_id = int(spec["cause_id"])
y_full = ((~censored) & (cause == cause_id) & (days <= float(horizon))).astype(np.int8)
observed = valid & ~(censored & (days <= float(horizon)))
row_idx, pos_idx = np.nonzero(observed)
y = y_full[row_idx, pos_idx]
else:
y_full = future_event_within(arrays, int(spec["event_index"]), float(horizon)).astype(np.int8)
valid_full = arrays["valid_mask"].astype(bool)
times = arrays["time_since_start_days"].astype(float)
seq_len = valid_full.shape[1]
current_len = seq_len - 1
any_future = np.zeros((valid_full.shape[0], current_len), dtype=bool)
for offset in range(1, seq_len):
cur_len = min(current_len, seq_len - offset)
if cur_len <= 0:
break
dt = times[:, offset : offset + cur_len] - times[:, :cur_len]
any_future[:, :cur_len] |= valid_full[:, offset : offset + cur_len] & (dt > 0.0) & (dt <= float(horizon))
last_time = np.where(valid_full, times, -np.inf).max(axis=1, keepdims=True)
followup = last_time - times[:, :-1]
observed = valid & any_future & ((y_full[:, :-1].astype(bool)) | (followup >= float(horizon)))
row_idx, pos_idx = np.nonzero(observed)
y = y_full[:, :-1][row_idx, pos_idx]
if row_idx.size and "patient_ids" in arrays:
keep = unique_landmark_filter(arrays, row_idx, pos_idx)
row_idx = row_idx[keep]
pos_idx = pos_idx[keep]
y = y[keep]
return row_idx, pos_idx, y.astype(np.int8), arrays["patient_ids"][row_idx].astype(str) if "patient_ids" in arrays else row_idx.astype(str)
def fit_classifier(X: np.ndarray, y: np.ndarray, model_kind: str, seed: int, max_train: int, require_lightgbm: bool) -> tuple[Any, str]:
rng = np.random.default_rng(seed)
if max_train > 0 and X.shape[0] > max_train:
idx = rng.choice(X.shape[0], size=max_train, replace=False)
X_fit = X[idx]
y_fit = y[idx]
else:
X_fit, y_fit = X, y
if np.unique(y_fit).size < 2:
rate = float(y_fit.mean()) if y_fit.size else 0.0
return rate, "constant_rate"
if model_kind == "clinical_landmark":
model = make_pipeline(
StandardScaler(with_mean=True, with_std=True),
LogisticRegression(max_iter=1000, C=0.5, solver="lbfgs"),
)
model.fit(X_fit, y_fit)
return model, "Clinical Landmark logistic (L2)"
if LGBMClassifier is not None:
model = LGBMClassifier(
n_estimators=400,
learning_rate=0.03,
num_leaves=31,
subsample=0.85,
colsample_bytree=0.85,
reg_lambda=1.0,
random_state=seed,
n_jobs=-1,
verbose=-1,
)
model.fit(X_fit, y_fit)
return model, "LightGBM Landmark"
if require_lightgbm:
raise RuntimeError("LightGBM backend is required but lightgbm could not be imported.")
model = HistGradientBoostingClassifier(
learning_rate=0.05,
max_iter=160,
max_leaf_nodes=31,
min_samples_leaf=80,
l2_regularization=0.01,
early_stopping=True,
random_state=seed,
)
model.fit(X_fit, y_fit)
return model, "HistGradientBoosting fallback for LightGBM"
def predict_classifier(model: Any, X: np.ndarray) -> np.ndarray:
if isinstance(model, float):
return np.full(X.shape[0], model, dtype=np.float64)
return model.predict_proba(X)[:, 1].astype(np.float64)
def risk_baseline_rows(
train: dict[str, np.ndarray],
eval_arrays: dict[str, np.ndarray],
split: str,
endpoints: list[str],
horizons: list[float],
seed: int,
max_train: int,
bootstrap_clusters: int,
require_lightgbm: bool,
model_cache: dict[tuple[str, float, str], tuple[Any, str]] | None = None,
) -> list[dict[str, Any]]:
rows = []
for endpoint in endpoints:
for horizon in horizons:
print(json.dumps({"stage": "risk_horizon_labels", "split": split, "endpoint": endpoint, "horizon_days": float(horizon)}), flush=True)
tr_row, tr_pos, y_train, _ = horizon_label(train, endpoint, horizon)
ev_row, ev_pos, y_eval, patient_ids = horizon_label(eval_arrays, endpoint, horizon)
if y_train.size == 0 or y_eval.size == 0:
continue
X_train, _ = pre_action_features(train, tr_row, tr_pos)
X_eval, _ = pre_action_features(eval_arrays, ev_row, ev_pos)
for kind in ("clinical_landmark", "lightgbm_landmark"):
print(
json.dumps(
{
"stage": "risk_model_fit",
"split": split,
"endpoint": endpoint,
"horizon_days": float(horizon),
"model_kind": kind,
"train_n": int(y_train.size),
"eval_n": int(y_eval.size),
"max_train": int(max_train),
}
),
flush=True,
)
cache_key = (endpoint, float(horizon), kind)
if model_cache is not None and cache_key in model_cache:
model, model_label = model_cache[cache_key]
else:
model, model_label = fit_classifier(X_train, y_train, kind, seed + int(horizon) + len(endpoint), max_train, require_lightgbm)
if model_cache is not None:
model_cache[cache_key] = (model, model_label)
print(
json.dumps(
{
"stage": "risk_model_predict",
"split": split,
"endpoint": endpoint,
"horizon_days": float(horizon),
"model": model_label,
}
),
flush=True,
)
p = predict_classifier(model, X_eval)
metrics = weighted_binary_metrics(y_eval, p)
boot = bootstrap_metric_ci(
y_eval,
p,
patient_ids,
bootstrap_clusters,
seed + int(horizon) * 13 + len(endpoint) * 101 + (0 if kind == "clinical_landmark" else 1),
)
base = {
"split": split,
"endpoint": endpoint,
"horizon_days": float(horizon),
"model": model_label,
"n": int(y_eval.size),
"events": int(y_eval.sum()),
**metrics,
**boot,
}
rows.append(base)
rows.extend(
calibration_curve_rows(
split=split,
endpoint=endpoint,
horizon_days=float(horizon),
model=model_label,
y=y_eval,
p=p,
)
)
for threshold in RISK_THRESHOLD_GRID.get(endpoint, RISK_THRESHOLD_GRID["default"]):
dca_boot = bootstrap_net_benefit_ci(
y_eval,
p,
patient_ids,
float(threshold),
bootstrap_clusters,
seed + int(horizon) * 17 + int(round(float(threshold) * 10000)) + (0 if kind == "clinical_landmark" else 1),
)
rows.append(
{
"split": split,
"endpoint": endpoint,
"horizon_days": float(horizon),
"model": model_label,
"metric_family": "decision_curve",
"threshold": float(threshold),
"net_benefit": net_benefit(y_eval, p, float(threshold)),
"bootstrap_n": int(bootstrap_clusters),
"n": int(y_eval.size),
"events": int(y_eval.sum()),
**dca_boot,
}
)
return rows
def propensity_baseline_rows(
train: dict[str, np.ndarray],
eval_arrays: dict[str, np.ndarray],
train_actions: dict[str, np.ndarray],
eval_actions: dict[str, np.ndarray],
ontology: dict[str, Any],
split: str,
seed: int,
max_train: int,
require_lightgbm: bool,
model_cache: dict[int, Any] | None = None,
) -> list[dict[str, Any]]:
rows = []
slot_names = [slot["slot_name"] for slot in ontology["slots"]]
local = ontology.get("local_label_to_action_value_id_by_slot", {})
valid_train = train["valid_mask"].astype(bool)
valid_eval = eval_arrays["valid_mask"].astype(bool)
for slot_id, slot_name in enumerate(slot_names):
present_id = local.get(slot_name, {}).get("present")
if present_id is None:
positive_label = next((label for label in ("是", "有", "面访", "住院", "持续", "社区转医院") if label in local.get(slot_name, {})), None)
if positive_label is None:
continue
present_id = int(local[slot_name][positive_label])
tr_mask = valid_train & train_actions["action_mask"][:, :, slot_id].astype(bool)
ev_mask = valid_eval & eval_actions["action_mask"][:, :, slot_id].astype(bool)
tr_row, tr_pos = np.nonzero(tr_mask)
ev_row, ev_pos = np.nonzero(ev_mask)
if tr_row.size == 0 or ev_row.size == 0:
continue
y_train = (train_actions["action_value_ids"][:, :, slot_id][tr_row, tr_pos] == int(present_id)).astype(np.int8)
y_eval = (eval_actions["action_value_ids"][:, :, slot_id][ev_row, ev_pos] == int(present_id)).astype(np.int8)
X_train, _ = pre_action_features(train, tr_row, tr_pos)
X_eval, _ = pre_action_features(eval_arrays, ev_row, ev_pos)
for kind, label in (("empirical_state_era", "Empirical service-state x era propensity"), ("lightgbm_landmark", "LightGBM propensity")):
if kind == "empirical_state_era":
key_train = service_era_key(train["service_state"][tr_row, tr_pos], train["visit_year"][tr_row, tr_pos])
key_eval = service_era_key(eval_arrays["service_state"][ev_row, ev_pos], eval_arrays["visit_year"][ev_row, ev_pos])
rates = {}
for k in np.unique(key_train):
m = key_train == k
rates[int(k)] = float((y_train[m].sum() + 1.0) / (m.sum() + 2.0))
fallback = float((y_train.sum() + 1.0) / (y_train.size + 2.0))
p = np.asarray([rates.get(int(k), fallback) for k in key_eval], dtype=np.float64)
else:
if model_cache is not None and slot_id in model_cache:
model = model_cache[slot_id]
else:
model, _ = fit_classifier(X_train, y_train, "lightgbm_landmark", seed + slot_id, max_train, require_lightgbm)
if model_cache is not None:
model_cache[slot_id] = model
p = predict_classifier(model, X_eval)
p = np.clip(p, 1.0e-4, 1.0 - 1.0e-4)
w = y_eval / p + (1 - y_eval) / (1.0 - p)
rows.append(
{
"split": split,
"model": label,
"slot_id": int(slot_id),
"slot_name": slot_name,
"n": int(y_eval.size),
"present_rate": float(y_eval.mean()) if y_eval.size else math.nan,
"mean_propensity": float(p.mean()) if p.size else math.nan,
"p01": float(np.quantile(p, 0.01)) if p.size else math.nan,
"p05": float(np.quantile(p, 0.05)) if p.size else math.nan,
"p50": float(np.quantile(p, 0.50)) if p.size else math.nan,
"p95": float(np.quantile(p, 0.95)) if p.size else math.nan,
"p99": float(np.quantile(p, 0.99)) if p.size else math.nan,
"extreme_propensity_rate": float(((p < 0.01) | (p > 0.99)).mean()) if p.size else math.nan,
"effective_sample_size_binary_ipw": effective_sample_size(w),
**weighted_binary_metrics(y_eval, p),
}
)
return rows
def load_actions(
split: str,
arrays: dict[str, np.ndarray],
action_dir: Path | None,
tensor_metadata: dict[str, Any],
ordinal_direction: dict[str, Any],
cat_value_vocab: dict[str, int],
ontology: dict[str, Any],
stage0_dir: Path,
max_windows: int,
) -> dict[str, np.ndarray]:
actions, _, _ = load_or_build_action_arrays(
split=split,
arrays=arrays,
action_dir=action_dir,
tensor_metadata=tensor_metadata,
ordinal_direction=ordinal_direction,
cat_value_vocab=cat_value_vocab,
ontology=ontology,
grace_days=(30.0, 60.0, 90.0),
stage0_dir=stage0_dir,
)
for key, value in actions.items():
if max_windows > 0 and value.shape[0] > max_windows:
actions[key] = value[:max_windows]
return actions
def main() -> None:
parser = argparse.ArgumentParser(description="Run matched non-neural baselines for SCTM-v5 top-journal tables.")
parser.add_argument("--tensor-dir", type=Path, required=True)
parser.add_argument("--stage0-dir", type=Path, required=True)
parser.add_argument("--action-dir", type=Path, default=None)
parser.add_argument("--out-dir", type=Path, required=True)
parser.add_argument("--splits", default="val,test")
parser.add_argument("--horizons", default="90,365")
parser.add_argument("--risk-endpoints", default="death,disengagement,primary_referral_relapse,service_escalation,clinical_deterioration,high_acuity_state")
parser.add_argument("--max-windows", type=int, default=0)
parser.add_argument("--max-risk-train", type=int, default=500000)
parser.add_argument("--bootstrap-clusters", type=int, default=200)
parser.add_argument("--seed", type=int, default=20260526)
parser.add_argument("--require-lightgbm", action="store_true")
args = parser.parse_args()
args.out_dir.mkdir(parents=True, exist_ok=True)
meta = load_json(args.tensor_dir / "tensor_metadata.json")
vocab = load_json(args.tensor_dir / "cat_value_vocab.json")
ordinal_direction = load_json(args.stage0_dir / "ordinal_direction_table.json")
ontology = build_action_ontology(meta, vocab)
train = load_npz_to_memory(resolve_split_path(args.tensor_dir, "train"))
if args.max_windows > 0:
train = {k: v[: args.max_windows] if getattr(v, "shape", (0,))[0] == train["valid_mask"].shape[0] else v for k, v in train.items()}
train_actions = load_actions("train", train, args.action_dir, meta, ordinal_direction, vocab, ontology, args.stage0_dir, args.max_windows)
print(json.dumps({"stage": "fit_semi_markov", "train_windows": int(train["valid_mask"].shape[0])}), flush=True)
semi = fit_semi_markov(train, smoothing=1.0)
print(json.dumps({"stage": "fit_empirical_grammar"}), flush=True)
grammar = fit_empirical_grammar(train, meta, vocab, smoothing=1.0)
print(json.dumps({"stage": "fit_baseline_models_start", "splits": args.splits}), flush=True)
horizons = [float(x) for x in args.horizons.split(",") if x.strip()]
endpoints = [x.strip() for x in args.risk_endpoints.split(",") if x.strip()]
semi_rows = []
grammar_rows = []
risk_rows = []
propensity_rows = []
risk_model_cache: dict[tuple[str, float, str], tuple[Any, str]] = {}
propensity_model_cache: dict[int, Any] = {}
for split in [s.strip() for s in args.splits.split(",") if s.strip()]:
print(json.dumps({"stage": "evaluate_split_start", "split": split}), flush=True)
arrays = load_npz_to_memory(resolve_split_path(args.tensor_dir, split))
if args.max_windows > 0:
arrays = {k: v[: args.max_windows] if getattr(v, "shape", (0,))[0] == arrays["valid_mask"].shape[0] else v for k, v in arrays.items()}
actions = load_actions(split, arrays, args.action_dir, meta, ordinal_direction, vocab, ontology, args.stage0_dir, args.max_windows)
print(json.dumps({"stage": "semi_markov_metrics", "split": split}), flush=True)
semi_rows.extend(semi_markov_metrics(semi, arrays, split))
print(json.dumps({"stage": "empirical_grammar_metrics", "split": split}), flush=True)
grammar_rows.extend(empirical_grammar_metrics(grammar, arrays, split, meta, vocab))
print(json.dumps({"stage": "risk_baseline_metrics", "split": split}), flush=True)
risk_rows.extend(
risk_baseline_rows(
train,
arrays,
split,
endpoints,
horizons,
args.seed,
args.max_risk_train,
args.bootstrap_clusters,
args.require_lightgbm,
risk_model_cache,
)
)
print(json.dumps({"stage": "propensity_baseline_metrics", "split": split}), flush=True)
propensity_rows.extend(propensity_baseline_rows(train, arrays, train_actions, actions, ontology, split, args.seed, args.max_risk_train, args.require_lightgbm, propensity_model_cache))
print(json.dumps({"stage": "evaluate_split_done", "split": split}), flush=True)
paths = {
"semi_markov": args.out_dir / "semi_markov_care_process_metrics.csv",
"empirical_grammar": args.out_dir / "empirical_state_year_grammar_metrics.csv",
"risk": args.out_dir / "clinical_lightgbm_risk_baseline_metrics.csv",
"propensity": args.out_dir / "empirical_lightgbm_propensity_baseline_metrics.csv",
"behrt_status": args.out_dir / "behrt_style_baseline_status.json",
}
pd.DataFrame(semi_rows).to_csv(paths["semi_markov"], index=False)
pd.DataFrame(grammar_rows).to_csv(paths["empirical_grammar"], index=False)
pd.DataFrame(risk_rows).to_csv(paths["risk"], index=False)
pd.DataFrame(propensity_rows).to_csv(paths["propensity"], index=False)
behrt_status = {
"status": "not_run_in_this_script",
"reason": "Exact BEHRT-style baseline requires separate event-token MLM pretraining and landmark fine-tuning; this non-neural matched baseline suite intentionally does not substitute a tabular transformer for that baseline.",
"required_spec": {
"architecture": "6-layer Transformer encoder, 12 heads, hidden size 768",
"pretraining": "15% masked language modeling over m1_event_token_v2 tokens",
"finetuning": "landmark risk heads for death, disengagement, event0 primary referral relapse, service escalation, clinical deterioration, and high acuity",
},
}
paths["behrt_status"].write_text(json.dumps(behrt_status, ensure_ascii=False, indent=2), encoding="utf-8")
summary = {
"status": "completed",
"script": "evaluate_v5_matched_baselines.py",
"splits": args.splits,
"horizons": horizons,
"risk_endpoints": endpoints,
"max_risk_train": int(args.max_risk_train),
"bootstrap_clusters": int(args.bootstrap_clusters),
"require_lightgbm": bool(args.require_lightgbm),
"lightgbm_backend": "lightgbm" if LGBMClassifier is not None else "unavailable",
"outputs": {k: str(v) for k, v in paths.items()},
"notes": (
"LightGBM rows use the lightgbm backend; use --require-lightgbm for fail-fast publication runs. "
"Clinical Landmark uses unweighted L2 logistic regression so predict_proba remains on the natural prevalence scale. "
"Event-horizon endpoints use the same teacher-forced observed-future-contact evaluable mask as evaluate_v5_topjournal.py; "
"PWE endpoints use the same next-transition CIF target convention. "
"Supported event endpoints include primary_referral_relapse, service_escalation, clinical_deterioration, and high_acuity_state."
),
}
(args.out_dir / "matched_baseline_summary.json").write_text(json.dumps(json_ready(summary), ensure_ascii=False, indent=2), encoding="utf-8")
print(json.dumps({"status": "completed", "out_dir": str(args.out_dir)}, ensure_ascii=False), flush=True)
if __name__ == "__main__":
main()