from __future__ import annotations import argparse import os import sys import time from pathlib import Path import numpy as np import torch PROJECT_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(PROJECT_ROOT)) sys.path.insert(0, str(Path(__file__).resolve().parent)) from common import ( # noqa: E402 load_config, project_path, resolve_device, resolve_dtype, seed_everything, ) from model.sa_pinn import build_model, loss_components, weighted_loss # noqa: E402 from problems import ( # noqa: E402 CASES, build_equation, generate_data, point_counts, relative_l2, to_tensors, ) DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Train a self-adaptive PINN") parser.add_argument( "--config", type=Path, default=Path(os.environ.get("SAPINN_CONFIG", DEFAULT_CONFIG)), help="YAML configuration file", ) parser.add_argument("--case", choices=CASES, default="laplace") parser.add_argument("--device", help="Override common.device, for example cpu or cuda:0") parser.add_argument("--seed", type=int, help="Override common.seed") parser.add_argument("--epochs", type=int, help="Override Adam iterations") parser.add_argument("--lr", type=float, help="Override network learning rate") parser.add_argument("--attention-lr", type=float, help="Override attention learning rate") parser.add_argument("--lbfgs-iters", type=int, help="Override L-BFGS iterations") parser.add_argument("--n-sol", type=int, help="Override data or initial-condition points") parser.add_argument("--n-pde", type=int, help="Override PDE collocation points") parser.add_argument("--n-bnd", type=int, help="Override boundary points") parser.add_argument("--test-res", type=int, help="Override test resolution per axis") parser.add_argument("--weight-dir", type=Path, help="Override checkpoint output directory") parser.add_argument("--result-dir", type=Path, help="Override history output directory") return parser.parse_args() def main() -> None: args = parse_args() config_path = args.config.expanduser().resolve() config = load_config(config_path) common = config["common"] case_config = config["cases"][args.case] device = resolve_device(args.device or str(common["device"])) dtype = resolve_dtype(str(common["dtype"])) seed = args.seed if args.seed is not None else int(common["seed"]) epochs = ( args.epochs if args.epochs is not None else int(case_config["training"]["epochs"]) ) lbfgs_iters = ( args.lbfgs_iters if args.lbfgs_iters is not None else int(case_config["training"]["lbfgs_iters"]) ) network_lr = ( args.lr if args.lr is not None else float(case_config["training"]["lr"]) ) attention_lr = ( args.attention_lr if args.attention_lr is not None else float(case_config["attention"]["lr"]) ) if min(epochs, lbfgs_iters) < 0 or epochs + lbfgs_iters == 0: raise ValueError("at least one optimizer iteration count must be positive") if min(network_lr, attention_lr) <= 0: raise ValueError("network and attention learning rates must be positive") data_config = dict(case_config["data"]) for key, value in ( ("n_sol", args.n_sol), ("n_pde", args.n_pde), ("n_bnd", args.n_bnd), ("test_res", args.test_res), ): if value is not None: data_config[key] = value weight_dir = project_path(args.weight_dir or common["weight_dir"], PROJECT_ROOT) result_dir = project_path(args.result_dir or common["result_dir"], PROJECT_ROOT) checkpoint_path = weight_dir / case_config["output"]["checkpoint_name"] history_path = result_dir / case_config["output"]["history_name"] seed_everything(seed) data = generate_data(args.case, data_config, seed) counts = point_counts(data) tensors = to_tensors(data, device, dtype) attention_enabled = bool(case_config["attention"]["enabled"]) model = build_model( case_config["model"], counts, attention_enabled, dtype ).to(device=device, dtype=dtype) equation = build_equation(args.case, data_config) parameter_count = sum(parameter.numel() for parameter in model.parameters()) print(f"Config: {config_path}") print(f"Case: {args.case}") print(f"Device: {device}") print(f"Points: {counts}") print(f"Parameters: {parameter_count:,}") def loss_value() -> tuple[torch.Tensor, dict[str, torch.Tensor]]: components = loss_components(model, equation, tensors) return weighted_loss(components, case_config["loss"]), components def evaluate() -> float | None: with torch.no_grad(): prediction = model(tensors["x_test"]).cpu().numpy() return relative_l2(prediction, data["u_exact"]) network_optimizer = torch.optim.Adam(model.network_parameters(), lr=network_lr) attention_parameters = list(model.attention_parameters()) attention_optimizer = ( torch.optim.Adam(attention_parameters, lr=attention_lr) if attention_parameters else None ) history = [] l2_history = [] log_interval = int(case_config["training"]["log_interval"]) started = time.time() for epoch in range(1, epochs + 1): network_optimizer.zero_grad(set_to_none=True) if attention_optimizer is not None: attention_optimizer.zero_grad(set_to_none=True) loss, components = loss_value() if not torch.isfinite(loss): raise FloatingPointError(f"SA-PINN loss became non-finite at epoch {epoch}") loss.backward() network_optimizer.step() if attention_optimizer is not None: for parameter in attention_parameters: if parameter.grad is not None: parameter.grad.neg_() attention_optimizer.step() history.append(loss.item()) if epoch == 1 or epoch % log_interval == 0 or epoch == epochs: error = evaluate() if error is not None: l2_history.append((epoch, error)) attention = model.att_pde() error_text = "N/A" if error is None else f"{error:.3e}" print( f"epoch={epoch:6d} loss={loss.item():.3e} " f"data={components['data'].item():.3e} " f"boundary={components['boundary'].item():.3e} " f"pde={components['pde'].item():.3e} l2={error_text} " f"attention_std={attention.std(unbiased=False).item():.3e}" ) if lbfgs_iters > 0: model.set_attention_trainable(False) lbfgs = torch.optim.LBFGS( model.network_parameters(), lr=0.8, max_iter=lbfgs_iters, max_eval=max(1, 2 * lbfgs_iters), tolerance_grad=1.0e-7, tolerance_change=1.0e-9, history_size=50, line_search_fn="strong_wolfe", ) def closure() -> torch.Tensor: lbfgs.zero_grad(set_to_none=True) closure_loss, _ = loss_value() if not torch.isfinite(closure_loss): raise FloatingPointError("SA-PINN L-BFGS loss became non-finite") closure_loss.backward() return closure_loss lbfgs.step(closure) model.set_attention_trainable(True) error = evaluate() if error is not None: l2_history.append((epochs + lbfgs_iters, error)) print(f"L-BFGS L2={'N/A' if error is None else f'{error:.6e}'}") final_l2 = evaluate() elapsed = time.time() - started weight_dir.mkdir(parents=True, exist_ok=True) result_dir.mkdir(parents=True, exist_ok=True) checkpoint = { "case": args.case, "architecture": "sa_pinn", "model_state": model.state_dict(), "model_config": case_config["model"], "attention_enabled": attention_enabled, "point_counts": counts, "data_config": data_config, "seed": seed, "epochs": epochs, "lbfgs_iters": lbfgs_iters, "final_l2": final_l2, } torch.save(checkpoint, checkpoint_path) l2_array = np.asarray(l2_history, dtype=np.float64).reshape(-1, 2) np.savez_compressed( history_path, loss=np.asarray(history, dtype=np.float64), l2_steps=l2_array[:, 0] if l2_array.size else np.array([]), l2_values=l2_array[:, 1] if l2_array.size else np.array([]), elapsed_seconds=elapsed, final_l2=np.nan if final_l2 is None else final_l2, ) print(f"Final L2={'N/A' if final_l2 is None else f'{final_l2:.6e}'}") print(f"Saved checkpoint: {checkpoint_path}") print(f"Saved history: {history_path}") if __name__ == "__main__": main()