852 lines
31 KiB
Python
852 lines
31 KiB
Python
"""Compute per-disease attribution to predicted mortality risk.
|
|
|
|
For each selected patient and landmark age, this script keeps only rows where
|
|
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
|
|
|
|
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,
|
|
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
|
|
from targets import CHECKUP_IDX, PAD_IDX
|
|
|
|
|
|
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",
|
|
]
|
|
|
|
SUMMARY_KEY_COLUMNS = [
|
|
"selected_disease_token_id",
|
|
"selected_disease_code",
|
|
"selected_disease_name",
|
|
"selected_disease_organ_system",
|
|
"selected_disease_organ_system_label",
|
|
"landmark_age",
|
|
"sex",
|
|
]
|
|
|
|
SUMMARY_MEAN_COLUMNS = [
|
|
"history_disease_count",
|
|
"selected_disease_history_count",
|
|
"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_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"),
|
|
}
|
|
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_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]],
|
|
summary_file: str,
|
|
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,
|
|
"summary_file": summary_file,
|
|
"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 update_summary_accumulator(
|
|
summary: dict[tuple[Any, ...], dict[str, float]],
|
|
table: pd.DataFrame,
|
|
) -> None:
|
|
if table.empty:
|
|
return
|
|
grouped = table.groupby(SUMMARY_KEY_COLUMNS, dropna=False, sort=False)
|
|
for key, group in grouped:
|
|
if not isinstance(key, tuple):
|
|
key = (key,)
|
|
acc = summary.setdefault(
|
|
key,
|
|
{"n": 0.0, **{column: 0.0 for column in SUMMARY_MEAN_COLUMNS}},
|
|
)
|
|
n = int(len(group))
|
|
acc["n"] += float(n)
|
|
for column in SUMMARY_MEAN_COLUMNS:
|
|
acc[column] += float(pd.to_numeric(group[column], errors="coerce").sum())
|
|
|
|
|
|
def write_summary_csv(
|
|
path: Path,
|
|
summary: dict[tuple[Any, ...], dict[str, float]],
|
|
) -> int:
|
|
rows: list[dict[str, Any]] = []
|
|
for key, acc in summary.items():
|
|
n = int(acc["n"])
|
|
out = {column: value for column, value in zip(SUMMARY_KEY_COLUMNS, key)}
|
|
out["n"] = n
|
|
for column in SUMMARY_MEAN_COLUMNS:
|
|
out[f"mean__{column}"] = acc[column] / n if n > 0 else np.nan
|
|
rows.append(out)
|
|
|
|
columns = [
|
|
*SUMMARY_KEY_COLUMNS,
|
|
"n",
|
|
*[f"mean__{column}" for column in SUMMARY_MEAN_COLUMNS],
|
|
]
|
|
pd.DataFrame(rows, columns=columns).sort_values(
|
|
["selected_disease_token_id", "landmark_age", "sex"],
|
|
kind="mergesort",
|
|
).to_csv(path, index=False)
|
|
return len(rows)
|
|
|
|
|
|
def concat_padded_tensor_batches(chunks: list[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
|
if not chunks:
|
|
raise ValueError("Cannot concatenate an empty chunk list")
|
|
|
|
fill_values = {
|
|
"event_seq": PAD_IDX,
|
|
"time_seq": 0.0,
|
|
"readout_mask": False,
|
|
"padding_mask": False,
|
|
"other_type": 0,
|
|
"other_value": 0.0,
|
|
"other_value_kind": 0,
|
|
"other_time": 0.0,
|
|
}
|
|
out: Dict[str, torch.Tensor] = {}
|
|
for key in chunks[0]:
|
|
tensors = [chunk[key] for chunk in chunks]
|
|
shapes = [tuple(t.shape) for t in tensors]
|
|
if len(set(shapes)) == 1:
|
|
out[key] = torch.cat(tensors, dim=0)
|
|
continue
|
|
|
|
if any(t.ndim == 0 for t in tensors):
|
|
raise ValueError(f"Cannot concatenate scalar tensor key={key!r} with mismatched shapes")
|
|
max_shape = list(shapes[0])
|
|
for shape in shapes[1:]:
|
|
if len(shape) != len(max_shape):
|
|
raise ValueError(f"Cannot concatenate key={key!r} with shapes {shapes}")
|
|
max_shape = [max(a, b) for a, b in zip(max_shape, shape)]
|
|
|
|
padded: list[torch.Tensor] = []
|
|
fill = fill_values.get(key, 0)
|
|
for tensor in tensors:
|
|
target_shape = [int(tensor.shape[0]), *max_shape[1:]]
|
|
padded_tensor = tensor.new_full(target_shape, fill)
|
|
slices = tuple(slice(0, int(size)) for size in tensor.shape)
|
|
padded_tensor[slices] = tensor
|
|
padded.append(padded_tensor)
|
|
out[key] = torch.cat(padded, dim=0)
|
|
return out
|
|
|
|
|
|
def build_disease_ablated_slice(
|
|
batch: Dict[str, torch.Tensor],
|
|
row_indices: torch.Tensor,
|
|
token_ids: torch.Tensor,
|
|
) -> Dict[str, torch.Tensor]:
|
|
"""Build an ablated slice for aligned (row, disease_token) pairs."""
|
|
event_seq = batch["event_seq"]
|
|
row_indices = row_indices.to(device=event_seq.device, dtype=torch.long)
|
|
token_ids = token_ids.to(device=event_seq.device, dtype=event_seq.dtype)
|
|
|
|
out: Dict[str, torch.Tensor] = {}
|
|
out["event_seq"] = event_seq[row_indices].clone()
|
|
out["time_seq"] = batch["time_seq"][row_indices]
|
|
out["readout_mask"] = batch["readout_mask"][row_indices].clone()
|
|
out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
|
|
out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()
|
|
|
|
seq_len = int(event_seq.shape[1])
|
|
positions = torch.arange(seq_len, device=event_seq.device)[None, :]
|
|
remove = (out["event_seq"] == token_ids[:, None]) & out["padding_mask"]
|
|
out["event_seq"] = torch.where(
|
|
remove,
|
|
torch.full_like(out["event_seq"], PAD_IDX),
|
|
out["event_seq"],
|
|
)
|
|
out["padding_mask"] &= ~remove
|
|
out["readout_mask"] &= ~remove
|
|
|
|
has_valid = out["padding_mask"].any(dim=1)
|
|
if not bool(has_valid.all().item()):
|
|
empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten()
|
|
out["event_seq"][empty_rows, 0] = CHECKUP_IDX
|
|
out["time_seq"][empty_rows, 0] = batch["t_query"][row_indices[empty_rows]].to(
|
|
dtype=out["time_seq"].dtype
|
|
)
|
|
out["padding_mask"][empty_rows, 0] = True
|
|
out["readout_mask"][empty_rows, 0] = True
|
|
out["landmark_pos"][empty_rows] = 0
|
|
|
|
has_readout = out["readout_mask"].any(dim=1)
|
|
if not bool(has_readout.all().item()):
|
|
rows = torch.nonzero(~has_readout, as_tuple=False).flatten()
|
|
local_valid = out["padding_mask"][rows]
|
|
last_pos = torch.where(
|
|
local_valid,
|
|
positions.expand(local_valid.shape[0], -1),
|
|
torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
|
|
).amax(dim=1)
|
|
out["readout_mask"][rows] = False
|
|
out["readout_mask"][rows, last_pos] = True
|
|
out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype)
|
|
|
|
repeated_keys = (
|
|
"sex",
|
|
"other_type",
|
|
"other_value",
|
|
"other_value_kind",
|
|
"other_time",
|
|
"t_query",
|
|
"patient_id",
|
|
"landmark_age",
|
|
"followup_end_time",
|
|
"death_time",
|
|
"row_idx",
|
|
)
|
|
for key in repeated_keys:
|
|
out[key] = batch[key][row_indices]
|
|
return out
|
|
|
|
|
|
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 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,
|
|
*,
|
|
eps: float,
|
|
) -> torch.Tensor:
|
|
return numerator / denominator.clamp_min(float(eps))
|
|
|
|
|
|
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 per-disease model attribution to mortality risk."
|
|
)
|
|
parser.add_argument("--run_path", type=str, required=True)
|
|
parser.add_argument(
|
|
"--disease",
|
|
type=str,
|
|
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("--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),
|
|
)
|
|
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]
|
|
|
|
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)
|
|
scanned_disease_tensor = torch.as_tensor(
|
|
scanned_disease_tokens,
|
|
dtype=torch.long,
|
|
device=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 * 32, 4096))
|
|
)
|
|
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 = (
|
|
Path(args.output_path)
|
|
if args.output_path
|
|
else output_name_for_run(
|
|
run_path,
|
|
eval_split,
|
|
tau,
|
|
all_diseases=args.disease is None or str(args.disease).strip() == "",
|
|
)
|
|
)
|
|
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}")
|
|
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}")
|
|
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 directory: {output_dir}")
|
|
|
|
written_rows = 0
|
|
shard_index = 0
|
|
shards: list[dict[str, Any]] = []
|
|
summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {}
|
|
row_base_cache: dict[int, dict[str, Any]] = {}
|
|
pending_batch_chunks: list[Dict[str, torch.Tensor]] = []
|
|
pending_meta_chunks: list[list[dict[str, Any]]] = []
|
|
pending_orig_risk_chunks: list[torch.Tensor] = []
|
|
pending_orig_hazard_chunks: list[torch.Tensor] = []
|
|
pending_n = 0
|
|
|
|
def flush_pending() -> None:
|
|
nonlocal written_rows, shard_index, pending_batch_chunks, pending_meta_chunks
|
|
nonlocal pending_orig_risk_chunks, pending_orig_hazard_chunks, pending_n
|
|
if pending_n == 0:
|
|
return
|
|
|
|
ablated_batch = concat_padded_tensor_batches(pending_batch_chunks)
|
|
meta_rows = [row for chunk in pending_meta_chunks for row in chunk]
|
|
orig_risk = torch.cat(pending_orig_risk_chunks, dim=0).to(device)
|
|
orig_hazard = torch.cat(pending_orig_hazard_chunks, dim=0).to(device)
|
|
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)
|
|
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))
|
|
value_block = torch.stack(
|
|
[
|
|
orig_risk,
|
|
orig_hazard,
|
|
ablated_risk,
|
|
ablated_hazard,
|
|
attr_prob,
|
|
attr_hazard,
|
|
ratio_prob,
|
|
ratio_hazard,
|
|
],
|
|
dim=1,
|
|
).detach().cpu().numpy()
|
|
|
|
for i, row in enumerate(meta_rows):
|
|
row["death_risk"] = float(value_block[i, 0])
|
|
row["death_hazard"] = float(value_block[i, 1])
|
|
row["ablated_death_risk"] = float(value_block[i, 2])
|
|
row["ablated_death_hazard"] = float(value_block[i, 3])
|
|
row["mortality_attribution_probability"] = float(value_block[i, 4])
|
|
row["mortality_attribution_hazard"] = float(value_block[i, 5])
|
|
row["mortality_attribution_probability_ratio"] = float(value_block[i, 6])
|
|
row["mortality_attribution_hazard_ratio"] = float(value_block[i, 7])
|
|
|
|
table = pd.DataFrame(meta_rows).reindex(columns=OUTPUT_COLUMNS)
|
|
update_summary_accumulator(summary_accumulator, table)
|
|
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_orig_risk_chunks = []
|
|
pending_orig_hazard_chunks = []
|
|
pending_n = 0
|
|
|
|
def get_row_base(row_idx: int) -> dict[str, Any]:
|
|
cached = row_base_cache.get(row_idx)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
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)
|
|
unique_tokens, token_counts = np.unique(hist_tokens, return_counts=True)
|
|
total_count, group_counts = historical_counts_by_group(
|
|
hist_tokens,
|
|
death_idx=death_idx,
|
|
token_to_group=token_to_group,
|
|
group_names=group_names,
|
|
)
|
|
cached = {
|
|
"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),
|
|
"_hist_tokens": hist_tokens,
|
|
"_token_counts": {
|
|
int(token): int(count)
|
|
for token, count in zip(unique_tokens.tolist(), token_counts.tolist())
|
|
},
|
|
"_group_counts": group_counts,
|
|
}
|
|
row_base_cache[row_idx] = cached
|
|
return cached
|
|
|
|
for batch in tqdm(loader, desc="Per-disease mortality attribution", dynamic_ncols=True):
|
|
batch_dev = {
|
|
k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
|
|
for k, v in batch.items()
|
|
}
|
|
hidden = infer_landmark_hidden(
|
|
model=model,
|
|
batch=batch_dev,
|
|
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_dev["event_seq"],
|
|
vocab_size=int(dataset.vocab_size),
|
|
device=device,
|
|
)
|
|
pair_indices = torch.nonzero(
|
|
occurred[:, scanned_disease_tensor],
|
|
as_tuple=False,
|
|
)
|
|
if pair_indices.numel() == 0:
|
|
continue
|
|
|
|
pair_offset = 0
|
|
while pair_offset < int(pair_indices.shape[0]):
|
|
capacity = int(attribution_batch_size) - pending_n
|
|
pair_stop = min(int(pair_indices.shape[0]), pair_offset + capacity)
|
|
pair_chunk = pair_indices[pair_offset:pair_stop]
|
|
local_rows = pair_chunk[:, 0].long()
|
|
disease_token_ids = scanned_disease_tensor[pair_chunk[:, 1]].long()
|
|
ablated_chunk = build_disease_ablated_slice(
|
|
batch=batch_dev,
|
|
row_indices=local_rows,
|
|
token_ids=disease_token_ids,
|
|
)
|
|
|
|
meta_chunk: list[dict[str, Any]] = []
|
|
row_ids = batch_dev["row_idx"][local_rows].detach().cpu().numpy().astype(np.int64)
|
|
disease_tokens_list = disease_token_ids.detach().cpu().numpy().astype(np.int64)
|
|
for local_pos, (row_idx, disease_token) in enumerate(
|
|
zip(row_ids.tolist(), disease_tokens_list.tolist())
|
|
):
|
|
disease_token = int(disease_token)
|
|
disease_meta = metadata[disease_token]
|
|
row_base = get_row_base(int(row_idx))
|
|
group_counts = row_base["_group_counts"]
|
|
disease_history_count = int(row_base["_token_counts"].get(disease_token, 0))
|
|
if disease_history_count <= 0:
|
|
raise RuntimeError(
|
|
"Internal mismatch: occurred mask selected disease "
|
|
f"{disease_token} for row {row_idx}, but cached history has count 0"
|
|
)
|
|
|
|
meta_chunk.append(
|
|
{
|
|
"patient_id": row_base["patient_id"],
|
|
"dataset_index": row_base["dataset_index"],
|
|
"eid": row_base["eid"],
|
|
"sex": row_base["sex"],
|
|
"landmark_age": row_base["landmark_age"],
|
|
"tau": row_base["tau"],
|
|
"followup_end_time": row_base["followup_end_time"],
|
|
"history_disease_count": row_base["history_disease_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)
|
|
),
|
|
}
|
|
)
|
|
|
|
if meta_chunk:
|
|
pending_batch_chunks.append(ablated_chunk)
|
|
pending_meta_chunks.append(meta_chunk)
|
|
pending_orig_risk_chunks.append(death_risk_tensor[local_rows].detach())
|
|
pending_orig_hazard_chunks.append(death_hazard_tensor[local_rows].detach())
|
|
pending_n += len(meta_chunk)
|
|
pair_offset = pair_stop
|
|
|
|
if pending_n >= int(attribution_batch_size):
|
|
flush_pending()
|
|
|
|
flush_pending()
|
|
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})
|
|
summary_path = output_dir / "summary_by_disease_age_sex.csv"
|
|
summary_rows = write_summary_csv(summary_path, summary_accumulator)
|
|
write_manifest(
|
|
output_dir,
|
|
rows=written_rows,
|
|
shards=shards,
|
|
summary_file=summary_path.name,
|
|
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}")
|
|
print(f"Wrote {summary_rows} summary rows to {summary_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|