#!/usr/bin/env python3 """4000 full-objective Adam updates; every frame contributes to every update.""" from __future__ import annotations import argparse import csv import json import os import platform import sys import time from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import numpy as np import torch from scripts.common import (ROOT, atomic_checkpoint, build_model, gather, load_dns, predict, read_config, relative_l2, resolve, sample_indices, save_json, setup, tensor) def objective(model, boundary_q, boundary_uv, residual_q, cfg, backward): """Exact sample means; chunking changes only summation order, not optimizer steps.""" chunk = cfg["training"]["gradient_chunk_size"] data_value, residual_value = 0.0, 0.0 flow = model.flow_taylor if cfg["training"]["derivative_backend"] == "taylor" else model.flow residual = model.residual_taylor if cfg["training"]["derivative_backend"] == "taylor" else model.residual_autograd for start in range(0, len(boundary_q), chunk): q = boundary_q[start:start + chunk] uv = boundary_uv[start:start + chunk] loss = ((flow(q)[:, :2] - uv)**2).sum() / len(boundary_q) if backward: loss.backward() data_value += float(loss.detach()) for start in range(0, len(residual_q), chunk): _, r, _ = residual(residual_q[start:start + chunk]) loss = (r**2).sum() / len(residual_q) if backward: loss.backward() residual_value += float(loss.detach()) regularization = cfg["training"]["slope_weight"] * model.slope_recovery() if backward: regularization.backward() reg_value = float(regularization.detach()) return {"loss": data_value + residual_value + reg_value, "loss_boundary": data_value, "loss_residual": residual_value, "loss_slope": reg_value} def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", default=str(ROOT / "config/config.yaml")) parser.add_argument("--device") parser.add_argument("--resume", type=Path) args = parser.parse_args() cfg = read_config(args.config) device, dtype = setup(cfg, args.device) dns = load_dns(cfg) out = resolve(cfg["outputs"]["results"]) out.mkdir(parents=True, exist_ok=True) indices = sample_indices(cfg, dns) np.savez_compressed(out / "sampling_indices.npz", **indices) bq, buv = gather(dns, indices["boundary"]) rq, _ = gather(dns, indices["residual"]) vq, vuv = gather(dns, indices["validation"]) bq, buv, rq = (tensor(x, device, dtype) for x in (bq, buv, rq)) model = build_model(cfg, device) tr = cfg["training"] optimizer = torch.optim.Adam(model.parameters(), lr=tr["learning_rate"], betas=tuple(tr["adam_betas"]), eps=tr["adam_eps"]) start_step, best_loss, best_step, elapsed_prior = 0, float("inf"), -1, 0.0 if args.resume: checkpoint = torch.load(args.resume, map_location=device, weights_only=False) if checkpoint["config"] != cfg or checkpoint["data_sha256"] != dns["sha256"]: raise ValueError("Resume requires identical config and DNS SHA256") model.load_state_dict(checkpoint["model_state_dict"]) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) start_step = checkpoint["completed_updates"] best_loss, best_step = checkpoint["best_loss"], checkpoint["best_step"] elapsed_prior = checkpoint["elapsed_seconds"] torch.set_rng_state(checkpoint["torch_rng_state"].cpu()) if device.type == "cuda" and checkpoint.get("cuda_rng_state") is not None: torch.cuda.set_rng_state(checkpoint["cuda_rng_state"].cpu()) run_start = time.perf_counter() parameter_count = sum(p.numel() for p in model.parameters()) manifest = {"status": "running", "config": cfg, "data_path": dns["path"], "data_sha256": dns["sha256"], "data_shapes": cfg["data"]["expected_shapes"], "time_range": [float(dns["t"][0]), float(dns["t"][-1])], "time_spacing": float(np.diff(dns["t"])[0]), "boundary_samples": len(bq), "residual_samples": len(rq), "sample_counts_per_time": {"residual": int(indices["residual"].shape[1]), "boundary": int(indices["boundary"].shape[1])}, "validation_samples": len(vq), "parameter_count": parameter_count, "framework": "PyTorch", "torch_version": torch.__version__, "python_version": platform.python_version(), "device": str(device), "accelerator": torch.cuda.get_device_name(device) if device.type == "cuda" else "CPU", "slurm_job_id": os.environ.get("SLURM_JOB_ID"), "dtype": str(dtype), "derivative_backend": tr["derivative_backend"], "selection": tr["selection"], "completed_updates": start_step, "started_unix": time.time()} save_json(out / "run_manifest.json", manifest) (out / "config_used.yaml").write_text(Path(args.config).read_text()) print(json.dumps(manifest, ensure_ascii=False), flush=True) def payload(step, score): return {"format_version": 1, "class_name": "TSAPINN", "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "config": cfg, "completed_updates": step, "training_objective": score, "best_loss": best_loss, "best_step": best_step, "data_sha256": dns["sha256"], "elapsed_seconds": elapsed_prior + time.perf_counter() - run_start, "torch_rng_state": torch.get_rng_state(), "cuda_rng_state": torch.cuda.get_rng_state() if device.type == "cuda" else None} history_path = out / "training_history.csv" evaluation_path = out / "evaluation_history.csv" history_fields = ["completed_updates", "loss", "loss_boundary", "loss_residual", "loss_slope", "learning_rate", "elapsed_seconds", "frequency_min", "frequency_mean", "frequency_max"] mode = "a" if start_step else "w" with history_path.open(mode, newline="", buffering=1) as hist, evaluation_path.open(mode, newline="", buffering=1) as evaluations: writer = csv.DictWriter(hist, fieldnames=history_fields) ew = csv.DictWriter(evaluations, fieldnames=["completed_updates", "heldout_u", "heldout_v", "snapshot_u", "snapshot_v"]) if not start_step: writer.writeheader() ew.writeheader() final_score = None for completed in range(start_step, tr["iterations"] + 1): # Score belongs to CURRENT weights, after exactly `completed` updates. is_final = completed == tr["iterations"] model.train() optimizer.zero_grad(set_to_none=True) with torch.set_grad_enabled(not is_final): score = objective(model, bq, buv, rq, cfg, backward=not is_final) if not all(np.isfinite(list(score.values()))): raise FloatingPointError(f"Nonfinite objective at update {completed}: {score}") if not is_final and any(p.grad is not None and not torch.isfinite(p.grad).all() for p in model.parameters()): raise FloatingPointError(f"Nonfinite gradients at update {completed}") if score["loss"] < best_loss: best_loss, best_step = score["loss"], completed atomic_checkpoint(resolve(cfg["outputs"]["checkpoint"]), payload(completed, score)) elapsed = elapsed_prior + time.perf_counter() - run_start frequencies = torch.cat([f.detach() for f in model.frequencies]) lr_exponent = completed // tr["decay_steps"] if tr["staircase"] else completed / tr["decay_steps"] lr = tr["learning_rate"] * tr["decay_rate"]**lr_exponent for group in optimizer.param_groups: group["lr"] = lr row = {"completed_updates": completed, **score, "learning_rate": lr, "elapsed_seconds": elapsed, "frequency_min": float(frequencies.min()), "frequency_mean": float(frequencies.mean()), "frequency_max": float(frequencies.max())} writer.writerow(row) if completed % tr["log_every"] == 0 or is_final: print(f"update={completed:4d}/4000 loss={score['loss']:.8e} boundary={score['loss_boundary']:.8e} residual={score['loss_residual']:.8e} slope={score['loss_slope']:.8e} lr={lr:.8e} elapsed={elapsed:.2f}s", flush=True) if completed % tr["eval_every"] == 0 or is_final: val_prediction = predict(model, vq, device, dtype, cfg["evaluation"]["batch_size"]) val_error = relative_l2(val_prediction[:, :2], vuv) snap = int(np.argmin(abs(dns["t"] - cfg["evaluation"]["snapshot_time"]))) sq = np.column_stack((dns["xy"], np.full(len(dns["xy"]), dns["t"][snap]))) spred = predict(model, sq, device, dtype, cfg["evaluation"]["batch_size"]) se = relative_l2(spred[:, :2], dns["uv"][snap]) ew.writerow({"completed_updates": completed, "heldout_u": val_error[0], "heldout_v": val_error[1], "snapshot_u": se[0], "snapshot_v": se[1]}) print(f"EVAL update={completed} heldout_relL2(u,v)=({val_error[0]:.8e},{val_error[1]:.8e}) snapshot_t={dns['t'][snap]:.3f} relL2(u,v)=({se[0]:.8e},{se[1]:.8e})", flush=True) manifest.update(completed_updates=completed, best_step=best_step, best_loss=best_loss, elapsed_seconds=elapsed) save_json(out / "run_manifest.json", manifest) if completed % tr["checkpoint_every"] == 0 or is_final: atomic_checkpoint(resolve(cfg["outputs"]["last_checkpoint"]), payload(completed, score)) if is_final: final_score = score break optimizer.step() elapsed = elapsed_prior + time.perf_counter() - run_start manifest.update(status="completed", completed_updates=tr["iterations"], elapsed_seconds=elapsed, best_step=best_step, best_loss=best_loss, final_objective=final_score, peak_device_memory_bytes=torch.cuda.max_memory_allocated() if device.type == "cuda" else 0) save_json(out / "run_manifest.json", manifest) print(f"TRAINING COMPLETE updates=4000 best_step={best_step} best_loss={best_loss:.8e} runtime_seconds={elapsed:.3f}", flush=True) if __name__ == "__main__": main()