Add single-disease mortality attribution script
This commit is contained in:
529
evaluate_single_disease_mortality_attribution.py
Normal file
529
evaluate_single_disease_mortality_attribution.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user