#!/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()