"""Export Weibull shape-parameter statistics on the test split. The script is intended for all_future checkpoints with dist_mode="weibull" or dist_mode="mixed". For full Weibull models it reads rho_head[Death]; for mixed models it reads rho_death_head. For full Weibull models it also exports disease token rho summaries, which are the main evidence for whether risk/hazard changes with horizon instead of following an exponential shape. """ from __future__ import annotations import argparse import contextlib import json from pathlib import Path from typing import Any, Dict, Iterable, List, Optional import numpy as np import pandas as pd import torch import torch.nn.functional as F import torch.multiprocessing as torch_mp from torch.utils.data import DataLoader from tqdm.auto import tqdm from eval_data import load_sequence_eval_dataset from evaluate_auc_v2 import ( LandmarkDataset, _build_first_occurrence_maps, _get_death_token_ids, build_model_from_dataset, cfg_get, collate_landmark_fn, load_checkpoint_state_dict, load_json_config, load_model_state, make_eval_indices, parse_float_list, parse_int_list, resolve_dist_mode_for_checkpoint, resolve_eval_device, validate_dataset_metadata, ) try: torch_mp.set_sharing_strategy("file_system") except RuntimeError: pass def quantile_summary(df: pd.DataFrame, group_cols: List[str], value_cols: List[str]) -> pd.DataFrame: probs = [0.01, 0.05, 0.25, 0.50, 0.75, 0.95, 0.99] rows: List[Dict[str, Any]] = [] grouped = [((), df)] if not group_cols else df.groupby(group_cols, dropna=False) for key, g in grouped: if not isinstance(key, tuple): key = (key,) base = {col: val for col, val in zip(group_cols, key)} base["n"] = int(len(g)) for col in value_cols: x = pd.to_numeric(g[col], errors="coerce").to_numpy(dtype=np.float64) x = x[np.isfinite(x)] if x.size == 0: continue row = dict(base) row["variable"] = col row["mean"] = float(np.mean(x)) row["std"] = float(np.std(x, ddof=1)) if x.size > 1 else 0.0 row["min"] = float(np.min(x)) row["max"] = float(np.max(x)) for p in probs: row[f"p{int(p * 100):02d}"] = float(np.quantile(x, p)) rows.append(row) return pd.DataFrame(rows) def load_labels_meta(path: Optional[str]) -> Optional[pd.DataFrame]: if path is None: return None fp = Path(path) if not fp.exists(): return None return pd.read_csv(fp) @torch.inference_mode() def infer_landmark_hidden_local( model, loader: DataLoader, device: torch.device, use_amp: bool, hidden_cache_dtype: str, ) -> tuple[np.ndarray, Dict[str, np.ndarray]]: """Minimal all_future landmark hidden inference for parameter export.""" out_dtype = np.float32 if str(hidden_cache_dtype).lower() == "float32" else np.float16 hidden_parts: List[np.ndarray] = [] arrays: Dict[str, List[np.ndarray]] = { "patient_id": [], "sex": [], "landmark_age": [], "followup_end_time": [], "death_time": [], } amp_enabled = bool(use_amp and device.type == "cuda") for batch in tqdm(loader, desc="Landmark hidden", dynamic_ncols=True): batch_dev = { k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) for k, v in batch.items() } amp_ctx = ( torch.autocast(device_type=device.type, dtype=torch.float16) if amp_enabled else contextlib.nullcontext() ) with amp_ctx: hidden = model( event_seq=batch_dev["event_seq"], time_seq=batch_dev["time_seq"], sex=batch_dev["sex"], padding_mask=batch_dev["padding_mask"], t_query=batch_dev["t_query"], other_type=batch_dev["other_type"], other_value=batch_dev["other_value"], other_value_kind=batch_dev["other_value_kind"], other_time=batch_dev["other_time"], target_mode="all_future", ) hidden_parts.append(hidden.detach().cpu().numpy().astype(out_dtype, copy=False)) for key in arrays: arrays[key].append(batch[key].cpu().numpy()) hidden_all = np.concatenate(hidden_parts, axis=0) row_arrays = {key: np.concatenate(parts, axis=0) for key, parts in arrays.items()} return hidden_all, row_arrays @torch.inference_mode() def project_death_params( model, hidden_all: np.ndarray, dist_mode: str, device: torch.device, batch_size: int, use_amp: bool, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", hidden_all.shape[0]) - 1)) if not hasattr(model, "vocab_size"): death_idx = int(model.risk_head.out_features - 1) compute_dtype = torch.float16 if (device.type == "cuda" and use_amp) else torch.float32 risk_w = model.risk_head.weight[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype) risk_b = None if model.risk_head.bias is not None: risk_b = model.risk_head.bias[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype) if dist_mode == "weibull": rho_w = model.rho_head.weight[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype) rho_b = model.rho_head.bias[death_idx: death_idx + 1].detach().to(device=device, dtype=compute_dtype) elif dist_mode == "mixed": rho_w = model.rho_death_head.weight.detach().to(device=device, dtype=compute_dtype) rho_b = model.rho_death_head.bias.detach().to(device=device, dtype=compute_dtype) else: raise ValueError("Death Weibull parameter export requires dist_mode='weibull' or 'mixed'.") logits_out: List[np.ndarray] = [] rate_out: List[np.ndarray] = [] rho_out: List[np.ndarray] = [] for start in tqdm(range(0, hidden_all.shape[0], batch_size), desc="Death eta/rho", dynamic_ncols=True): end = min(start + batch_size, hidden_all.shape[0]) h = torch.from_numpy(hidden_all[start:end]).to(device=device, dtype=compute_dtype, non_blocking=True) logits = F.linear(h, risk_w, risk_b).squeeze(-1) rate = F.softplus(logits) + 1e-8 rho = F.softplus(F.linear(h, rho_w, rho_b).squeeze(-1)) + 1e-6 logits_out.append(logits.float().cpu().numpy()) rate_out.append(rate.float().cpu().numpy()) rho_out.append(rho.float().cpu().numpy()) del h, logits, rate, rho return ( np.concatenate(logits_out).astype(np.float32, copy=False), np.concatenate(rate_out).astype(np.float32, copy=False), np.concatenate(rho_out).astype(np.float32, copy=False), ) @torch.inference_mode() def export_all_token_rho_summary( model, hidden_all: np.ndarray, dataset, device: torch.device, output_dir: Path, token_chunk_size: int, row_batch_size: int, use_amp: bool, horizons: np.ndarray, ) -> None: if not hasattr(model, "rho_head"): print("[INFO] Skipping all-token rho summary because this is not a full Weibull model.") return special = {0, 1, 2} token_ids = [ int(t) for t, code in dataset.label_id_to_code.items() if int(t) not in special and not str(code).startswith("<") ] token_ids = sorted(set(token_ids)) death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", len(token_ids)) - 1)) if not hasattr(model, "vocab_size"): death_idx = int(model.risk_head.out_features - 1) compute_dtype = torch.float16 if (device.type == "cuda" and use_amp) else torch.float32 rows: List[Dict[str, Any]] = [] for chunk_start in tqdm(range(0, len(token_ids), token_chunk_size), desc="All-token rho chunks", dynamic_ncols=True): chunk = token_ids[chunk_start: chunk_start + token_chunk_size] w = model.rho_head.weight[chunk].detach().to(device=device, dtype=compute_dtype) b = model.rho_head.bias[chunk].detach().to(device=device, dtype=compute_dtype) vals_parts: List[np.ndarray] = [] for row_start in range(0, hidden_all.shape[0], row_batch_size): row_end = min(row_start + row_batch_size, hidden_all.shape[0]) h = torch.from_numpy(hidden_all[row_start:row_end]).to(device=device, dtype=compute_dtype, non_blocking=True) rho = F.softplus(F.linear(h, w, b)) + 1e-6 vals_parts.append(rho.float().cpu().numpy()) del h, rho vals = np.concatenate(vals_parts, axis=0) for j, token in enumerate(chunk): x = vals[:, j].astype(np.float64, copy=False) row = { "token": int(token), "label_code": dataset.label_id_to_code.get(int(token), ""), "endpoint_type": "death" if int(token) == int(death_idx) else "disease", "n_landmark_rows": int(x.size), "rho_mean": float(np.mean(x)), "rho_std": float(np.std(x, ddof=1)) if x.size > 1 else 0.0, "rho_minus_one_mean": float(np.mean(x - 1.0)), "frac_rho_gt_1": float(np.mean(x > 1.0)), "frac_rho_lt_1": float(np.mean(x < 1.0)), "frac_rho_gt_1_1": float(np.mean(x > 1.1)), "frac_rho_lt_0_9": float(np.mean(x < 0.9)), "rho_p01": float(np.quantile(x, 0.01)), "rho_p05": float(np.quantile(x, 0.05)), "rho_p25": float(np.quantile(x, 0.25)), "rho_p50": float(np.quantile(x, 0.50)), "rho_p75": float(np.quantile(x, 0.75)), "rho_p95": float(np.quantile(x, 0.95)), "rho_p99": float(np.quantile(x, 0.99)), } for horizon in horizons.tolist(): h = float(horizon) if h <= 0: continue # Shape-only time scaling. For rho=1 this equals 1, i.e. an # exponential model with constant instantaneous hazard. inst_scale = np.power(h, x - 1.0) cumhaz_scale = np.power(h, x) row[f"instant_hazard_scale_h{h:g}y_vs_1y_mean"] = float(np.mean(inst_scale)) row[f"instant_hazard_scale_h{h:g}y_vs_1y_p50"] = float(np.quantile(inst_scale, 0.50)) row[f"cumhaz_scale_h{h:g}y_mean"] = float(np.mean(cumhaz_scale)) row[f"cumhaz_scale_h{h:g}y_p50"] = float(np.quantile(cumhaz_scale, 0.50)) rows.append(row) del vals, vals_parts out = pd.DataFrame(rows) out.to_csv(output_dir / "all_token_weibull_shape_summary.csv", index=False) out[out["endpoint_type"] == "disease"].to_csv( output_dir / "disease_token_weibull_shape_summary.csv", index=False ) out[out["endpoint_type"] == "death"].to_csv( output_dir / "death_token_weibull_shape_summary.csv", index=False ) disease = out[out["endpoint_type"] == "disease"].copy() if not disease.empty: pd.DataFrame([ { "n_tokens": int(len(disease)), "rho_mean_across_tokens": float(disease["rho_mean"].mean()), "rho_median_across_tokens": float(disease["rho_p50"].median()), "tokens_with_mean_rho_gt_1": int((disease["rho_mean"] > 1.0).sum()), "tokens_with_mean_rho_lt_1": int((disease["rho_mean"] < 1.0).sum()), "frac_tokens_with_mean_rho_gt_1": float((disease["rho_mean"] > 1.0).mean()), "frac_tokens_with_mean_rho_lt_1": float((disease["rho_mean"] < 1.0).mean()), "tokens_with_mean_rho_gt_1_1": int((disease["rho_mean"] > 1.1).sum()), "tokens_with_mean_rho_lt_0_9": int((disease["rho_mean"] < 0.9).sum()), } ]).to_csv(output_dir / "disease_weibull_shape_overall_summary.csv", index=False) def main() -> None: parser = argparse.ArgumentParser(description="Export test-split Weibull shape parameter statistics.") parser.add_argument("--run_path", type=str, required=True) parser.add_argument("--output_path", type=str, default=None) parser.add_argument("--eval_split", type=str, default="test", choices=["test", "val", "valid", "validation", "train", "all"]) parser.add_argument("--landmark_start", type=float, default=None) parser.add_argument("--landmark_stop", type=float, default=None) parser.add_argument("--landmark_step", type=float, default=None) parser.add_argument("--horizons", type=str, default=None) parser.add_argument("--batch_size", type=int, default=None) parser.add_argument( "--num_workers", type=int, default=0, help=( "DataLoader workers. Default 0 avoids Linux multiprocessing " "'received 0 items of ancdata' failures on shared filesystems." ), ) parser.add_argument("--device", type=str, default=None) parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None) parser.add_argument("--hidden_cache_dtype", type=str, default="float32", choices=["float16", "float32"]) parser.add_argument( "--include_all_token_rho_summary", action=argparse.BooleanOptionalAction, default=True, help=( "For full Weibull models, export disease/death token rho summaries. " "Use --no-include_all_token_rho_summary to skip the heavier token projection." ), ) parser.add_argument("--token_chunk_size", type=int, default=32) parser.add_argument("--row_batch_size", type=int, default=512) args = parser.parse_args() run_path = Path(args.run_path) config_path = run_path / "train_config.json" ckpt_path = run_path / "best_model.pt" if not config_path.exists(): raise FileNotFoundError(config_path) if not ckpt_path.exists(): raise FileNotFoundError(ckpt_path) cfg = load_json_config(config_path) model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower() if model_target_mode != "all_future": raise ValueError("This export is intended for all_future checkpoints.") data_prefix = cfg.get("data_prefix", "ukb") labels_file = cfg.get("labels_file", "labels.csv") no_event_interval_years = cfg.get("no_event_interval_years", 5.0) include_no_event_in_uts_target = cfg.get("include_no_event_in_uts_target", False) dataset = load_sequence_eval_dataset( model_target_mode=model_target_mode, data_prefix=data_prefix, labels_file=labels_file, no_event_interval_years=float(no_event_interval_years), include_no_event_in_uts_target=bool(include_no_event_in_uts_target), min_history_events=int(cfg.get("all_future_min_history_events", 1)), min_future_events=int(cfg.get("all_future_min_future_events", 1)), extra_info_types=parse_int_list(cfg.get("extra_info_types", None)), ) validate_dataset_metadata(dataset, cfg) subset_indices = make_eval_indices(dataset, args, cfg) first_occurrence_by_token, _, _, _ = _build_first_occurrence_maps(dataset, subset_indices) landmark_start = float(cfg_get(args, cfg, "landmark_start", 40.0)) landmark_stop = float(cfg_get(args, cfg, "landmark_stop", 80.0)) landmark_step = float(cfg_get(args, cfg, "landmark_step", 5.0)) landmark_ages = np.arange(landmark_start, landmark_stop, landmark_step, dtype=np.float32) if landmark_ages.size == 0: raise ValueError("No landmark ages produced.") horizons = np.asarray( parse_float_list(cfg_get(args, cfg, "horizons", "1,5,10")) or [1.0, 5.0, 10.0], dtype=np.float32, ) if horizons.size == 0: raise ValueError("No horizons provided.") state_dict = load_checkpoint_state_dict(ckpt_path, map_location="cpu") dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict) if dist_mode not in {"weibull", "mixed"}: raise ValueError( f"Resolved dist_mode={dist_mode!r}; expected 'weibull' or 'mixed' for Weibull shape export." ) cfg_model = dict(cfg) cfg_model["dist_mode"] = dist_mode device = resolve_eval_device(args.device) model = build_model_from_dataset(args, cfg_model, dataset).to(device) load_model_state(model, state_dict) model.eval() death_token_ids = _get_death_token_ids(dataset, None) death_idx = int(death_token_ids[0]) attn_mask_mode = str(cfg.get("attn_mask_mode", "target_aware")) landmark_dataset = LandmarkDataset( dataset=dataset, subset_indices=subset_indices, landmark_ages=landmark_ages, attn_mask_mode=attn_mask_mode, model_target_mode=model_target_mode, min_history_events=int(cfg_get(args, cfg, "min_history_events", 1)), first_occurrence_by_token=first_occurrence_by_token, death_token_ids=death_token_ids, ) batch_size = int(cfg_get(args, cfg, "batch_size", 128)) num_workers = int(cfg_get(args, cfg, "num_workers", 0)) loader_kwargs = { "batch_size": batch_size, "shuffle": False, "collate_fn": collate_landmark_fn, "num_workers": num_workers, "pin_memory": device.type == "cuda", } if num_workers > 0: loader_kwargs["persistent_workers"] = True loader_kwargs["prefetch_factor"] = 2 loader = DataLoader(landmark_dataset, **loader_kwargs) use_amp = bool(cfg_get(args, cfg, "use_amp", False)) hidden_all, row_arrays = infer_landmark_hidden_local( model=model, loader=loader, device=device, use_amp=use_amp, hidden_cache_dtype=str(args.hidden_cache_dtype), ) eta, rate, rho = project_death_params( model=model, hidden_all=hidden_all, dist_mode=dist_mode, device=device, batch_size=int(args.row_batch_size), use_amp=use_amp, ) rows = pd.DataFrame({ "patient_id": row_arrays["patient_id"].astype(np.int64), "sex": row_arrays["sex"].astype(np.int64), "sex_label": np.where(row_arrays["sex"].astype(np.int64) == 0, "female", "male"), "landmark_age": row_arrays["landmark_age"].astype(np.float32), "followup_end_time": row_arrays["followup_end_time"].astype(np.float32), "death_time": row_arrays["death_time"].astype(np.float32), "death_eta": eta, "death_rate": rate, "death_rho": rho, }) for horizon in horizons.tolist(): h = float(horizon) cumulative_hazard = rows["death_rate"].to_numpy(dtype=np.float64) * np.power(h, rows["death_rho"].to_numpy(dtype=np.float64)) rows[f"death_cumhaz_h{h:g}y"] = cumulative_hazard rows[f"death_risk_h{h:g}y"] = -np.expm1(-cumulative_hazard) rows[f"death_observed_h{h:g}y"] = ( (rows["death_time"].to_numpy(dtype=np.float64) > rows["landmark_age"].to_numpy(dtype=np.float64)) & (rows["death_time"].to_numpy(dtype=np.float64) <= rows["landmark_age"].to_numpy(dtype=np.float64) + h) ).astype(np.int8) output_dir = Path(args.output_path) if args.output_path else run_path / "weibull_death_parameter_stats_test" output_dir.mkdir(parents=True, exist_ok=True) rows.to_csv(output_dir / "death_weibull_parameters_by_landmark.csv", index=False) value_cols = ["death_eta", "death_rate", "death_rho"] for horizon in horizons.tolist(): h = float(horizon) value_cols.extend([f"death_cumhaz_h{h:g}y", f"death_risk_h{h:g}y"]) quantile_summary(rows, [], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_overall.csv", index=False) quantile_summary(rows, ["landmark_age"], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_by_landmark_age.csv", index=False) quantile_summary(rows, ["sex_label"], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_by_sex.csv", index=False) quantile_summary(rows, ["sex_label", "landmark_age"], value_cols).to_csv(output_dir / "death_weibull_parameter_summary_by_sex_landmark_age.csv", index=False) metadata = { "run_path": str(run_path), "config_path": str(config_path), "checkpoint_path": str(ckpt_path), "eval_split": str(args.eval_split), "model_target_mode": model_target_mode, "time_mode": str(cfg.get("time_mode")), "dist_mode_config": str(cfg.get("dist_mode")), "dist_mode_resolved": dist_mode, "extra_info_types": cfg.get("extra_info_types"), "death_token_id": death_idx, "death_label_code": dataset.label_id_to_code.get(death_idx, "Death"), "n_selected_patients": int(len(subset_indices)), "n_landmark_rows": int(len(rows)), "landmark_ages": [float(x) for x in landmark_ages.tolist()], "horizons": [float(x) for x in horizons.tolist()], } with (output_dir / "metadata.json").open("w", encoding="utf-8") as f: json.dump(metadata, f, indent=2) if args.include_all_token_rho_summary and dist_mode == "weibull": export_all_token_rho_summary( model=model, hidden_all=hidden_all, dataset=dataset, device=device, output_dir=output_dir, token_chunk_size=int(args.token_chunk_size), row_batch_size=int(args.row_batch_size), use_amp=use_amp, horizons=horizons, ) elif dist_mode == "mixed": pd.DataFrame([ { "dist_mode": dist_mode, "disease_shape_available": False, "death_shape_available": True, "note": ( "The mixed model uses Weibull rho only for Death. " "Non-death disease hazards are exponential, equivalent to fixed rho=1." ), } ]).to_csv(output_dir / "disease_shape_not_available_for_mixed_model.csv", index=False) print(f"Wrote Weibull shape parameter statistics to: {output_dir}") if __name__ == "__main__": main()