Remove legacy event and mixed distribution paths

evaluate_calibration.py

@@ -23,12 +23,11 @@ from __future__ import annotations
 import argparse
 import contextlib
 import json
-import math
 import os
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 import numpy as np
 import pandas as pd
@@ -67,7 +66,7 @@ from evaluate_auc_v2 import (
 )
 from losses import build_loss
 from model_architectures import resolve_model_architecture
-from targets import CHECKUP_IDX, PAD_IDX
+from targets import PAD_IDX, RESERVED_IDX
 from train_util import load_eid_file
 
 
@@ -971,8 +970,6 @@ def _risk_probability_matrix(
     rho: Optional[np.ndarray],
     horizons: np.ndarray,
     dist_mode: str,
-    token: int,
-    death_idx: int,
 ) -> np.ndarray:
     """Convert one token's logits to all horizon risks at once."""
     logits = np.asarray(logits, dtype=np.float32)
@@ -982,13 +979,7 @@ def _risk_probability_matrix(
         + np.maximum(logits, np.float32(0.0))
         + np.float32(1e-8)
     )
-    use_weibull = (
-        str(dist_mode).lower() == "weibull"
-        or (
-            str(dist_mode).lower() == "mixed"
-            and int(token) == int(death_idx)
-        )
-    )
+    use_weibull = str(dist_mode).lower() == "weibull"
     if use_weibull:
         if rho is None:
             raise RuntimeError(
@@ -1047,7 +1038,6 @@ def _evaluate_calibration_token(
     label_id_to_code: Dict[int, str],
     dist_mode: str,
     horizons: np.ndarray,
-    death_index: int,
     min_cases: int,
     min_controls: int,
     max_ipcw_weight: float,
@@ -1114,8 +1104,6 @@ def _evaluate_calibration_token(
             ),
             horizons=horizons,
             dist_mode=dist_mode,
-            token=token,
-            death_idx=death_index,
         )
         results = compute_ipcw_horizons(
             probabilities=probabilities,
@@ -1200,9 +1188,6 @@ def evaluate_landmark_calibration(
 
     patient_count = len(landmark_dataset.subset_indices)
     death_tokens = set(int(value) for value in landmark_dataset.death_token_ids)
-    death_index = int(
-        getattr(model, "death_idx", getattr(model, "vocab_size", 1) - 1)
-    )
     strata = _build_calibration_strata(
         row_arrays["sex"],
         row_arrays["landmark_age"],
@@ -1265,7 +1250,6 @@ def evaluate_landmark_calibration(
                     ),
                     "dist_mode": dist_mode,
                     "horizons": horizons,
-                    "death_index": death_index,
                     "min_cases": min_cases,
                     "min_controls": min_controls,
                     "max_ipcw_weight": max_ipcw_weight,
@@ -1498,19 +1482,12 @@ def build_calibration_summary(metrics: pd.DataFrame) -> pd.DataFrame:
 
 def _build_point_process_criterion(
     dist_mode: str,
-    death_index: int,
 ) -> Any:
-    ignored = {PAD_IDX, CHECKUP_IDX}
+    ignored = {PAD_IDX, RESERVED_IDX}
     if dist_mode == "exponential":
         return build_loss("exponential", ignored_idx=ignored)
     if dist_mode == "weibull":
         return build_loss("weibull", ignored_idx=ignored)
-    if dist_mode == "mixed":
-        return build_loss(
-            "mixed",
-            death_idx=death_index,
-            ignored_idx=ignored,
-        )
     raise ValueError(f"Unsupported dist_mode: {dist_mode!r}")
 
 
@@ -1523,16 +1500,7 @@ def evaluate_point_process_nll(
     device: torch.device,
     use_amp: bool,
 ) -> Dict[str, Any]:
-    criterion = _build_point_process_criterion(
-        dist_mode,
-        death_index=int(
-            getattr(
-                model,
-                "death_idx",
-                int(getattr(model, "vocab_size", 1)) - 1,
-            )
-        ),
-    )
+    criterion = _build_point_process_criterion(dist_mode)
     model.eval().to(device)
     total_nll = 0.0
     query_count = 0
@@ -1587,13 +1555,7 @@ def evaluate_point_process_nll(
                     exposure=batch_device["exposure"],
                 )
             else:
-                loss = criterion(
-                    logits=logits,
-                    death_rho=model.calc_death_rho(hidden),
-                    targets=batch_device["future_targets"],
-                    dt=batch_device["future_dt"],
-                    exposure=batch_device["exposure"],
-                )
+                raise ValueError(f"Unsupported dist_mode: {dist_mode!r}")
 
         if not torch.isfinite(loss):
             raise RuntimeError("Non-finite point-process NLL encountered.")
@@ -1601,7 +1563,7 @@ def evaluate_point_process_nll(
         total_nll += float(loss.detach().cpu()) * batch_size
         query_count += batch_size
         valid_targets = batch["future_targets"] > PAD_IDX
-        valid_targets &= batch["future_targets"] != CHECKUP_IDX
+        valid_targets &= batch["future_targets"] != RESERVED_IDX
         future_event_count += int(valid_targets.sum().item())
         exposure_sum += float(batch["exposure"].sum().item())