Add Weibull shape export scripts

export_weibull_death_parameter_stats.py

@@ -0,0 +1,502 @@
+"""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
+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,
+)
+
+
+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=None)
+    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", 4))
+    loader = DataLoader(
+        landmark_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        collate_fn=collate_landmark_fn,
+        num_workers=num_workers,
+        pin_memory=device.type == "cuda",
+        persistent_workers=num_workers > 0,
+        prefetch_factor=2 if num_workers > 0 else None,
+    )
+
+    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()

export_weibull_shape_parameter_stats.py

@@ -0,0 +1,7 @@
+"""Compatibility entry point for Weibull shape-parameter export."""
+
+from export_weibull_death_parameter_stats import main
+
+
+if __name__ == "__main__":
+    main()

plot_next_token_to_all_future_auc.R

@@ -0,0 +1,553 @@
+#!/usr/bin/env Rscript
+
+# Paper-grade single-panel figures supporting the conclusion that fixed-landmark
+# horizon evaluation favors all_future over next_token.
+#
+# Outputs are written as separate panel files. This script intentionally does not
+# combine panels with plot_grid().
+
+suppressPackageStartupMessages({
+  library(cowplot)
+  library(dplyr)
+  library(ggplot2)
+  library(jsonlite)
+  library(readr)
+  library(stringr)
+  library(tibble)
+  library(tidyr)
+})
+
+root_dir <- "."
+runs_dir <- file.path(root_dir, "runs")
+out_dir <- file.path(root_dir, "figures_next_token_to_all_future_absolute_smoking_alcohol_bmi")
+dir.create(out_dir, showWarnings = FALSE, recursive = TRUE)
+
+required_time_mode <- "absolute"
+required_extra_info_types <- c(11L, 66L, 67L)
+required_extra_info_signature <- paste(sort(required_extra_info_types), collapse = ",")
+
+theme_set(
+  theme_cowplot(font_size = 9) +
+    theme(
+      plot.background = element_rect(fill = "white", color = NA),
+      panel.background = element_rect(fill = "white", color = NA),
+      legend.background = element_rect(fill = "white", color = NA),
+      legend.key = element_rect(fill = "white", color = NA)
+    )
+)
+
+target_cols <- c(
+  "next_token" = "#B54A3A",
+  "all_future" = "#2C7FB8"
+)
+
+dist_shapes <- c(
+  "exponential" = 16,
+  "weibull" = 17,
+  "mixed" = 15
+)
+
+read_run_config <- function(run_path) {
+  cfg_path <- file.path(run_path, "train_config.json")
+  if (!file.exists(cfg_path)) return(NULL)
+  cfg <- jsonlite::read_json(cfg_path, simplifyVector = TRUE)
+  extra_info_types <- cfg$extra_info_types %||% integer(0)
+  extra_info_signature <- paste(sort(as.integer(extra_info_types)), collapse = ",")
+  tibble(
+    run = basename(run_path),
+    model_target_mode = as.character(cfg$model_target_mode %||% NA_character_),
+    target_mode = as.character(cfg$target_mode %||% NA_character_),
+    dist_mode = as.character(cfg$dist_mode %||% NA_character_),
+    time_mode = as.character(cfg$time_mode %||% NA_character_),
+    readout_name = as.character(cfg$readout_name %||% NA_character_),
+    attn_mask_mode = as.character(cfg$attn_mask_mode %||% NA_character_),
+    extra_info_signature = extra_info_signature
+  )
+}
+
+`%||%` <- function(x, y) {
+  if (is.null(x) || length(x) == 0) y else x
+}
+
+load_one_result <- function(run_path, file_name, eval_family) {
+  cfg <- read_run_config(run_path)
+  if (is.null(cfg)) return(NULL)
+  fp <- file.path(run_path, file_name)
+  if (!file.exists(fp)) return(NULL)
+
+  df <- suppressMessages(readr::read_csv(fp, show_col_types = FALSE))
+  if (!("auc" %in% names(df)) || nrow(df) == 0) return(NULL)
+
+  out <- df %>%
+    mutate(
+      run = basename(run_path),
+      eval_family = eval_family,
+      auc = as.numeric(auc)
+    ) %>%
+    left_join(cfg, by = "run", suffix = c("", "_cfg"))
+
+  coalesce_joined <- function(data, col) {
+    cfg_col <- paste0(col, "_cfg")
+    if (col %in% names(data) && cfg_col %in% names(data)) {
+      dplyr::coalesce(data[[col]], data[[cfg_col]])
+    } else if (col %in% names(data)) {
+      data[[col]]
+    } else if (cfg_col %in% names(data)) {
+      data[[cfg_col]]
+    } else {
+      rep(NA_character_, nrow(data))
+    }
+  }
+
+  for (col in c("model_target_mode", "target_mode", "dist_mode", "time_mode", "readout_name", "attn_mask_mode")) {
+    out[[col]] <- coalesce_joined(out, col)
+  }
+
+  out %>%
+    select(-any_of(c(
+      "model_target_mode_cfg", "target_mode_cfg", "dist_mode_cfg",
+      "time_mode_cfg", "readout_name_cfg", "attn_mask_mode_cfg"
+    )))
+}
+
+run_paths <- list.dirs(runs_dir, recursive = FALSE, full.names = TRUE)
+
+landmark_auc <- bind_rows(lapply(
+  run_paths,
+  load_one_result,
+  file_name = "df_auc_landmark.csv",
+  eval_family = "Fixed landmark + horizon"
+)) %>%
+  filter(time_mode == "absolute")
+
+token_auc <- bind_rows(lapply(
+  run_paths,
+  load_one_result,
+  file_name = "df_both.csv",
+  eval_family = "Delphi2M-style token"
+)) %>%
+  filter(time_mode == "absolute")
+
+if (nrow(landmark_auc) == 0) {
+  stop("No landmark AUC files found under runs/*/df_auc_landmark.csv")
+}
+if (nrow(token_auc) == 0) {
+  stop("No token AUC files found under runs/*/df_both.csv")
+}
+
+landmark_auc <- landmark_auc %>%
+  filter(
+    time_mode == required_time_mode,
+    extra_info_signature == required_extra_info_signature
+  )
+
+token_auc <- token_auc %>%
+  filter(
+    time_mode == required_time_mode,
+    extra_info_signature == required_extra_info_signature
+  )
+
+if (nrow(landmark_auc) == 0 || nrow(token_auc) == 0) {
+  stop(
+    "No AUC rows remain after filtering for time_mode='",
+    required_time_mode,
+    "' and extra_info_types='",
+    required_extra_info_signature,
+    "'."
+  )
+}
+
+message(
+  "Using runs with time_mode='", required_time_mode,
+  "' and extra_info_types='", required_extra_info_signature, "':"
+)
+print(sort(unique(landmark_auc$run)))
+
+classify_endpoint <- function(data) {
+  data %>%
+    mutate(
+      endpoint_type = if_else(
+        str_to_lower(as.character(label_code)) == "death",
+        "Death",
+        "Non-death disease"
+      ),
+      endpoint_type = factor(endpoint_type, levels = c("Non-death disease", "Death"))
+    )
+}
+
+landmark_auc <- classify_endpoint(landmark_auc)
+token_auc <- classify_endpoint(token_auc)
+
+landmark_auc_disease <- landmark_auc %>% filter(endpoint_type == "Non-death disease")
+token_auc_disease <- token_auc %>% filter(endpoint_type == "Non-death disease")
+landmark_auc_death <- landmark_auc %>% filter(endpoint_type == "Death")
+token_auc_death <- token_auc %>% filter(endpoint_type == "Death")
+
+if (nrow(landmark_auc_death) == 0 || nrow(token_auc_death) == 0) {
+  warning("Death rows were not found in one or both AUC tables.")
+}
+
+auc_all <- bind_rows(
+  landmark_auc_disease %>% mutate(horizon = as.numeric(horizon), offset = NA_real_),
+  token_auc_disease %>% mutate(horizon = NA_real_, offset = as.numeric(offset))
+) %>%
+  mutate(
+    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
+    eval_family = factor(eval_family, levels = c("Delphi2M-style token", "Fixed landmark + horizon")),
+    dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed")),
+    model_label = recode(
+      as.character(model_target_mode),
+      "next_token" = "next-token objective",
+      "all_future" = "all-future objective"
+    )
+  )
+
+mean_ci <- function(x) {
+  x <- x[is.finite(x)]
+  n <- length(x)
+  m <- mean(x)
+  se <- sd(x) / sqrt(n)
+  tibble(mean = m, ymin = m - 1.96 * se, ymax = m + 1.96 * se, n = n)
+}
+
+save_panel <- function(plot, name, width = 3.6, height = 3.0) {
+  pdf_path <- file.path(out_dir, paste0(name, ".pdf"))
+  png_path <- file.path(out_dir, paste0(name, ".png"))
+  cowplot::save_plot(pdf_path, plot, base_width = width, base_height = height, bg = "white")
+  cowplot::save_plot(png_path, plot, base_width = width, base_height = height, dpi = 600, bg = "white")
+  message("Wrote: ", pdf_path)
+  message("Wrote: ", png_path)
+}
+
+# Panel 1: run-level mean AUC under the clinically aligned landmark/horizon task.
+# Death is excluded here and plotted separately below.
+landmark_run <- landmark_auc_disease %>%
+  mutate(model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))) %>%
+  group_by(run, model_target_mode, dist_mode, time_mode, target_mode) %>%
+  summarise(mean_auc = mean(auc, na.rm = TRUE), median_auc = median(auc, na.rm = TRUE), .groups = "drop")
+
+landmark_summary <- landmark_run %>%
+  group_by(model_target_mode) %>%
+  summarise(mean_ci(mean_auc), .groups = "drop")
+
+p1 <- ggplot(landmark_run, aes(x = model_target_mode, y = mean_auc)) +
+  geom_point(
+    aes(color = model_target_mode, shape = dist_mode),
+    position = position_jitter(width = 0.09, height = 0, seed = 1),
+    size = 2.2,
+    alpha = 0.88
+  ) +
+  geom_errorbar(
+    data = landmark_summary,
+    aes(x = model_target_mode, y = mean, ymin = ymin, ymax = ymax, color = model_target_mode),
+    width = 0.12,
+    linewidth = 0.55,
+    inherit.aes = FALSE
+  ) +
+  geom_point(
+    data = landmark_summary,
+    aes(x = model_target_mode, y = mean, color = model_target_mode),
+    size = 3.4,
+    inherit.aes = FALSE
+  ) +
+  scale_color_manual(values = target_cols, guide = "none") +
+  scale_shape_manual(values = dist_shapes, na.translate = FALSE) +
+  scale_x_discrete(labels = c("next_token", "all_future")) +
+  coord_cartesian(ylim = c(0.58, 0.78)) +
+  labs(
+    x = NULL,
+    y = "Mean AUC per run",
+    shape = "Risk head",
+    title = "Non-death landmark AUC (absolute time)"
+  ) +
+  theme(
+    plot.title = element_text(face = "bold", size = 10),
+    axis.text.x = element_text(size = 9),
+    legend.position = c(0.72, 0.20),
+    legend.background = element_blank()
+  )
+
+save_panel(p1, "panel_01_landmark_overall")
+
+# Panel 2: landmark AUC by prediction horizon.
+landmark_horizon_run <- landmark_auc_disease %>%
+  mutate(
+    horizon = as.numeric(horizon),
+    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))
+  ) %>%
+  group_by(run, model_target_mode, horizon) %>%
+  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
+
+landmark_horizon_summary <- landmark_horizon_run %>%
+  group_by(model_target_mode, horizon) %>%
+  summarise(mean_ci(mean_auc), .groups = "drop")
+
+p2 <- ggplot(landmark_horizon_run, aes(x = horizon, y = mean_auc, color = model_target_mode)) +
+  geom_line(aes(group = run), alpha = 0.18, linewidth = 0.35) +
+  geom_point(alpha = 0.32, size = 1.1) +
+  geom_ribbon(
+    data = landmark_horizon_summary,
+    aes(x = horizon, y = mean, ymin = ymin, ymax = ymax, fill = model_target_mode, group = model_target_mode),
+    alpha = 0.13,
+    color = NA,
+    inherit.aes = FALSE
+  ) +
+  geom_line(data = landmark_horizon_summary, aes(y = mean), linewidth = 0.85) +
+  geom_point(data = landmark_horizon_summary, aes(y = mean), size = 2.0) +
+  scale_color_manual(
+    values = target_cols,
+    labels = c("next_token", "all_future"),
+    name = NULL
+  ) +
+  scale_fill_manual(values = target_cols, guide = "none") +
+  scale_x_continuous(breaks = c(1, 5, 10)) +
+  coord_cartesian(ylim = c(0.58, 0.78)) +
+  labs(
+    x = "Prediction horizon, years",
+    y = "Mean AUC per run",
+    title = "Non-death landmark AUC across horizons"
+  ) +
+  theme(
+    plot.title = element_text(face = "bold", size = 10),
+    legend.position = c(0.31, 0.20),
+    legend.background = element_blank()
+  )
+
+save_panel(p2, "panel_02_landmark_by_horizon", width = 3.8, height = 3.0)
+
+# Panel 3: Delphi2M-style token AUC by offset. This documents why the old
+# evaluation can make next_token look competitive, especially near the event.
+token_offset_run <- token_auc_disease %>%
+  mutate(
+    offset = as.numeric(offset),
+    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))
+  ) %>%
+  group_by(run, model_target_mode, offset) %>%
+  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
+
+token_offset_summary <- token_offset_run %>%
+  group_by(model_target_mode, offset) %>%
+  summarise(mean_ci(mean_auc), .groups = "drop")
+
+p3 <- ggplot(token_offset_run, aes(x = offset, y = mean_auc, color = model_target_mode)) +
+  geom_line(aes(group = run), alpha = 0.18, linewidth = 0.35) +
+  geom_point(alpha = 0.32, size = 1.1) +
+  geom_ribbon(
+    data = token_offset_summary,
+    aes(x = offset, y = mean, ymin = ymin, ymax = ymax, fill = model_target_mode, group = model_target_mode),
+    alpha = 0.13,
+    color = NA,
+    inherit.aes = FALSE
+  ) +
+  geom_line(data = token_offset_summary, aes(y = mean), linewidth = 0.85) +
+  geom_point(data = token_offset_summary, aes(y = mean), size = 2.0) +
+  scale_color_manual(
+    values = target_cols,
+    labels = c("next_token", "all_future"),
+    name = NULL
+  ) +
+  scale_fill_manual(values = target_cols, guide = "none") +
+  scale_x_continuous(breaks = c(0.1, 1, 5, 10), trans = "log10") +
+  coord_cartesian(ylim = c(0.55, 0.82)) +
+  labs(
+    x = "Minimum offset before event, years",
+    y = "Mean AUC per run",
+    title = "Non-death token AUC by offset"
+  ) +
+  theme(
+    plot.title = element_text(face = "bold", size = 10),
+    legend.position = c(0.31, 0.20),
+    legend.background = element_blank()
+  )
+
+save_panel(p3, "panel_03_token_auc_by_offset", width = 3.8, height = 3.0)
+
+# Panel 4: within-run contrast between old token evaluation and landmark
+# evaluation. Each run contributes one point per evaluation family.
+run_eval_contrast <- auc_all %>%
+  group_by(run, model_target_mode, dist_mode, eval_family) %>%
+  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
+
+p4 <- ggplot(run_eval_contrast, aes(x = eval_family, y = mean_auc, color = model_target_mode)) +
+  geom_line(aes(group = run), alpha = 0.34, linewidth = 0.45) +
+  geom_point(aes(shape = dist_mode), size = 2.0, alpha = 0.84) +
+  stat_summary(
+    aes(group = model_target_mode),
+    fun = mean,
+    geom = "point",
+    size = 3.3,
+    shape = 18,
+    position = position_dodge(width = 0.16)
+  ) +
+  scale_color_manual(
+    values = target_cols,
+    labels = c("next_token", "all_future"),
+    name = NULL
+  ) +
+  scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
+  coord_cartesian(ylim = c(0.58, 0.78)) +
+  labs(
+    x = NULL,
+    y = "Mean AUC per run",
+    title = "Evaluation choice changes the conclusion (absolute time)"
+  ) +
+  theme(
+    plot.title = element_text(face = "bold", size = 10),
+    axis.text.x = element_text(angle = 18, hjust = 1),
+    legend.position = "right"
+  )
+
+save_panel(p4, "panel_04_evaluation_contrast", width = 4.3, height = 3.1)
+
+# Panel 5: disease-level distribution for the landmark task, pooled over
+# horizons and runs. This shows the shift without hiding heterogeneity.
+landmark_density <- landmark_auc_disease %>%
+  mutate(model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))) %>%
+  filter(is.finite(auc))
+
+p5 <- ggplot(landmark_density, aes(x = auc, fill = model_target_mode, color = model_target_mode)) +
+  geom_density(alpha = 0.20, linewidth = 0.65, adjust = 1.1) +
+  geom_vline(
+    data = landmark_density %>%
+      group_by(model_target_mode) %>%
+      summarise(mean_auc = mean(auc), .groups = "drop"),
+    aes(xintercept = mean_auc, color = model_target_mode),
+    linewidth = 0.75,
+    linetype = "22"
+  ) +
+  scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
+  scale_fill_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
+  coord_cartesian(xlim = c(0.35, 1.0)) +
+  labs(
+    x = "AUC",
+    y = "Density",
+    title = "Non-death landmark AUC distribution"
+  ) +
+  theme(
+    plot.title = element_text(face = "bold", size = 10),
+    legend.position = c(0.24, 0.82),
+    legend.background = element_blank()
+  )
+
+save_panel(p5, "panel_05_landmark_auc_distribution", width = 3.8, height = 3.0)
+
+# Panel 6: death-only fixed landmark + horizon AUC. Death has one endpoint token,
+# so each line is a run trajectory across horizons.
+death_landmark_run <- landmark_auc_death %>%
+  mutate(
+    horizon = as.numeric(horizon),
+    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
+    dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed"))
+  ) %>%
+  group_by(run, model_target_mode, dist_mode, horizon) %>%
+  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
+
+death_landmark_summary <- death_landmark_run %>%
+  group_by(model_target_mode, horizon) %>%
+  summarise(mean_ci(mean_auc), .groups = "drop")
+
+p6 <- ggplot(death_landmark_run, aes(x = horizon, y = mean_auc, color = model_target_mode)) +
+  geom_line(aes(group = run), alpha = 0.42, linewidth = 0.45) +
+  geom_point(aes(shape = dist_mode), alpha = 0.9, size = 2.0) +
+  geom_line(data = death_landmark_summary, aes(y = mean, group = model_target_mode), linewidth = 0.9) +
+  geom_point(data = death_landmark_summary, aes(y = mean), size = 2.2) +
+  scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
+  scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
+  scale_x_continuous(breaks = c(1, 5, 10)) +
+  coord_cartesian(ylim = c(0.58, 0.95)) +
+  labs(
+    x = "Prediction horizon, years",
+    y = "AUC",
+    title = "Death-only landmark AUC"
+  ) +
+  theme(
+    plot.title = element_text(face = "bold", size = 10),
+    legend.position = "right"
+  )
+
+save_panel(p6, "panel_06_death_landmark_by_horizon", width = 3.9, height = 3.0)
+
+# Panel 7: death-only Delphi2M-style token AUC by offset.
+death_token_run <- token_auc_death %>%
+  mutate(
+    offset = as.numeric(offset),
+    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
+    dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed"))
+  ) %>%
+  group_by(run, model_target_mode, dist_mode, offset) %>%
+  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
+
+death_token_summary <- death_token_run %>%
+  group_by(model_target_mode, offset) %>%
+  summarise(mean_ci(mean_auc), .groups = "drop")
+
+p7 <- ggplot(death_token_run, aes(x = offset, y = mean_auc, color = model_target_mode)) +
+  geom_line(aes(group = run), alpha = 0.42, linewidth = 0.45) +
+  geom_point(aes(shape = dist_mode), alpha = 0.9, size = 2.0) +
+  geom_line(data = death_token_summary, aes(y = mean, group = model_target_mode), linewidth = 0.9) +
+  geom_point(data = death_token_summary, aes(y = mean), size = 2.2) +
+  scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
+  scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
+  scale_x_continuous(breaks = c(0.1, 1, 5, 10), trans = "log10") +
+  coord_cartesian(ylim = c(0.58, 0.95)) +
+  labs(
+    x = "Minimum offset before event, years",
+    y = "AUC",
+    title = "Death-only token AUC"
+  ) +
+  theme(
+    plot.title = element_text(face = "bold", size = 10),
+    legend.position = "right"
+  )
+
+save_panel(p7, "panel_07_death_token_auc_by_offset", width = 3.9, height = 3.0)
+
+# Panel 8: death-only contrast between the two evaluation families.
+death_eval_contrast <- bind_rows(
+  landmark_auc_death %>% mutate(horizon = as.numeric(horizon), offset = NA_real_),
+  token_auc_death %>% mutate(horizon = NA_real_, offset = as.numeric(offset))
+) %>%
+  mutate(
+    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
+    eval_family = factor(eval_family, levels = c("Delphi2M-style token", "Fixed landmark + horizon")),
+    dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed"))
+  ) %>%
+  group_by(run, model_target_mode, dist_mode, eval_family) %>%
+  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")
+
+p8 <- ggplot(death_eval_contrast, aes(x = eval_family, y = mean_auc, color = model_target_mode)) +
+  geom_line(aes(group = run), alpha = 0.38, linewidth = 0.5) +
+  geom_point(aes(shape = dist_mode), size = 2.2, alpha = 0.9) +
+  stat_summary(
+    aes(group = model_target_mode),
+    fun = mean,
+    geom = "point",
+    size = 3.4,
+    shape = 18,
+    position = position_dodge(width = 0.16)
+  ) +
+  scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
+  scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
+  coord_cartesian(ylim = c(0.58, 0.95)) +
+  labs(
+    x = NULL,
+    y = "Mean AUC per run",
+    title = "Death endpoint evaluated separately"
+  ) +
+  theme(
+    plot.title = element_text(face = "bold", size = 10),
+    axis.text.x = element_text(angle = 18, hjust = 1),
+    legend.position = "right"
+  )
+
+save_panel(p8, "panel_08_death_evaluation_contrast", width = 4.3, height = 3.1)
+
+# Export the exact run-level summaries used by the figures.
+readr::write_csv(landmark_run, file.path(out_dir, "landmark_run_summary.csv"))
+readr::write_csv(token_offset_run, file.path(out_dir, "token_offset_run_summary.csv"))
+readr::write_csv(run_eval_contrast, file.path(out_dir, "run_evaluation_contrast.csv"))
+readr::write_csv(death_landmark_run, file.path(out_dir, "death_landmark_run_summary.csv"))
+readr::write_csv(death_token_run, file.path(out_dir, "death_token_offset_run_summary.csv"))
+readr::write_csv(death_eval_contrast, file.path(out_dir, "death_evaluation_contrast.csv"))
+
+message("Done. Panels are in: ", normalizePath(out_dir, winslash = "/"))

run_weibull_shape_exports.sh

@@ -0,0 +1,76 @@
+#!/usr/bin/env bash
+
+# Export Weibull shape-parameter summaries for the all_future models trained
+# with smoking/alcohol/BMI extra information.
+#
+# Bash 4.2 compatible. Run from the DeepHealth repository root on the Linux
+# server, for example:
+#
+#   bash run_weibull_shape_exports.sh
+#
+# Optional overrides:
+#   PYTHON=python3 DEVICE=cuda BATCH_SIZE=128 NUM_WORKERS=4 ROW_BATCH_SIZE=512 \
+#     bash run_weibull_shape_exports.sh
+
+set -euo pipefail
+
+PYTHON="${PYTHON:-python}"
+DEVICE="${DEVICE:-cuda}"
+BATCH_SIZE="${BATCH_SIZE:-128}"
+NUM_WORKERS="${NUM_WORKERS:-4}"
+ROW_BATCH_SIZE="${ROW_BATCH_SIZE:-512}"
+LANDMARK_START="${LANDMARK_START:-40}"
+LANDMARK_STOP="${LANDMARK_STOP:-80}"
+LANDMARK_STEP="${LANDMARK_STEP:-5}"
+HORIZONS="${HORIZONS:-1,5,10}"
+
+RUNS=(
+  "runs/relative_weibull_all_future_pure_disease_20260620_095229"
+  "runs/relative_mixed_all_future_pure_disease_20260620_132415"
+  "runs/absolute_weibull_all_future_pure_disease_20260620_114816"
+  "runs/absolute_mixed_all_future_pure_disease_20260620_161804"
+)
+
+echo "Python: ${PYTHON}"
+echo "Device: ${DEVICE}"
+echo "Batch size: ${BATCH_SIZE}"
+echo "Workers: ${NUM_WORKERS}"
+echo "Row batch size: ${ROW_BATCH_SIZE}"
+echo "Horizons: ${HORIZONS}"
+echo
+
+for run_path in "${RUNS[@]}"; do
+  if [[ ! -d "${run_path}" ]]; then
+    echo "[ERROR] Missing run directory: ${run_path}" >&2
+    exit 1
+  fi
+  if [[ ! -f "${run_path}/best_model.pt" ]]; then
+    echo "[ERROR] Missing checkpoint: ${run_path}/best_model.pt" >&2
+    exit 1
+  fi
+  if [[ ! -f "${run_path}/train_config.json" ]]; then
+    echo "[ERROR] Missing config: ${run_path}/train_config.json" >&2
+    exit 1
+  fi
+
+  output_path="${run_path}/weibull_shape_parameter_stats_test"
+  echo "=== Exporting Weibull shape stats: ${run_path} ==="
+  "${PYTHON}" export_weibull_shape_parameter_stats.py \
+    --run_path "${run_path}" \
+    --output_path "${output_path}" \
+    --eval_split test \
+    --device "${DEVICE}" \
+    --batch_size "${BATCH_SIZE}" \
+    --num_workers "${NUM_WORKERS}" \
+    --row_batch_size "${ROW_BATCH_SIZE}" \
+    --hidden_cache_dtype float32 \
+    --landmark_start "${LANDMARK_START}" \
+    --landmark_stop "${LANDMARK_STOP}" \
+    --landmark_step "${LANDMARK_STEP}" \
+    --horizons "${HORIZONS}" \
+    --include_all_token_rho_summary
+  echo "Wrote: ${output_path}"
+  echo
+done
+
+echo "All Weibull shape exports completed."