Batch mortality attribution queries

evaluate_event_free_survival.py

@@ -288,14 +288,21 @@ def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[st
     }
 
 
-def ablate_event_history_for_tokens(
+def build_group_ablated_batch(
     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
+    group_names: Sequence[str],
+    organ_groups: dict[str, list[int]],
+    occurred: torch.Tensor,
+) -> tuple[Dict[str, torch.Tensor], list[str]]:
+    """Build a group-major batch for organ/system mortality attribution."""
+    active_groups: list[str] = []
+    for group in group_names:
+        ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device)
+        if ids.numel() > 0 and bool(occurred[:, ids].any().item()):
+            active_groups.append(group)
+
+    if not active_groups:
+        return {}, []
 
     event_rows: list[torch.Tensor] = []
     time_rows: list[torch.Tensor] = []
@@ -306,6 +313,8 @@ def ablate_event_history_for_tokens(
     time_seq = batch["time_seq"]
     readout_mask = batch["readout_mask"]
     padding_mask = batch["padding_mask"].bool()
+    for group in active_groups:
+        selected = {int(token) for token in organ_groups[group]}
         for i in range(event_seq.shape[0]):
             valid = padding_mask[i]
             events = event_seq[i, valid]
@@ -346,7 +355,8 @@ def ablate_event_history_for_tokens(
             readout_rows.append(kept_reads)
             landmark_positions.append(landmark_pos.to(dtype=batch["landmark_pos"].dtype))
 
-    out = dict(batch)
+    repeat_count = len(active_groups)
+    out: Dict[str, torch.Tensor] = {}
     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(
@@ -354,7 +364,33 @@ def ablate_event_history_for_tokens(
     )
     out["padding_mask"] = out["event_seq"] > PAD_IDX
     out["landmark_pos"] = torch.stack(landmark_positions)
-    return out
+
+    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:
+        value = batch[key]
+        repeats = (repeat_count,) + (1,) * (value.ndim - 1)
+        out[key] = value.repeat(repeats)
+    return out, active_groups
+
+
+def slice_tensor_batch(
+    batch: Dict[str, torch.Tensor],
+    start: int,
+    stop: int,
+) -> Dict[str, torch.Tensor]:
+    return {key: value[start:stop] for key, value in batch.items()}
 
 
 @torch.no_grad()
@@ -503,6 +539,12 @@ def parse_args() -> argparse.Namespace:
     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 expanded organ/system 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)
@@ -578,6 +620,11 @@ def main() -> None:
     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 * 4, batch_size))
+    )
+    if attribution_batch_size <= 0:
+        raise ValueError("attribution_batch_size must be positive")
     num_workers = int(cfg_get(args, cfg, "num_workers", 4))
     loader = DataLoader(
         IndexedLandmarkDataset(landmark_dataset),
@@ -602,6 +649,7 @@ def main() -> None:
     print(f"Death token: {death_idx}")
     print(f"Organ/system groups: {len(group_names)}")
     print(f"Landmark rows: {len(landmark_dataset)}")
+    print(f"Attribution batch size: {attribution_batch_size}")
     print(f"Output: {output_path}")
 
     rows: list[dict[str, Any]] = []
@@ -644,18 +692,31 @@ def main() -> None:
 
         group_mortality_attr_prob: dict[str, np.ndarray] = {}
         group_mortality_attr_hazard: dict[str, np.ndarray] = {}
+        batch_n = int(batch["event_seq"].shape[0])
+        zeros = np.zeros(batch_n, dtype=np.float32)
         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
+            group_mortality_attr_prob[group] = zeros.copy()
+            group_mortality_attr_hazard[group] = zeros.copy()
 
-            ablated_batch = ablate_event_history_for_tokens(batch, organ_groups[group])
-            ablated_death_risk = death_risk_for_batch(
+        ablated_batch, active_groups = build_group_ablated_batch(
+            batch=batch,
+            group_names=group_names,
+            organ_groups=organ_groups,
+            occurred=occurred,
+        )
+        if active_groups:
+            ablated_risk_chunks: list[torch.Tensor] = []
+            expanded_n = int(ablated_batch["event_seq"].shape[0])
+            for start in range(0, expanded_n, attribution_batch_size):
+                chunk = slice_tensor_batch(
+                    ablated_batch,
+                    start,
+                    min(start + attribution_batch_size, expanded_n),
+                )
+                ablated_risk_chunks.append(
+                    death_risk_for_batch(
                         model=model,
-                batch=ablated_batch,
+                        batch=chunk,
                         device=device,
                         model_target_mode=model_target_mode,
                         readout_name=readout_name,
@@ -663,13 +724,21 @@ def main() -> None:
                         dist_mode=dist_mode,
                         tau=tau,
                     )
+                )
+
+            ablated_death_risk = torch.cat(ablated_risk_chunks, dim=0).view(
+                len(active_groups), batch_n
+            )
             ablated_death_hazard = mortality_hazard_from_risk(ablated_death_risk)
-            group_mortality_attr_prob[group] = (
-                death_risk_tensor - ablated_death_risk
+            active_attr_hazard = (
+                death_hazard_tensor[None, :] - ablated_death_hazard
             ).detach().cpu().numpy()
-            group_mortality_attr_hazard[group] = (
-                death_hazard_tensor - ablated_death_hazard
+            active_attr_prob = (
+                death_risk_tensor[None, :] - ablated_death_risk
             ).detach().cpu().numpy()
+            for group_idx, group in enumerate(active_groups):
+                group_mortality_attr_prob[group] = active_attr_prob[group_idx]
+                group_mortality_attr_hazard[group] = active_attr_hazard[group_idx]
 
         row_indices = batch["row_idx"].cpu().numpy().astype(np.int64)
         for j, row_idx in enumerate(row_indices.tolist()):