DeepHealth/export_weibull_death_parameter_stats.py

"""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()