Add mortality attribution evaluation

evaluate_event_free_survival.py

@@ -1,9 +1,12 @@
-"""Compute landmark future event-free survival summaries for DeepHealth.
+"""Compute landmark future death and incident system-disease risks.
 
 For each selected patient and landmark age, this script computes:
 
-* P(alive and no new modeled disease within tau years);
-* P(alive and no new disease in each ICD-10 chapter-derived system);
+* future death risk within tau years;
+* future incident disease risk for each ICD-10 chapter-derived system;
+* model attribution of each historical organ/system disease set to predicted
+  mortality risk, computed by deleting that system's historical disease tokens
+  and re-querying the model;
 * historical modeled-disease count;
 * historical modeled-disease count within each ICD-10 chapter-derived system.
 
@@ -38,8 +41,9 @@ from evaluate_auc_v2 import (
     resolve_eval_device,
     validate_dataset_metadata,
 )
-from future_event_free_survival import (
-    future_event_free_survival_from_probabilities,
+from future_risk import (
+    death_risk_from_probabilities,
+    new_disease_risk_from_probabilities,
     probabilities_from_logits,
 )
 from models import DeepHealth
@@ -284,6 +288,75 @@ def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[st
     }
 
 
+def ablate_event_history_for_tokens(
+    batch: Dict[str, torch.Tensor],
+    token_ids: Sequence[int],
+) -> Dict[str, torch.Tensor]:
+    """Return a batch with selected disease tokens removed from event history."""
+    selected = {int(token) for token in token_ids}
+    if not selected:
+        return batch
+
+    event_rows: list[torch.Tensor] = []
+    time_rows: list[torch.Tensor] = []
+    readout_rows: list[torch.Tensor] = []
+    landmark_positions: list[torch.Tensor] = []
+
+    event_seq = batch["event_seq"]
+    time_seq = batch["time_seq"]
+    readout_mask = batch["readout_mask"]
+    padding_mask = batch["padding_mask"].bool()
+    for i in range(event_seq.shape[0]):
+        valid = padding_mask[i]
+        events = event_seq[i, valid]
+        times = time_seq[i, valid]
+        reads = readout_mask[i, valid]
+        keep = torch.ones_like(events, dtype=torch.bool)
+        for token in selected:
+            keep &= events != int(token)
+
+        kept_events = events[keep]
+        kept_times = times[keep]
+        kept_reads = reads[keep]
+        if kept_events.numel() == 0:
+            kept_events = torch.tensor(
+                [CHECKUP_IDX],
+                dtype=event_seq.dtype,
+                device=event_seq.device,
+            )
+            kept_times = batch["t_query"][i : i + 1].to(
+                dtype=time_seq.dtype,
+                device=time_seq.device,
+            )
+            kept_reads = torch.ones(1, dtype=torch.bool, device=readout_mask.device)
+
+        if bool(kept_reads.any()):
+            landmark_pos = torch.nonzero(kept_reads, as_tuple=False)[-1, 0]
+        else:
+            landmark_pos = torch.tensor(
+                int(kept_events.numel() - 1),
+                dtype=batch["landmark_pos"].dtype,
+                device=batch["landmark_pos"].device,
+            )
+            kept_reads = torch.zeros_like(kept_events, dtype=torch.bool)
+            kept_reads[int(landmark_pos.item())] = True
+
+        event_rows.append(kept_events)
+        time_rows.append(kept_times)
+        readout_rows.append(kept_reads)
+        landmark_positions.append(landmark_pos.to(dtype=batch["landmark_pos"].dtype))
+
+    out = dict(batch)
+    out["event_seq"] = pad_sequence(event_rows, batch_first=True, padding_value=PAD_IDX)
+    out["time_seq"] = pad_sequence(time_rows, batch_first=True, padding_value=0.0)
+    out["readout_mask"] = pad_sequence(
+        readout_rows, batch_first=True, padding_value=False
+    )
+    out["padding_mask"] = out["event_seq"] > PAD_IDX
+    out["landmark_pos"] = torch.stack(landmark_positions)
+    return out
+
+
 @torch.no_grad()
 def infer_landmark_hidden(
     *,
@@ -351,6 +424,42 @@ def make_occurred_mask(
     return occurred
 
 
+def mortality_hazard_from_risk(risk: torch.Tensor, eps: float = 1e-7) -> torch.Tensor:
+    return -torch.log1p(-risk.clamp(0.0, 1.0 - float(eps)))
+
+
+def death_risk_for_batch(
+    *,
+    model: DeepHealth,
+    batch: Dict[str, torch.Tensor],
+    device: torch.device,
+    model_target_mode: str,
+    readout_name: str,
+    readout_reduce: str,
+    dist_mode: str,
+    tau: float,
+) -> torch.Tensor:
+    hidden = infer_landmark_hidden(
+        model=model,
+        batch=batch,
+        device=device,
+        model_target_mode=model_target_mode,
+        readout_name=readout_name,
+        readout_reduce=readout_reduce,
+    )
+    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,
+    )
+    return death_risk_from_probabilities(probabilities)
+
+
 def historical_counts_by_group(
     tokens: np.ndarray,
     *,
@@ -373,12 +482,12 @@ def historical_counts_by_group(
 
 
 def output_name_for_run(run_path: Path, eval_split: str, tau: float) -> Path:
-    return run_path / f"event_free_survival_{eval_split}_tau{tau:g}y.csv"
+    return run_path / f"future_risk_{eval_split}_tau{tau:g}y.csv"
 
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
-        description="Compute landmark event-free survival summaries."
+        description="Compute landmark death and incident system-disease risks."
     )
     parser.add_argument("--run_path", type=str, required=True)
     parser.add_argument("--output_path", type=str, default=None)
@@ -496,7 +605,7 @@ def main() -> None:
     print(f"Output: {output_path}")
 
     rows: list[dict[str, Any]] = []
-    for batch in tqdm(loader, desc="Event-free survival", dynamic_ncols=True):
+    for batch in tqdm(loader, desc="Future risks", dynamic_ncols=True):
         hidden = infer_landmark_hidden(
             model=model,
             batch=batch,
@@ -521,20 +630,45 @@ def main() -> None:
             device=device,
         )
 
-        all_survival = future_event_free_survival_from_probabilities(
-            probabilities,
-            occurred,
-            disease_ids=None,
-            vocab_size=int(dataset.vocab_size),
-        ).detach().cpu().numpy()
+        death_risk_tensor = death_risk_from_probabilities(probabilities)
+        death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor)
+        death_risk = death_risk_tensor.detach().cpu().numpy()
 
-        group_survival: dict[str, np.ndarray] = {}
+        group_risk: dict[str, np.ndarray] = {}
         for group in group_names:
-            group_survival[group] = future_event_free_survival_from_probabilities(
+            group_risk[group] = new_disease_risk_from_probabilities(
                 probabilities,
                 occurred,
-                disease_ids=organ_groups[group],
-                vocab_size=int(dataset.vocab_size),
+                organ_groups[group],
+            ).detach().cpu().numpy()
+
+        group_mortality_attr_prob: dict[str, np.ndarray] = {}
+        group_mortality_attr_hazard: dict[str, np.ndarray] = {}
+        for group in group_names:
+            ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=device)
+            if ids.numel() == 0 or not bool(occurred[:, ids].any().item()):
+                zeros = np.zeros(batch["event_seq"].shape[0], dtype=np.float32)
+                group_mortality_attr_prob[group] = zeros
+                group_mortality_attr_hazard[group] = zeros
+                continue
+
+            ablated_batch = ablate_event_history_for_tokens(batch, organ_groups[group])
+            ablated_death_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_death_hazard = mortality_hazard_from_risk(ablated_death_risk)
+            group_mortality_attr_prob[group] = (
+                death_risk_tensor - ablated_death_risk
+            ).detach().cpu().numpy()
+            group_mortality_attr_hazard[group] = (
+                death_hazard_tensor - ablated_death_hazard
             ).detach().cpu().numpy()
 
         row_indices = batch["row_idx"].cpu().numpy().astype(np.int64)
@@ -559,11 +693,17 @@ def main() -> None:
                 "tau": tau,
                 "followup_end_time": float(meta["followup_end_time"]),
                 "history_disease_count": int(total_count),
-                "event_free_survival_all": float(all_survival[j]),
+                "death_risk": float(death_risk[j]),
             }
             for group in group_names:
                 out[f"history_count__{group}"] = int(group_counts[group])
-                out[f"event_free_survival__{group}"] = float(group_survival[group][j])
+                out[f"new_disease_risk__{group}"] = float(group_risk[group][j])
+                out[f"mortality_attribution_probability__{group}"] = float(
+                    group_mortality_attr_prob[group][j]
+                )
+                out[f"mortality_attribution_hazard__{group}"] = float(
+                    group_mortality_attr_hazard[group][j]
+                )
             rows.append(out)
 
     df = pd.DataFrame(rows)