From 00ccedf82f144e95d90526dedec1ed9b4e761190 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Mon, 29 Jun 2026 10:42:32 +0800 Subject: [PATCH] Report death distribution parameters for disease attribution --- ...te_single_disease_mortality_attribution.py | 203 ++++++++++-------- 1 file changed, 110 insertions(+), 93 deletions(-) diff --git a/evaluate_single_disease_mortality_attribution.py b/evaluate_single_disease_mortality_attribution.py index e9a7527..61ff7fc 100644 --- a/evaluate_single_disease_mortality_attribution.py +++ b/evaluate_single_disease_mortality_attribution.py @@ -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),