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

Report death distribution parameters for disease attribution

Browse Source
This commit is contained in:
Jiarui Li
2026-06-29 10:42:32 +08:00
parent fe826c4052
commit 00ccedf82f
1 changed files with 110 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files

203
evaluate_single_disease_mortality_attribution.py
View File

@@ -1,10 +1,10 @@
"""Compute per-disease attribution to predicted mortality risk.
"""Compute per-disease attribution to predicted mortality distribution parameters.
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.
that historical disease token, re-queries the model, and reports the original
and ablated fitted death distribution parameters. If --disease is omitted, all
disease tokens in the mapping are scanned.
Death is always token vocab_size - 1.
"""
@@ -18,6 +18,7 @@ from typing import Any, Dict
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
@@ -36,15 +37,12 @@ from evaluate_event_free_survival import (
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,
mortality_hazard_from_risk,
)
from future_risk import death_risk_from_probabilities, probabilities_from_logits
from targets import CHECKUP_IDX, PAD_IDX
@@ -54,7 +52,6 @@ OUTPUT_COLUMNS = [
"eid",
"sex",
"landmark_age",
"tau",
"followup_end_time",
"history_disease_count",
"selected_disease_history_count",
@@ -64,14 +61,13 @@ OUTPUT_COLUMNS = [
"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",
"death_distribution",
"original_death_lambda",
"ablated_death_lambda",
"original_death_scale",
"ablated_death_scale",
"original_death_shape",
"ablated_death_shape",
]
SUMMARY_KEY_COLUMNS = [
@@ -88,14 +84,15 @@ 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",
]
SUMMARY_PARAMETER_COLUMNS = [
"original_death_lambda",
"ablated_death_lambda",
"original_death_scale",
"ablated_death_scale",
"original_death_shape",
"ablated_death_shape",
]
@@ -128,7 +125,7 @@ def write_manifest(
summary_file: str,
scanned_diseases: list[dict[str, Any]],
eval_split: str,
tau: float,
dist_mode: str,
landmark_start: float,
landmark_stop: float,
landmark_step: float,
@@ -141,7 +138,7 @@ def write_manifest(
"summary_file": summary_file,
"scanned_diseases": scanned_diseases,
"eval_split": eval_split,
"tau": float(tau),
"dist_mode": str(dist_mode),
"landmark_start": float(landmark_start),
"landmark_stop": float(landmark_stop),
"landmark_step": float(landmark_step),
@@ -162,12 +159,23 @@ def update_summary_accumulator(
key = (key,)
acc = summary.setdefault(
key,
{"n": 0.0, **{column: 0.0 for column in SUMMARY_MEAN_COLUMNS}},
{
"n": 0.0,
**{column: 0.0 for column in SUMMARY_MEAN_COLUMNS},
**{f"count__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
**{f"sum__{column}": 0.0 for column in SUMMARY_PARAMETER_COLUMNS},
**{f"sumsq__{column}": 0.0 for column in SUMMARY_PARAMETER_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())
for column in SUMMARY_PARAMETER_COLUMNS:
values = pd.to_numeric(group[column], errors="coerce").dropna()
acc[f"count__{column}"] += float(len(values))
acc[f"sum__{column}"] += float(values.sum())
acc[f"sumsq__{column}"] += float((values * values).sum())
def write_summary_csv(
@@ -181,12 +189,23 @@ def write_summary_csv(
out["n"] = n
for column in SUMMARY_MEAN_COLUMNS:
out[f"mean__{column}"] = acc[column] / n if n > 0 else np.nan
for column in SUMMARY_PARAMETER_COLUMNS:
count = int(acc[f"count__{column}"])
mean = acc[f"sum__{column}"] / count if count > 0 else np.nan
second = acc[f"sumsq__{column}"] / count if count > 0 else np.nan
out[f"mean__{column}"] = mean
out[f"var__{column}"] = second - mean * mean if count > 0 else np.nan
rows.append(out)
columns = [
*SUMMARY_KEY_COLUMNS,
"n",
*[f"mean__{column}" for column in SUMMARY_MEAN_COLUMNS],
*[
name
for column in SUMMARY_PARAMETER_COLUMNS
for name in (f"mean__{column}", f"var__{column}")
],
]
pd.DataFrame(rows, columns=columns).sort_values(
["selected_disease_token_id", "landmark_age", "sex"],
@@ -365,23 +384,57 @@ def resolve_disease_tokens(
return out
def safe_ratio(
numerator: torch.Tensor,
denominator: torch.Tensor,
def death_distribution_parameters(
model,
hidden: torch.Tensor,
*,
eps: float,
) -> torch.Tensor:
return numerator / denominator.clamp_min(float(eps))
dist_mode: str,
eps: float = 1e-8,
) -> tuple[str, torch.Tensor]:
"""Return death distribution parameters with columns matching PARAMETER_VALUE_COLUMNS."""
logits = model.calc_risk(hidden)
death_idx = int(logits.shape[1]) - 1
death_lambda = F.softplus(logits[:, death_idx]) + float(eps)
if dist_mode == "exponential":
nan = torch.full_like(death_lambda, float("nan"))
return "exponential", torch.stack([death_lambda, nan, nan], dim=1)
if dist_mode == "weibull":
rho = model.calc_weibull_rho(hidden)[:, death_idx].to(dtype=death_lambda.dtype)
elif dist_mode == "mixed":
rho = model.calc_death_rho(hidden).to(dtype=death_lambda.dtype)
else:
raise ValueError(f"Unsupported dist_mode={dist_mode!r}")
shape = rho.clamp_min(float(eps))
scale = torch.pow(death_lambda.clamp_min(float(eps)), -1.0 / shape)
nan = torch.full_like(death_lambda, float("nan"))
return "weibull", torch.stack([nan, scale, shape], dim=1)
def output_name_for_run(run_path: Path, eval_split: str, tau: float, *, all_diseases: bool) -> Path:
def parameter_pair_block(original: torch.Tensor, ablated: torch.Tensor) -> torch.Tensor:
return torch.stack(
[
original[:, 0],
ablated[:, 0],
original[:, 1],
ablated[:, 1],
original[:, 2],
ablated[:, 2],
],
dim=1,
)
def output_name_for_run(run_path: Path, eval_split: str, *, 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"
return run_path / f"single_disease_mortality_parameters_{eval_split}_{scope}"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Compute per-disease model attribution to mortality risk."
description="Compute per-disease model attribution to mortality distribution parameters."
)
parser.add_argument("--run_path", type=str, required=True)
parser.add_argument(
@@ -408,7 +461,6 @@ def parse_args() -> argparse.Namespace:
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(
@@ -420,12 +472,6 @@ def parse_args() -> argparse.Namespace:
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.",
)
parser.add_argument(
"--shard_rows",
type=int,
@@ -476,9 +522,6 @@ def main() -> None:
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,
@@ -504,6 +547,7 @@ def main() -> None:
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)
death_distribution_name = "exponential" if dist_mode == "exponential" else "weibull"
cfg_model = dict(cfg)
cfg_model["dist_mode"] = dist_mode
device = resolve_eval_device(args.device)
@@ -519,8 +563,6 @@ def main() -> None:
)
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")
if int(args.shard_rows) <= 0:
raise ValueError("--shard_rows must be positive")
@@ -542,7 +584,6 @@ def main() -> None:
else output_name_for_run(
run_path,
eval_split,
tau,
all_diseases=args.disease is None or str(args.disease).strip() == "",
)
)
@@ -553,7 +594,6 @@ def main() -> None:
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}")
@@ -598,7 +638,6 @@ def main() -> None:
"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,
@@ -625,18 +664,11 @@ def main() -> None:
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,
_death_distribution, original_params = death_distribution_parameters(
model,
hidden,
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)
event_np = batch["event_seq"].numpy()
valid_event = (event_np >= 0) & (event_np < int(dataset.vocab_size))
selected_event = np.zeros_like(valid_event, dtype=bool)
@@ -659,35 +691,22 @@ def main() -> None:
token_ids=disease_token_ids,
)
with torch.no_grad():
ablated_risk = death_risk_for_batch(
ablated_hidden = infer_landmark_hidden(
model=model,
batch=ablated_chunk,
device=device,
model_target_mode=model_target_mode,
readout_name=readout_name,
readout_reduce=readout_reduce,
dist_mode=dist_mode,
tau=tau,
)
orig_risk = death_risk_tensor[local_rows]
orig_hazard = death_hazard_tensor[local_rows]
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,
_ablated_distribution, ablated_params = death_distribution_parameters(
model,
ablated_hidden,
dist_mode=dist_mode,
)
value_block = parameter_pair_block(
original_params[local_rows],
ablated_params,
).detach().cpu().numpy()
row_ids = batch["row_idx"][local_rows_np].numpy().astype(np.int64, copy=False)
disease_tokens_list = disease_tokens_np
@@ -726,7 +745,6 @@ def main() -> None:
"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,
@@ -740,14 +758,13 @@ def main() -> None:
"history_count__selected_organ_system": int(
group_counts.get(str(disease_meta.get("organ_system", "")), 0)
),
"death_risk": float(all_values[i, 0]),
"death_hazard": float(all_values[i, 1]),
"ablated_death_risk": float(all_values[i, 2]),
"ablated_death_hazard": float(all_values[i, 3]),
"mortality_attribution_probability": float(all_values[i, 4]),
"mortality_attribution_hazard": float(all_values[i, 5]),
"mortality_attribution_probability_ratio": float(all_values[i, 6]),
"mortality_attribution_hazard_ratio": float(all_values[i, 7]),
"death_distribution": death_distribution_name,
"original_death_lambda": float(all_values[i, 0]),
"ablated_death_lambda": float(all_values[i, 1]),
"original_death_scale": float(all_values[i, 2]),
"ablated_death_scale": float(all_values[i, 3]),
"original_death_shape": float(all_values[i, 4]),
"ablated_death_shape": float(all_values[i, 5]),
}
)
@@ -787,7 +804,7 @@ def main() -> None:
for token, meta in scanned_disease_items
],
eval_split=eval_split,
tau=tau,
dist_mode=dist_mode,
landmark_start=float(args.landmark_start),
landmark_stop=float(args.landmark_stop),
landmark_step=float(args.landmark_step),