From 34353d33d5fa37761c422b296a909560645d6e10 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Sat, 27 Jun 2026 15:30:38 +0800 Subject: [PATCH] Scan all diseases for mortality attribution --- ...te_single_disease_mortality_attribution.py | 303 ++++++++++++------ 1 file changed, 201 insertions(+), 102 deletions(-) diff --git a/evaluate_single_disease_mortality_attribution.py b/evaluate_single_disease_mortality_attribution.py index a1173d5..38fe9f5 100644 --- a/evaluate_single_disease_mortality_attribution.py +++ b/evaluate_single_disease_mortality_attribution.py @@ -1,15 +1,17 @@ -"""Compute single-disease attribution to predicted mortality risk. +"""Compute per-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. +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 @@ -74,13 +76,8 @@ OUTPUT_COLUMNS = [ ] -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) - +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"), } @@ -94,10 +91,38 @@ def write_compressed_npz(path: Path, frames: list[pd.DataFrame]) -> int: 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 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( @@ -184,6 +209,22 @@ def resolve_disease_token( ) +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, @@ -193,22 +234,31 @@ def safe_ratio( 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 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 single-disease model attribution to mortality risk." + description="Compute per-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.", + 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("--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) @@ -270,7 +320,11 @@ def main() -> None: Path(args.organ_mapping_path), vocab_size=int(dataset.vocab_size), ) - disease_token, disease_meta = resolve_disease_token(args.disease, metadata) + 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), @@ -333,12 +387,18 @@ def main() -> None: prefetch_factor=2 if num_workers > 0 else None, ) - output_path = normalize_npz_output_path( + output_path = ( Path(args.output_path) if args.output_path - else output_name_for_run(run_path, eval_split, disease_token, tau) + else output_name_for_run( + run_path, + eval_split, + tau, + all_diseases=args.disease is None or str(args.disease).strip() == "", + ) ) - output_path.parent.mkdir(parents=True, exist_ok=True) + 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}") @@ -348,22 +408,26 @@ def main() -> None: 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')}" - ) + 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: {output_path}") + print(f"Output directory: {output_dir}") written_rows = 0 - output_frames: list[pd.DataFrame] = [] + 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, pending_batch_chunks, pending_meta_chunks, pending_n + nonlocal written_rows, shard_index, pending_batch_chunks, pending_meta_chunks, pending_n if pending_n == 0: return @@ -414,13 +478,18 @@ def main() -> None: ratio_hazard[i].detach().cpu() ) - output_frames.append(pd.DataFrame(meta_rows).reindex(columns=OUTPUT_COLUMNS)) + 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="Single-disease mortality attribution", dynamic_ncols=True): + for batch in tqdm(loader, desc="Per-disease mortality attribution", dynamic_ncols=True): hidden = infer_landmark_hidden( model=model, batch=batch, @@ -447,82 +516,112 @@ def main() -> None: 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: + 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 - 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, - ) + 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 - 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()), - } + 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, ) - if meta_chunk: - pending_batch_chunks.append(ablated_chunk) - pending_meta_chunks.append(meta_chunk) - pending_n += len(meta_chunk) - row_offset = row_stop + 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 - if pending_n >= int(attribution_batch_size): - flush_pending() + 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 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__":