from __future__ import annotations import argparse import csv import os import sys from pathlib import Path from time import time import numpy as np import torch PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from model.ContinuousFlowOperator import ( # noqa: E402 CFOBatch, ContinuousFlowOperator, quintic_spline_coefficients_np, rmse, trajectory_relative_l2, ) from scripts.common import ensure_dirs, load_burgers_splits, load_config, write_json # noqa: E402 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Train Quintic CFO on 1D Burgers") parser.add_argument("--config", default=str(PROJECT_ROOT / "config" / "config.yaml")) parser.add_argument("--steps", type=int, default=None) parser.add_argument("--batch-size", type=int, default=None) parser.add_argument("--eval-interval", type=int, default=None) parser.add_argument("--print-interval", type=int, default=None) parser.add_argument("--smoke", action="store_true") return parser.parse_args() def device_from_runtime() -> torch.device: if torch.cuda.is_available(): return torch.device("cuda") return torch.device("cpu") def prepare_splines(cfg: dict, train_data: np.ndarray, tspan: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: results_dir = Path(cfg["paths"]["results_dir"]) cache_path = results_dir / "quintic_spline_train_seed43_ratio100.npz" if cfg["training"].get("spline_cache", True) and cache_path.exists(): cached = np.load(cache_path) return cached["coeffs"], cached["starts"], cached["ends"] print("Preparing quintic spline coefficients for train trajectories...", flush=True) coeffs, starts, ends = quintic_spline_coefficients_np(train_data, tspan) if cfg["training"].get("spline_cache", True): np.savez_compressed(cache_path, coeffs=coeffs, starts=starts, ends=ends) return coeffs, starts.astype(np.float32), ends.astype(np.float32) def make_batch( coeffs: np.ndarray, starts: np.ndarray, ends: np.ndarray, batch_size: int, device: torch.device, generator: torch.Generator, ) -> CFOBatch: idx = torch.randint(0, len(coeffs), (batch_size,), generator=generator) idx_np = idx.cpu().numpy() spline_coef = torch.from_numpy(coeffs[idx_np]).to(device=device, dtype=torch.float32) t_start = torch.from_numpy(starts[idx_np]).to(device=device, dtype=torch.float32) t_end = torch.from_numpy(ends[idx_np]).to(device=device, dtype=torch.float32) dt = t_end - t_start delta_t = torch.rand(batch_size, generator=generator).to(device=device, dtype=torch.float32) * dt eps = torch.randn(tuple(spline_coef[:, 0].shape), generator=generator).to(device=device, dtype=torch.float32) return CFOBatch(spline_coef=spline_coef, t_start=t_start, t_end=t_end, delta_t=delta_t, eps=eps) @torch.no_grad() def evaluate(model: ContinuousFlowOperator, data: np.ndarray, device: torch.device, steps_per_segment: int, batch_size: int) -> dict: preds = [] for start in range(0, len(data), batch_size): chunk = torch.from_numpy(data[start : start + batch_size, 0]).to(device=device, dtype=torch.float32) pred = model.rollout(chunk, trajectory_points=data.shape[1], steps_per_segment=steps_per_segment) preds.append(pred.numpy()) pred_arr = np.concatenate(preds, axis=0) return { "trajectory_relative_l2": trajectory_relative_l2(data, pred_arr), "rmse": rmse(data, pred_arr), } def save_checkpoint(path: Path, model: ContinuousFlowOperator, optimizer: torch.optim.Optimizer, step: int, metrics: dict, cfg: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_suffix(".tmp") torch.save( { "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "step": step, "metrics": metrics, "config": cfg, }, tmp, ) tmp.replace(path) def main() -> None: args = parse_args() cfg = load_config(args.config) ensure_dirs(cfg) train_cfg = cfg["training"] infer_cfg = cfg["inference"] if args.steps is not None: train_cfg["steps"] = args.steps if args.batch_size is not None: train_cfg["batch_size"] = args.batch_size if args.eval_interval is not None: train_cfg["eval_interval"] = args.eval_interval if args.print_interval is not None: train_cfg["print_interval"] = args.print_interval if args.smoke: train_cfg["steps"] = min(int(train_cfg["steps"]), 2) train_cfg["batch_size"] = min(int(train_cfg["batch_size"]), 4) train_cfg["eval_interval"] = 1 train_cfg["print_interval"] = 1 np.random.seed(int(train_cfg["seed"])) torch.manual_seed(int(train_cfg["seed"])) device = device_from_runtime() print(f"Runtime device={device} torch={torch.__version__} cuda_available={torch.cuda.is_available()}", flush=True) splits, tspan = load_burgers_splits(cfg) train_data, val_data, test_data = splits["train"], splits["val"], splits["test"] print(f"Data split shapes train={train_data.shape} val={val_data.shape} test={test_data.shape}", flush=True) coeffs, starts, ends = prepare_splines(cfg, train_data, tspan) print(f"Spline dataset coeffs={coeffs.shape} starts={starts.shape} ends={ends.shape}", flush=True) model = ContinuousFlowOperator(input_size=cfg["model"]["input_size"], gamma=cfg["model"]["gamma"]).to(device) optimizer = torch.optim.Adam( model.parameters(), lr=float(train_cfg["learning_rate"]), betas=tuple(float(x) for x in train_cfg["betas"]), ) gen = torch.Generator(device="cpu") gen.manual_seed(int(train_cfg["seed"])) history_path = Path(cfg["paths"]["results_dir"]) / "train_history.csv" best_ckpt = Path(train_cfg["best_checkpoint"]) latest_ckpt = Path(train_cfg["latest_checkpoint"]) best_val = float("inf") best_step = -1 start_time = time() with history_path.open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=["step", "loss", "val_rel_l2", "val_rmse", "elapsed_sec"]) writer.writeheader() for step in range(1, int(train_cfg["steps"]) + 1): model.train() batch = make_batch(coeffs, starts, ends, int(train_cfg["batch_size"]), device, gen) optimizer.zero_grad(set_to_none=True) loss = model.loss_from_batch(batch) loss.backward() optimizer.step() row = { "step": step, "loss": float(loss.detach().cpu()), "val_rel_l2": "", "val_rmse": "", "elapsed_sec": time() - start_time, } do_eval = (step % int(train_cfg["eval_interval"]) == 0) or args.smoke if do_eval: metrics = evaluate( model, val_data, device, steps_per_segment=int(infer_cfg["steps_per_segment"]), batch_size=int(infer_cfg["eval_batch_size"]), ) row["val_rel_l2"] = metrics["trajectory_relative_l2"] row["val_rmse"] = metrics["rmse"] if metrics["trajectory_relative_l2"] < best_val: best_val = metrics["trajectory_relative_l2"] best_step = step save_checkpoint(best_ckpt, model, optimizer, step, metrics, cfg) print( f"step={step} loss={row['loss']:.8e} val_rel_l2={metrics['trajectory_relative_l2']:.8e} " f"val_rmse={metrics['rmse']:.8e} best_step={best_step} best_val={best_val:.8e}", flush=True, ) elif step % int(train_cfg["print_interval"]) == 0 or step == 1: print(f"step={step} loss={row['loss']:.8e} elapsed_sec={row['elapsed_sec']:.1f}", flush=True) writer.writerow(row) f.flush() final_metrics = evaluate( model, test_data, device, steps_per_segment=int(infer_cfg["steps_per_segment"]), batch_size=int(infer_cfg["eval_batch_size"]), ) save_checkpoint(latest_ckpt, model, optimizer, int(train_cfg["steps"]), final_metrics, cfg) summary = { "status": "success", "steps": int(train_cfg["steps"]), "best_step": best_step, "best_val_trajectory_relative_l2": best_val, "latest_test_trajectory_relative_l2": final_metrics["trajectory_relative_l2"], "latest_test_rmse": final_metrics["rmse"], "best_checkpoint": str(best_ckpt), "latest_checkpoint": str(latest_ckpt), } write_json(Path(cfg["paths"]["results_dir"]) / "train_summary.json", summary) print("Training complete.", flush=True) print(summary, flush=True) if __name__ == "__main__": main()