Scan all diseases for mortality attribution

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
+each scanned disease token has already occurred in the history. It then deletes
-deletes that historical disease token, re-queries the model, and reports both
+that historical disease token, re-queries the model, and reports both
-differences and ratios on probability and cumulative-hazard scales.
+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:
+def write_compressed_npz_table(path: Path, table: pd.DataFrame) -> int:
-    if frames:
+    table = table.reindex(columns=OUTPUT_COLUMNS)
-        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"),
     }
@@ -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:
+def normalize_output_dir(path: Path) -> Path:
-    if path.suffix.lower() == ".npz":
+    if path.suffix:
-        return path
+        return path.with_suffix(path.suffix + "_shards")
-    return path.with_suffix(".npz")
+    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:
+def output_name_for_run(run_path: Path, eval_split: str, tau: float, *, all_diseases: bool) -> Path:
-    return run_path / f"single_disease_mortality_attribution_{eval_split}_token{token_id}_tau{tau:g}y.npz"
+    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,
+        default=None,
-        help="Disease token_id, ICD-10 code, exact name, or unambiguous name substring.",
+        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(
+    if len(scanned_disease_items) == len(metadata):
-        "Disease: "
+        print(f"Diseases: all mapped diseases ({len(scanned_disease_items)})")
-        f"token={disease_token}, code={disease_meta.get('code')}, name={disease_meta.get('name')}"
+    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(
+        event_values = batch["event_seq"].detach().cpu().numpy().reshape(-1)
-            occurred[:, disease_token].to(dtype=torch.bool),
+        batch_disease_tokens = sorted(
-            as_tuple=False,
+            {
-        ).flatten()
+                int(token)
-        if active_rows.numel() == 0:
+                for token in event_values.tolist()
+                if int(token) in scanned_disease_token_set
+            }
+        )
+        if not batch_disease_tokens:
             continue
 
-        row_offset = 0
+        for disease_token in batch_disease_tokens:
-        while row_offset < int(active_rows.numel()):
+            disease_meta = metadata[disease_token]
-            capacity = int(attribution_batch_size) - pending_n
+            active_rows = torch.nonzero(
-            row_stop = min(int(active_rows.numel()), row_offset + capacity)
+                occurred[:, disease_token].to(dtype=torch.bool),
-            row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
+                as_tuple=False,
-            ablated_chunk = build_group_ablated_slice(
+            ).flatten()
-                batch=batch,
+            if active_rows.numel() == 0:
-                token_ids=[disease_token],
+                continue
-                row_indices=row_indices,
-            )
 
-            meta_chunk: list[dict[str, Any]] = []
+            row_offset = 0
-            row_ids = batch["row_idx"][row_indices.cpu()].cpu().numpy().astype(np.int64)
+            while row_offset < int(active_rows.numel()):
-            for local_pos, row_idx in enumerate(row_ids.tolist()):
+                capacity = int(attribution_batch_size) - pending_n
-                meta = landmark_dataset.rows[int(row_idx)]
+                row_stop = min(int(active_rows.numel()), row_offset + capacity)
-                dataset_index = int(meta["dataset_index"])
+                row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
-                sample = dataset.samples[dataset_index]
+                ablated_chunk = build_group_ablated_slice(
-                hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
+                    batch=batch,
-                total_count, group_counts = historical_counts_by_group(
+                    token_ids=[disease_token],
-                    hist_tokens,
+                    row_indices=row_indices,
-                    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:
+                meta_chunk: list[dict[str, Any]] = []
-                pending_batch_chunks.append(ablated_chunk)
+                row_ids = batch["row_idx"][row_indices.cpu()].cpu().numpy().astype(np.int64)
-                pending_meta_chunks.append(meta_chunk)
+                for local_pos, row_idx in enumerate(row_ids.tolist()):
-                pending_n += len(meta_chunk)
+                    meta = landmark_dataset.rows[int(row_idx)]
-            row_offset = row_stop
+                    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):
+                    orig_row = int(row_indices[local_pos].detach().cpu())
-                flush_pending()
+                    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)
+    if not shards:
-    print(f"Wrote {written_rows} rows to {output_path}")
+        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__":