JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove legacy event and mixed distribution paths

Browse Source
This commit is contained in:
Jiarui Li
2026-08-01 14:23:18 +08:00
parent dfb22adf2d
commit de6f9b75b9
22 changed files with 370 additions and 463 deletions
Download Patch File Download Diff File Expand all files Collapse all files

89
evaluate_auc_v2.py
View File

@@ -1,7 +1,7 @@
"""Evaluate landmark fixed-horizon incident disease AUC for DeepHealth.
This script supports DeepHealth fixed-horizon risk scores for exponential,
Weibull, and mixed all-future distributions.
This script supports DeepHealth fixed-horizon risk scores for exponential and
Weibull all-future distributions.
The default horizons are 0.1, 1, 5, and 10 years. As in Delphi2M, 0.1 years
is reported as the no-gap evaluation.
@@ -53,9 +53,9 @@ from eval_data import (
)
from model_architectures import resolve_model_architecture
from models import DeepHealth
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
from targets import NO_EVENT_IDX, PAD_IDX, RESERVED_IDX
SPECIAL_TOKENS = {PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX}
SPECIAL_TOKENS = {PAD_IDX, RESERVED_IDX, NO_EVENT_IDX}
def parse_int_list(value: Any) -> Optional[List[int]]:
@@ -137,29 +137,19 @@ def load_checkpoint_state_dict(checkpoint_path: Path, map_location: str | torch.
def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, Any]) -> str:
mode = str(cfg_dist_mode).lower()
if mode not in {"exponential", "weibull"}:
raise ValueError(
f"Unsupported dist_mode={mode!r}; expected exponential or weibull."
)
has_rho_head = any(str(k).startswith("rho_head.")
for k in state_dict.keys())
has_rho_death_head = any(str(k).startswith("rho_death_head.")
for k in state_dict.keys())
if has_rho_head:
if mode != "weibull":
print(
"[WARN] Checkpoint contains rho_head weights; overriding dist_mode to 'weibull' for evaluation.")
return "weibull"
if has_rho_death_head:
if mode != "mixed":
print(
"[WARN] Checkpoint contains rho_death_head weights; overriding dist_mode to 'mixed' for evaluation.")
return "mixed"
if mode == "weibull":
print(
"[WARN] dist_mode is 'weibull' but checkpoint has no rho_head weights; overriding dist_mode to 'exponential'.")
return "exponential"
if mode == "mixed":
print(
"[WARN] dist_mode is 'mixed' but checkpoint has no rho_death_head weights; overriding dist_mode to 'exponential'.")
return "exponential"
return mode if mode in {"exponential", "weibull", "mixed"} else "exponential"
if mode == "weibull" and not has_rho_head:
raise RuntimeError("Weibull checkpoint is missing rho_head parameters.")
if mode == "exponential" and has_rho_head:
raise RuntimeError(
"Exponential checkpoint unexpectedly contains rho_head parameters."
)
return mode
def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
@@ -290,7 +280,7 @@ def _metadata_count_map(dataset: HealthDataset, labels_meta: Optional[pd.DataFra
return out
def _get_death_token_ids(dataset: HealthDataset, labels_meta: Optional[pd.DataFrame]) -> List[int]:
def _get_death_token_ids(dataset: HealthDataset) -> List[int]:
return [int(dataset.vocab_size) - 1]
@@ -418,7 +408,7 @@ class LandmarkDataset(Dataset):
prefix_times = full_time[prefix_mask]
valid_history_mask = ~np.isin(prefix_events, np.array(
[PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX], dtype=np.int64))
[PAD_IDX, RESERVED_IDX, NO_EVENT_IDX], dtype=np.int64))
if valid_history_mask.sum() < self.min_history_events:
continue
@@ -671,24 +661,11 @@ def project_distribution_chunk(
device=device, dtype=compute_dtype)
rho_weight = None
rho_bias = None
death_rho_weight = None
death_rho_bias = None
mixed_death_cols: List[int] = []
death_idx = int(getattr(model, "death_idx", getattr(model, "vocab_size", 0) - 1))
if dist_mode == "weibull":
rho_weight = model.rho_head.weight[disease_ids].detach().to(
device=device, dtype=compute_dtype)
rho_bias = model.rho_head.bias[disease_ids].detach().to(
device=device, dtype=compute_dtype)
elif dist_mode == "mixed":
mixed_death_cols = [j for j, token in enumerate(disease_ids)
if int(token) == death_idx]
if mixed_death_cols:
death_rho_weight = model.rho_death_head.weight.detach().to(
device=device, dtype=compute_dtype)
death_rho_bias = model.rho_death_head.bias.detach().to(
device=device, dtype=compute_dtype)
out_parts: List[np.ndarray] = []
rho_parts: List[np.ndarray] = []
@@ -703,14 +680,6 @@ def project_distribution_chunk(
if dist_mode == "weibull":
assert rho_weight is not None and rho_bias is not None
rho = F.softplus(torch.matmul(h, rho_weight.t()) + rho_bias) + 1e-6
elif dist_mode == "mixed" and mixed_death_cols:
assert death_rho_weight is not None and death_rho_bias is not None
rho = torch.ones_like(logits)
death_rho = F.softplus(
torch.matmul(h, death_rho_weight.t()).squeeze(-1) + death_rho_bias.squeeze(0)
) + 1e-6
for col in mixed_death_cols:
rho[:, int(col)] = death_rho
out_parts.append(logits.float().cpu(
).numpy().astype(np.float32, copy=False))
@@ -747,7 +716,6 @@ def _init_worker(
exclude_death_competing: bool,
death_token_ids: np.ndarray,
dist_mode: str,
model_death_idx: int,
) -> None:
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")
@@ -772,7 +740,6 @@ def _init_worker(
"exclude_death_competing": bool(exclude_death_competing),
"death_token_ids": set(int(x) for x in np.asarray(death_token_ids, dtype=np.int64).tolist()),
"dist_mode": str(dist_mode).lower(),
"model_death_idx": int(model_death_idx),
"first_time_cache": {},
}
)
@@ -798,8 +765,6 @@ def _score_to_probability(
score_mode: str,
horizon: float,
dist_mode: str,
token: int,
death_idx: int,
) -> np.ndarray:
if score_mode == "eta":
return logits.astype(np.float64, copy=False)
@@ -811,11 +776,6 @@ def _score_to_probability(
raise RuntimeError("Weibull risk scoring requires rho parameters.")
exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False))
return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False)
if dist_mode == "mixed" and int(token) == int(death_idx):
if rho is None:
raise RuntimeError("Mixed death risk scoring requires death rho parameters.")
exposure = np.power(np.float32(horizon), rho.astype(np.float32, copy=False))
return (-np.expm1(-rate * exposure)).astype(np.float64, copy=False)
return (-np.expm1(-rate * np.float32(horizon))).astype(np.float64, copy=False)
@@ -832,7 +792,6 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
rho_chunk = _WORKER["rho_chunk"]
rho_token = None if rho_chunk is None else rho_chunk[:, int(j)]
dist_mode = _WORKER["dist_mode"]
model_death_idx = int(_WORKER["model_death_idx"])
first_time_patient = _first_time_by_patient(token)
is_death_target = token in _WORKER["death_token_ids"]
@@ -891,8 +850,6 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
score_mode=score_mode,
horizon=horizon,
dist_mode=dist_mode,
token=token,
death_idx=model_death_idx,
)
control_scores = _score_to_probability(
logits_token[idx[control_idx]],
@@ -900,8 +857,6 @@ def _eval_token(task: Tuple[int, int, str]) -> List[Dict[str, Any]]:
score_mode=score_mode,
horizon=horizon,
dist_mode=dist_mode,
token=token,
death_idx=model_death_idx,
)
auc, auc_var = get_auc_delong_var(case_scores, control_scores)
@@ -1019,8 +974,6 @@ def evaluate_landmark_auc(
death_token_ids=np.asarray(
landmark_dataset.death_token_ids, dtype=np.int64),
dist_mode=dist_mode,
model_death_idx=int(getattr(
model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
)
nested = [_eval_token(t) for t in tqdm(
tasks, desc=f"AUC chunk {chunk_idx}", leave=False, dynamic_ncols=True)]
@@ -1048,8 +1001,6 @@ def evaluate_landmark_auc(
np.asarray(landmark_dataset.death_token_ids,
dtype=np.int64),
dist_mode,
int(getattr(
model, "death_idx", getattr(model, "vocab_size", 1) - 1)),
),
) as ex:
nested = list(
@@ -1253,10 +1204,6 @@ def main() -> None:
state_dict = load_checkpoint_state_dict(model_ckpt_path, map_location="cpu")
dist_mode = resolve_dist_mode_for_checkpoint(dist_mode_cfg, state_dict)
if dist_mode not in {"exponential", "weibull", "mixed"}:
raise ValueError(
f"Unsupported dist_mode={dist_mode!r}; expected exponential, weibull, or mixed."
)
if score_mode == "eta":
print(
@@ -1302,7 +1249,7 @@ def main() -> None:
"Please use a checkpoint trained with the same no-event vocabulary configuration."
)
death_token_ids = _get_death_token_ids(dataset, labels_meta)
death_token_ids = _get_death_token_ids(dataset)
min_history_events = int(cfg_get(args, cfg, "min_history_events", 1))
landmark_dataset = LandmarkDataset(
dataset=dataset,