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

Add disease history ablation modes

Browse Source
This commit is contained in:
Jiarui Li
2026-07-29 13:53:15 +08:00
parent 31f129a7dc
commit c622ec50f7
10 changed files with 1178 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
AGENTS.md Normal file
View File

@@ -0,0 +1,7 @@
# 项目协作原则
- 不回答或执行超出用户问题或请求范围的事项。
- 如果确有必要超出范围，必须先停止并询问用户；得到明确同意后才能继续。
- 回答应简单直接、逻辑清晰。
- 永远不要使用 Node.js 读取、处理、分析或转换数据；涉及数据任务时必须使用其他工具。
- Python 与数据处理统一使用本机 Miniconda：`C:\ProgramData\miniconda3`；不要使用 Codex 自带或其他 Python/Conda 运行时。

25
README.md
View File

@@ -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`：用于评估脚本重建模型和输入方式

195
dataset.py
View File

@@ -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),

53
eval_data.py
View File

@@ -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
<NO_EVENT> gap tokens.
All-future training uses the observed history without reusing the
next-step view that contains imputed <NO_EVENT> 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}")

47
evaluate_auc.py
View File

@@ -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)),

44
evaluate_auc_v2.py
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"

57
tests/test_dataset_checkup.py Normal file
View File

@@ -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()

381
tests/test_disease_history_ablation.py Normal file
View File

@@ -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()

53
train_all_future.py
View File

@@ -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(

337
train_disease_history_ablation_linux.sh Normal file
View File

@@ -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/<campaign>/seed_<seed>/traj_mixer_v5/...
batch_logs/<campaign>/seed_<seed>/<mode>.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