Add disease history ablation modes

This commit is contained in:
2026-07-29 13:53:15 +08:00
parent 31f129a7dc
commit c622ec50f7
10 changed files with 1178 additions and 21 deletions

View File

@@ -30,7 +30,12 @@ from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset
from tqdm.auto import tqdm
from dataset import HealthDataset
from dataset import (
DISEASE_HISTORY_MODE_TIMED,
HealthDataset,
normalize_disease_history_mode,
transform_disease_history,
)
from delphi2m_auc_report import (
DEFAULT_DELPHI2M_PERIODS_YEARS,
build_delphi2m_auc_report,
@@ -338,6 +343,7 @@ class LandmarkDataset(Dataset):
min_history_events: int,
first_occurrence_by_token: Dict[int, Tuple[np.ndarray, np.ndarray]],
death_token_ids: Sequence[int],
disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED,
) -> None:
self.dataset = dataset
self.subset_indices = np.asarray(subset_indices, dtype=np.int64)
@@ -348,6 +354,16 @@ class LandmarkDataset(Dataset):
"model_target_mode must be next_token or all_future, got "
f"{self.model_target_mode!r}"
)
self.disease_history_mode = normalize_disease_history_mode(
disease_history_mode
)
if (
self.model_target_mode != "all_future"
and self.disease_history_mode != DISEASE_HISTORY_MODE_TIMED
):
raise ValueError(
"ordered/set disease history is only supported for all_future models"
)
self.min_history_events = int(min_history_events)
self.first_occurrence_by_token = first_occurrence_by_token
@@ -428,10 +444,16 @@ class LandmarkDataset(Dataset):
readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
readout_mask[-1] = True
else:
event_seq_landmark = prefix_events.astype(
np.int64, copy=False)
time_seq_landmark = prefix_times.astype(
np.float32, copy=False)
(
event_seq_landmark,
time_seq_landmark,
model_t_query,
) = transform_disease_history(
event_seq=prefix_events,
actual_time_seq=prefix_times,
actual_t_query=landmark_age,
disease_history_mode=self.disease_history_mode,
)
landmark_pos = int(len(event_seq_landmark) - 1)
readout_mask = np.zeros(len(event_seq_landmark), dtype=bool)
@@ -444,7 +466,11 @@ class LandmarkDataset(Dataset):
"followup_end_time": np.float32(followup_end),
"death_time": np.float32(self.patient_death_time[patient_id]),
"landmark_pos": landmark_pos,
"t_query": np.float32(landmark_age),
"t_query": (
np.float32(landmark_age)
if self.model_target_mode == "next_token"
else model_t_query
),
"event_seq": event_seq_landmark,
"time_seq": time_seq_landmark,
"readout_mask": readout_mask,
@@ -1137,6 +1163,9 @@ def main() -> None:
f"got {model_target_mode!r}"
)
dist_mode_cfg = str(cfg.get("dist_mode", "exponential"))
disease_history_mode = normalize_disease_history_mode(
cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED)
)
output_path = Path(
cfg_get(args, cfg, "output_path", None)
@@ -1162,6 +1191,7 @@ def main() -> None:
min_history_events=int(cfg.get("all_future_min_history_events", 1)),
min_future_events=int(cfg.get("all_future_min_future_events", 1)),
extra_info_types=parse_int_list(cfg.get("extra_info_types", None)),
disease_history_mode=disease_history_mode,
)
validate_dataset_metadata(dataset, cfg)
@@ -1282,6 +1312,7 @@ def main() -> None:
min_history_events=min_history_events,
first_occurrence_by_token=first_occurrence_by_token,
death_token_ids=death_token_ids,
disease_history_mode=disease_history_mode,
)
batch_size = int(cfg_get(args, cfg, "batch_size", 128))
@@ -1322,6 +1353,7 @@ def main() -> None:
print(f"Number of selected patients: {len(subset_indices)}")
print(f"No-event support: {bool(has_no_event)}")
print(f"Model target mode: {model_target_mode}")
print(f"Disease history mode: {disease_history_mode}")
print(f"Landmark query mode: {landmark_query_mode}")
print(
"Landmark token mode: no_event"