| from __future__ import annotations |
|
|
| import argparse |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| from model.gpinn import ( |
| exact_poisson1d, |
| exact_poisson2d, |
| gPINN, |
| gPINNBurgers, |
| gPINNPoisson2D, |
| ) |
|
|
|
|
| DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml" |
| CASES = ("1d", "2d", "burgers") |
|
|
|
|
| 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 load_checkpoint(weight_dir: Path, filename: str) -> dict: |
| checkpoint_path = weight_dir / filename |
| if not checkpoint_path.is_file(): |
| raise FileNotFoundError( |
| f"checkpoint not found: {checkpoint_path}. Run scripts/train.py first." |
| ) |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) |
| if not isinstance(checkpoint, dict) or "model_state" not in checkpoint: |
| raise ValueError(f"invalid checkpoint format: {checkpoint_path}") |
| return checkpoint |
|
|
|
|
| def infer_poisson1d( |
| weight_dir: Path, |
| result_dir: Path, |
| device: torch.device, |
| dtype: torch.dtype, |
| ) -> float: |
| checkpoint = load_checkpoint(weight_dir, "gpinn_poisson1d.pt") |
| layers = checkpoint.get("layers", [1, 30, 30, 30, 30, 1]) |
| model = gPINN(layers).to(device=device, dtype=dtype) |
| model.load_state_dict(checkpoint["model_state"], strict=True) |
| model.eval() |
|
|
| x = torch.linspace(0.0, np.pi, 1000, dtype=dtype, device=device).unsqueeze(-1) |
| exact = exact_poisson1d(x.cpu().numpy().reshape(-1)) |
| with torch.no_grad(): |
| prediction = model(x).cpu().numpy().reshape(-1) |
| relative_l2 = float(np.linalg.norm(prediction - exact) / np.linalg.norm(exact)) |
|
|
| figure, axis = plt.subplots(figsize=(8, 4)) |
| axis.plot(x.cpu().numpy(), exact, "k-", label="Exact") |
| axis.plot(x.cpu().numpy(), prediction, "r--", label="gPINN") |
| axis.set_xlabel("x") |
| axis.set_ylabel("u") |
| axis.legend() |
| figure.tight_layout() |
| output_path = result_dir / "gpinn_poisson1d.png" |
| figure.savefig(output_path, dpi=150) |
| plt.close(figure) |
| print(f"[1D] relative L2={relative_l2:.6e}, plot={output_path}") |
| return relative_l2 |
|
|
|
|
| def infer_poisson2d( |
| weight_dir: Path, |
| result_dir: Path, |
| device: torch.device, |
| dtype: torch.dtype, |
| default_exponent: float, |
| ) -> float: |
| checkpoint = load_checkpoint(weight_dir, "gpinn_poisson2d.pt") |
| layers = checkpoint.get("layers", [2, 30, 30, 30, 1]) |
| architecture = checkpoint.get("architecture", "gpinn") |
| if architecture == "poisson2d_hard_bc": |
| model = gPINNPoisson2D(layers) |
| elif architecture == "gpinn": |
| |
| model = gPINN(layers) |
| else: |
| raise ValueError(f"unsupported 2D Poisson architecture: {architecture}") |
| model = model.to(device=device, dtype=dtype) |
| model.load_state_dict(checkpoint["model_state"], strict=True) |
| model.eval() |
|
|
| exponent = float(checkpoint.get("a", default_exponent)) |
| axis = np.linspace(0.0, 1.0, 100) |
| grid_x, grid_y = np.meshgrid(axis, axis) |
| coordinates = torch.as_tensor( |
| np.column_stack([grid_x.ravel(), grid_y.ravel()]), dtype=dtype, device=device |
| ) |
| with torch.no_grad(): |
| prediction = model(coordinates).cpu().numpy().reshape(grid_x.shape) |
| exact = exact_poisson2d(grid_x, grid_y, exponent) |
| absolute_error = np.abs(prediction - exact) |
| relative_l2 = float( |
| np.linalg.norm(prediction.ravel() - exact.ravel()) / np.linalg.norm(exact.ravel()) |
| ) |
|
|
| figure, axes = plt.subplots(1, 3, figsize=(15, 4)) |
| for plot_axis, title, field in zip( |
| axes, |
| ("Exact", "gPINN", "Absolute error"), |
| (exact, prediction, absolute_error), |
| strict=True, |
| ): |
| image = plot_axis.imshow(field, extent=(0, 1, 0, 1), origin="lower", cmap="jet") |
| plot_axis.set_title(title) |
| plot_axis.set_xlabel("x") |
| plot_axis.set_ylabel("y") |
| figure.colorbar(image, ax=plot_axis) |
| figure.tight_layout() |
| output_path = result_dir / "gpinn_poisson2d.png" |
| figure.savefig(output_path, dpi=150) |
| plt.close(figure) |
| print(f"[2D] relative L2={relative_l2:.6e}, plot={output_path}") |
| return relative_l2 |
|
|
|
|
| def load_burgers_grid(data_path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]: |
| if not data_path.is_file(): |
| raise FileNotFoundError(f"Burgers data not found: {data_path}") |
| with np.load(data_path) as data: |
| missing = {"t", "x", "usol"}.difference(data.files) |
| if missing: |
| raise ValueError(f"Burgers data is missing arrays: {sorted(missing)}") |
| time_axis = np.asarray(data["t"]).reshape(-1) |
| space_axis = np.asarray(data["x"]).reshape(-1) |
| exact = np.asarray(data["usol"]) |
| expected_shape = (space_axis.size, time_axis.size) |
| if exact.shape != expected_shape: |
| raise ValueError(f"usol shape must be {expected_shape}, got {exact.shape}") |
| return time_axis, space_axis, exact |
|
|
|
|
| def infer_burgers( |
| weight_dir: Path, |
| result_dir: Path, |
| data_path: Path, |
| device: torch.device, |
| dtype: torch.dtype, |
| ) -> float: |
| checkpoint = load_checkpoint(weight_dir, "gpinn_burgers.pt") |
| layers = checkpoint.get("layers", [2, 32, 32, 32, 1]) |
| model = gPINNBurgers(layers).to(device=device, dtype=dtype) |
| model.load_state_dict(checkpoint["model_state"], strict=True) |
| model.eval() |
|
|
| time_axis, space_axis, exact = load_burgers_grid(data_path) |
| time_grid, space_grid = np.meshgrid(time_axis, space_axis) |
| coordinates = torch.as_tensor( |
| np.column_stack([space_grid.ravel(), time_grid.ravel()]), |
| dtype=dtype, |
| device=device, |
| ) |
| with torch.no_grad(): |
| prediction = model(coordinates).cpu().numpy().reshape(exact.shape) |
| absolute_error = np.abs(prediction - exact) |
| relative_l2 = float( |
| np.linalg.norm(prediction.ravel() - exact.ravel()) / np.linalg.norm(exact.ravel()) |
| ) |
|
|
| figure, axes = plt.subplots(1, 3, figsize=(15, 4)) |
| for plot_axis, title, field in zip( |
| axes, |
| ("Exact", "gPINN", "Absolute error"), |
| (exact, prediction, absolute_error), |
| strict=True, |
| ): |
| image = plot_axis.contourf(time_axis, space_axis, field, 100, cmap="jet") |
| plot_axis.set_title(title) |
| plot_axis.set_xlabel("t") |
| plot_axis.set_ylabel("x") |
| figure.colorbar(image, ax=plot_axis) |
| figure.tight_layout() |
| output_path = result_dir / "gpinn_burgers.png" |
| figure.savefig(output_path, dpi=150) |
| plt.close(figure) |
| print(f"[Burgers] relative L2={relative_l2:.6e}, plot={output_path}") |
| return relative_l2 |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Run gPINN inference and visualization") |
| parser.add_argument( |
| "--config", |
| type=Path, |
| default=Path(os.environ.get("GPINN_CONFIG", DEFAULT_CONFIG)), |
| help="YAML configuration file", |
| ) |
| parser.add_argument("--case", choices=(*CASES, "all"), default="all") |
| parser.add_argument("--device", help="Override common.device, for example cpu or cuda:0") |
| parser.add_argument("--data", type=Path, help="Override Burgers.npz path") |
| parser.add_argument("--weight-dir", type=Path, help="Override checkpoint input directory") |
| parser.add_argument("--result-dir", type=Path, help="Override plot 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"] |
| device = resolve_device(args.device or str(common["device"])) |
| dtype = resolve_dtype(str(common["dtype"])) |
| weight_dir = project_path(args.weight_dir or common["weight_dir"]) |
| result_dir = project_path(args.result_dir or common["result_dir"]) |
| data_path = project_path(args.data or config["burgers"]["data"]) |
| result_dir.mkdir(parents=True, exist_ok=True) |
|
|
| print(f"Config: {config_path}") |
| print(f"Device: {device}") |
| selected_cases = CASES if args.case == "all" else (args.case,) |
| if "1d" in selected_cases: |
| infer_poisson1d(weight_dir, result_dir, device, dtype) |
| if "2d" in selected_cases: |
| infer_poisson2d( |
| weight_dir, |
| result_dir, |
| device, |
| dtype, |
| float(config["poisson2d"]["a"]), |
| ) |
| if "burgers" in selected_cases: |
| infer_burgers(weight_dir, result_dir, data_path, device, dtype) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
