DeepHealth/evaluate_single_disease_mortality_attribution.py

"""Compute per-disease attribution to predicted mortality risk.

For each selected patient and landmark age, this script keeps only rows where
each scanned disease token has already occurred in the history. It then deletes
that historical disease token, re-queries the model, and reports both
differences and ratios on probability and cumulative-hazard scales. If
--disease is omitted, all disease tokens in the mapping are scanned.

Death is always token vocab_size - 1.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader
from tqdm.auto import tqdm

from evaluate_auc_v2 import (
    build_model_from_dataset,
    cfg_get,
    load_checkpoint_state_dict,
    load_json_config,
    load_model_state,
    resolve_dist_mode_for_checkpoint,
    resolve_eval_device,
    validate_dataset_metadata,
)
from evaluate_event_free_survival import (
    IndexedLandmarkDataset,
    LandmarkDataset,
    build_first_occurrence_maps_for_landmarks,
    build_group_ablated_slice,
    collate_indexed_landmark_fn,
    death_risk_for_batch,
    historical_counts_by_group,
    infer_landmark_hidden,
    load_eval_sequence_dataset,
    load_organ_groups,
    make_landmark_ages,
    make_occurred_mask,
    mortality_hazard_from_risk,
)
from future_risk import death_risk_from_probabilities, probabilities_from_logits


OUTPUT_COLUMNS = [
    "patient_id",
    "dataset_index",
    "eid",
    "sex",
    "landmark_age",
    "tau",
    "followup_end_time",
    "history_disease_count",
    "selected_disease_history_count",
    "selected_disease_token_id",
    "selected_disease_code",
    "selected_disease_name",
    "selected_disease_organ_system",
    "selected_disease_organ_system_label",
    "history_count__selected_organ_system",
    "death_risk",
    "death_hazard",
    "ablated_death_risk",
    "ablated_death_hazard",
    "mortality_attribution_probability",
    "mortality_attribution_hazard",
    "mortality_attribution_probability_ratio",
    "mortality_attribution_hazard_ratio",
]


def write_compressed_npz_table(path: Path, table: pd.DataFrame) -> int:
    table = table.reindex(columns=OUTPUT_COLUMNS)
    arrays: dict[str, np.ndarray] = {
        "__columns__": np.asarray(OUTPUT_COLUMNS, dtype="U"),
    }
    for column in OUTPUT_COLUMNS:
        values = table[column] if column in table else pd.Series([], dtype=object)
        if values.dtype == object:
            arrays[column] = values.fillna("").astype(str).to_numpy(dtype="U")
        else:
            arrays[column] = values.to_numpy()
    np.savez_compressed(path, **arrays)
    return int(len(table))


def normalize_output_dir(path: Path) -> Path:
    if path.suffix:
        return path.with_suffix(path.suffix + "_shards")
    return path


def write_manifest(
    output_dir: Path,
    *,
    rows: int,
    shards: list[dict[str, Any]],
    scanned_diseases: list[dict[str, Any]],
    eval_split: str,
    tau: float,
    landmark_start: float,
    landmark_stop: float,
    landmark_step: float,
) -> None:
    payload = {
        "format": "compressed_npz_shards",
        "columns": OUTPUT_COLUMNS,
        "rows": int(rows),
        "shards": shards,
        "scanned_diseases": scanned_diseases,
        "eval_split": eval_split,
        "tau": float(tau),
        "landmark_start": float(landmark_start),
        "landmark_stop": float(landmark_stop),
        "landmark_step": float(landmark_step),
    }
    with (output_dir / "manifest.json").open("w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)


def load_disease_metadata(
    mapping_path: Path,
    *,
    vocab_size: int,
) -> dict[int, dict[str, Any]]:
    if not mapping_path.exists():
        raise FileNotFoundError(f"Disease mapping file not found: {mapping_path}")
    table = pd.read_csv(mapping_path)
    required = {"token_id", "code", "name", "is_death"}
    missing = required - set(table.columns)
    if missing:
        raise ValueError(f"{mapping_path} is missing columns: {sorted(missing)}")

    death_idx = int(vocab_size) - 1
    out: dict[int, dict[str, Any]] = {}
    for row in table.itertuples(index=False):
        token = int(getattr(row, "token_id"))
        if token < 0 or token >= int(vocab_size) or token == death_idx:
            continue
        if int(getattr(row, "is_death")) == 1:
            continue
        meta = {
            "token_id": token,
            "code": str(getattr(row, "code")),
            "name": str(getattr(row, "name")),
        }
        for column in (
            "icd10_chapter",
            "icd10_chapter_title",
            "organ_system",
            "organ_system_label",
        ):
            if hasattr(row, column):
                meta[column] = str(getattr(row, column))
        out[token] = meta
    return out


def resolve_disease_token(
    value: str,
    metadata: dict[int, dict[str, Any]],
) -> tuple[int, dict[str, Any]]:
    text = str(value).strip()
    if text == "":
        raise ValueError("--disease must not be empty")

    if text.isdigit() or (text.startswith("-") and text[1:].isdigit()):
        token = int(text)
        if token not in metadata:
            raise ValueError(f"Disease token_id {token} was not found in the mapping")
        return token, metadata[token]

    lower = text.lower()
    exact = [
        (token, meta)
        for token, meta in metadata.items()
        if str(meta.get("code", "")).lower() == lower
        or str(meta.get("name", "")).lower() == lower
    ]
    if len(exact) == 1:
        return exact[0]
    if len(exact) > 1:
        raise ValueError(f"--disease={value!r} matched multiple diseases exactly")

    contains = [
        (token, meta)
        for token, meta in metadata.items()
        if lower in str(meta.get("code", "")).lower()
        or lower in str(meta.get("name", "")).lower()
    ]
    if len(contains) == 1:
        return contains[0]
    if not contains:
        raise ValueError(f"--disease={value!r} did not match any disease token")
    preview = ", ".join(
        f"{token}:{meta.get('code')} ({meta.get('name')})"
        for token, meta in contains[:10]
    )
    raise ValueError(
        f"--disease={value!r} matched {len(contains)} diseases; use token_id or code. "
        f"First matches: {preview}"
    )


def resolve_disease_tokens(
    value: str | None,
    metadata: dict[int, dict[str, Any]],
) -> list[tuple[int, dict[str, Any]]]:
    if value is None or str(value).strip() == "":
        return [(token, metadata[token]) for token in sorted(metadata)]
    out: list[tuple[int, dict[str, Any]]] = []
    seen: set[int] = set()
    for part in str(value).split(","):
        token, meta = resolve_disease_token(part, metadata)
        if token not in seen:
            out.append((token, meta))
            seen.add(token)
    return out


def safe_ratio(
    numerator: torch.Tensor,
    denominator: torch.Tensor,
    *,
    eps: float,
) -> torch.Tensor:
    return numerator / denominator.clamp_min(float(eps))


def output_name_for_run(run_path: Path, eval_split: str, tau: float, *, all_diseases: bool) -> Path:
    scope = "all_diseases" if all_diseases else "selected_diseases"
    return run_path / f"single_disease_mortality_attribution_{eval_split}_{scope}_tau{tau:g}y"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Compute per-disease model attribution to mortality risk."
    )
    parser.add_argument("--run_path", type=str, required=True)
    parser.add_argument(
        "--disease",
        type=str,
        default=None,
        help=(
            "Optional disease token_id, ICD-10 code, exact name, unambiguous name "
            "substring, or comma-separated list. If omitted, scan all disease tokens."
        ),
    )
    parser.add_argument(
        "--output_path",
        type=str,
        default=None,
        help="Output directory for compressed .npz shards.",
    )
    parser.add_argument("--organ_mapping_path", type=str, default="icd10_chapter_organ_mapping.csv")
    parser.add_argument("--eval_split", type=str, default=None)
    parser.add_argument("--dataset_subset_size", type=int, default=None)
    parser.add_argument("--train_eid_file", type=str, default=None)
    parser.add_argument("--val_eid_file", type=str, default=None)
    parser.add_argument("--test_eid_file", type=str, default=None)
    parser.add_argument("--landmark_start", type=float, default=40.0)
    parser.add_argument("--landmark_stop", type=float, default=80.0)
    parser.add_argument("--landmark_step", type=float, default=5.0)
    parser.add_argument("--tau", type=float, default=5.0)
    parser.add_argument("--min_history_events", type=int, default=None)
    parser.add_argument("--batch_size", type=int, default=None)
    parser.add_argument(
        "--attribution_batch_size",
        type=int,
        default=None,
        help="Forward batch size for disease-token ablation queries.",
    )
    parser.add_argument("--num_workers", type=int, default=None)
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument("--extra_info_types", type=str, default=None)
    parser.add_argument(
        "--ratio_eps",
        type=float,
        default=1e-7,
        help="Small lower bound for ratio denominators.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    run_path = Path(args.run_path)
    config_path = run_path / "train_config.json"
    checkpoint_path = run_path / "best_model.pt"
    if not config_path.exists():
        raise FileNotFoundError(f"train_config.json not found: {config_path}")
    if not checkpoint_path.exists():
        raise FileNotFoundError(f"best_model.pt not found: {checkpoint_path}")

    cfg = load_json_config(config_path)
    model_target_mode = str(cfg.get("model_target_mode", "next_token")).lower()
    if model_target_mode not in {"next_token", "all_future"}:
        raise ValueError(f"Unsupported model_target_mode: {model_target_mode!r}")

    target_mode = str(cfg.get("target_mode", "uts"))
    attn_mask_mode = str(
        cfg.get("attn_mask_mode", "non_strict_time" if target_mode == "uts" else "target_aware")
    )
    readout_name = str(
        cfg.get("readout_name", "same_time_group_end" if target_mode == "uts" else "token")
    )
    readout_reduce = str(cfg.get("readout_reduce", "mean"))

    dataset, subset_indices, eval_split, split_source = load_eval_sequence_dataset(args, cfg)
    validate_dataset_metadata(dataset, cfg)

    metadata = load_disease_metadata(
        Path(args.organ_mapping_path),
        vocab_size=int(dataset.vocab_size),
    )
    scanned_disease_items = resolve_disease_tokens(args.disease, metadata)
    if not scanned_disease_items:
        raise ValueError("No diseases selected for attribution")
    scanned_disease_tokens = [token for token, _meta in scanned_disease_items]
    scanned_disease_token_set = set(scanned_disease_tokens)

    landmark_ages = make_landmark_ages(
        float(args.landmark_start),
        float(args.landmark_stop),
        float(args.landmark_step),
    )
    tau = float(args.tau)
    if tau < 0:
        raise ValueError("tau must be non-negative")

    first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
        dataset,
        subset_indices,
    )
    death_idx = int(dataset.vocab_size) - 1
    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_idx],
    )

    organ_groups, _organ_labels, token_to_group = load_organ_groups(
        Path(args.organ_mapping_path),
        vocab_size=int(dataset.vocab_size),
    )
    group_names = sorted(organ_groups)

    state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu")
    dist_mode = resolve_dist_mode_for_checkpoint(str(cfg.get("dist_mode", "exponential")), state_dict)
    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()

    batch_size = int(cfg_get(args, cfg, "batch_size", 128))
    attribution_batch_size = int(
        cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 4, batch_size))
    )
    if attribution_batch_size <= 0:
        raise ValueError("attribution_batch_size must be positive")
    if float(args.ratio_eps) <= 0:
        raise ValueError("--ratio_eps must be positive")

    num_workers = int(cfg_get(args, cfg, "num_workers", 4))
    loader = DataLoader(
        IndexedLandmarkDataset(landmark_dataset),
        batch_size=batch_size,
        shuffle=False,
        collate_fn=collate_indexed_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,
    )

    output_path = (
        Path(args.output_path)
        if args.output_path
        else output_name_for_run(
            run_path,
            eval_split,
            tau,
            all_diseases=args.disease is None or str(args.disease).strip() == "",
        )
    )
    output_dir = normalize_output_dir(output_path)
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Eval split: {eval_split}")
    print(f"Split source: {split_source}")
    print(f"Selected patients: {len(subset_indices)}")
    print(f"Landmark ages: {landmark_ages.tolist()}")
    print(f"Tau: {tau:g} years")
    print(f"Dist mode: {dist_mode}")
    print(f"Device: {device}")
    print(f"Death token: {death_idx}")
    if len(scanned_disease_items) == len(metadata):
        print(f"Diseases: all mapped diseases ({len(scanned_disease_items)})")
    else:
        preview = ", ".join(
            f"{token}:{meta.get('code')}" for token, meta in scanned_disease_items[:10]
        )
        print(f"Diseases: {len(scanned_disease_items)} selected ({preview})")
    print(f"Landmark rows: {len(landmark_dataset)}")
    print(f"Attribution batch size: {attribution_batch_size}")
    print(f"Output directory: {output_dir}")

    written_rows = 0
    shard_index = 0
    shards: list[dict[str, Any]] = []
    pending_batch_chunks: list[Dict[str, torch.Tensor]] = []
    pending_meta_chunks: list[list[dict[str, Any]]] = []
    pending_n = 0

    def flush_pending() -> None:
        nonlocal written_rows, shard_index, pending_batch_chunks, pending_meta_chunks, pending_n
        if pending_n == 0:
            return

        ablated_batch = {
            key: torch.cat([chunk[key] for chunk in pending_batch_chunks], dim=0)
            for key in pending_batch_chunks[0]
        }
        meta_rows = [row for chunk in pending_meta_chunks for row in chunk]
        with torch.no_grad():
            ablated_risk = death_risk_for_batch(
                model=model,
                batch=ablated_batch,
                device=device,
                model_target_mode=model_target_mode,
                readout_name=readout_name,
                readout_reduce=readout_reduce,
                dist_mode=dist_mode,
                tau=tau,
            )
        ablated_hazard = mortality_hazard_from_risk(ablated_risk)
        orig_risk = torch.as_tensor(
            [row.pop("_death_risk") for row in meta_rows],
            dtype=ablated_risk.dtype,
            device=ablated_risk.device,
        )
        orig_hazard = torch.as_tensor(
            [row.pop("_death_hazard") for row in meta_rows],
            dtype=ablated_hazard.dtype,
            device=ablated_hazard.device,
        )

        attr_prob = orig_risk - ablated_risk
        attr_hazard = orig_hazard - ablated_hazard
        ratio_prob = safe_ratio(orig_risk, ablated_risk, eps=float(args.ratio_eps))
        ratio_hazard = safe_ratio(orig_hazard, ablated_hazard, eps=float(args.ratio_eps))

        for i, row in enumerate(meta_rows):
            row["death_risk"] = float(orig_risk[i].detach().cpu())
            row["death_hazard"] = float(orig_hazard[i].detach().cpu())
            row["ablated_death_risk"] = float(ablated_risk[i].detach().cpu())
            row["ablated_death_hazard"] = float(ablated_hazard[i].detach().cpu())
            row["mortality_attribution_probability"] = float(attr_prob[i].detach().cpu())
            row["mortality_attribution_hazard"] = float(attr_hazard[i].detach().cpu())
            row["mortality_attribution_probability_ratio"] = float(
                ratio_prob[i].detach().cpu()
            )
            row["mortality_attribution_hazard_ratio"] = float(
                ratio_hazard[i].detach().cpu()
            )

        table = pd.DataFrame(meta_rows).reindex(columns=OUTPUT_COLUMNS)
        shard_name = f"part-{shard_index:06d}.npz"
        shard_path = output_dir / shard_name
        shard_rows = write_compressed_npz_table(shard_path, table)
        shards.append({"file": shard_name, "rows": int(shard_rows)})
        shard_index += 1
        written_rows += len(meta_rows)
        pending_batch_chunks = []
        pending_meta_chunks = []
        pending_n = 0

    for batch in tqdm(loader, desc="Per-disease mortality attribution", dynamic_ncols=True):
        hidden = infer_landmark_hidden(
            model=model,
            batch=batch,
            device=device,
            model_target_mode=model_target_mode,
            readout_name=readout_name,
            readout_reduce=readout_reduce,
        )
        with torch.no_grad():
            logits = model.calc_risk(hidden)
            rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None
            death_rho = model.calc_death_rho(hidden) if dist_mode == "mixed" else None
            probabilities = probabilities_from_logits(
                logits,
                tau,
                dist_mode=dist_mode,
                rho=rho,
                death_rho=death_rho,
            )
        death_risk_tensor = death_risk_from_probabilities(probabilities)
        death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor)
        occurred = make_occurred_mask(
            batch["event_seq"].to(device),
            vocab_size=int(dataset.vocab_size),
            device=device,
        )
        event_values = batch["event_seq"].detach().cpu().numpy().reshape(-1)
        batch_disease_tokens = sorted(
            {
                int(token)
                for token in event_values.tolist()
                if int(token) in scanned_disease_token_set
            }
        )
        if not batch_disease_tokens:
            continue

        for disease_token in batch_disease_tokens:
            disease_meta = metadata[disease_token]
            active_rows = torch.nonzero(
                occurred[:, disease_token].to(dtype=torch.bool),
                as_tuple=False,
            ).flatten()
            if active_rows.numel() == 0:
                continue

            row_offset = 0
            while row_offset < int(active_rows.numel()):
                capacity = int(attribution_batch_size) - pending_n
                row_stop = min(int(active_rows.numel()), row_offset + capacity)
                row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
                ablated_chunk = build_group_ablated_slice(
                    batch=batch,
                    token_ids=[disease_token],
                    row_indices=row_indices,
                )

                meta_chunk: list[dict[str, Any]] = []
                row_ids = batch["row_idx"][row_indices.cpu()].cpu().numpy().astype(np.int64)
                for local_pos, row_idx in enumerate(row_ids.tolist()):
                    meta = landmark_dataset.rows[int(row_idx)]
                    dataset_index = int(meta["dataset_index"])
                    sample = dataset.samples[dataset_index]
                    hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
                    total_count, group_counts = historical_counts_by_group(
                        hist_tokens,
                        death_idx=death_idx,
                        token_to_group=token_to_group,
                        group_names=group_names,
                    )
                    disease_history_count = int((hist_tokens == int(disease_token)).sum())
                    if disease_history_count <= 0:
                        continue

                    orig_row = int(row_indices[local_pos].detach().cpu())
                    meta_chunk.append(
                        {
                            "patient_id": int(meta["patient_id"]),
                            "dataset_index": dataset_index,
                            "eid": int(sample.get("eid", -1)),
                            "sex": int(meta["sex"]),
                            "landmark_age": float(meta["landmark_age"]),
                            "tau": tau,
                            "followup_end_time": float(meta["followup_end_time"]),
                            "history_disease_count": int(total_count),
                            "selected_disease_history_count": disease_history_count,
                            "selected_disease_token_id": int(disease_token),
                            "selected_disease_code": str(disease_meta.get("code", "")),
                            "selected_disease_name": str(disease_meta.get("name", "")),
                            "selected_disease_organ_system": str(
                                disease_meta.get("organ_system", "")
                            ),
                            "selected_disease_organ_system_label": str(
                                disease_meta.get("organ_system_label", "")
                            ),
                            "history_count__selected_organ_system": int(
                                group_counts.get(str(disease_meta.get("organ_system", "")), 0)
                            ),
                            "_death_risk": float(death_risk_tensor[orig_row].detach().cpu()),
                            "_death_hazard": float(death_hazard_tensor[orig_row].detach().cpu()),
                        }
                    )

                if meta_chunk:
                    pending_batch_chunks.append(ablated_chunk)
                    pending_meta_chunks.append(meta_chunk)
                    pending_n += len(meta_chunk)
                row_offset = row_stop

                if pending_n >= int(attribution_batch_size):
                    flush_pending()

    flush_pending()
    if not shards:
        empty_path = output_dir / "part-000000.npz"
        write_compressed_npz_table(empty_path, pd.DataFrame(columns=OUTPUT_COLUMNS))
        shards.append({"file": empty_path.name, "rows": 0})
    write_manifest(
        output_dir,
        rows=written_rows,
        shards=shards,
        scanned_diseases=[
            {"token_id": int(token), **{k: v for k, v in meta.items() if k != "token_id"}}
            for token, meta in scanned_disease_items
        ],
        eval_split=eval_split,
        tau=tau,
        landmark_start=float(args.landmark_start),
        landmark_stop=float(args.landmark_stop),
        landmark_step=float(args.landmark_step),
    )
    print(f"Wrote {written_rows} rows in {len(shards)} shard(s) to {output_dir}")


if __name__ == "__main__":
    main()