DeepHealth/train_all_future.py

"""
Train DeepHealth with query-conditioned all-future supervision.

Training samples are patient-level. For each patient and each __getitem__ call,
AllFutureHealthDataset randomly samples a query time t_query, uses events at or
before t_query as history, and uses events after t_query as the future target set.

Validation/test samples are deterministic query points built from future event
times, then split by patient.
"""
from __future__ import annotations

import argparse
import json
import logging
import math
import time
from pathlib import Path
from typing import Any, Dict

import numpy as np
import torch
from torch.nn.utils import clip_grad_norm_
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 losses import build_loss
from models import DeepHealth
from targets import CHECKUP_IDX, PAD_IDX
from train_util import (
    configure_torch_for_training,
    create_unique_run_dir,
    load_extra_info_types_file,
    resolve_device,
    save_checkpoint,
    save_config,
    set_optimizer_lr,
    set_seed,
    setup_logging,
    split_all_future_datasets,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Train DeepHealth with all-future supervision")

    parser.add_argument("--data_prefix", type=str, default="ukb")
    parser.add_argument("--labels_file", type=str, default="labels.csv")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--extra_info_types_file", type=str, default=None)

    parser.add_argument("--train_ratio", type=float, default=0.7)
    parser.add_argument("--val_ratio", type=float, default=0.15)
    parser.add_argument("--test_ratio", type=float, default=0.15)
    parser.add_argument("--min_history_events", type=int, default=1)
    parser.add_argument("--min_future_events", type=int, default=1)
    parser.add_argument("--validation_query_seed", type=int, default=None)

    parser.add_argument("--n_embd", type=int, default=120)
    parser.add_argument("--n_head", type=int, default=10)
    parser.add_argument("--n_hist_layer", type=int, default=12)
    parser.add_argument("--n_tab_layer", type=int, default=4)
    parser.add_argument("--n_bins", type=int, default=16)
    parser.add_argument("--time_mode", type=str, default="relative",
                        choices=["relative", "absolute"])
    parser.add_argument("--dist_mode", type=str, default="exponential",
                        choices=["exponential", "weibull", "mixed"])
    parser.add_argument("--dropout", type=float, default=0.0)

    parser.add_argument("--batch_size", type=int, default=128)
    parser.add_argument("--base_lr", type=float, default=3e-4)
    parser.add_argument("--weight_decay", type=float, default=0.1)
    parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99))
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--max_epochs", type=int, default=200)
    parser.add_argument("--warmup_epochs", type=int, default=10)
    parser.add_argument("--patience", type=int, default=15)
    parser.add_argument("--min_lr_ratio", type=float, default=0.1)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--device", type=str, default="cuda")

    args = parser.parse_args()
    if args.min_history_events < 1:
        raise ValueError("min_history_events must be >= 1")
    if args.min_future_events < 1:
        raise ValueError("min_future_events must be >= 1")
    if not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0):
        raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
    if args.validation_query_seed is None:
        args.validation_query_seed = int(args.seed)
    args.extra_info_types = (
        load_extra_info_types_file(args.extra_info_types_file)
        if args.extra_info_types_file is not None
        else None
    )
    return args


def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float:
    if epoch < args.warmup_epochs:
        return adaptive_lr * (epoch + 1) / args.warmup_epochs
    progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs)
    cosine = 0.5 * (1 + math.cos(math.pi * progress))
    return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio))


def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
    non_blocking = device.type == "cuda"
    return {
        key: value.to(device, non_blocking=non_blocking)
        if isinstance(value, torch.Tensor)
        else value
        for key, value in batch.items()
    }


def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth:
    return DeepHealth(
        vocab_size=dataset.vocab_size,
        n_embd=args.n_embd,
        n_head=args.n_head,
        n_hist_layer=args.n_hist_layer,
        n_tab_layer=args.n_tab_layer,
        n_types=dataset.n_types,
        n_cont_types=dataset.n_cont_types,
        n_categories=dataset.n_categories,
        cont_type_ids=dataset.cont_type_ids,
        n_bins=args.n_bins,
        target_mode="all_future",
        time_mode=args.time_mode,
        dist_mode=args.dist_mode,
        dropout=args.dropout,
    )


def build_criterion(args: argparse.Namespace, dataset: AllFutureHealthDataset):
    ignored_idx = {PAD_IDX, CHECKUP_IDX}
    if args.dist_mode == "exponential":
        return build_loss("exponential", ignored_idx=ignored_idx)
    if args.dist_mode == "weibull":
        return build_loss("weibull", ignored_idx=ignored_idx)
    if args.dist_mode == "mixed":
        return build_loss(
            "mixed",
            death_idx=dataset.vocab_size - 1,
            ignored_idx=ignored_idx,
        )
    raise ValueError(f"Unknown dist_mode: {args.dist_mode}")


def compute_all_future_loss(
    args: argparse.Namespace,
    model: DeepHealth,
    criterion,
    batch: Dict[str, torch.Tensor],
    device: torch.device,
) -> torch.Tensor:
    batch = move_batch_to_device(batch, device)

    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"],
        target_mode="all_future",
    )
    logits = model.calc_risk(hidden)

    if args.dist_mode == "exponential":
        loss = criterion(
            logits=logits,
            targets=batch["future_targets"],
            exposure=batch["exposure"],
        )
    elif args.dist_mode == "weibull":
        loss = criterion(
            logits=logits,
            weibull_rho=model.calc_weibull_rho(hidden),
            targets=batch["future_targets"],
            dt=batch["future_dt"],
            exposure=batch["exposure"],
        )
    else:
        loss = criterion(
            logits=logits,
            death_rho=model.calc_death_rho(hidden),
            targets=batch["future_targets"],
            dt=batch["future_dt"],
            exposure=batch["exposure"],
        )

    if not torch.isfinite(loss):
        raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
    return loss


def run_epoch(
    logger: logging.Logger,
    args: argparse.Namespace,
    model: DeepHealth,
    criterion,
    loader: DataLoader,
    optimizer: AdamW | None,
    device: torch.device,
    is_train: bool,
) -> float:
    model.train(is_train)
    total = 0.0
    n_batches = 0
    skipped = 0
    desc = "train" if is_train else "val"

    progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
    for batch_idx, batch in enumerate(progress):
        try:
            loss = compute_all_future_loss(args, model, criterion, batch, device)
            if is_train:
                if optimizer is None:
                    raise ValueError("optimizer is required for training")
                optimizer.zero_grad(set_to_none=True)
                loss.backward()
                if args.grad_clip > 0:
                    clip_grad_norm_(model.parameters(), args.grad_clip)
                optimizer.step()

            total += float(loss.detach().cpu())
            n_batches += 1
            avg = total / max(1, n_batches)
            progress.set_postfix(loss=f"{float(loss.detach().cpu()):.4f}", avg=f"{avg:.4f}", skipped=skipped)
        except RuntimeError as exc:
            if "Loss is not finite" not in str(exc):
                raise
            skipped += 1
            logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}")

    if skipped:
        logger.info(f"Skipped {skipped} batches due to non-finite loss")
    return total / max(1, n_batches) if n_batches else float("inf")


def build_metadata(
    args: argparse.Namespace,
    dataset: AllFutureHealthDataset,
    run_name: str,
    train_subset,
    val_subset,
    test_subset,
) -> Dict[str, Any]:
    return {
        "run_name": run_name,
        "dataset_class": "AllFutureHealthDataset",
        "collate_fn": "all_future_collate_fn",
        "model_class": "DeepHealth",
        "model_target_mode": "all_future",
        "target_mode": "all_future",
        "dist_mode": args.dist_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),
        "extra_info_types_file": (
            Path(args.extra_info_types_file).name
            if args.extra_info_types_file is not None
            else None
        ),
        "extra_info_types": [int(x) for x in dataset.extra_info_types],
        "dataset_metadata": {
            "vocab_size": int(dataset.vocab_size),
            "n_types": int(dataset.n_types),
            "n_cont_types": int(dataset.n_cont_types),
            "n_categories": int(dataset.n_categories),
            "cont_type_ids": [int(x) for x in dataset.cont_type_ids],
            "extra_info_types": [int(x) for x in dataset.extra_info_types],
        },
        "split_sizes": {
            "train": int(len(train_subset)),
            "val": int(len(val_subset)),
            "test": int(len(test_subset)),
        },
        "resolved_readout_name": "none",
        "resolved_loss_name": args.dist_mode,
    }


def main() -> None:
    args = parse_args()
    set_seed(args.seed)
    device = resolve_device(args.device)
    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}"
    )
    logger = setup_logging(run_dir)

    logger.info(f"Starting all-future training run: {run_name}")
    logger.info(f"Device: {device}")
    logger.info(f"extra_info_types: {args.extra_info_types or 'all'}")

    logger.info("Loading all-future datasets...")
    train_dataset = AllFutureHealthDataset(
        data_prefix=args.data_prefix,
        labels_file=args.labels_file,
        split="train",
        min_history_events=args.min_history_events,
        min_future_events=args.min_future_events,
        validation_query_seed=args.validation_query_seed,
        extra_info_types=args.extra_info_types,
    )
    val_dataset = AllFutureHealthDataset(
        data_prefix=args.data_prefix,
        labels_file=args.labels_file,
        split="valid",
        min_history_events=args.min_history_events,
        min_future_events=args.min_future_events,
        validation_query_seed=args.validation_query_seed,
        extra_info_types=args.extra_info_types,
    )
    test_dataset = AllFutureHealthDataset(
        data_prefix=args.data_prefix,
        labels_file=args.labels_file,
        split="test",
        min_history_events=args.min_history_events,
        min_future_events=args.min_future_events,
        validation_query_seed=args.validation_query_seed,
        extra_info_types=args.extra_info_types,
    )
    train_subset, val_subset, test_subset = split_all_future_datasets(
        train_dataset=train_dataset,
        val_dataset=val_dataset,
        test_dataset=test_dataset,
        train_ratio=args.train_ratio,
        val_ratio=args.val_ratio,
        test_ratio=args.test_ratio,
        seed=args.seed,
    )
    logger.info(
        f"Patients/queries: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
    )

    train_loader = DataLoader(
        train_subset,
        batch_size=args.batch_size,
        sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)),
        collate_fn=all_future_collate_fn,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
        prefetch_factor=2 if args.num_workers > 0 else None,
    )
    val_loader = DataLoader(
        val_subset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=all_future_collate_fn,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
        prefetch_factor=2 if args.num_workers > 0 else None,
    )
    test_loader = DataLoader(
        test_subset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=all_future_collate_fn,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
        prefetch_factor=2 if args.num_workers > 0 else None,
    )

    model = build_model(args, train_dataset).to(device)
    optimizer = AdamW(
        model.parameters(),
        lr=args.base_lr,
        betas=tuple(args.betas),
        weight_decay=args.weight_decay,
    )
    criterion = build_criterion(args, train_dataset)
    adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)

    save_config(
        args,
        run_dir / "train_config.json",
        extra=build_metadata(args, train_dataset, run_name, train_subset, val_subset, test_subset),
    )

    best_val = float("inf")
    patience = 0
    history = []
    best_model_path = run_dir / "best_model.pt"
    start = time.time()

    for epoch in range(args.max_epochs):
        lr = get_lr(epoch, args, adaptive_lr)
        set_optimizer_lr(optimizer, lr)

        train_loss = run_epoch(logger, args, model, criterion, train_loader, optimizer, device, True)
        with torch.no_grad():
            val_loss = run_epoch(logger, args, model, criterion, val_loader, None, device, False)

        is_best = val_loss < best_val
        if is_best:
            best_val = val_loss
            patience = 0
            save_checkpoint(model, best_model_path)
        else:
            patience += 1

        logger.info(
            f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | "
            f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | "
            f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | "
            f"elapsed={time.time() - start:.1f}s"
        )
        history.append({
            "epoch": epoch + 1,
            "lr": lr,
            "train_loss": train_loss,
            "val_loss": val_loss,
            "best_val_loss": best_val,
            "is_best": int(is_best),
        })
        if patience >= args.patience:
            logger.info(f"Early stopping triggered at epoch {epoch + 1}")
            break

    with (run_dir / "history.json").open("w", encoding="utf-8") as f:
        json.dump(history, f, indent=2)

    logger.info("Evaluating best model on all-future test queries...")
    model.load_state_dict(torch.load(best_model_path, map_location=device))
    with torch.no_grad():
        test_loss = run_epoch(logger, args, model, criterion, test_loader, None, device, False)
    logger.info(f"Test loss: {test_loss:.6f}")
    logger.info(f"Best checkpoint: {best_model_path}")


if __name__ == "__main__":
    main()