Vectorize per-disease attribution batches

evaluate_single_disease_mortality_attribution.py

@@ -35,7 +35,6 @@ from evaluate_event_free_survival import (
     IndexedLandmarkDataset,
     LandmarkDataset,
     build_first_occurrence_maps_for_landmarks,
-    build_group_ablated_slice,
     collate_indexed_landmark_fn,
     death_risk_for_batch,
     historical_counts_by_group,
@@ -47,7 +46,7 @@ from evaluate_event_free_survival import (
     mortality_hazard_from_risk,
 )
 from future_risk import death_risk_from_probabilities, probabilities_from_logits
-from targets import PAD_IDX
+from targets import CHECKUP_IDX, PAD_IDX
 
 
 OUTPUT_COLUMNS = [
@@ -239,6 +238,76 @@ def concat_padded_tensor_batches(chunks: list[Dict[str, torch.Tensor]]) -> Dict[
     return out
 
 
+def build_disease_ablated_slice(
+    batch: Dict[str, torch.Tensor],
+    row_indices: torch.Tensor,
+    token_ids: torch.Tensor,
+) -> Dict[str, torch.Tensor]:
+    """Build an ablated slice for aligned (row, disease_token) pairs."""
+    event_seq = batch["event_seq"]
+    row_indices = row_indices.to(device=event_seq.device, dtype=torch.long)
+    token_ids = token_ids.to(device=event_seq.device, dtype=event_seq.dtype)
+
+    out: Dict[str, torch.Tensor] = {}
+    out["event_seq"] = event_seq[row_indices].clone()
+    out["time_seq"] = batch["time_seq"][row_indices]
+    out["readout_mask"] = batch["readout_mask"][row_indices].clone()
+    out["padding_mask"] = batch["padding_mask"][row_indices].bool().clone()
+    out["landmark_pos"] = batch["landmark_pos"][row_indices].clone()
+
+    seq_len = int(event_seq.shape[1])
+    positions = torch.arange(seq_len, device=event_seq.device)[None, :]
+    remove = (out["event_seq"] == token_ids[:, None]) & out["padding_mask"]
+    out["event_seq"] = torch.where(
+        remove,
+        torch.full_like(out["event_seq"], PAD_IDX),
+        out["event_seq"],
+    )
+    out["padding_mask"] &= ~remove
+    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_indices[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",
+        "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:
+        out[key] = batch[key][row_indices]
+    return out
+
+
 def load_disease_metadata(
     mapping_path: Path,
     *,
@@ -486,7 +555,7 @@ def main() -> None:
 
     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))
+        cfg_get(args, cfg, "attribution_batch_size", max(batch_size * 8, batch_size))
     )
     if attribution_batch_size <= 0:
         raise ValueError("attribution_batch_size must be positive")
@@ -616,6 +685,7 @@ def main() -> None:
         dataset_index = int(meta["dataset_index"])
         sample = dataset.samples[dataset_index]
         hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
+        unique_tokens, token_counts = np.unique(hist_tokens, return_counts=True)
         total_count, group_counts = historical_counts_by_group(
             hist_tokens,
             death_idx=death_idx,
@@ -632,15 +702,23 @@ def main() -> None:
             "followup_end_time": float(meta["followup_end_time"]),
             "history_disease_count": int(total_count),
             "_hist_tokens": hist_tokens,
+            "_token_counts": {
+                int(token): int(count)
+                for token, count in zip(unique_tokens.tolist(), token_counts.tolist())
+            },
             "_group_counts": group_counts,
         }
         row_base_cache[row_idx] = cached
         return cached
 
     for batch in tqdm(loader, desc="Per-disease mortality attribution", dynamic_ncols=True):
+        batch_dev = {
+            k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v)
+            for k, v in batch.items()
+        }
         hidden = infer_landmark_hidden(
             model=model,
-            batch=batch,
+            batch=batch_dev,
             device=device,
             model_target_mode=model_target_mode,
             readout_name=readout_name,
@@ -660,48 +738,48 @@ def main() -> None:
         death_risk_tensor = death_risk_from_probabilities(probabilities)
         death_hazard_tensor = mortality_hazard_from_risk(death_risk_tensor)
         occurred = make_occurred_mask(
-            batch["event_seq"].to(device),
+            batch_dev["event_seq"],
             vocab_size=int(dataset.vocab_size),
             device=device,
         )
-        active_token_mask = occurred[:, scanned_disease_tensor].any(dim=0)
+        pair_indices = torch.nonzero(
-        batch_disease_tokens = scanned_disease_tensor[
+            occurred[:, scanned_disease_tensor],
-            active_token_mask
-        ].detach().cpu().tolist()
-        if not batch_disease_tokens:
-            continue
-
-        for disease_token in batch_disease_tokens:
-            disease_meta = metadata[disease_token]
-            active_rows = torch.nonzero(
-                occurred[:, disease_token].to(dtype=torch.bool),
             as_tuple=False,
-            ).flatten()
+        )
-            if active_rows.numel() == 0:
+        if pair_indices.numel() == 0:
             continue
 
-            row_offset = 0
+        pair_offset = 0
-            while row_offset < int(active_rows.numel()):
+        while pair_offset < int(pair_indices.shape[0]):
             capacity = int(attribution_batch_size) - pending_n
-                row_stop = min(int(active_rows.numel()), row_offset + capacity)
+            pair_stop = min(int(pair_indices.shape[0]), pair_offset + capacity)
-                row_indices = active_rows[row_offset:row_stop].to(device=batch["event_seq"].device)
+            pair_chunk = pair_indices[pair_offset:pair_stop]
-                ablated_chunk = build_group_ablated_slice(
+            local_rows = pair_chunk[:, 0].long()
-                    batch=batch,
+            disease_token_ids = scanned_disease_tensor[pair_chunk[:, 1]].long()
-                    token_ids=[disease_token],
+            ablated_chunk = build_disease_ablated_slice(
-                    row_indices=row_indices,
+                batch=batch_dev,
+                row_indices=local_rows,
+                token_ids=disease_token_ids,
             )
 
             meta_chunk: list[dict[str, Any]] = []
-                row_ids = batch["row_idx"][row_indices.cpu()].cpu().numpy().astype(np.int64)
+            row_ids = batch_dev["row_idx"][local_rows].detach().cpu().numpy().astype(np.int64)
-                for local_pos, row_idx in enumerate(row_ids.tolist()):
+            disease_tokens_list = disease_token_ids.detach().cpu().numpy().astype(np.int64)
+            for local_pos, (row_idx, disease_token) in enumerate(
+                zip(row_ids.tolist(), disease_tokens_list.tolist())
+            ):
+                disease_token = int(disease_token)
+                disease_meta = metadata[disease_token]
                 row_base = get_row_base(int(row_idx))
-                    hist_tokens = row_base["_hist_tokens"]
                 group_counts = row_base["_group_counts"]
-                    disease_history_count = int((hist_tokens == int(disease_token)).sum())
+                disease_history_count = int(row_base["_token_counts"].get(disease_token, 0))
                 if disease_history_count <= 0:
-                        continue
+                    raise RuntimeError(
+                        "Internal mismatch: occurred mask selected disease "
+                        f"{disease_token} for row {row_idx}, but cached history has count 0"
+                    )
 
-                    orig_row = int(row_indices[local_pos].detach().cpu())
+                orig_row = int(local_rows[local_pos].detach().cpu().item())
                 meta_chunk.append(
                     {
                         "patient_id": row_base["patient_id"],
@@ -734,7 +812,7 @@ def main() -> None:
                 pending_batch_chunks.append(ablated_chunk)
                 pending_meta_chunks.append(meta_chunk)
                 pending_n += len(meta_chunk)
-                row_offset = row_stop
+            pair_offset = pair_stop
 
             if pending_n >= int(attribution_batch_size):
                 flush_pending()