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 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 each scanned disease token has already occurred in the history. It then deletes
that historical disease token, re-queries the model, and reports both that historical disease token, re-queries the model, and reports the original
differences and ratios on probability and cumulative-hazard scales. If and ablated fitted death distribution parameters. If --disease is omitted, all
--disease is omitted, all disease tokens in the mapping are scanned. disease tokens in the mapping are scanned.
Death is always token vocab_size - 1. Death is always token vocab_size - 1.
""" """
@@ -18,6 +18,7 @@ from typing import Any, Dict
import numpy as np import numpy as np
import pandas as pd import pandas as pd
import torch import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
from tqdm.auto import tqdm from tqdm.auto import tqdm
@@ -36,15 +37,12 @@ from evaluate_event_free_survival import (
LandmarkDataset, LandmarkDataset,
build_first_occurrence_maps_for_landmarks, build_first_occurrence_maps_for_landmarks,
collate_indexed_landmark_fn, collate_indexed_landmark_fn,
death_risk_for_batch,
historical_counts_by_group, historical_counts_by_group,
infer_landmark_hidden, infer_landmark_hidden,
load_eval_sequence_dataset, load_eval_sequence_dataset,
load_organ_groups, load_organ_groups,
make_landmark_ages, 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 from targets import CHECKUP_IDX, PAD_IDX
@@ -54,7 +52,6 @@ OUTPUT_COLUMNS = [
"eid", "eid",
"sex", "sex",
"landmark_age", "landmark_age",
"tau",
"followup_end_time", "followup_end_time",
"history_disease_count", "history_disease_count",
"selected_disease_history_count", "selected_disease_history_count",
@@ -64,14 +61,13 @@ OUTPUT_COLUMNS = [
"selected_disease_organ_system", "selected_disease_organ_system",
"selected_disease_organ_system_label", "selected_disease_organ_system_label",
"history_count__selected_organ_system", "history_count__selected_organ_system",
"death_risk", "death_distribution",
"death_hazard", "original_death_lambda",
"ablated_death_risk", "ablated_death_lambda",
"ablated_death_hazard", "original_death_scale",
"mortality_attribution_probability", "ablated_death_scale",
"mortality_attribution_hazard", "original_death_shape",
"mortality_attribution_probability_ratio", "ablated_death_shape",
"mortality_attribution_hazard_ratio",
] ]
SUMMARY_KEY_COLUMNS = [ SUMMARY_KEY_COLUMNS = [
@@ -88,14 +84,15 @@ SUMMARY_MEAN_COLUMNS = [
"history_disease_count", "history_disease_count",
"selected_disease_history_count", "selected_disease_history_count",
"history_count__selected_organ_system", "history_count__selected_organ_system",
"death_risk", ]
"death_hazard",
"ablated_death_risk", SUMMARY_PARAMETER_COLUMNS = [
"ablated_death_hazard", "original_death_lambda",
"mortality_attribution_probability", "ablated_death_lambda",
"mortality_attribution_hazard", "original_death_scale",
"mortality_attribution_probability_ratio", "ablated_death_scale",
"mortality_attribution_hazard_ratio", "original_death_shape",
"ablated_death_shape",
] ]
@@ -128,7 +125,7 @@ def write_manifest(
summary_file: str, summary_file: str,
scanned_diseases: list[dict[str, Any]], scanned_diseases: list[dict[str, Any]],
eval_split: str, eval_split: str,
tau: float, dist_mode: str,
landmark_start: float, landmark_start: float,
landmark_stop: float, landmark_stop: float,
landmark_step: float, landmark_step: float,
@@ -141,7 +138,7 @@ def write_manifest(
"summary_file": summary_file, "summary_file": summary_file,
"scanned_diseases": scanned_diseases, "scanned_diseases": scanned_diseases,
"eval_split": eval_split, "eval_split": eval_split,
"tau": float(tau), "dist_mode": str(dist_mode),
"landmark_start": float(landmark_start), "landmark_start": float(landmark_start),
"landmark_stop": float(landmark_stop), "landmark_stop": float(landmark_stop),
"landmark_step": float(landmark_step), "landmark_step": float(landmark_step),
@@ -162,12 +159,23 @@ def update_summary_accumulator(
key = (key,) key = (key,)
acc = summary.setdefault( acc = summary.setdefault(
key, 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)) n = int(len(group))
acc["n"] += float(n) acc["n"] += float(n)
for column in SUMMARY_MEAN_COLUMNS: for column in SUMMARY_MEAN_COLUMNS:
acc[column] += float(pd.to_numeric(group[column], errors="coerce").sum()) 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( def write_summary_csv(
@@ -181,12 +189,23 @@ def write_summary_csv(
out["n"] = n out["n"] = n
for column in SUMMARY_MEAN_COLUMNS: for column in SUMMARY_MEAN_COLUMNS:
out[f"mean__{column}"] = acc[column] / n if n > 0 else np.nan 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) rows.append(out)
columns = [ columns = [
*SUMMARY_KEY_COLUMNS, *SUMMARY_KEY_COLUMNS,
"n", "n",
*[f"mean__{column}" for column in SUMMARY_MEAN_COLUMNS], *[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( pd.DataFrame(rows, columns=columns).sort_values(
["selected_disease_token_id", "landmark_age", "sex"], ["selected_disease_token_id", "landmark_age", "sex"],
@@ -365,23 +384,57 @@ def resolve_disease_tokens(
return out return out
def safe_ratio( def death_distribution_parameters(
numerator: torch.Tensor, model,
denominator: torch.Tensor, hidden: torch.Tensor,
*, *,
eps: float, dist_mode: str,
) -> torch.Tensor: eps: float = 1e-8,
return numerator / denominator.clamp_min(float(eps)) ) -> 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" 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: def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser( 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("--run_path", type=str, required=True)
parser.add_argument( 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_start", type=float, default=40.0)
parser.add_argument("--landmark_stop", type=float, default=80.0) parser.add_argument("--landmark_stop", type=float, default=80.0)
parser.add_argument("--landmark_step", type=float, default=5.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("--min_history_events", type=int, default=None)
parser.add_argument("--batch_size", type=int, default=None) parser.add_argument("--batch_size", type=int, default=None)
parser.add_argument( parser.add_argument(
@@ -420,12 +472,6 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--num_workers", type=int, default=None) parser.add_argument("--num_workers", type=int, default=None)
parser.add_argument("--device", type=str, default=None) parser.add_argument("--device", type=str, default=None)
parser.add_argument("--extra_info_types", 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( parser.add_argument(
"--shard_rows", "--shard_rows",
type=int, type=int,
@@ -476,9 +522,6 @@ def main() -> None:
float(args.landmark_stop), float(args.landmark_stop),
float(args.landmark_step), 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( first_occurrence_by_token = build_first_occurrence_maps_for_landmarks(
dataset, dataset,
@@ -504,6 +547,7 @@ def main() -> None:
state_dict = load_checkpoint_state_dict(checkpoint_path, map_location="cpu") 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) 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 = dict(cfg)
cfg_model["dist_mode"] = dist_mode cfg_model["dist_mode"] = dist_mode
device = resolve_eval_device(args.device) device = resolve_eval_device(args.device)
@@ -519,8 +563,6 @@ def main() -> None:
) )
if attribution_batch_size <= 0: if attribution_batch_size <= 0:
raise ValueError("attribution_batch_size must be positive") 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: if int(args.shard_rows) <= 0:
raise ValueError("--shard_rows must be positive") raise ValueError("--shard_rows must be positive")
@@ -542,7 +584,6 @@ def main() -> None:
else output_name_for_run( else output_name_for_run(
run_path, run_path,
eval_split, eval_split,
tau,
all_diseases=args.disease is None or str(args.disease).strip() == "", 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"Split source: {split_source}")
print(f"Selected patients: {len(subset_indices)}") print(f"Selected patients: {len(subset_indices)}")
print(f"Landmark ages: {landmark_ages.tolist()}") print(f"Landmark ages: {landmark_ages.tolist()}")
print(f"Tau: {tau:g} years")
print(f"Dist mode: {dist_mode}") print(f"Dist mode: {dist_mode}")
print(f"Device: {device}") print(f"Device: {device}")
print(f"Death token: {death_idx}") print(f"Death token: {death_idx}")
@@ -598,7 +638,6 @@ def main() -> None:
"eid": int(sample.get("eid", -1)), "eid": int(sample.get("eid", -1)),
"sex": int(meta["sex"]), "sex": int(meta["sex"]),
"landmark_age": float(meta["landmark_age"]), "landmark_age": float(meta["landmark_age"]),
"tau": tau,
"followup_end_time": float(meta["followup_end_time"]), "followup_end_time": float(meta["followup_end_time"]),
"history_disease_count": int(total_count), "history_disease_count": int(total_count),
"_hist_tokens": hist_tokens, "_hist_tokens": hist_tokens,
@@ -625,18 +664,11 @@ def main() -> None:
readout_reduce=readout_reduce, readout_reduce=readout_reduce,
) )
with torch.no_grad(): with torch.no_grad():
logits = model.calc_risk(hidden) _death_distribution, original_params = death_distribution_parameters(
rho = model.calc_weibull_rho(hidden) if dist_mode == "weibull" else None model,
death_rho = model.calc_death_rho(hidden) if dist_mode == "mixed" else None hidden,
probabilities = probabilities_from_logits(
logits,
tau,
dist_mode=dist_mode, 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() event_np = batch["event_seq"].numpy()
valid_event = (event_np >= 0) & (event_np < int(dataset.vocab_size)) valid_event = (event_np >= 0) & (event_np < int(dataset.vocab_size))
selected_event = np.zeros_like(valid_event, dtype=bool) selected_event = np.zeros_like(valid_event, dtype=bool)
@@ -659,35 +691,22 @@ def main() -> None:
token_ids=disease_token_ids, token_ids=disease_token_ids,
) )
with torch.no_grad(): with torch.no_grad():
ablated_risk = death_risk_for_batch( ablated_hidden = infer_landmark_hidden(
model=model, model=model,
batch=ablated_chunk, batch=ablated_chunk,
device=device, device=device,
model_target_mode=model_target_mode, model_target_mode=model_target_mode,
readout_name=readout_name, readout_name=readout_name,
readout_reduce=readout_reduce, readout_reduce=readout_reduce,
dist_mode=dist_mode,
tau=tau,
) )
orig_risk = death_risk_tensor[local_rows] _ablated_distribution, ablated_params = death_distribution_parameters(
orig_hazard = death_hazard_tensor[local_rows] model,
ablated_hazard = mortality_hazard_from_risk(ablated_risk) ablated_hidden,
attr_prob = orig_risk - ablated_risk dist_mode=dist_mode,
attr_hazard = orig_hazard - ablated_hazard )
ratio_prob = safe_ratio(orig_risk, ablated_risk, eps=float(args.ratio_eps)) value_block = parameter_pair_block(
ratio_hazard = safe_ratio(orig_hazard, ablated_hazard, eps=float(args.ratio_eps)) original_params[local_rows],
value_block = torch.stack( ablated_params,
[
orig_risk,
orig_hazard,
ablated_risk,
ablated_hazard,
attr_prob,
attr_hazard,
ratio_prob,
ratio_hazard,
],
dim=1,
).detach().cpu().numpy() ).detach().cpu().numpy()
row_ids = batch["row_idx"][local_rows_np].numpy().astype(np.int64, copy=False) row_ids = batch["row_idx"][local_rows_np].numpy().astype(np.int64, copy=False)
disease_tokens_list = disease_tokens_np disease_tokens_list = disease_tokens_np
@@ -726,7 +745,6 @@ def main() -> None:
"eid": row_base["eid"], "eid": row_base["eid"],
"sex": row_base["sex"], "sex": row_base["sex"],
"landmark_age": row_base["landmark_age"], "landmark_age": row_base["landmark_age"],
"tau": row_base["tau"],
"followup_end_time": row_base["followup_end_time"], "followup_end_time": row_base["followup_end_time"],
"history_disease_count": row_base["history_disease_count"], "history_disease_count": row_base["history_disease_count"],
"selected_disease_history_count": disease_history_count, "selected_disease_history_count": disease_history_count,
@@ -740,14 +758,13 @@ def main() -> None:
"history_count__selected_organ_system": int( "history_count__selected_organ_system": int(
group_counts.get(str(disease_meta.get("organ_system", "")), 0) group_counts.get(str(disease_meta.get("organ_system", "")), 0)
), ),
"death_risk": float(all_values[i, 0]), "death_distribution": death_distribution_name,
"death_hazard": float(all_values[i, 1]), "original_death_lambda": float(all_values[i, 0]),
"ablated_death_risk": float(all_values[i, 2]), "ablated_death_lambda": float(all_values[i, 1]),
"ablated_death_hazard": float(all_values[i, 3]), "original_death_scale": float(all_values[i, 2]),
"mortality_attribution_probability": float(all_values[i, 4]), "ablated_death_scale": float(all_values[i, 3]),
"mortality_attribution_hazard": float(all_values[i, 5]), "original_death_shape": float(all_values[i, 4]),
"mortality_attribution_probability_ratio": float(all_values[i, 6]), "ablated_death_shape": float(all_values[i, 5]),
"mortality_attribution_hazard_ratio": float(all_values[i, 7]),
} }
) )
@@ -787,7 +804,7 @@ def main() -> None:
for token, meta in scanned_disease_items for token, meta in scanned_disease_items
], ],
eval_split=eval_split, eval_split=eval_split,
tau=tau, dist_mode=dist_mode,
landmark_start=float(args.landmark_start), landmark_start=float(args.landmark_start),
landmark_stop=float(args.landmark_stop), landmark_stop=float(args.landmark_stop),
landmark_step=float(args.landmark_step), landmark_step=float(args.landmark_step),