from __future__ import annotations import json import logging import sys import time from datetime import datetime from pathlib import Path from typing import Any, Dict, Iterable, Tuple import numpy as np import torch from torch.optim import AdamW from torch.utils.data import Subset from dataset import AllFutureHealthDataset, HealthDataset from models import DeepHealth def create_unique_run_dir(name_fn, runs_root: Path = Path("runs")) -> tuple[Path, str]: while True: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") run_name = name_fn(timestamp) run_dir = runs_root / run_name try: run_dir.mkdir(parents=True, exist_ok=False) return run_dir, run_name except FileExistsError: time.sleep(1.0) def setup_logging(run_dir: Path) -> logging.Logger: run_dir.mkdir(parents=True, exist_ok=True) logger = logging.getLogger("DeepHealth") logger.setLevel(logging.INFO) logger.handlers.clear() formatter = logging.Formatter( "%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) console_handler = logging.StreamHandler(sys.stdout) console_handler.setLevel(logging.INFO) console_handler.setFormatter(formatter) logger.addHandler(console_handler) file_handler = logging.FileHandler(run_dir / "train.log", mode="w") file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) return logger def set_seed(seed: int) -> None: np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) def load_extra_info_types_file(path: str) -> list[int]: file_path = Path(path) if not file_path.is_file(): raise FileNotFoundError(f"extra_info_types_file not found: {path}") text = file_path.read_text(encoding="utf-8").strip() if not text: return [] if text.startswith("["): raw_items = json.loads(text) if not isinstance(raw_items, list): raise ValueError("extra_info_types_file JSON must be a list") else: raw_items = [] for line in text.splitlines(): line = line.split("#", 1)[0].strip() if line: raw_items.extend(line.replace(",", " ").replace(";", " ").split()) try: return [int(x) for x in raw_items] except (TypeError, ValueError) as exc: raise ValueError(f"Invalid extra info type id in {path}") from exc def configure_torch_for_training(device: torch.device) -> None: if device.type != "cuda": return torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True if hasattr(torch, "set_float32_matmul_precision"): torch.set_float32_matmul_precision("high") def resolve_device(device_arg: str) -> torch.device: requested = device_arg.strip().lower() if requested == "cpu": return torch.device("cpu") if requested == "cuda": return torch.device("cuda" if torch.cuda.is_available() else "cpu") if requested.startswith("cuda:"): if not torch.cuda.is_available(): return torch.device("cpu") index = int(requested.split(":", 1)[1]) if index < 0 or index >= torch.cuda.device_count(): raise ValueError(f"Requested CUDA device is out of range: {device_arg}") return torch.device(f"cuda:{index}") raise ValueError(f"Unsupported device: {device_arg}") def split_dataset( dataset: HealthDataset, train_ratio: float, val_ratio: float, test_ratio: float, seed: int, ) -> Tuple[Subset, Subset, Subset]: total = train_ratio + val_ratio + test_ratio if not np.isclose(total, 1.0, atol=1e-6): raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}") indices = np.random.RandomState(seed).permutation(len(dataset)) n_train = int(len(dataset) * train_ratio) n_val = int(len(dataset) * val_ratio) return ( Subset(dataset, indices[:n_train]), Subset(dataset, indices[n_train:n_train + n_val]), Subset(dataset, indices[n_train + n_val:]), ) def split_all_future_datasets( train_dataset: AllFutureHealthDataset, val_dataset: AllFutureHealthDataset, test_dataset: AllFutureHealthDataset, train_ratio: float, val_ratio: float, test_ratio: float, seed: int, ) -> Tuple[Subset, Subset, Subset]: total = train_ratio + val_ratio + test_ratio if not np.isclose(total, 1.0, atol=1e-6): raise ValueError(f"train/val/test ratios must sum to 1.0, got {total}") patient_indices = np.random.RandomState(seed).permutation(len(train_dataset.patients)) n_train = int(len(patient_indices) * train_ratio) n_val = int(len(patient_indices) * val_ratio) train_patient_idx = patient_indices[:n_train] val_patient_set = set(int(x) for x in patient_indices[n_train:n_train + n_val]) test_patient_set = set(int(x) for x in patient_indices[n_train + n_val:]) val_query_idx = [ i for i, (pidx, _t_query) in enumerate(val_dataset.valid_queries) if int(pidx) in val_patient_set ] test_query_idx = [ i for i, (pidx, _t_query) in enumerate(test_dataset.valid_queries) if int(pidx) in test_patient_set ] if not val_query_idx: raise ValueError("All-future validation split has no valid query samples.") if not test_query_idx: raise ValueError("All-future test split has no valid query samples.") return ( Subset(train_dataset, train_patient_idx), Subset(val_dataset, np.asarray(val_query_idx, dtype=np.int64)), Subset(test_dataset, np.asarray(test_query_idx, dtype=np.int64)), ) def build_optimizer(args: Any, model: DeepHealth) -> AdamW: return AdamW( model.parameters(), lr=args.base_lr, betas=tuple(args.betas), weight_decay=args.weight_decay, ) def set_optimizer_lr(optimizer: AdamW, lr: float) -> None: for param_group in optimizer.param_groups: param_group["lr"] = lr def save_checkpoint(model: DeepHealth, checkpoint_path: Path) -> None: torch.save(model.state_dict(), checkpoint_path) def save_config( args: Any, config_path: Path, extra: Dict[str, Any] | None = None, ) -> None: config: Dict[str, Any] = {} for key, value in vars(args).items(): if isinstance(value, tuple): config[key] = list(value) elif isinstance(value, list): config[key] = value elif isinstance(value, (int, float, str, bool, type(None))): config[key] = value else: config[key] = str(value) if extra: config.update(extra) config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")