| from __future__ import annotations |
|
|
| import sys |
| from collections.abc import Mapping |
| from pathlib import Path |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import matplotlib.tri as tri |
| import numpy as np |
| import torch |
| import yaml |
| from matplotlib.patches import Polygon |
|
|
|
|
| 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 ( |
| build_evaluation_points, |
| combined_coordinates, |
| combined_exact_solution, |
| load_mat_data, |
| ) |
| from model.xpinn import XPINNPoisson2D |
|
|
|
|
| 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 load_model( |
| checkpoint_path: Path, |
| fallback_model_config: dict, |
| device: torch.device, |
| dtype: torch.dtype, |
| ) -> XPINNPoisson2D: |
| 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, Mapping): |
| raise ValueError(f"invalid XPINN checkpoint: {checkpoint_path}") |
| if "model_state" in checkpoint: |
| state = checkpoint["model_state"] |
| model_config = checkpoint.get("model_config", fallback_model_config) |
| elif checkpoint and all(torch.is_tensor(value) for value in checkpoint.values()): |
| state = checkpoint |
| model_config = fallback_model_config |
| else: |
| raise ValueError(f"checkpoint contains no model state: {checkpoint_path}") |
| model = XPINNPoisson2D(model_config, dtype=dtype).to(device=device, dtype=dtype) |
| model.load_state_dict(state, strict=True) |
| model.eval() |
| return model |
|
|
|
|
| def predict( |
| model: XPINNPoisson2D, evaluation_points: dict[str, torch.Tensor] |
| ) -> np.ndarray: |
| with torch.no_grad(): |
| predictions = model.predict( |
| evaluation_points["xy1"], |
| evaluation_points["xy2"], |
| evaluation_points["xy3"], |
| ) |
| return torch.cat(predictions).cpu().numpy().reshape(-1) |
|
|
|
|
| def mask_polygon(data: Mapping) -> np.ndarray: |
| boundary_x = np.asarray(data["xb"]).reshape(-1) |
| boundary_y = np.asarray(data["yb"]).reshape(-1) |
| return np.vstack( |
| ( |
| np.column_stack((boundary_x, boundary_y)), |
| np.array( |
| [ |
| [1.8, boundary_y[-1]], |
| [1.8, -1.7], |
| [-1.6, -1.7], |
| [-1.6, 1.55], |
| [1.8, 1.55], |
| [1.8, boundary_y[-1]], |
| ] |
| ), |
| np.array([[boundary_x[-1], boundary_y[-1]]]), |
| ) |
| ) |
|
|
|
|
| def save_figure( |
| data: Mapping, |
| exact: np.ndarray, |
| prediction: np.ndarray, |
| output_path: Path, |
| ) -> None: |
| x, y = combined_coordinates(data) |
| if exact.size != x.size or prediction.size != x.size: |
| raise ValueError("plot coordinates, exact values, and predictions must have equal size") |
| triangulation = tri.Triangulation(x, y) |
| polygon = mask_polygon(data) |
| fields = ( |
| ("Exact", exact), |
| ("XPINN", prediction), |
| ("Absolute error", np.abs(exact - prediction)), |
| ) |
| figure, axes = plt.subplots(1, 3, figsize=(18, 5)) |
| for axis, (title, field) in zip(axes, fields, strict=True): |
| contour = axis.tricontourf(triangulation, field, 100, cmap="jet") |
| axis.add_patch( |
| Polygon(polygon, closed=True, facecolor="white", edgecolor="white") |
| ) |
| axis.plot(data["xi1"].reshape(-1), data["yi1"].reshape(-1), "w-", linewidth=0.5) |
| axis.plot(data["xi2"].reshape(-1), data["yi2"].reshape(-1), "w-", linewidth=0.5) |
| axis.set_title(title) |
| axis.set_xlabel("x") |
| axis.set_ylabel("y") |
| axis.set_aspect("equal") |
| figure.colorbar(contour, ax=axis) |
| figure.tight_layout() |
| figure.savefig(output_path, dpi=150) |
| plt.close(figure) |
|
|
|
|
| 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"])) |
| data_path = project_path(config["data"]["mat_file"]) |
| weight_dir = project_path(common["weight_dir"]) |
| result_dir = project_path(common["result_dir"]) |
| checkpoint_path = weight_dir / config["training"]["checkpoint_name"] |
| output_path = result_dir / config["inference"]["figure_name"] |
| result_dir.mkdir(parents=True, exist_ok=True) |
|
|
| print(f"Config: {config_path}") |
| print(f"Data: {data_path}") |
| print(f"Checkpoint: {checkpoint_path}") |
| print(f"Device: {device}") |
| data = load_mat_data(data_path) |
| model = load_model(checkpoint_path, config["model"], device, dtype) |
| evaluation_points = build_evaluation_points(data, device, dtype) |
| prediction = predict(model, evaluation_points) |
| exact = combined_exact_solution(data) |
| relative_l2 = float( |
| np.linalg.norm(exact - prediction) / np.linalg.norm(exact) |
| ) |
| save_figure(data, exact, prediction, output_path) |
| print(f"Relative L2={relative_l2:.6e}") |
| print(f"Plot: {output_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
