from __future__ import annotations import random import sys import time from pathlib import Path import numpy as np import torch import yaml PROJECT_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(PROJECT_ROOT)) sys.path.insert(0, str(Path(__file__).resolve().parent)) from data_utils import ( # noqa: E402 build_evaluation_points, build_training_batch, exact_subdomain_values, load_mat_data, ) from model.xpinn import XPINNPoisson2D # noqa: E402 DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml" def load_config(path: Path) -> dict: with path.open("r", encoding="utf-8") as stream: config = yaml.safe_load(stream) if not isinstance(config, dict) or "root" not in config: raise ValueError(f"config must contain a 'root' mapping: {path}") return config["root"] def project_path(value: str | Path) -> Path: path = Path(value).expanduser() return path if path.is_absolute() else PROJECT_ROOT / path def resolve_device(requested: str) -> torch.device: if requested == "auto": return torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device(requested) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false") return device def resolve_dtype(name: str) -> torch.dtype: try: return {"float32": torch.float32, "float64": torch.float64}[name] except KeyError as error: raise ValueError(f"unsupported dtype: {name}") from error def seed_everything(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def compute_loss( outputs: dict[str, torch.Tensor], batch: dict[str, torch.Tensor], weights: dict, ) -> tuple[torch.Tensor, dict[str, float]]: boundary_loss = torch.mean( (batch["ub"] - outputs["boundary_prediction"]).square() ) pde_loss = sum( torch.mean(outputs[f"residual{domain}"].square()) for domain in (1, 2, 3) ) interface_residual_loss = torch.mean( outputs["interface1_residual"].square() ) + torch.mean(outputs["interface2_residual"].square()) interface_value_loss = ( torch.mean( (outputs["interface1_domain1"] - outputs["interface1_average"]).square() ) + torch.mean( (outputs["interface1_domain2"] - outputs["interface1_average"]).square() ) + torch.mean( (outputs["interface2_domain1"] - outputs["interface2_average"]).square() ) + torch.mean( (outputs["interface2_domain3"] - outputs["interface2_average"]).square() ) ) total = ( float(weights["boundary"]) * boundary_loss + float(weights["pde"]) * pde_loss + float(weights["interface_residual"]) * interface_residual_loss + float(weights["interface_value"]) * interface_value_loss ) return total, { "boundary": boundary_loss.item(), "pde": pde_loss.item(), "interface_residual": interface_residual_loss.item(), "interface_value": interface_value_loss.item(), } def clear_coordinate_grads(batch: dict[str, torch.Tensor]) -> None: for value in batch.values(): if value.requires_grad and value.grad is not None: value.grad = None def evaluate( model: XPINNPoisson2D, points: dict[str, torch.Tensor], references: tuple[torch.Tensor, torch.Tensor, torch.Tensor], ) -> tuple[float, float, float, float]: with torch.no_grad(): predictions = model.predict(points["xy1"], points["xy2"], points["xy3"]) errors = tuple( ( torch.linalg.vector_norm(prediction - reference) / torch.linalg.vector_norm(reference) ).item() for prediction, reference in zip(predictions, references, strict=True) ) combined_prediction = torch.cat(predictions) combined_reference = torch.cat(references) combined_error = ( torch.linalg.vector_norm(combined_prediction - combined_reference) / torch.linalg.vector_norm(combined_reference) ).item() return errors[0], errors[1], errors[2], combined_error def main() -> None: config_path = DEFAULT_CONFIG.resolve() config = load_config(config_path) common = config["common"] device = resolve_device(str(common["device"])) dtype = resolve_dtype(str(common["dtype"])) seed = int(common["seed"]) data_path = project_path(config["data"]["mat_file"]) weight_dir = project_path(common["weight_dir"]) checkpoint_path = weight_dir / config["training"]["checkpoint_name"] steps = int(config["training"]["steps"]) if steps <= 0: raise ValueError("training steps must be positive") counts = {key: int(value) for key, value in config["data"]["samples"].items()} seed_everything(seed) print(f"Config: {config_path}") print(f"Data: {data_path}") print(f"Device: {device}") print(f"Samples: {counts}") data = load_mat_data(data_path) batch = build_training_batch(data, counts, seed, device, dtype) evaluation_points = build_evaluation_points(data, device, dtype) references = exact_subdomain_values(data, device, dtype) model = XPINNPoisson2D(config["model"], dtype=dtype).to( device=device, dtype=dtype ) optimizer = torch.optim.Adam( model.parameters(), lr=float(config["training"]["lr"]) ) log_interval = int(config["training"]["log_interval"]) started = time.time() for step in range(1, steps + 1): clear_coordinate_grads(batch) outputs = model.training_outputs(batch) loss, parts = compute_loss(outputs, batch, config["loss"]) if not torch.isfinite(loss): raise FloatingPointError(f"XPINN loss became non-finite at step {step}") optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() if step == 1 or step % log_interval == 0 or step == steps: error1, error2, error3, combined_error = evaluate( model, evaluation_points, references ) print( f"step={step:4d} loss={loss.item():.3e} " f"boundary={parts['boundary']:.3e} pde={parts['pde']:.3e} " f"interface_r={parts['interface_residual']:.3e} " f"interface_u={parts['interface_value']:.3e} " f"l2=({error1:.3e}, {error2:.3e}, {error3:.3e}) " f"combined={combined_error:.3e}" ) weight_dir.mkdir(parents=True, exist_ok=True) checkpoint = { "case": "poisson2d", "architecture": "xpinn_poisson2d", "model_state": model.state_dict(), "model_config": config["model"], "sample_counts": counts, "step": steps, } torch.save(checkpoint, checkpoint_path) print(f"Training finished in {time.time() - started:.1f}s") print(f"Saved checkpoint: {checkpoint_path}") if __name__ == "__main__": main()