scripts/train.py

#!/usr/bin/env python3
"""Train the paper-faithful FAST-DeepONet Darcy experiment."""

from __future__ import annotations

import argparse
import csv
import math
import os
import sys
import time
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch import Tensor
from torch.utils.data import DataLoader, TensorDataset

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from model.FASTDeepONet import FASTDeepONet, compute_spectral_basis
from scripts.common import (
    atomic_write_json,
    balanced_train_queries,
    fit_target_transform,
    fixed_tune_queries,
    load_config,
    load_split,
    make_coordinate_features,
    relative_l2,
    seed_everything,
    select_device,
)


def validate_config(config: dict[str, Any]) -> None:
    data = config["data"]
    model = config["model"]
    training = config["training"]
    regularization = config["regularization"]
    required = {
        "data.grid_size": (data["grid_size"], 64),
        "data.points_per_field": (data["points_per_field"], 4096),
        "model.branch_input_dim": (model["branch_input_dim"], 4096),
        "model.requested_factor_rank": (model["requested_factor_rank"], 16),
        "model.effective_factor_rank": (model["effective_factor_rank"], 16),
        "model.residual_width": (model["residual_width"], 16),
        "model.factor_scale": (float(model["factor_scale"]), 1.0),
        "model.branch_width": (model["branch_width"], 128),
        "model.trunk_width": (model["trunk_width"], 128),
        "model.latent_dim": (model["latent_dim"], 128),
        "model.hidden_layers": (model["hidden_layers"], 3),
        "model.fourier_order": (model["fourier_order"], 4),
        "model.trunk_input_dim": (model["trunk_input_dim"], 18),
        "regularization.directional_lambda": (float(regularization["directional_lambda"]), 1.0),
        "training.epochs": (training["epochs"], 300),
        "training.batch_size": (training["batch_size"], 2048),
        "training.train_queries_per_epoch": (training["train_queries_per_epoch"], 8192),
        "training.tune_queries": (training["tune_queries"], 4096),
        "training.model_seed": (training["model_seed"], 50000000),
        "training.schedule_seed": (training["schedule_seed"], 20260715),
    }
    mismatches = [f"{key}: {actual!r} != {expected!r}" for key, (actual, expected) in required.items() if actual != expected]
    if model["raw_coordinate_dim"] + 4 * model["fourier_order"] != model["trunk_input_dim"]:
        mismatches.append("raw/Fourier feature count does not match trunk_input_dim")
    if training["amp"] or training["early_stopping"]:
        mismatches.append("AMP and early stopping must both be disabled")
    if config["data"]["split_counts"] != {"train": 60, "tune": 10, "calib": 20, "test": 10}:
        mismatches.append("split_counts must be train/tune/calib/test=60/10/20/10")
    if mismatches:
        raise ValueError("Configuration is inconsistent:\n- " + "\n- ".join(mismatches))


def build_model(config: dict[str, Any], spectral_basis: np.ndarray) -> FASTDeepONet:
    model_config = config["model"]
    return FASTDeepONet(
        spectral_basis=spectral_basis,
        branch_input_dim=model_config["branch_input_dim"],
        requested_factor_rank=model_config["requested_factor_rank"],
        residual_width=model_config["residual_width"],
        factor_scale=model_config["factor_scale"],
        branch_width=model_config["branch_width"],
        trunk_input_dim=model_config["trunk_input_dim"],
        trunk_width=model_config["trunk_width"],
        latent_dim=model_config["latent_dim"],
        hidden_layers=model_config["hidden_layers"],
    )


def tau_for_epoch(epoch: int, epochs: int, initial: float, final: float) -> float:
    return max(final, initial - (initial - final) * epoch / epochs)


def query_loader(
    field_indices: np.ndarray,
    coordinate_indices: np.ndarray,
    batch_size: int,
    shuffle: bool,
    generator_seed: int | None,
    num_workers: int,
) -> DataLoader:
    dataset = TensorDataset(
        torch.from_numpy(field_indices.astype(np.int64, copy=False)),
        torch.from_numpy(coordinate_indices.astype(np.int64, copy=False)),
    )
    generator = None
    if generator_seed is not None:
        generator = torch.Generator()
        generator.manual_seed(generator_seed)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        generator=generator,
        num_workers=num_workers,
        drop_last=False,
    )


def evaluate_queries(
    model: FASTDeepONet,
    inputs: Tensor,
    normalized_targets: Tensor,
    coordinates: Tensor,
    loader: DataLoader,
    device: torch.device,
) -> float:
    model.eval()
    squared_error_sum = 0.0
    sample_count = 0
    with torch.inference_mode():
        for fields_cpu, coordinates_cpu in loader:
            fields = fields_cpu.to(device=device, non_blocking=False)
            coordinate_ids = coordinates_cpu.to(device=device, non_blocking=False)
            prediction = model(inputs[fields], coordinates[coordinate_ids])
            target = normalized_targets[fields, coordinate_ids]
            squared_error_sum += float(torch.sum((prediction - target) ** 2).item())
            sample_count += int(target.numel())
    if sample_count == 0:
        raise RuntimeError("Evaluation loader produced no samples")
    return squared_error_sum / sample_count


def write_history_csv(path: Path, history: list[dict[str, float | int]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    fieldnames = [
        "epoch",
        "tau",
        "learning_rate",
        "train_total_loss",
        "train_mse",
        "train_regularizer",
        "tune_mse",
        "best_tune_mse",
    ]
    with temporary.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(history)
    os.replace(temporary, path)


def atomic_torch_save(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    torch.save(payload, temporary)
    os.replace(temporary, path)


def run_formal_training(config: dict[str, Any], config_path: Path) -> None:
    validate_config(config)
    if config["runtime"]["formal_training_requires_slurm"] and not os.environ.get("SLURM_JOB_ID"):
        raise RuntimeError("Formal 300-epoch training is only allowed inside a SLURM compute-node job")
    device = select_device(require_accelerator=True)
    training = config["training"]
    model_config = config["model"]
    data_config = config["data"]
    regularization = config["regularization"]
    paths = config["paths"]
    seed_everything(int(training["model_seed"]))

    print(
        f"formal_training=true slurm_job_id={os.environ.get('SLURM_JOB_ID')} "
        f"device={device} device_name={torch.cuda.get_device_name(device)}",
        flush=True,
    )
    split_counts = data_config["split_counts"]
    train_names, train_inputs_np, train_targets_np, _ = load_split(
        paths["data_root"], "train", split_counts["train"], data_config["points_per_field"]
    )
    tune_names, tune_inputs_np, tune_targets_np, _ = load_split(
        paths["data_root"], "tune", split_counts["tune"], data_config["points_per_field"]
    )
    basis, singular_values, numerical_rank, effective_rank = compute_spectral_basis(
        train_inputs_np,
        model_config["requested_factor_rank"],
        model_config["svd_relative_tolerance"],
    )
    if effective_rank != model_config["effective_factor_rank"]:
        raise RuntimeError(
            f"Effective factor rank is {effective_rank}, expected {model_config['effective_factor_rank']}"
        )
    target_mean_np, target_std = fit_target_transform(train_targets_np)
    train_targets_normalized_np = (train_targets_np - target_mean_np) / target_std
    tune_targets_normalized_np = (tune_targets_np - target_mean_np) / target_std
    coordinate_features_np = make_coordinate_features(
        data_config["grid_size"], model_config["fourier_order"]
    )
    if coordinate_features_np.shape != (data_config["points_per_field"], model_config["trunk_input_dim"]):
        raise RuntimeError(f"Unexpected coordinate feature shape: {coordinate_features_np.shape}")

    model = build_model(config, basis).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=float(training["learning_rate"]))
    scheduler = torch.optim.lr_scheduler.ExponentialLR(
        optimizer, gamma=float(training["scheduler_gamma"])
    )
    train_inputs = torch.from_numpy(train_inputs_np).to(device)
    train_targets = torch.from_numpy(train_targets_normalized_np.astype(np.float32)).to(device)
    tune_inputs = torch.from_numpy(tune_inputs_np).to(device)
    tune_targets = torch.from_numpy(tune_targets_normalized_np.astype(np.float32)).to(device)
    coordinates = torch.from_numpy(coordinate_features_np).to(device)

    tune_fields_np, tune_coordinates_np = fixed_tune_queries(
        split_counts["tune"],
        data_config["points_per_field"],
        training["tune_queries"],
        training["tune_seed"],
    )
    tune_loader = query_loader(
        tune_fields_np,
        tune_coordinates_np,
        training["batch_size"],
        False,
        None,
        training["num_workers"],
    )

    checkpoint_path = Path(paths["checkpoint"])
    results_dir = Path(paths["results_dir"])
    results_dir.mkdir(parents=True, exist_ok=True)
    history: list[dict[str, float | int]] = []
    best_tune_mse = math.inf
    best_epoch = 0
    start_time = time.time()

    for zero_based_epoch in range(training["epochs"]):
        epoch = zero_based_epoch + 1
        tau = tau_for_epoch(
            epoch,
            training["epochs"],
            float(regularization["tau_initial"]),
            float(regularization["tau_final"]),
        )
        train_fields_np, train_coordinates_np = balanced_train_queries(
            split_counts["train"],
            data_config["points_per_field"],
            training["train_queries_per_epoch"],
            training["schedule_seed"],
            zero_based_epoch,
        )
        train_loader = query_loader(
            train_fields_np,
            train_coordinates_np,
            training["batch_size"],
            True,
            int(training["model_seed"] + 1009 * epoch),
            training["num_workers"],
        )

        model.train()
        totals = {"total": 0.0, "mse": 0.0, "regularizer": 0.0, "queries": 0}
        for fields_cpu, coordinates_cpu in train_loader:
            fields = fields_cpu.to(device=device, non_blocking=False)
            coordinate_ids = coordinates_cpu.to(device=device, non_blocking=False)
            optimizer.zero_grad(set_to_none=True)
            prediction = model(train_inputs[fields], coordinates[coordinate_ids])
            target = train_targets[fields, coordinate_ids]
            mse = torch.mean((prediction - target) ** 2)
            directional = model.directional_regularizer(tau)
            total_loss = mse + float(regularization["directional_lambda"]) * directional
            total_loss.backward()
            optimizer.step()
            batch_queries = int(target.numel())
            totals["total"] += float(total_loss.detach().item()) * batch_queries
            totals["mse"] += float(mse.detach().item()) * batch_queries
            totals["regularizer"] += float(directional.detach().item()) * batch_queries
            totals["queries"] += batch_queries
        if totals["queries"] != training["train_queries_per_epoch"]:
            raise RuntimeError(
                f"Epoch {epoch} consumed {totals['queries']} queries, expected {training['train_queries_per_epoch']}"
            )
        scheduler.step()
        tune_mse = evaluate_queries(
            model, tune_inputs, tune_targets, coordinates, tune_loader, device
        )
        train_total = totals["total"] / totals["queries"]
        train_mse = totals["mse"] / totals["queries"]
        train_regularizer = totals["regularizer"] / totals["queries"]
        learning_rate = float(optimizer.param_groups[0]["lr"])

        if not all(math.isfinite(value) for value in (train_total, train_mse, train_regularizer, tune_mse)):
            raise FloatingPointError(f"Non-finite metric detected at epoch {epoch}")
        if tune_mse < best_tune_mse:
            best_tune_mse = tune_mse
            best_epoch = epoch
            checkpoint = {
                "model_state_dict": model.state_dict(),
                "model_arguments": model.model_arguments(),
                "spectral_basis": torch.from_numpy(basis),
                "singular_values": torch.from_numpy(singular_values.copy()),
                "numerical_rank": numerical_rank,
                "effective_factor_rank": effective_rank,
                "target_mean": torch.from_numpy(target_mean_np),
                "target_std": target_std,
                "best_epoch": best_epoch,
                "best_tune_mse": best_tune_mse,
                "config": config,
                "config_path": str(config_path),
                "train_filenames": train_names,
                "tune_filenames": tune_names,
                "training_completed_epochs_at_save": epoch,
            }
            atomic_torch_save(checkpoint_path, checkpoint)

        record: dict[str, float | int] = {
            "epoch": epoch,
            "tau": tau,
            "learning_rate": learning_rate,
            "train_total_loss": train_total,
            "train_mse": train_mse,
            "train_regularizer": train_regularizer,
            "tune_mse": tune_mse,
            "best_tune_mse": best_tune_mse,
        }
        history.append(record)
        atomic_write_json(results_dir / "training_history.json", history)
        write_history_csv(results_dir / "training_history.csv", history)
        print(
            f"epoch={epoch:03d}/{training['epochs']} tau={tau:.6f} "
            f"train_loss={train_total:.8e} train_mse={train_mse:.8e} "
            f"train_reg={train_regularizer:.8e} tune_loss={tune_mse:.8e} "
            f"best_tune_loss={best_tune_mse:.8e} lr={learning_rate:.8e}",
            flush=True,
        )

    elapsed_seconds = time.time() - start_time
    summary = {
        "status": "completed",
        "formal_training": True,
        "completed_epochs": len(history),
        "requested_epochs": training["epochs"],
        "best_epoch": best_epoch,
        "best_tune_mse": best_tune_mse,
        "checkpoint": str(checkpoint_path),
        "effective_factor_rank": effective_rank,
        "numerical_rank": numerical_rank,
        "elapsed_seconds": elapsed_seconds,
        "slurm_job_id": os.environ.get("SLURM_JOB_ID"),
        "device": str(device),
        "device_name": torch.cuda.get_device_name(device),
    }
    atomic_write_json(results_dir / "training_summary.json", summary)
    print(
        f"training_complete epochs={len(history)} best_epoch={best_epoch} "
        f"best_tune_loss={best_tune_mse:.8e} checkpoint={checkpoint_path}",
        flush=True,
    )


def run_smoke_test(config: dict[str, Any]) -> None:
    validate_config(config)
    device = select_device(require_accelerator=False)
    training = config["training"]
    model_config = config["model"]
    data_config = config["data"]
    paths = config["paths"]
    split_counts = data_config["split_counts"]
    seed_everything(int(training["model_seed"]))
    _, train_inputs_np, train_targets_np, _ = load_split(
        paths["data_root"], "train", split_counts["train"], data_config["points_per_field"]
    )
    _, tune_inputs_np, tune_targets_np, _ = load_split(
        paths["data_root"], "tune", split_counts["tune"], data_config["points_per_field"]
    )
    _, _, test_targets_np, _ = load_split(
        paths["data_root"], "test", split_counts["test"], data_config["points_per_field"]
    )
    basis, _, _, effective_rank = compute_spectral_basis(
        train_inputs_np, model_config["requested_factor_rank"], model_config["svd_relative_tolerance"]
    )
    if effective_rank != 16:
        raise AssertionError(f"smoke effective rank {effective_rank} != 16")
    target_mean_np, target_std = fit_target_transform(train_targets_np)
    train_targets_normalized_np = (train_targets_np - target_mean_np) / target_std
    tune_targets_normalized_np = (tune_targets_np - target_mean_np) / target_std
    coordinates_np = make_coordinate_features(data_config["grid_size"], model_config["fourier_order"])
    model = build_model(config, basis).to(device)
    train_inputs = torch.from_numpy(train_inputs_np).to(device)
    train_targets = torch.from_numpy(train_targets_normalized_np.astype(np.float32)).to(device)
    tune_inputs = torch.from_numpy(tune_inputs_np).to(device)
    tune_targets = torch.from_numpy(tune_targets_normalized_np.astype(np.float32)).to(device)
    coordinates = torch.from_numpy(coordinates_np).to(device)
    fields_np, coordinate_ids_np = balanced_train_queries(60, 4096, 64, 20260715, 0)
    fields = torch.from_numpy(fields_np).to(device)
    coordinate_ids = torch.from_numpy(coordinate_ids_np).to(device)
    report: dict[str, Any] = {"formal_training": False, "device": str(device), "checks": {}}

    prediction = model(train_inputs[fields], coordinates[coordinate_ids])
    if prediction.shape != (64,) or not torch.isfinite(prediction).all():
        raise AssertionError(f"forward output failed: shape={prediction.shape}")
    report["checks"]["forward_pass"] = "pass"
    print("smoke forward_pass=PASS", flush=True)

    model.zero_grad(set_to_none=True)
    mse = torch.mean((prediction - train_targets[fields, coordinate_ids]) ** 2)
    total = mse + model.directional_regularizer(0.1)
    total.backward()
    gradient_parameters = [parameter.grad for parameter in model.parameters() if parameter.requires_grad]
    if not gradient_parameters or any(gradient is None or not torch.isfinite(gradient).all() for gradient in gradient_parameters):
        raise AssertionError("backward produced missing or non-finite gradients")
    report["checks"]["backward_pass"] = "pass"
    print("smoke backward_pass=PASS", flush=True)

    optimizer = torch.optim.Adam(model.parameters(), lr=1.0e-3)
    losses: list[float] = []
    for _ in range(5):
        optimizer.zero_grad(set_to_none=True)
        prediction = model(train_inputs[fields], coordinates[coordinate_ids])
        loss = torch.mean((prediction - train_targets[fields, coordinate_ids]) ** 2) + model.directional_regularizer(0.1)
        loss.backward()
        optimizer.step()
        losses.append(float(loss.detach().item()))
    if not all(math.isfinite(value) for value in losses) or losses[-1] > losses[0]:
        raise AssertionError(f"tiny train loss did not decrease: {losses}")
    report["checks"]["train_loop_dry_run"] = {"status": "pass", "losses": losses}
    print(f"smoke train_loop_dry_run=PASS initial={losses[0]:.6e} final={losses[-1]:.6e}", flush=True)

    tune_fields_np, tune_coordinates_np = fixed_tune_queries(10, 4096, 64, 0)
    validation_loader = query_loader(tune_fields_np, tune_coordinates_np, 32, False, None, 0)
    validation_mse = evaluate_queries(
        model, tune_inputs, tune_targets, coordinates, validation_loader, device
    )
    if not math.isfinite(validation_mse):
        raise AssertionError("validation MSE is not finite")
    report["checks"]["validation_loop_dry_run"] = {"status": "pass", "mse": validation_mse}
    print(f"smoke validation_loop_dry_run=PASS mse={validation_mse:.6e}", flush=True)

    sample_metrics = [relative_l2(target, np.zeros_like(target)) for target in test_targets_np[:2]]
    mean_metric = float(np.mean(sample_metrics))
    if not all(math.isfinite(value) for value in sample_metrics + [mean_metric]):
        raise AssertionError("relative L2 smoke metric is not finite")
    report["checks"]["relative_l2_dry_run"] = {
        "status": "pass",
        "per_sample": sample_metrics,
        "mean": mean_metric,
    }
    print(f"smoke relative_l2_dry_run=PASS mean={mean_metric:.6e}", flush=True)

    expected_branch_feature_dim = model.effective_factor_rank + model.residual_width
    if expected_branch_feature_dim != 32 or coordinates_np.shape != (4096, 18):
        raise AssertionError("model/config shape consistency failed")
    report["checks"]["config_consistency"] = "pass"
    report["status"] = "passed"
    print("smoke config_consistency=PASS", flush=True)
    atomic_write_json(Path(paths["results_dir"]) / "smoke_test_report.json", report)
    print("smoke_test_complete checks=6/6 status=PASS", flush=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--config",
        type=Path,
        default=PROJECT_ROOT / "config" / "config.yaml",
        help="Experiment YAML path",
    )
    parser.add_argument(
        "--smoke-test",
        action="store_true",
        help="Run the six non-formal small-tensor checks; never writes the formal checkpoint",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    config_path = args.config.resolve()
    config = load_config(config_path)
    if args.smoke_test:
        run_smoke_test(config)
    else:
        run_formal_training(config, config_path)


if __name__ == "__main__":
    main()