Stream mortality attribution batches

evaluate_event_free_survival.py

@@ -288,82 +288,57 @@ def collate_indexed_landmark_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[st
     }
 
 
-def build_group_ablated_batch(
+def build_group_ablated_slice(
     batch: Dict[str, torch.Tensor],
-    group_names: Sequence[str],
+    token_ids: Sequence[int],
-    organ_groups: dict[str, list[int]],
+    row_start: int,
-    occurred: torch.Tensor,
+    row_stop: int,
-) -> tuple[Dict[str, torch.Tensor], list[str]]:
+) -> Dict[str, torch.Tensor]:
-    """Build a group-major batch for organ/system mortality attribution."""
+    """Build one fixed-width ablated slice without rebuilding variable-length rows."""
-    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] = []
-    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 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]
-            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))
-
-    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["event_seq"] = event_seq[row_start:row_stop].clone()
-    out["time_seq"] = pad_sequence(time_rows, batch_first=True, padding_value=0.0)
+    out["time_seq"] = batch["time_seq"][row_start:row_stop]
-    out["readout_mask"] = pad_sequence(
+    out["readout_mask"] = batch["readout_mask"][row_start:row_stop].clone()
-        readout_rows, batch_first=True, padding_value=False
+    out["padding_mask"] = batch["padding_mask"][row_start:row_stop].bool().clone()
+    out["landmark_pos"] = batch["landmark_pos"][row_start:row_stop].clone()
+
+    seq_len = int(event_seq.shape[1])
+    positions = torch.arange(seq_len, device=event_seq.device)[None, :]
+    ids = torch.as_tensor(token_ids, dtype=event_seq.dtype, device=event_seq.device)
+    remove = torch.isin(out["event_seq"], ids) & out["padding_mask"]
+    out["event_seq"] = torch.where(
+        remove,
+        torch.full_like(out["event_seq"], PAD_IDX),
+        out["event_seq"],
     )
-    out["padding_mask"] = out["event_seq"] > PAD_IDX
+    out["padding_mask"] &= ~remove
-    out["landmark_pos"] = torch.stack(landmark_positions)
+    out["readout_mask"] &= ~remove
+
+    has_valid = out["padding_mask"].any(dim=1)
+    if not bool(has_valid.all().item()):
+        empty_rows = torch.nonzero(~has_valid, as_tuple=False).flatten()
+        out["event_seq"][empty_rows, 0] = CHECKUP_IDX
+        out["time_seq"][empty_rows, 0] = batch["t_query"][
+            row_start + empty_rows
+        ].to(dtype=out["time_seq"].dtype)
+        out["padding_mask"][empty_rows, 0] = True
+        out["readout_mask"][empty_rows, 0] = True
+        out["landmark_pos"][empty_rows] = 0
+
+    has_readout = out["readout_mask"].any(dim=1)
+    if not bool(has_readout.all().item()):
+        rows = torch.nonzero(~has_readout, as_tuple=False).flatten()
+        local_valid = out["padding_mask"][rows]
+        last_pos = torch.where(
+            local_valid,
+            positions.expand(local_valid.shape[0], -1),
+            torch.zeros_like(positions.expand(local_valid.shape[0], -1)),
+        ).amax(dim=1)
+        out["readout_mask"][rows] = False
+        out["readout_mask"][rows, last_pos] = True
+        out["landmark_pos"][rows] = last_pos.to(dtype=out["landmark_pos"].dtype)
 
     repeated_keys = (
         "sex",
@@ -379,18 +354,62 @@ def build_group_ablated_batch(
         "row_idx",
     )
     for key in repeated_keys:
-        value = batch[key]
+        out[key] = batch[key][row_start:row_stop]
-        repeats = (repeat_count,) + (1,) * (value.ndim - 1)
+    return out
-        out[key] = value.repeat(repeats)
-    return out, active_groups
 
 
-def slice_tensor_batch(
+def concat_tensor_batches(chunks: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+    return {
+        key: torch.cat([chunk[key] for chunk in chunks], dim=0)
+        for key in chunks[0]
+    }
+
+
+def iter_group_ablated_batches(
     batch: Dict[str, torch.Tensor],
-    start: int,
+    group_names: Sequence[str],
-    stop: int,
+    organ_groups: dict[str, list[int]],
-) -> Dict[str, torch.Tensor]:
+    occurred: torch.Tensor,
-    return {key: value[start:stop] for key, value in batch.items()}
+    max_batch_size: int,
+):
+    """Yield ablated chunks as soon as enough rows are available for a forward pass."""
+    batch_n = int(batch["event_seq"].shape[0])
+    pending_batches: list[Dict[str, torch.Tensor]] = []
+    pending_groups: list[str] = []
+    pending_rows: list[int] = []
+    pending_n = 0
+
+    for group in group_names:
+        ids = torch.as_tensor(organ_groups[group], dtype=torch.long, device=occurred.device)
+        if ids.numel() == 0 or not bool(occurred[:, ids].any().item()):
+            continue
+
+        row_start = 0
+        while row_start < batch_n:
+            capacity = int(max_batch_size) - pending_n
+            row_stop = min(batch_n, row_start + capacity)
+            chunk = build_group_ablated_slice(
+                batch=batch,
+                token_ids=organ_groups[group],
+                row_start=row_start,
+                row_stop=row_stop,
+            )
+            chunk_n = int(row_stop - row_start)
+            pending_batches.append(chunk)
+            pending_groups.extend([group] * chunk_n)
+            pending_rows.extend(range(row_start, row_stop))
+            pending_n += chunk_n
+            row_start = row_stop
+
+            if pending_n >= int(max_batch_size):
+                yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
+                pending_batches = []
+                pending_groups = []
+                pending_rows = []
+                pending_n = 0
+
+    if pending_batches:
+        yield concat_tensor_batches(pending_batches), pending_groups, pending_rows
 
 
 @torch.no_grad()
@@ -646,6 +665,7 @@ def main() -> None:
     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}")
     print(f"Organ/system groups: {len(group_names)}")
     print(f"Landmark rows: {len(landmark_dataset)}")
@@ -698,47 +718,34 @@ def main() -> None:
             group_mortality_attr_prob[group] = zeros.copy()
             group_mortality_attr_hazard[group] = zeros.copy()
 
-        ablated_batch, active_groups = build_group_ablated_batch(
+        for ablated_chunk, chunk_groups, chunk_rows in iter_group_ablated_batches(
             batch=batch,
             group_names=group_names,
             organ_groups=organ_groups,
             occurred=occurred,
-        )
+            max_batch_size=attribution_batch_size,
-        if active_groups:
+        ):
-            ablated_risk_chunks: list[torch.Tensor] = []
+            ablated_death_risk = death_risk_for_batch(
-            expanded_n = int(ablated_batch["event_seq"].shape[0])
+                model=model,
-            for start in range(0, expanded_n, attribution_batch_size):
+                batch=ablated_chunk,
-                chunk = slice_tensor_batch(
+                device=device,
-                    ablated_batch,
+                model_target_mode=model_target_mode,
-                    start,
+                readout_name=readout_name,
-                    min(start + attribution_batch_size, expanded_n),
+                readout_reduce=readout_reduce,
-                )
+                dist_mode=dist_mode,
-                ablated_risk_chunks.append(
+                tau=tau,
-                    death_risk_for_batch(
-                        model=model,
-                        batch=chunk,
-                        device=device,
-                        model_target_mode=model_target_mode,
-                        readout_name=readout_name,
-                        readout_reduce=readout_reduce,
-                        dist_mode=dist_mode,
-                        tau=tau,
-                    )
-                )
-
-            ablated_death_risk = torch.cat(ablated_risk_chunks, dim=0).view(
-                len(active_groups), batch_n
             )
+            row_tensor = torch.as_tensor(chunk_rows, dtype=torch.long, device=device)
             ablated_death_hazard = mortality_hazard_from_risk(ablated_death_risk)
-            active_attr_hazard = (
+            attr_prob = (
-                death_hazard_tensor[None, :] - ablated_death_hazard
+                death_risk_tensor[row_tensor] - ablated_death_risk
             ).detach().cpu().numpy()
-            active_attr_prob = (
+            attr_hazard = (
-                death_risk_tensor[None, :] - ablated_death_risk
+                death_hazard_tensor[row_tensor] - ablated_death_hazard
             ).detach().cpu().numpy()
-            for group_idx, group in enumerate(active_groups):
+            for local_idx, (group, row_idx) in enumerate(zip(chunk_groups, chunk_rows)):
-                group_mortality_attr_prob[group] = active_attr_prob[group_idx]
+                group_mortality_attr_prob[group][row_idx] = attr_prob[local_idx]
-                group_mortality_attr_hazard[group] = active_attr_hazard[group_idx]
+                group_mortality_attr_hazard[group][row_idx] = attr_hazard[local_idx]
 
         row_indices = batch["row_idx"].cpu().numpy().astype(np.int64)
         for j, row_idx in enumerate(row_indices.tolist()):