scripts/train.py

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