diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..effe788 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,7 @@ +# 项目协作原则 + +- 不回答或执行超出用户问题或请求范围的事项。 +- 如果确有必要超出范围，必须先停止并询问用户；得到明确同意后才能继续。 +- 回答应简单直接、逻辑清晰。 +- 永远不要使用 Node.js 读取、处理、分析或转换数据；涉及数据任务时必须使用其他工具。 +- Python 与数据处理统一使用本机 Miniconda：`C:\ProgramData\miniconda3`；不要使用 Codex 自带或其他 Python/Conda 运行时。 diff --git a/README.md b/README.md index a0ec3ad..bd9aa53 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ 疾病序列 stream + 统一的额外信息 token stream ``` -疾病、死亡、checkup 事件仍然保存在事件序列里；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。 +疾病、死亡、checkup 事件保存在预处理事件文件中；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。dataset 仅在实验选择了至少一种 extra-info type 时保留 checkup；显式传入空列表时，模型输入是没有 checkup 的纯疾病历史。 ## 数据准备 @@ -267,6 +267,28 @@ python train_all_future.py \ --extra_pool_reduce mean ``` +纯疾病历史的 T/O/S 消融仅用于 +`TrajMixer + all_future + relative + Weibull + extra_info_types=[]`： + +```bash +python train_all_future.py \ + --model_architecture traj_mixer_v5 \ + --time_mode relative \ + --dist_mode weibull \ + --extra_info_types_file extra_info_types_none.txt \ + --disease_history_mode ordered +``` + +`--disease_history_mode` 的含义： + +- `timed`：保留疾病事件顺序和真实患病时间（T）。 +- `ordered`：保留疾病事件顺序，以 `0, 1, ..., G-1` 的首发日期顺序组替代真实时间；同日首发疾病共享一个位置，query 位置为 `G`（O）。 +- `set`：疾病代码去重并排序，疾病和 query 的模型时间全部为 `0`，不保留顺序或患病时间（S）。 + +查询点、未来疾病、`future_dt`、exposure、Landmark 年龄与 AUC +病例/对照始终使用真实时间。训练配置会记录 `disease_history_mode`； +旧配置缺少该字段时按 `timed` 处理。 + 选择额外信息变量： ```bash @@ -279,6 +301,7 @@ python train_next_step.py --extra_info_types_file extra_info_types_smoking_alcoh - `extra_info_types_file`：训练时使用的列表文件名 - `extra_info_types`：解析后的实际 type id 列表，用于评估脚本复现变量选择 +- `disease_history_mode`：all-future 模型使用的 T/O/S 疾病历史表示 - `extra_pool_reduce`：同一 `other_time` 的 extra-info tokens 池化方式，默认为 `mean` - `model_target_mode`、`time_mode`、`dist_mode`、`dataset_class`、`collate_fn`、`resolved_loss_name`：用于评估脚本重建模型和输入方式 diff --git a/dataset.py b/dataset.py index bfe4e3f..e8c66ca 100644 --- a/dataset.py +++ b/dataset.py @@ -20,6 +20,167 @@ from targets import ( ONE_DAY_YEARS = 1.0 / DAYS_PER_YEAR +DISEASE_HISTORY_MODE_TIMED = "timed" +DISEASE_HISTORY_MODE_ORDERED = "ordered" +DISEASE_HISTORY_MODE_SET = "set" +DISEASE_HISTORY_MODES = ( + DISEASE_HISTORY_MODE_TIMED, + DISEASE_HISTORY_MODE_ORDERED, + DISEASE_HISTORY_MODE_SET, +) + + +def normalize_disease_history_mode(mode: str | None) -> str: + value = DISEASE_HISTORY_MODE_TIMED if mode is None else str(mode).lower() + if value not in DISEASE_HISTORY_MODES: + raise ValueError( + "disease_history_mode must be one of " + f"{list(DISEASE_HISTORY_MODES)}, got {mode!r}" + ) + return value + + +def transform_disease_history( + event_seq: np.ndarray, + actual_time_seq: np.ndarray, + actual_t_query: float, + disease_history_mode: str, +) -> Tuple[np.ndarray, np.ndarray, np.float32]: + """ + Convert an already-truncated disease history into its model representation. + + ``timed`` keeps the real event/query times. ``ordered`` preserves the + chronological event order but replaces calendar time with ordinal event-time + groups. Diseases first recorded on the same day share one ordinal position. + ``set`` removes both time and order by sorting the unique disease codes and + assigning every disease and the query the same model time. + """ + mode = normalize_disease_history_mode(disease_history_mode) + events = np.asarray(event_seq, dtype=np.int64) + times = np.asarray(actual_time_seq, dtype=np.float32) + if events.ndim != 1 or times.ndim != 1 or events.shape != times.shape: + raise ValueError( + "event_seq and actual_time_seq must be aligned 1D arrays, got " + f"{events.shape} and {times.shape}" + ) + + if mode == DISEASE_HISTORY_MODE_TIMED: + return events, times, np.float32(actual_t_query) + + special = events <= NO_EVENT_IDX + if np.any(special): + raise ValueError( + f"{mode} disease history must contain only disease events; " + f"found special token ids {np.unique(events[special]).tolist()}" + ) + + if mode == DISEASE_HISTORY_MODE_ORDERED: + _, ordinal_groups = np.unique(times, return_inverse=True) + model_times = ordinal_groups.astype(np.float32, copy=False) + n_groups = int(model_times.max()) + 1 if model_times.size else 0 + return events, model_times, np.float32(n_groups) + + set_events = np.unique(events) + model_times = np.zeros(set_events.size, dtype=np.float32) + return set_events, model_times, np.float32(0.0) + + +def transform_disease_history_batch_at_position( + event_seq: torch.Tensor, + actual_time_seq: torch.Tensor, + padding_mask: torch.Tensor, + query_position: int, + disease_history_mode: str, + vocab_size: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Build a model-visible prefix for token-position all-future evaluation. + + Actual event times remain outside this return value for AUC bookkeeping. + For ordered/set modes, events after ``query_position`` are explicitly + masked so collapsing time cannot expose future diseases. + """ + mode = normalize_disease_history_mode(disease_history_mode) + if event_seq.ndim != 2 or actual_time_seq.shape != event_seq.shape: + raise ValueError( + "event_seq and actual_time_seq must be aligned 2D tensors, got " + f"{tuple(event_seq.shape)} and {tuple(actual_time_seq.shape)}" + ) + if padding_mask.shape != event_seq.shape: + raise ValueError( + "padding_mask must match event_seq, got " + f"{tuple(padding_mask.shape)} and {tuple(event_seq.shape)}" + ) + if query_position < 0 or query_position >= event_seq.size(1): + raise ValueError( + f"query_position={query_position} is outside sequence length " + f"{event_seq.size(1)}" + ) + + padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool) + if not torch.all(padding_mask[:, query_position]): + raise ValueError("query_position must be valid for every batch row") + + if mode == DISEASE_HISTORY_MODE_TIMED: + return ( + event_seq, + actual_time_seq, + padding_mask, + actual_time_seq[:, query_position], + ) + + positions = torch.arange( + event_seq.size(1), + device=event_seq.device, + )[None, :] + history_mask = padding_mask & (positions <= query_position) + visible_events = event_seq.masked_select(history_mask) + if torch.any(visible_events <= NO_EVENT_IDX): + special_ids = torch.unique( + visible_events[visible_events <= NO_EVENT_IDX] + ).detach().cpu().tolist() + raise ValueError( + f"{mode} disease history must contain only disease events; " + f"found special token ids {special_ids}" + ) + + if mode == DISEASE_HISTORY_MODE_ORDERED: + model_times = torch.zeros_like(actual_time_seq) + model_t_query = torch.zeros( + event_seq.size(0), + device=actual_time_seq.device, + dtype=actual_time_seq.dtype, + ) + for row_idx in range(event_seq.size(0)): + row_mask = history_mask[row_idx] + _, ordinal_groups = torch.unique( + actual_time_seq[row_idx, row_mask], + sorted=True, + return_inverse=True, + ) + model_times[row_idx, row_mask] = ordinal_groups.to( + dtype=actual_time_seq.dtype + ) + model_t_query[row_idx] = float( + int(ordinal_groups.max().item()) + 1 + if ordinal_groups.numel() + else 0 + ) + return event_seq, model_times, history_mask, model_t_query + + sentinel = torch.full_like(event_seq, int(vocab_size)) + sortable = torch.where(history_mask, event_seq, sentinel) + set_events = torch.sort(sortable, dim=1).values + set_mask = set_events != int(vocab_size) + set_events = set_events.masked_fill(~set_mask, PAD_IDX) + model_times = torch.zeros_like(actual_time_seq) + model_t_query = torch.zeros( + event_seq.size(0), + device=actual_time_seq.device, + dtype=actual_time_seq.dtype, + ) + return set_events, model_times, set_mask, model_t_query + def load_label_vocab( labels_file: str, @@ -255,6 +416,15 @@ class _ExpoBaseDataset(Dataset): times_days_raw = rows[:, 1].astype(np.float32) labels_raw = rows[:, 2].astype(np.int64) + # CHECKUP is the assessment landmark for selected extra-info tokens. + # An explicitly empty selection represents a disease-only history, + # so retaining CHECKUP in that case would introduce an empty + # landmark token that is not part of the disease sequence. + if not self.extra_info_types: + keep = labels_raw != CHECKUP_IDX + times_days_raw = times_days_raw[keep] + labels_raw = labels_raw[keep] + if len(labels_raw) == 0: yield eid, times_days_raw, labels_raw continue @@ -368,6 +538,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset): min_future_events: int = 1, validation_query_seed: int = 42, extra_info_types: Iterable[int] | None = None, + disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED, ) -> None: if split not in {"train", "valid", "test"}: raise ValueError(f"split must be train/valid/test, got {split!r}") @@ -379,6 +550,18 @@ class AllFutureHealthDataset(_ExpoBaseDataset): extra_info_types=extra_info_types, ) + self.disease_history_mode = normalize_disease_history_mode( + disease_history_mode + ) + if ( + self.disease_history_mode != DISEASE_HISTORY_MODE_TIMED + and self.extra_info_types + ): + raise ValueError( + f"disease_history_mode={self.disease_history_mode!r} is only " + "supported with an explicitly empty extra-info selection" + ) + self.split = split self.min_history_events = int(min_history_events) self.min_future_events = int(min_future_events) @@ -504,11 +687,17 @@ class AllFutureHealthDataset(_ExpoBaseDataset): labels = patient["labels"] hist = times <= t_query fut = times > t_query + event_seq, model_time_seq, model_t_query = transform_disease_history( + event_seq=labels[hist], + actual_time_seq=times[hist], + actual_t_query=t_query, + disease_history_mode=self.disease_history_mode, + ) return { - "event_seq": torch.from_numpy(labels[hist]).long(), - "time_seq": torch.from_numpy(times[hist]).float(), - "t_query": torch.tensor(t_query, dtype=torch.float32), + "event_seq": torch.from_numpy(event_seq).long(), + "time_seq": torch.from_numpy(model_time_seq).float(), + "t_query": torch.tensor(model_t_query, dtype=torch.float32), "future_targets": torch.from_numpy(labels[fut]).long(), "future_dt": torch.from_numpy(times[fut] - np.float32(t_query)).float(), "exposure": torch.tensor(np.float32(patient["t_obs"] - t_query), dtype=torch.float32), diff --git a/eval_data.py b/eval_data.py index 5e70b9f..a201644 100644 --- a/eval_data.py +++ b/eval_data.py @@ -9,7 +9,12 @@ import numpy as np import torch from torch.nn.utils.rnn import pad_sequence -from dataset import AllFutureHealthDataset, HealthDataset +from dataset import ( + DISEASE_HISTORY_MODE_TIMED, + AllFutureHealthDataset, + HealthDataset, + normalize_disease_history_mode, +) from model_architectures import resolve_model_architecture from models import DeepHealth from targets import PAD_IDX @@ -61,6 +66,40 @@ def validate_training_mode_config(cfg: Dict[str, Any]) -> None: "model_target_mode must be next_token or all_future, got " f"{model_target_mode!r}" ) + + disease_history_mode = normalize_disease_history_mode( + cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED) + ) + if disease_history_mode != DISEASE_HISTORY_MODE_TIMED: + expected = { + "model_target_mode": "all_future", + "time_mode": "relative", + "dist_mode": "weibull", + "model_architecture": "traj_mixer_v5", + } + actual = { + "model_target_mode": model_target_mode, + "time_mode": str(cfg.get("time_mode", "")).lower(), + "dist_mode": str(cfg.get("dist_mode", "")).lower(), + "model_architecture": str( + cfg.get("model_architecture", "") + ).lower(), + } + mismatches = [ + f"{name}={actual[name]!r} (expected {value!r})" + for name, value in expected.items() + if actual[name] != value + ] + extra_info_types = cfg.get("extra_info_types", None) + if extra_info_types != []: + mismatches.append("extra_info_types must be []") + if mismatches: + raise ValueError( + f"disease_history_mode={disease_history_mode!r} is only valid " + "for the no-extra TrajMixer + all_future + relative + Weibull " + "ablation; " + "; ".join(mismatches) + ) + if model_target_mode != "next_token": return @@ -208,9 +247,10 @@ class AllFutureSequenceEvalDataset: """ Eval-only sequence view for all-future checkpoints. - All-future training uses the observed history, including CHECKUP state - tokens, without reusing the next-step view that contains imputed - gap tokens. + All-future training uses the observed history without reusing the + next-step view that contains imputed gap tokens. CHECKUP is + retained only when the experiment selects at least one extra-info type; + an explicitly empty selection is a disease-only history. """ def __init__( @@ -220,6 +260,7 @@ class AllFutureSequenceEvalDataset: min_history_events: int = 1, min_future_events: int = 1, extra_info_types: Iterable[int] | None = None, + disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED, ) -> None: base = AllFutureHealthDataset( data_prefix=data_prefix, @@ -228,6 +269,7 @@ class AllFutureSequenceEvalDataset: min_history_events=min_history_events, min_future_events=min_future_events, extra_info_types=extra_info_types, + disease_history_mode=disease_history_mode, ) self.base = base @@ -239,6 +281,7 @@ class AllFutureSequenceEvalDataset: self.n_categories = base.n_categories self.cont_type_ids = base.cont_type_ids self.extra_info_types = base.extra_info_types + self.disease_history_mode = base.disease_history_mode self.samples: List[Dict[str, Any]] = [] for patient in base.patients: @@ -288,6 +331,7 @@ def load_sequence_eval_dataset( min_history_events: int, min_future_events: int, extra_info_types: Iterable[int] | None, + disease_history_mode: str = DISEASE_HISTORY_MODE_TIMED, ): mode = str(model_target_mode).lower() if mode == "next_token": @@ -304,6 +348,7 @@ def load_sequence_eval_dataset( min_history_events=min_history_events, min_future_events=min_future_events, extra_info_types=extra_info_types, + disease_history_mode=disease_history_mode, ) raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}") diff --git a/evaluate_auc.py b/evaluate_auc.py index b2a8047..a3b36f5 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -39,7 +39,12 @@ import torch from torch.utils.data import DataLoader, Subset from tqdm.auto import tqdm -from dataset import HealthDataset +from dataset import ( + DISEASE_HISTORY_MODE_TIMED, + HealthDataset, + normalize_disease_history_mode, + transform_disease_history_batch_at_position, +) from delphi2m_auc_report import ( DEFAULT_DELPHI2M_PERIODS_YEARS, build_delphi2m_auc_report, @@ -333,6 +338,7 @@ def infer_readout_hidden( loader: DataLoader, device: torch.device, model_target_mode: str, + disease_history_mode: str, use_amp: bool, hidden_cache_dtype: str = "float16", ) -> Tuple[np.ndarray, Dict[str, np.ndarray]]: @@ -342,6 +348,16 @@ def infer_readout_hidden( raise ValueError( f"model_target_mode must be next_token or all_future, got {model_target_mode!r}" ) + disease_history_mode = normalize_disease_history_mode( + disease_history_mode + ) + if ( + model_target_mode != "all_future" + and disease_history_mode != DISEASE_HISTORY_MODE_TIMED + ): + raise ValueError( + "ordered/set disease history is only supported for all_future models" + ) hidden_parts: List[np.ndarray] = [] arrays: Dict[str, List[np.ndarray]] = { @@ -385,12 +401,25 @@ def infer_readout_hidden( active = padding_mask[:, pos].bool() if not active.any(): continue - hidden_pos = model( + ( + model_event_seq, + model_time_seq, + model_padding_mask, + model_t_query, + ) = transform_disease_history_batch_at_position( event_seq=event_seq[active], - time_seq=time_seq[active], - sex=batch_dev["sex"][active], + actual_time_seq=time_seq[active], padding_mask=padding_mask[active], - t_query=time_seq[active, pos], + query_position=pos, + disease_history_mode=disease_history_mode, + vocab_size=model.vocab_size, + ) + hidden_pos = model( + event_seq=model_event_seq, + time_seq=model_time_seq, + sex=batch_dev["sex"][active], + padding_mask=model_padding_mask, + t_query=model_t_query, other_type=batch_dev["other_type"][active], other_value=batch_dev["other_value"][active], other_value_kind=batch_dev["other_value_kind"][active], @@ -899,6 +928,7 @@ def evaluate_auc_pipeline( offsets: Sequence[float], device: torch.device, model_target_mode: str, + disease_history_mode: str, num_workers_auc: int, use_amp: bool, auc_task_chunk_size: int = 0, @@ -957,6 +987,7 @@ def evaluate_auc_pipeline( loader=loader, device=device, model_target_mode=model_target_mode, + disease_history_mode=disease_history_mode, use_amp=use_amp, hidden_cache_dtype=hidden_cache_dtype, ) @@ -1165,6 +1196,9 @@ def main() -> None: f"got {model_target_mode!r}" ) dist_mode_cfg = cfg.get("dist_mode", "exponential") + disease_history_mode = normalize_disease_history_mode( + cfg.get("disease_history_mode", DISEASE_HISTORY_MODE_TIMED) + ) device = resolve_eval_device(args.device) if device.type == "cuda": @@ -1179,6 +1213,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) @@ -1209,6 +1244,7 @@ def main() -> None: print(f"Resolved dist_mode for evaluation: {dist_mode}") print(f"Resolved model architecture: {model_architecture}") print(f"Model target mode for AUC: {model_target_mode}") + print(f"Disease history mode for AUC: {disease_history_mode}") print( "AUC score semantics: evaluate_auc.py uses disease-specific eta/logit scores; " "dist_mode affects model loading but is not converted to horizon-specific risk probability." @@ -1253,6 +1289,7 @@ def main() -> None: offsets=auc_offsets, device=device, model_target_mode=model_target_mode, + disease_history_mode=disease_history_mode, num_workers_auc=int(cfg_get(args, cfg, "num_workers_auc", max( 1, (os.cpu_count() or 2) - 1))), use_amp=bool(cfg_get(args, cfg, "use_amp", False)), diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index fdc29de..05b7386 100644 --- a/evaluate_auc_v2.py +++ b/evaluate_auc_v2.py @@ -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" diff --git a/tests/test_dataset_checkup.py b/tests/test_dataset_checkup.py new file mode 100644 index 0000000..4bea4d0 --- /dev/null +++ b/tests/test_dataset_checkup.py @@ -0,0 +1,57 @@ +import unittest +from pathlib import Path + +import numpy as np + +from dataset import _ExpoBaseDataset +from targets import CHECKUP_IDX +from train_util import load_extra_info_types_file + + +class CheckupSelectionTests(unittest.TestCase): + @staticmethod + def _base(extra_info_types): + dataset = _ExpoBaseDataset.__new__(_ExpoBaseDataset) + dataset.extra_info_types = list(extra_info_types) + dataset.event_data = np.asarray( + [ + [101, 10, CHECKUP_IDX], + [101, 20, 2], + [101, 30, 3], + ], + dtype=np.float64, + ) + return dataset + + def test_explicit_empty_extra_info_removes_checkup(self): + project_root = Path(__file__).resolve().parents[1] + selected_types = load_extra_info_types_file( + str(project_root / "extra_info_types_none.txt") + ) + self.assertEqual(selected_types, []) + dataset = self._base(selected_types) + + rows = list(dataset._iter_patient_events(impute_no_event_gaps=False)) + + self.assertEqual(len(rows), 1) + eid, times, labels = rows[0] + self.assertEqual(eid, 101) + np.testing.assert_array_equal(times, np.asarray([20, 30], dtype=np.float32)) + self.assertNotIn(CHECKUP_IDX, labels.tolist()) + + def test_selected_extra_info_keeps_checkup(self): + dataset = self._base([11]) + + rows = list(dataset._iter_patient_events(impute_no_event_gaps=False)) + + self.assertEqual(len(rows), 1) + _, times, labels = rows[0] + np.testing.assert_array_equal( + times, + np.asarray([10, 20, 30], dtype=np.float32), + ) + self.assertEqual(labels[0], CHECKUP_IDX) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_disease_history_ablation.py b/tests/test_disease_history_ablation.py new file mode 100644 index 0000000..03eda65 --- /dev/null +++ b/tests/test_disease_history_ablation.py @@ -0,0 +1,381 @@ +import unittest +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import torch + +from dataset import ( + AllFutureHealthDataset, + all_future_collate_fn, + transform_disease_history, + transform_disease_history_batch_at_position, +) +from eval_data import validate_training_mode_config +from losses import build_loss +from models import DeepHealth +from train_all_future import parse_args + + +class DiseaseHistoryTransformTests(unittest.TestCase): + def test_timed_ordered_and_set_representations(self): + events = np.asarray([9, 4, 7], dtype=np.int64) + times = np.asarray([50.0, 60.0, 65.0], dtype=np.float32) + + timed_events, timed_times, timed_query = transform_disease_history( + events, times, 70.0, "timed" + ) + np.testing.assert_array_equal(timed_events, events) + np.testing.assert_array_equal(timed_times, times) + self.assertEqual(float(timed_query), 70.0) + + ordered_events, ordered_times, ordered_query = transform_disease_history( + events, times, 70.0, "ordered" + ) + np.testing.assert_array_equal(ordered_events, events) + np.testing.assert_array_equal( + ordered_times, + np.asarray([0.0, 1.0, 2.0], dtype=np.float32), + ) + self.assertEqual(float(ordered_query), 3.0) + + set_events, set_times, set_query = transform_disease_history( + events, times, 70.0, "set" + ) + np.testing.assert_array_equal( + set_events, + np.asarray([4, 7, 9], dtype=np.int64), + ) + np.testing.assert_array_equal( + set_times, + np.zeros(3, dtype=np.float32), + ) + self.assertEqual(float(set_query), 0.0) + + def test_ordered_removes_calendar_time_but_keeps_order(self): + events = np.asarray([9, 4, 7], dtype=np.int64) + first = transform_disease_history( + events, + np.asarray([20.0, 21.0, 70.0], dtype=np.float32), + 75.0, + "ordered", + ) + second = transform_disease_history( + events, + np.asarray([50.0, 60.0, 65.0], dtype=np.float32), + 70.0, + "ordered", + ) + np.testing.assert_array_equal(first[0], second[0]) + np.testing.assert_array_equal(first[1], second[1]) + self.assertEqual(float(first[2]), float(second[2])) + + reversed_events = transform_disease_history( + events[::-1], + np.asarray([50.0, 60.0, 65.0], dtype=np.float32), + 70.0, + "ordered", + )[0] + self.assertFalse(np.array_equal(first[0], reversed_events)) + + def test_ordered_keeps_same_day_diseases_in_one_order_group(self): + events, model_times, model_query = transform_disease_history( + np.asarray([9, 4, 7], dtype=np.int64), + np.asarray([50.0, 50.0, 65.0], dtype=np.float32), + 70.0, + "ordered", + ) + np.testing.assert_array_equal( + events, + np.asarray([9, 4, 7], dtype=np.int64), + ) + np.testing.assert_array_equal( + model_times, + np.asarray([0.0, 0.0, 1.0], dtype=np.float32), + ) + self.assertEqual(float(model_query), 2.0) + + def test_set_removes_order_and_calendar_time(self): + first = transform_disease_history( + np.asarray([9, 4, 7], dtype=np.int64), + np.asarray([50.0, 60.0, 65.0], dtype=np.float32), + 70.0, + "set", + ) + second = transform_disease_history( + np.asarray([7, 9, 4], dtype=np.int64), + np.asarray([20.0, 21.0, 70.0], dtype=np.float32), + 75.0, + "set", + ) + np.testing.assert_array_equal(first[0], second[0]) + np.testing.assert_array_equal(first[1], second[1]) + self.assertEqual(float(first[2]), float(second[2])) + + def test_all_future_targets_stay_on_actual_time(self): + patient = { + "times": np.asarray([50.0, 60.0, 65.0, 75.0], dtype=np.float32), + "labels": np.asarray([9, 4, 7, 12], dtype=np.int64), + "t_obs": 75.0, + "sex": 0, + "other_type": np.zeros(0, dtype=np.int64), + "other_value": np.zeros(0, dtype=np.float32), + "other_value_kind": np.zeros(0, dtype=np.int64), + "other_time": np.zeros(0, dtype=np.float32), + } + + items = {} + for mode in ("timed", "ordered", "set"): + dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset) + dataset.disease_history_mode = mode + items[mode] = dataset._build_item(patient, 70.0) + + for item in items.values(): + torch.testing.assert_close( + item["future_targets"], + torch.tensor([12], dtype=torch.long), + ) + torch.testing.assert_close( + item["future_dt"], + torch.tensor([5.0]), + ) + torch.testing.assert_close( + item["exposure"], + torch.tensor(5.0), + ) + + torch.testing.assert_close( + items["timed"]["time_seq"], + torch.tensor([50.0, 60.0, 65.0]), + ) + self.assertEqual(float(items["timed"]["t_query"]), 70.0) + torch.testing.assert_close( + items["ordered"]["event_seq"], + torch.tensor([9, 4, 7]), + ) + torch.testing.assert_close( + items["ordered"]["time_seq"], + torch.tensor([0.0, 1.0, 2.0]), + ) + self.assertEqual(float(items["ordered"]["t_query"]), 3.0) + torch.testing.assert_close( + items["set"]["event_seq"], + torch.tensor([4, 7, 9]), + ) + torch.testing.assert_close( + items["set"]["time_seq"], + torch.zeros(3), + ) + self.assertEqual(float(items["set"]["t_query"]), 0.0) + + def test_batch_prefix_transform_masks_future_events(self): + events = torch.tensor( + [ + [9, 4, 7, 12], + [8, 5, 11, 0], + ], + dtype=torch.long, + ) + actual_times = torch.tensor( + [ + [50.0, 60.0, 65.0, 75.0], + [45.0, 55.0, 80.0, 0.0], + ] + ) + mask = events > 0 + + timed = transform_disease_history_batch_at_position( + events, actual_times, mask, 1, "timed", vocab_size=20 + ) + torch.testing.assert_close(timed[0], events) + torch.testing.assert_close(timed[1], actual_times) + torch.testing.assert_close(timed[2], mask) + torch.testing.assert_close(timed[3], torch.tensor([60.0, 55.0])) + + ordered = transform_disease_history_batch_at_position( + events, actual_times, mask, 1, "ordered", vocab_size=20 + ) + torch.testing.assert_close( + ordered[2], + torch.tensor( + [ + [True, True, False, False], + [True, True, False, False], + ] + ), + ) + torch.testing.assert_close(ordered[3], torch.tensor([2.0, 2.0])) + + disease_set = transform_disease_history_batch_at_position( + events, actual_times, mask, 1, "set", vocab_size=20 + ) + torch.testing.assert_close( + disease_set[0], + torch.tensor( + [ + [4, 9, 0, 0], + [5, 8, 0, 0], + ] + ), + ) + torch.testing.assert_close( + disease_set[2], + torch.tensor( + [ + [True, True, False, False], + [True, True, False, False], + ] + ), + ) + torch.testing.assert_close(disease_set[3], torch.zeros(2)) + + +class DiseaseSetModelTests(unittest.TestCase): + @staticmethod + def _model(): + return DeepHealth( + vocab_size=16, + n_embd=24, + n_head=4, + n_layer=2, + n_types=1, + n_cont_types=0, + n_categories=1, + cont_type_ids=[], + n_bins=4, + target_mode="all_future", + time_mode="relative", + dist_mode="weibull", + dropout=0.0, + model_architecture="traj_mixer_v5", + ) + + def test_equal_time_query_is_permutation_invariant(self): + torch.manual_seed(7) + model = self._model().eval() + event_seq = torch.tensor( + [ + [4, 7, 9], + [9, 4, 7], + ], + dtype=torch.long, + ) + time_seq = torch.zeros(2, 3) + empty_long = torch.zeros(2, 0, dtype=torch.long) + empty_float = torch.zeros(2, 0) + + with torch.inference_mode(): + hidden = model( + event_seq=event_seq, + time_seq=time_seq, + sex=torch.zeros(2, dtype=torch.long), + padding_mask=torch.ones(2, 3, dtype=torch.bool), + t_query=torch.zeros(2), + other_type=empty_long, + other_value=empty_float, + other_value_kind=empty_long, + other_time=empty_float, + ) + + torch.testing.assert_close(hidden[0], hidden[1], atol=1e-6, rtol=1e-6) + + def test_ordered_and_set_support_finite_weibull_loss(self): + patient = { + "times": np.asarray([50.0, 60.0, 65.0, 75.0], dtype=np.float32), + "labels": np.asarray([9, 4, 7, 12], dtype=np.int64), + "t_obs": 75.0, + "sex": 0, + "other_type": np.zeros(0, dtype=np.int64), + "other_value": np.zeros(0, dtype=np.float32), + "other_value_kind": np.zeros(0, dtype=np.int64), + "other_time": np.zeros(0, dtype=np.float32), + } + criterion = build_loss("weibull", ignored_idx={0, 1}) + + for mode in ("ordered", "set"): + dataset = AllFutureHealthDataset.__new__(AllFutureHealthDataset) + dataset.disease_history_mode = mode + item = dataset._build_item(patient, 70.0) + batch = all_future_collate_fn([item, item]) + model = self._model() + hidden = model( + event_seq=batch["event_seq"], + time_seq=batch["time_seq"], + sex=batch["sex"], + padding_mask=batch["padding_mask"], + t_query=batch["t_query"], + other_type=batch["other_type"], + other_value=batch["other_value"], + other_value_kind=batch["other_value_kind"], + other_time=batch["other_time"], + ) + loss = criterion( + logits=model.calc_risk(hidden), + weibull_rho=model.calc_weibull_rho(hidden), + targets=batch["future_targets"], + dt=batch["future_dt"], + exposure=batch["exposure"], + ) + self.assertTrue(torch.isfinite(loss), msg=f"{mode} loss={loss}") + + +class DiseaseHistoryConfigTests(unittest.TestCase): + def test_training_cli_accepts_ordered_ablation(self): + project_root = Path(__file__).resolve().parents[1] + with patch( + "sys.argv", + [ + "train_all_future.py", + "--disease_history_mode", + "ordered", + "--model_architecture", + "traj_mixer_v5", + "--time_mode", + "relative", + "--dist_mode", + "weibull", + "--extra_info_types_file", + str(project_root / "extra_info_types_none.txt"), + ], + ): + args = parse_args() + self.assertEqual(args.disease_history_mode, "ordered") + self.assertEqual(args.extra_info_types, []) + + def test_legacy_config_defaults_to_timed(self): + validate_training_mode_config( + { + "model_target_mode": "all_future", + "time_mode": "relative", + "dist_mode": "weibull", + "model_architecture": "traj_mixer_v5", + "extra_info_types": [11, 66, 67], + } + ) + + def test_ordered_config_requires_exact_ablation_setup(self): + validate_training_mode_config( + { + "model_target_mode": "all_future", + "time_mode": "relative", + "dist_mode": "weibull", + "model_architecture": "traj_mixer_v5", + "extra_info_types": [], + "disease_history_mode": "ordered", + } + ) + with self.assertRaises(ValueError): + validate_training_mode_config( + { + "model_target_mode": "all_future", + "time_mode": "relative", + "dist_mode": "weibull", + "model_architecture": "traj_mixer_v5", + "extra_info_types": [11], + "disease_history_mode": "ordered", + } + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/train_all_future.py b/train_all_future.py index 6bc5200..e6b110e 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -25,7 +25,12 @@ from torch.optim import AdamW from torch.utils.data import DataLoader, RandomSampler from tqdm.auto import tqdm -from dataset import AllFutureHealthDataset, all_future_collate_fn +from dataset import ( + DISEASE_HISTORY_MODES, + DISEASE_HISTORY_MODE_TIMED, + AllFutureHealthDataset, + all_future_collate_fn, +) from losses import build_loss from model_architectures import ( DEFAULT_MODEL_ARCHITECTURE, @@ -74,6 +79,16 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--runs_root", type=str, default="runs") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--extra_info_types_file", type=str, default=None) + parser.add_argument( + "--disease_history_mode", + type=str, + default=DISEASE_HISTORY_MODE_TIMED, + choices=DISEASE_HISTORY_MODES, + help=( + "timed=real disease times; ordered=chronological disease order with " + "ordinal positions; set=unordered disease set with no disease time" + ), + ) parser.add_argument("--train_ratio", type=float, default=0.7) parser.add_argument("--val_ratio", type=float, default=0.15) @@ -134,6 +149,27 @@ def parse_args() -> argparse.Namespace: if args.extra_info_types_file is not None else None ) + if args.disease_history_mode != DISEASE_HISTORY_MODE_TIMED: + expected = { + "model_architecture": "traj_mixer_v5", + "time_mode": "relative", + "dist_mode": "weibull", + } + mismatches = [ + f"{name}={getattr(args, name)!r} (expected {value!r})" + for name, value in expected.items() + if getattr(args, name) != value + ] + if args.extra_info_types != []: + mismatches.append( + "extra_info_types must be [] via extra_info_types_none.txt" + ) + if mismatches: + raise ValueError( + f"disease_history_mode={args.disease_history_mode!r} is reserved " + "for the no-extra TrajMixer + all_future + relative + Weibull " + "ablation; " + "; ".join(mismatches) + ) return args @@ -296,6 +332,7 @@ def build_metadata( "model_target_mode": "all_future", "target_mode": "all_future", "dist_mode": args.dist_mode, + "disease_history_mode": args.disease_history_mode, "all_future_min_history_events": int(args.min_history_events), "all_future_min_future_events": int(args.min_future_events), "all_future_validation_query_seed": int(args.validation_query_seed), @@ -330,7 +367,15 @@ def main() -> None: configure_torch_for_training(device) run_dir, run_name = create_unique_run_dir( - lambda timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}", + lambda timestamp: ( + ( + "" + if args.disease_history_mode == DISEASE_HISTORY_MODE_TIMED + else f"{args.disease_history_mode}_" + ) + + f"{args.time_mode}_{args.dist_mode}_" + f"all_future_pure_disease_{timestamp}" + ), runs_root=Path(args.runs_root) / args.model_architecture, ) logger = setup_logging(run_dir) @@ -338,6 +383,7 @@ def main() -> None: logger.info(f"Starting all-future training run: {run_name}") logger.info(f"Device: {device}") logger.info(f"Model architecture: {args.model_architecture}") + logger.info(f"Disease history mode: {args.disease_history_mode}") logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}") logger.info("Loading all-future datasets...") @@ -349,6 +395,7 @@ def main() -> None: min_future_events=args.min_future_events, validation_query_seed=args.validation_query_seed, extra_info_types=args.extra_info_types, + disease_history_mode=args.disease_history_mode, ) val_dataset = AllFutureHealthDataset( data_prefix=args.data_prefix, @@ -358,6 +405,7 @@ def main() -> None: min_future_events=args.min_future_events, validation_query_seed=args.validation_query_seed, extra_info_types=args.extra_info_types, + disease_history_mode=args.disease_history_mode, ) test_dataset = AllFutureHealthDataset( data_prefix=args.data_prefix, @@ -367,6 +415,7 @@ def main() -> None: min_future_events=args.min_future_events, validation_query_seed=args.validation_query_seed, extra_info_types=args.extra_info_types, + disease_history_mode=args.disease_history_mode, ) if args.train_eid_file and args.val_eid_file and args.test_eid_file: train_subset, val_subset, test_subset = split_all_future_datasets_by_eid_files( diff --git a/train_disease_history_ablation_linux.sh b/train_disease_history_ablation_linux.sh new file mode 100644 index 0000000..258932a --- /dev/null +++ b/train_disease_history_ablation_linux.sh @@ -0,0 +1,337 @@ +#!/usr/bin/env bash +# +# Train the no-extra-information T/O/S disease-history ablation. +# +# Fixed model: +# TrajMixer + all_future + relative + Weibull +# +# History modes: +# timed (T): disease identities, order, and real first-onset times +# ordered (O): disease identities and chronological order only +# set (S): unordered disease set only +# +# Jobs assigned to one GPU run sequentially. Different GPUs run in parallel. +# +# Examples: +# bash train_disease_history_ablation_linux.sh --gpus 0,1,2 +# bash train_disease_history_ablation_linux.sh --gpus 0 --seeds 42 --modes ordered,set +# bash train_disease_history_ablation_linux.sh --gpus 0,1 --dry-run +# + +set -uo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +GPU_CSV="" +SEED_CSV="42,43,44" +MODE_CSV="timed,ordered,set" +NUM_WORKERS=4 +BATCH_SIZE=256 +PYTHON_BIN="${PYTHON_BIN:-python}" +CAMPAIGN_NAME="disease_history_ablation_no_extra" +DRY_RUN=0 + +EXTRA_INFO_TYPES_FILE="$SCRIPT_DIR/extra_info_types_none.txt" +ENTRYPOINT="$SCRIPT_DIR/train_all_future.py" + +usage() { + cat <<'EOF' +Usage: + bash train_disease_history_ablation_linux.sh --gpus GPU_LIST [options] + +Required: + --gpus LIST Comma-separated GPU ids, for example 0,1,2. + +Options: + --seeds LIST Comma-separated seeds (default: 42,43,44). + --modes LIST Subset of timed,ordered,set (default: all three). + --batch-size N Batch size per task (default: 256). + --num-workers N DataLoader workers per task (default: 4). + --python PATH Python executable (default: $PYTHON_BIN or python). + --campaign NAME Output campaign name. + --dry-run Print commands without creating files or training. + -h, --help Show this help message. + +Fixed experiment settings: + architecture traj_mixer_v5 + target all_future + time mode relative + distribution weibull + extra information extra_info_types_none.txt + +Outputs: + runs//seed_/traj_mixer_v5/... + batch_logs//seed_/.log +EOF +} + +while (($# > 0)); do + case "$1" in + --gpus) + [[ $# -ge 2 ]] || { + echo "ERROR: --gpus requires a value." >&2 + exit 2 + } + GPU_CSV="$2" + shift 2 + ;; + --seeds) + [[ $# -ge 2 ]] || { + echo "ERROR: --seeds requires a value." >&2 + exit 2 + } + SEED_CSV="$2" + shift 2 + ;; + --modes) + [[ $# -ge 2 ]] || { + echo "ERROR: --modes requires a value." >&2 + exit 2 + } + MODE_CSV="$2" + shift 2 + ;; + --batch-size) + [[ $# -ge 2 ]] || { + echo "ERROR: --batch-size requires a value." >&2 + exit 2 + } + BATCH_SIZE="$2" + shift 2 + ;; + --num-workers) + [[ $# -ge 2 ]] || { + echo "ERROR: --num-workers requires a value." >&2 + exit 2 + } + NUM_WORKERS="$2" + shift 2 + ;; + --python) + [[ $# -ge 2 ]] || { + echo "ERROR: --python requires a value." >&2 + exit 2 + } + PYTHON_BIN="$2" + shift 2 + ;; + --campaign) + [[ $# -ge 2 ]] || { + echo "ERROR: --campaign requires a value." >&2 + exit 2 + } + CAMPAIGN_NAME="$2" + shift 2 + ;; + --dry-run) + DRY_RUN=1 + shift + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "ERROR: unknown argument: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +[[ -n "$GPU_CSV" ]] || { + echo "ERROR: --gpus is required." >&2 + usage >&2 + exit 2 +} +[[ -n "$SEED_CSV" ]] || { + echo "ERROR: --seeds must not be empty." >&2 + exit 2 +} +[[ -n "$MODE_CSV" ]] || { + echo "ERROR: --modes must not be empty." >&2 + exit 2 +} +[[ "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || { + echo "ERROR: --batch-size must be a positive integer." >&2 + exit 2 +} +[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || { + echo "ERROR: --num-workers must be a non-negative integer." >&2 + exit 2 +} +[[ "$CAMPAIGN_NAME" =~ ^[A-Za-z0-9._-]+$ ]] || { + echo "ERROR: --campaign may contain only letters, numbers, ., _, and -." >&2 + exit 2 +} +[[ -f "$EXTRA_INFO_TYPES_FILE" ]] || { + echo "ERROR: missing extra-info file: $EXTRA_INFO_TYPES_FILE" >&2 + exit 2 +} +[[ -f "$ENTRYPOINT" ]] || { + echo "ERROR: missing training entrypoint: $ENTRYPOINT" >&2 + exit 2 +} +command -v "$PYTHON_BIN" >/dev/null 2>&1 || { + echo "ERROR: Python executable not found: $PYTHON_BIN" >&2 + exit 2 +} + +IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV" +declare -A SEEN_GPUS=() +for gpu in "${GPU_IDS[@]}"; do + [[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || { + echo "ERROR: invalid GPU id: $gpu" >&2 + exit 2 + } + [[ -z "${SEEN_GPUS[$gpu]+x}" ]] || { + echo "ERROR: duplicate GPU id: $gpu" >&2 + exit 2 + } + SEEN_GPUS["$gpu"]=1 +done + +IFS=',' read -r -a SEEDS <<< "$SEED_CSV" +declare -A SEEN_SEEDS=() +for seed in "${SEEDS[@]}"; do + [[ "$seed" =~ ^[0-9]+$ ]] || { + echo "ERROR: invalid seed: $seed" >&2 + exit 2 + } + [[ -z "${SEEN_SEEDS[$seed]+x}" ]] || { + echo "ERROR: duplicate seed: $seed" >&2 + exit 2 + } + SEEN_SEEDS["$seed"]=1 +done + +IFS=',' read -r -a MODES <<< "$MODE_CSV" +declare -A SEEN_MODES=() +for mode in "${MODES[@]}"; do + case "$mode" in + timed|ordered|set) + ;; + *) + echo "ERROR: invalid mode: $mode (expected timed, ordered, or set)." >&2 + exit 2 + ;; + esac + [[ -z "${SEEN_MODES[$mode]+x}" ]] || { + echo "ERROR: duplicate mode: $mode" >&2 + exit 2 + } + SEEN_MODES["$mode"]=1 +done + +RUNS_ROOT="$SCRIPT_DIR/runs/$CAMPAIGN_NAME" +LOG_ROOT="$SCRIPT_DIR/batch_logs/$CAMPAIGN_NAME" +if ((!DRY_RUN)); then + mkdir -p "$RUNS_ROOT" "$LOG_ROOT" +fi + +declare -a JOB_SEEDS=() +declare -a JOB_MODES=() +for seed in "${SEEDS[@]}"; do + for mode in "${MODES[@]}"; do + JOB_SEEDS+=("$seed") + JOB_MODES+=("$mode") + done +done + +print_command() { + printf '%q ' "$@" + printf '\n' +} + +run_job() { + local job_index="$1" + local gpu="$2" + local seed="${JOB_SEEDS[$job_index]}" + local mode="${JOB_MODES[$job_index]}" + local seed_runs_root="$RUNS_ROOT/seed_$seed" + local seed_log_root="$LOG_ROOT/seed_$seed" + local log_file="$seed_log_root/$mode.log" + local -a command=( + "$PYTHON_BIN" + -u + "$ENTRYPOINT" + --runs_root "$seed_runs_root" + --seed "$seed" + --batch_size "$BATCH_SIZE" + --num_workers "$NUM_WORKERS" + --device cuda + --model_architecture traj_mixer_v5 + --time_mode relative + --dist_mode weibull + --extra_info_types_file "$EXTRA_INFO_TYPES_FILE" + --disease_history_mode "$mode" + ) + + if ((!DRY_RUN)); then + mkdir -p "$seed_runs_root" "$seed_log_root" + fi + + echo "[$(date '+%F %T')] START seed=$seed mode=$mode gpu=$gpu" + echo " log=$log_file" + if ((DRY_RUN)); then + printf ' CUDA_VISIBLE_DEVICES=%q ' "$gpu" + print_command "${command[@]}" + return 0 + fi + + if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \ + "${command[@]}" >"$log_file" 2>&1; then + echo "[$(date '+%F %T')] DONE seed=$seed mode=$mode gpu=$gpu" + return 0 + else + local exit_code=$? + echo "[$(date '+%F %T')] FAIL seed=$seed mode=$mode gpu=$gpu exit=$exit_code" >&2 + echo " See: $log_file" >&2 + return "$exit_code" + fi +} + +worker() { + local slot="$1" + local gpu="${GPU_IDS[$slot]}" + local job_index + local failed=0 + + for ((job_index = slot; job_index < ${#JOB_SEEDS[@]}; job_index += ${#GPU_IDS[@]})); do + run_job "$job_index" "$gpu" || failed=1 + done + return "$failed" +} + +echo "Campaign: $CAMPAIGN_NAME" +echo "Seeds: ${SEEDS[*]}" +echo "Modes: ${MODES[*]}" +echo "GPUs: ${GPU_IDS[*]}" +echo "Configurations per seed: ${#MODES[@]}" +echo "Total tasks: ${#JOB_SEEDS[@]}" +echo "Runs root: $RUNS_ROOT" +echo "Log root: $LOG_ROOT" +echo + +declare -a WORKER_PIDS=() +for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do + worker "$slot" & + WORKER_PIDS+=("$!") +done + +overall_status=0 +for pid in "${WORKER_PIDS[@]}"; do + wait "$pid" || overall_status=1 +done + +if ((overall_status != 0)); then + echo "One or more training tasks failed. Inspect logs under: $LOG_ROOT" >&2 + exit 1 +fi + +if ((DRY_RUN)); then + echo "Dry run completed successfully." +else + echo "All disease-history ablation tasks completed successfully." +fi