Optimize per-disease attribution scanning

evaluate_single_disease_mortality_attribution.py

@@ -438,7 +438,6 @@ def main() -> None:
     if not scanned_disease_items:
         raise ValueError("No diseases selected for attribution")
     scanned_disease_tokens = [token for token, _meta in scanned_disease_items]
-    scanned_disease_token_set = set(scanned_disease_tokens)
 
     landmark_ages = make_landmark_ages(
         float(args.landmark_start),
@@ -476,6 +475,11 @@ def main() -> None:
     cfg_model = dict(cfg)
     cfg_model["dist_mode"] = dist_mode
     device = resolve_eval_device(args.device)
+    scanned_disease_tensor = torch.as_tensor(
+        scanned_disease_tokens,
+        dtype=torch.long,
+        device=device,
+    )
     model = build_model_from_dataset(args, cfg_model, dataset).to(device)
     load_model_state(model, state_dict)
     model.eval()
@@ -537,6 +541,7 @@ def main() -> None:
     shard_index = 0
     shards: list[dict[str, Any]] = []
     summary_accumulator: dict[tuple[Any, ...], dict[str, float]] = {}
+    row_base_cache: dict[int, dict[str, Any]] = {}
     pending_batch_chunks: list[Dict[str, torch.Tensor]] = []
     pending_meta_chunks: list[list[dict[str, Any]]] = []
     pending_n = 0
@@ -602,6 +607,36 @@ def main() -> None:
         pending_meta_chunks = []
         pending_n = 0
 
+    def get_row_base(row_idx: int) -> dict[str, Any]:
+        cached = row_base_cache.get(row_idx)
+        if cached is not None:
+            return cached
+
+        meta = landmark_dataset.rows[int(row_idx)]
+        dataset_index = int(meta["dataset_index"])
+        sample = dataset.samples[dataset_index]
+        hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
+        total_count, group_counts = historical_counts_by_group(
+            hist_tokens,
+            death_idx=death_idx,
+            token_to_group=token_to_group,
+            group_names=group_names,
+        )
+        cached = {
+            "patient_id": int(meta["patient_id"]),
+            "dataset_index": dataset_index,
+            "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,
+            "_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):
         hidden = infer_landmark_hidden(
             model=model,
@@ -629,14 +664,10 @@ def main() -> None:
             vocab_size=int(dataset.vocab_size),
             device=device,
         )
-        event_values = batch["event_seq"].detach().cpu().numpy().reshape(-1)
-        batch_disease_tokens = sorted(
-            {
-                int(token)
-                for token in event_values.tolist()
-                if int(token) in scanned_disease_token_set
-            }
-        )
+        active_token_mask = occurred[:, scanned_disease_tensor].any(dim=0)
+        batch_disease_tokens = scanned_disease_tensor[
+            active_token_mask
+        ].detach().cpu().tolist()
         if not batch_disease_tokens:
             continue
 
@@ -663,16 +694,9 @@ def main() -> None:
                 meta_chunk: list[dict[str, Any]] = []
                 row_ids = batch["row_idx"][row_indices.cpu()].cpu().numpy().astype(np.int64)
                 for local_pos, row_idx in enumerate(row_ids.tolist()):
-                    meta = landmark_dataset.rows[int(row_idx)]
-                    dataset_index = int(meta["dataset_index"])
-                    sample = dataset.samples[dataset_index]
-                    hist_tokens = np.asarray(meta["event_seq"], dtype=np.int64)
-                    total_count, group_counts = historical_counts_by_group(
-                        hist_tokens,
-                        death_idx=death_idx,
-                        token_to_group=token_to_group,
-                        group_names=group_names,
-                    )
+                    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())
                     if disease_history_count <= 0:
                         continue
@@ -680,14 +704,14 @@ def main() -> None:
                     orig_row = int(row_indices[local_pos].detach().cpu())
                     meta_chunk.append(
                         {
-                            "patient_id": int(meta["patient_id"]),
-                            "dataset_index": dataset_index,
-                            "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),
+                            "patient_id": row_base["patient_id"],
+                            "dataset_index": row_base["dataset_index"],
+                            "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,
                             "selected_disease_token_id": int(disease_token),
                             "selected_disease_code": str(disease_meta.get("code", "")),