DeepHealthExpo/train_util.py

from __future__ import annotations

import logging
import sys
import time
import csv
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, 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_eid_file(path: str | Path) -> set[int]:
    file_path = Path(path)
    if not file_path.is_file():
        raise FileNotFoundError(f"eid split file not found: {file_path}")
    with file_path.open(newline="", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f)
        if reader.fieldnames is None or "eid" not in reader.fieldnames:
            raise ValueError(
                f"eid split file must contain an 'eid' column: {file_path}"
            )
        out: set[int] = set()
        for row in reader:
            raw = (row.get("eid") or "").strip()
            if raw:
                out.add(int(raw))
    if not out:
        raise ValueError(f"eid split file is empty: {file_path}")
    return out


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_dataset_by_eid_files(
    dataset: HealthDataset,
    train_eid_file: str | Path,
    val_eid_file: str | Path,
    test_eid_file: str | Path,
) -> Tuple[Subset, Subset, Subset]:
    split_sets = {
        "train": load_eid_file(train_eid_file),
        "val": load_eid_file(val_eid_file),
        "test": load_eid_file(test_eid_file),
    }
    overlaps = (
        split_sets["train"] & split_sets["val"],
        split_sets["train"] & split_sets["test"],
        split_sets["val"] & split_sets["test"],
    )
    if any(overlaps):
        raise ValueError("eid split files must be disjoint")

    split_indices: Dict[str, list[int]] = {"train": [], "val": [], "test": []}
    for idx, sample in enumerate(dataset.samples):
        eid = int(sample["eid"])
        for split_name, eid_set in split_sets.items():
            if eid in eid_set:
                split_indices[split_name].append(idx)
                break

    missing = [name for name, indices in split_indices.items() if not indices]
    if missing:
        raise ValueError(f"Empty dataset split(s) after eid filtering: {missing}")

    return (
        Subset(dataset, np.asarray(split_indices["train"], dtype=np.int64)),
        Subset(dataset, np.asarray(split_indices["val"], dtype=np.int64)),
        Subset(dataset, np.asarray(split_indices["test"], dtype=np.int64)),
    )


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 split_all_future_datasets_by_eid_files(
    train_dataset: AllFutureHealthDataset,
    val_dataset: AllFutureHealthDataset,
    test_dataset: AllFutureHealthDataset,
    train_eid_file: str | Path,
    val_eid_file: str | Path,
    test_eid_file: str | Path,
) -> Tuple[Subset, Subset, Subset]:
    split_sets = {
        "train": load_eid_file(train_eid_file),
        "val": load_eid_file(val_eid_file),
        "test": load_eid_file(test_eid_file),
    }
    overlaps = (
        split_sets["train"] & split_sets["val"],
        split_sets["train"] & split_sets["test"],
        split_sets["val"] & split_sets["test"],
    )
    if any(overlaps):
        raise ValueError("eid split files must be disjoint")

    train_patient_idx = [
        idx
        for idx, patient in enumerate(train_dataset.patients)
        if int(patient["eid"]) in split_sets["train"]
    ]
    val_query_idx = [
        idx
        for idx, (pidx, _t_query) in enumerate(val_dataset.valid_queries)
        if int(val_dataset.patients[int(pidx)]["eid"]) in split_sets["val"]
    ]
    test_query_idx = [
        idx
        for idx, (pidx, _t_query) in enumerate(test_dataset.valid_queries)
        if int(test_dataset.patients[int(pidx)]["eid"]) in split_sets["test"]
    ]

    if not train_patient_idx:
        raise ValueError("All-future training eid split has no patients.")
    if not val_query_idx:
        raise ValueError("All-future validation eid split has no valid query samples.")
    if not test_query_idx:
        raise ValueError("All-future test eid split has no valid query samples.")

    return (
        Subset(train_dataset, np.asarray(train_patient_idx, dtype=np.int64)),
        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")