JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Scan all diseases for mortality attribution

Browse Source
This commit is contained in:
Jiarui Li
2026-06-27 15:30:38 +08:00
parent 128a8ce655
commit 34353d33d5
1 changed files with 201 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

165
evaluate_single_disease_mortality_attribution.py
View File

@@ -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)
def write_compressed_npz_table(path: Path, table: pd.DataFrame) -> int:
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:
if path.suffix.lower() == ".npz":
def normalize_output_dir(path: Path) -> Path:
if path.suffix:
return path.with_suffix(path.suffix + "_shards")
return path
return path.with_suffix(".npz")
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,6 +516,19 @@ def main() -> None:
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,
@@ -521,8 +603,25 @@ def main() -> None:
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__":