diff --git a/evaluate_single_disease_mortality_attribution.py b/evaluate_single_disease_mortality_attribution.py new file mode 100644 index 0000000..a1173d5 --- /dev/null +++ b/evaluate_single_disease_mortality_attribution.py @@ -0,0 +1,529 @@ +"""Compute single-disease attribution to predicted mortality risk. + +For each selected patient and landmark age, this script keeps only rows where +the requested 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. + +Death is always token vocab_size - 1. +""" +from __future__ import annotations + +import argparse +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(path: Path, frames: list[pd.DataFrame]) -> int: + if frames: + table = pd.concat(frames, ignore_index=True) + table = table.reindex(columns=OUTPUT_COLUMNS) + else: + table = pd.DataFrame(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_npz_output_path(path: Path) -> Path: + if path.suffix.lower() == ".npz": + return path + return path.with_suffix(".npz") + + +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 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, token_id: int, tau: float) -> Path: + return run_path / f"single_disease_mortality_attribution_{eval_split}_token{token_id}_tau{tau:g}y.npz" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Compute single-disease model attribution to mortality risk." + ) + parser.add_argument("--run_path", type=str, required=True) + parser.add_argument( + "--disease", + type=str, + required=True, + help="Disease token_id, ICD-10 code, exact name, or unambiguous name substring.", + ) + parser.add_argument("--output_path", type=str, default=None, help="Compressed .npz output path.") + 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), + ) + disease_token, disease_meta = resolve_disease_token(args.disease, metadata) + + 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 = normalize_npz_output_path( + Path(args.output_path) + if args.output_path + else output_name_for_run(run_path, eval_split, disease_token, tau) + ) + output_path.parent.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}") + print( + "Disease: " + f"token={disease_token}, code={disease_meta.get('code')}, name={disease_meta.get('name')}" + ) + print(f"Landmark rows: {len(landmark_dataset)}") + print(f"Attribution batch size: {attribution_batch_size}") + print(f"Output: {output_path}") + + written_rows = 0 + output_frames: list[pd.DataFrame] = [] + 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, 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() + ) + + output_frames.append(pd.DataFrame(meta_rows).reindex(columns=OUTPUT_COLUMNS)) + written_rows += len(meta_rows) + pending_batch_chunks = [] + pending_meta_chunks = [] + pending_n = 0 + + for batch in tqdm(loader, desc="Single-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, + ) + 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() + written_rows = write_compressed_npz(output_path, output_frames) + print(f"Wrote {written_rows} rows to {output_path}") + + +if __name__ == "__main__": + main()