DeepHealthExpo/train_exposure_autoencoder.py

"""Pretrain a TimesNet + ConvNeXtV2 autoencoder on training-set exposure."""
from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import numpy as np
import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

from backbones import TimesNetExposureAutoencoder
from dataset import ExposureCache
from train_util import (
    configure_torch_for_training,
    create_unique_run_dir,
    load_eid_file,
    resolve_device,
    set_seed,
    setup_logging,
)


class ExposureWindowDataset(Dataset):
    def __init__(self, cache: ExposureCache, row_indices: np.ndarray):
        self.cache = cache
        self.row_indices = np.asarray(row_indices, dtype=np.int64)

    def __len__(self) -> int:
        return len(self.row_indices)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        row = int(self.row_indices[index])
        return {
            "daily": torch.from_numpy(
                np.array(self.cache.daily[row], dtype=np.float32, copy=True)
            ),
            "monthly": torch.from_numpy(
                np.array(self.cache.monthly[row], dtype=np.float32, copy=True)
            ),
        }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Pretrain a TimesNet + ConvNeXtV2 exposure autoencoder"
    )
    parser.add_argument("--exposure_cache_dir", required=True)
    parser.add_argument("--train_eid_file", default="ukb_train_eid.csv")
    parser.add_argument("--runs_root", default="runs")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--val_fraction", type=float, default=0.05)
    parser.add_argument("--n_embd", type=int, default=120)
    parser.add_argument("--d_model", type=int, default=None)
    parser.add_argument("--n_layers", type=int, default=2)
    parser.add_argument("--top_k", type=int, default=3)
    parser.add_argument("--n_convnext_blocks", type=int, default=2)
    parser.add_argument("--conv_kernel_size", type=int, default=7)
    parser.add_argument("--mlp_ratio", type=float, default=4.0)
    parser.add_argument("--dropout", type=float, default=0.0)
    parser.add_argument("--mask_ratio", type=float, default=0.25)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--base_lr", type=float, default=3e-4)
    parser.add_argument("--weight_decay", type=float, default=0.05)
    parser.add_argument("--max_epochs", type=int, default=100)
    parser.add_argument("--warmup_epochs", type=int, default=5)
    parser.add_argument("--patience", type=int, default=12)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True)
    args = parser.parse_args()
    if not 0.0 < args.val_fraction < 1.0:
        parser.error("--val_fraction must be between 0 and 1")
    if not 0.0 <= args.mask_ratio < 1.0:
        parser.error("--mask_ratio must be in [0, 1)")
    return args


def select_rows(
    cache: ExposureCache, train_eids: set[int], val_fraction: float, seed: int
) -> tuple[np.ndarray, np.ndarray]:
    candidate_eids = np.asarray(
        sorted(set(map(int, cache.eids)) & train_eids), dtype=np.int64
    )
    if len(candidate_eids) < 2:
        raise ValueError("Need at least two training EIDs with cached exposure")
    rng = np.random.default_rng(seed)
    rng.shuffle(candidate_eids)
    n_val = max(1, int(round(len(candidate_eids) * val_fraction)))
    val_eids = candidate_eids[:n_val]
    fit_eids = candidate_eids[n_val:]
    valid_row = np.asarray(cache.row_index, dtype=np.int64) >= 0
    fit_event_rows = valid_row & np.isin(cache.eids, fit_eids)
    val_event_rows = valid_row & np.isin(cache.eids, val_eids)
    fit_rows = np.unique(np.asarray(cache.row_index[fit_event_rows], dtype=np.int64))
    val_rows = np.unique(np.asarray(cache.row_index[val_event_rows], dtype=np.int64))
    if len(fit_rows) == 0 or len(val_rows) == 0:
        raise ValueError("Training/validation exposure rows are empty after filtering")
    return fit_rows, val_rows


def channel_stats(
    cache: ExposureCache, rows: np.ndarray, chunk_size: int = 256
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    results = []
    for source in (cache.daily, cache.monthly):
        sums = np.zeros(source.shape[-1], dtype=np.float64)
        squares = np.zeros_like(sums)
        counts = np.zeros_like(sums)
        for start in tqdm(range(0, len(rows), chunk_size), desc="Channel statistics"):
            values = np.asarray(source[rows[start:start + chunk_size]], dtype=np.float64)
            finite = np.isfinite(values)
            clean = np.where(finite, values, 0.0)
            sums += clean.sum(axis=(0, 1))
            squares += np.square(clean).sum(axis=(0, 1))
            counts += finite.sum(axis=(0, 1))
        mean = sums / np.maximum(counts, 1.0)
        variance = squares / np.maximum(counts, 1.0) - np.square(mean)
        std = np.sqrt(np.maximum(variance, 1e-12))
        results.extend([mean.astype(np.float32), std.astype(np.float32)])
    return tuple(results)


def masked_mse(
    prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor
) -> torch.Tensor:
    error = (prediction - target).square() * mask
    return error.sum() / mask.sum().clamp_min(1.0)


def run_epoch(
    model: TimesNetExposureAutoencoder,
    loader: DataLoader,
    device: torch.device,
    stats: tuple[torch.Tensor, ...],
    mask_ratio: float,
    optimizer: AdamW | None,
    scaler: torch.amp.GradScaler,
    grad_clip: float,
    amp_enabled: bool,
) -> float:
    training = optimizer is not None
    model.train(training)
    total_loss = 0.0
    total_samples = 0
    daily_mean, daily_std, monthly_mean, monthly_std = stats
    context = torch.enable_grad if training else torch.no_grad
    with context():
        for batch in tqdm(loader, desc="train" if training else "val", leave=False):
            daily = batch["daily"].to(device, non_blocking=True)
            monthly = batch["monthly"].to(device, non_blocking=True)
            daily_observed = torch.isfinite(daily)
            monthly_observed = torch.isfinite(monthly)
            daily = (torch.nan_to_num(daily) - daily_mean) / daily_std
            monthly = (torch.nan_to_num(monthly) - monthly_mean) / monthly_std
            daily = daily * daily_observed
            monthly = monthly * monthly_observed

            if training and mask_ratio > 0:
                daily_input_mask = daily_observed & (
                    torch.rand_like(daily) >= mask_ratio
                )
                monthly_input_mask = monthly_observed & (
                    torch.rand_like(monthly) >= mask_ratio
                )
            else:
                daily_input_mask = daily_observed
                monthly_input_mask = monthly_observed
            daily_input = daily * daily_input_mask
            monthly_input = monthly * monthly_input_mask

            if training:
                optimizer.zero_grad(set_to_none=True)
            with torch.autocast(
                device_type=device.type, dtype=torch.float16,
                enabled=amp_enabled,
            ):
                daily_hat, monthly_hat, _ = model(
                    daily_input, monthly_input,
                    daily_input_mask, monthly_input_mask,
                )
                loss = (
                    masked_mse(daily_hat, daily, daily_observed)
                    + masked_mse(monthly_hat, monthly, monthly_observed)
                )
            if training:
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                scaler.step(optimizer)
                scaler.update()
            batch_size = daily.size(0)
            total_loss += float(loss.detach()) * batch_size
            total_samples += batch_size
    return total_loss / max(total_samples, 1)


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


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 stamp: f"exposure_ae_{stamp}", Path(args.runs_root)
    )
    logger = setup_logging(run_dir)
    cache = ExposureCache(args.exposure_cache_dir)
    train_rows, val_rows = select_rows(
        cache, load_eid_file(args.train_eid_file), args.val_fraction, args.seed
    )
    raw_stats = channel_stats(cache, train_rows)
    stats = tuple(
        torch.as_tensor(value, device=device).view(1, 1, -1)
        for value in raw_stats
    )
    loader_kwargs = dict(
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
    )
    train_loader = DataLoader(
        ExposureWindowDataset(cache, train_rows), shuffle=True, **loader_kwargs
    )
    val_loader = DataLoader(
        ExposureWindowDataset(cache, val_rows), shuffle=False, **loader_kwargs
    )
    model = TimesNetExposureAutoencoder(
        n_embd=args.n_embd, d_model=args.d_model, n_layers=args.n_layers,
        top_k=args.top_k, n_convnext_blocks=args.n_convnext_blocks,
        conv_kernel_size=args.conv_kernel_size, mlp_ratio=args.mlp_ratio,
        dropout=args.dropout,
    ).to(device)
    optimizer = AdamW(
        model.parameters(), lr=args.base_lr,
        weight_decay=args.weight_decay, betas=(0.9, 0.95),
    )
    amp_enabled = bool(args.amp and device.type == "cuda")
    scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)
    logger.info(
        f"Run {run_name}: device={device}, train_rows={len(train_rows):,}, "
        f"val_rows={len(val_rows):,}"
    )
    config = vars(args) | {
        "train_rows": len(train_rows),
        "val_rows": len(val_rows),
        "daily_mean": raw_stats[0].tolist(),
        "daily_std": raw_stats[1].tolist(),
        "monthly_mean": raw_stats[2].tolist(),
        "monthly_std": raw_stats[3].tolist(),
    }
    (run_dir / "train_config.json").write_text(
        json.dumps(config, indent=2), encoding="utf-8"
    )

    best_loss = float("inf")
    stale_epochs = 0
    history = []
    for epoch in range(args.max_epochs):
        lr = learning_rate(epoch, args)
        for group in optimizer.param_groups:
            group["lr"] = lr
        train_loss = run_epoch(
            model, train_loader, device, stats, args.mask_ratio, optimizer,
            scaler, args.grad_clip, amp_enabled,
        )
        val_loss = run_epoch(
            model, val_loader, device, stats, 0.0, None,
            scaler, args.grad_clip, amp_enabled,
        )
        logger.info(
            f"Epoch {epoch + 1:03d} | lr={lr:.3e} | "
            f"train={train_loss:.6f} | val={val_loss:.6f}"
        )
        history.append(
            {"epoch": epoch + 1, "lr": lr,
             "train_loss": train_loss, "val_loss": val_loss}
        )
        if val_loss < best_loss:
            best_loss = val_loss
            stale_epochs = 0
            torch.save(
                {
                    "model_state_dict": model.state_dict(),
                    "encoder_state_dict": model.encoder.state_dict(),
                    "model_config": {
                        key: config[key] for key in (
                            "n_embd", "d_model", "n_layers", "top_k",
                            "n_convnext_blocks", "conv_kernel_size",
                            "mlp_ratio", "dropout",
                        )
                    },
                    "normalization": {
                        "daily_mean": raw_stats[0],
                        "daily_std": raw_stats[1],
                        "monthly_mean": raw_stats[2],
                        "monthly_std": raw_stats[3],
                    },
                    "epoch": epoch + 1,
                    "val_loss": val_loss,
                },
                run_dir / "best.pt",
            )
        else:
            stale_epochs += 1
        (run_dir / "history.json").write_text(
            json.dumps(history, indent=2), encoding="utf-8"
        )
        if stale_epochs >= args.patience:
            logger.info(f"Early stopping after {epoch + 1} epochs")
            break
    logger.info(f"Best validation loss: {best_loss:.6f}")
    logger.info(f"Checkpoint: {run_dir / 'best.pt'}")


if __name__ == "__main__":
    main()