Add exposure autoencoder pretraining

README.md

@@ -58,3 +58,15 @@ disease event token + pre-onset exposure embedding -> same next-token Transforme
 ```
 
 The key constraint is that a disease event's own pre-onset exposure must not be used to predict that same disease event.
+
+Pretrain the exposure encoder as a denoising autoencoder using training-set EIDs:
+
+```bash
+python train_exposure_autoencoder.py \
+  --exposure_cache_dir ukb_exposure_cache \
+  --train_eid_file ukb_train_eid.csv
+```
+
+The best checkpoint contains both `model_state_dict`, an `encoder_state_dict`
+compatible with the default gated `TimesNetExposureEncoder`, and the channel
+normalization statistics needed when the encoder is attached to DeepHealth.

backbones.py

@@ -474,3 +474,117 @@ class TimesNetExposureEncoder(nn.Module):
         if self.gate is not None:
             h = torch.sigmoid(self.gate) * h
         return h
+
+
+class TimesNetSequenceDecoder(nn.Module):
+    """Decode a fixed-size latent vector into a multivariate time series."""
+
+    def __init__(
+        self,
+        output_dim: int,
+        latent_dim: int,
+        d_model: int,
+        n_layers: int = 2,
+        top_k: int = 3,
+        n_convnext_blocks: int = 2,
+        conv_kernel_size: int = 7,
+        mlp_ratio: float = 4.0,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.latent_proj = nn.Linear(latent_dim, d_model)
+        self.position_proj = nn.Linear(3, d_model)
+        self.blocks = nn.ModuleList([
+            TimesNetBlock(
+                d_model=d_model,
+                top_k=top_k,
+                n_convnext_blocks=n_convnext_blocks,
+                conv_kernel_size=conv_kernel_size,
+                mlp_ratio=mlp_ratio,
+                dropout=dropout,
+            )
+            for _ in range(n_layers)
+        ])
+        self.final_ln = nn.LayerNorm(d_model)
+        self.output_proj = nn.Linear(d_model, output_dim)
+
+    def forward(self, latent: torch.Tensor, length: int) -> torch.Tensor:
+        if latent.dim() != 2:
+            raise ValueError(
+                f"latent must have shape (B, D), got {tuple(latent.shape)}"
+            )
+        position = torch.linspace(
+            0.0, 1.0, length, device=latent.device, dtype=latent.dtype
+        )
+        position = torch.stack(
+            [position, torch.sin(2 * torch.pi * position),
+             torch.cos(2 * torch.pi * position)],
+            dim=-1,
+        )
+        h = self.latent_proj(latent).unsqueeze(1)
+        h = h + self.position_proj(position).unsqueeze(0)
+        for block in self.blocks:
+            h = block(h)
+        return self.output_proj(self.final_ln(h))
+
+
+class TimesNetExposureAutoencoder(nn.Module):
+    """Dual-resolution exposure autoencoder with a reusable event encoder."""
+
+    def __init__(
+        self,
+        n_embd: int = 120,
+        daily_input_dim: int = 4,
+        monthly_input_dim: int = 2,
+        d_model: int | None = None,
+        n_layers: int = 2,
+        top_k: int = 3,
+        n_convnext_blocks: int = 2,
+        conv_kernel_size: int = 7,
+        mlp_ratio: float = 4.0,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        d_model = n_embd if d_model is None else d_model
+        encoder_kwargs = dict(
+            n_embd=n_embd,
+            daily_input_dim=daily_input_dim,
+            monthly_input_dim=monthly_input_dim,
+            d_model=d_model,
+            n_layers=n_layers,
+            top_k=top_k,
+            n_convnext_blocks=n_convnext_blocks,
+            conv_kernel_size=conv_kernel_size,
+            mlp_ratio=mlp_ratio,
+            dropout=dropout,
+            use_gate=True,
+        )
+        decoder_kwargs = dict(
+            latent_dim=n_embd,
+            d_model=d_model,
+            n_layers=n_layers,
+            top_k=top_k,
+            n_convnext_blocks=n_convnext_blocks,
+            conv_kernel_size=conv_kernel_size,
+            mlp_ratio=mlp_ratio,
+            dropout=dropout,
+        )
+        self.encoder = TimesNetExposureEncoder(**encoder_kwargs)
+        self.daily_decoder = TimesNetSequenceDecoder(
+            output_dim=daily_input_dim, **decoder_kwargs
+        )
+        self.monthly_decoder = TimesNetSequenceDecoder(
+            output_dim=monthly_input_dim, **decoder_kwargs
+        )
+
+    def forward(
+        self,
+        daily: torch.Tensor,
+        monthly: torch.Tensor,
+        daily_mask: torch.Tensor | None = None,
+        monthly_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        latent = self.encoder(daily, monthly, daily_mask, monthly_mask)
+        daily_reconstruction = self.daily_decoder(latent, daily.size(1))
+        monthly_reconstruction = self.monthly_decoder(latent, monthly.size(1))
+        return daily_reconstruction, monthly_reconstruction, latent

dataset.py

@@ -42,7 +42,7 @@ def _monthly_exposure_columns() -> list[str]:
 
 
 def _load_readonly_npy(path: Path) -> np.ndarray:
-    arr = np.load(path)
+    arr = np.load(path, mmap_mode="r")
     arr.setflags(write=False)
     return arr
 

train_exposure_autoencoder.py

@@ -0,0 +1,330 @@
+"""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()