#!/usr/bin/env python3 """Test a trained AFEM model on alternative scatterer geometries. Supports: square, multi-circle, and the original circle. Usage: python src/test_media.py # uses src/test_config.yaml python src/test_media.py --k-test 30.0 --geometry circle python src/test_media.py --config my_test.yaml # custom config All test parameters live in the YAML config. CLI args serve as overrides. """ import argparse import copy import os import sys import time from pathlib import Path from typing import Optional import numpy as np import torch from torch_geometric.data import Batch _project_root = Path(__file__).resolve().parent.parent if str(_project_root) not in sys.path: sys.path.insert(0, str(_project_root)) from src.network import create_model from src.utils import load_checkpoint, load_config, setup_helmholtz_config from src.helmholtz_alt import ( HelmholtzProblemSquare, HelmholtzProblemMultiCircle, create_helmholtz_problem_square, create_helmholtz_problem_multi_circle, ) # ═══════════════════════════════════════════════════════════════════════ # Geometry factory mapping # ═══════════════════════════════════════════════════════════════════════ _GEOMETRY_FACTORIES = { "square": create_helmholtz_problem_square, "multi_circle": create_helmholtz_problem_multi_circle, "circle": None, # default HelmholtzProblem } # ═══════════════════════════════════════════════════════════════════════ # Epsilon_r property patching # ═══════════════════════════════════════════════════════════════════════ def _patch_epsilon_r(env): inner_fp = env.fem_problem.fem_problem if hasattr(inner_fp, "eps_r_at_midpoints"): def _eps_r(self): return inner_fp.eps_r_at_midpoints(self.mesh) type(env)._epsilon_r_elements = property(_eps_r) # ═══════════════════════════════════════════════════════════════════════ # Fine FEM reference (computed once, interpolated later) # ═══════════════════════════════════════════════════════════════════════ def _compute_fine_fem_reference(env, n_refine: int = 2): """Compute fine-FEM reference on initial mesh + n_refine uniform refinement.""" from skfem import Basis, ElementTriP1 fp = env.fem_problem.fem_problem ref_mesh = copy.deepcopy(env.mesh) for _ in range(n_refine): ref_mesh = ref_mesh.refined(np.arange(ref_mesh.t.shape[1])) ref_basis = Basis(ref_mesh, ElementTriP1()) ref_sol = fp.calculate_solution(ref_basis, cache=False) # Interpolate to coarse mesh vertices pts = env.mesh.p.T finder = ref_mesh.element_finder() cells = finder(*pts.T) cells = np.clip(cells, 0, ref_mesh.t.shape[1] - 1) i0, i1, i2 = ref_mesh.t[0, cells], ref_mesh.t[1, cells], ref_mesh.t[2, cells] p = ref_mesh.p x, y = pts[:, 0], pts[:, 1] x0, y0 = p[0, i0], p[1, i0] x1, y1 = p[0, i1], p[1, i1] x2, y2 = p[0, i2], p[1, i2] denom = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0) denom = np.where(np.abs(denom) < 1e-15, 1.0, denom) w0 = ((x1 - x) * (y2 - y) - (x2 - x) * (y1 - y)) / denom w1 = ((x2 - x) * (y0 - y) - (x0 - x) * (y2 - y)) / denom w2 = 1.0 - w0 - w1 u_ref_on_coarse = w0 * ref_sol[i0] + w1 * ref_sol[i1] + w2 * ref_sol[i2] return u_ref_on_coarse, ref_mesh, ref_sol def _interpolate_ref_to_mesh(target_pts, ref_mesh, ref_sol): """Interpolate cached reference solution to arbitrary mesh vertices.""" finder = ref_mesh.element_finder() cells = finder(*target_pts.T) cells = np.clip(cells, 0, ref_mesh.t.shape[1] - 1) i0, i1, i2 = ref_mesh.t[0, cells], ref_mesh.t[1, cells], ref_mesh.t[2, cells] p = ref_mesh.p x, y = target_pts[:, 0], target_pts[:, 1] x0, y0 = p[0, i0], p[1, i0] x1, y1 = p[0, i1], p[1, i1] x2, y2 = p[0, i2], p[1, i2] denom = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0) denom = np.where(np.abs(denom) < 1e-15, 1.0, denom) w0 = ((x1 - x) * (y2 - y) - (x2 - x) * (y1 - y)) / denom w1 = ((x2 - x) * (y0 - y) - (x0 - x) * (y2 - y)) / denom w2 = 1.0 - w0 - w1 return w0 * ref_sol[i0] + w1 * ref_sol[i1] + w2 * ref_sol[i2] def _compute_ref_grid(env, n_refine: int = 3, resolution: int = 200): """Compute fine reference on a regular grid for smooth heatmaps.""" from skfem import Basis, ElementTriP1 fp = env.fem_problem.fem_problem ref_mesh = copy.deepcopy(env.mesh) for _ in range(n_refine): ref_mesh = ref_mesh.refined(np.arange(ref_mesh.t.shape[1])) ref_basis = Basis(ref_mesh, ElementTriP1()) ref_sol = fp.calculate_solution(ref_basis, cache=False) boundary = fp._domain._boundary x_vec = np.linspace(boundary[0], boundary[2], resolution) y_vec = np.linspace(boundary[1], boundary[3], resolution) X, Y = np.meshgrid(x_vec, y_vec) grid_pts = np.column_stack([X.ravel(), Y.ravel()]) U_grid = np.zeros(len(grid_pts), dtype=np.complex128) batch_size = 4096 for start in range(0, len(grid_pts), batch_size): end = min(start + batch_size, len(grid_pts)) batch = grid_pts[start:end] finder = ref_mesh.element_finder() cells = finder(*batch.T) cells = np.clip(cells, 0, ref_mesh.t.shape[1] - 1) i0, i1, i2 = ref_mesh.t[0, cells], ref_mesh.t[1, cells], ref_mesh.t[2, cells] p = ref_mesh.p x, y = batch[:, 0], batch[:, 1] x0, y0 = p[0, i0], p[1, i0] x1, y1 = p[0, i1], p[1, i1] x2, y2 = p[0, i2], p[1, i2] denom = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0) denom = np.where(np.abs(denom) < 1e-15, 1.0, denom) w0 = ((x1 - x) * (y2 - y) - (x2 - x) * (y1 - y)) / denom w1 = ((x2 - x) * (y0 - y) - (x0 - x) * (y2 - y)) / denom w2 = 1.0 - w0 - w1 U_grid[start:end] = w0 * ref_sol[i0] + w1 * ref_sol[i1] + w2 * ref_sol[i2] return {"X": X, "Y": Y, "E_scat": U_grid.reshape(resolution, resolution)} def _compute_step_error(scalar, u_ref) -> float: if u_ref is None: return float("nan") diff = np.abs(scalar - u_ref) denom = np.linalg.norm(np.abs(u_ref)) if denom < 1e-12: denom = 1.0 return float(np.linalg.norm(diff) / denom) # ═══════════════════════════════════════════════════════════════════════ # Visualization # ═══════════════════════════════════════════════════════════════════════ def _render_field(ax, triang, values, title, vmin, vmax, show_mesh=True): tcf = ax.tripcolor(triang, values, shading="gouraud", cmap="jet", vmin=vmin, vmax=vmax) if show_mesh and triang is not None: n = triang.triangles.shape[0] ax.triplot(triang, lw=(0.5 if n < 500 else 0.3), color="black", alpha=(0.7 if n < 2000 else 0.5)) ax.set_aspect("equal") ax.set_title(title, fontsize=9) ax.set_xticks([]) ax.set_yticks([]) return tcf def _draw_scatterer(ax, geometry: str, env): fp = env.fem_problem.fem_problem if geometry == "square": sq = getattr(fp, "_sq_cx", 0.5), getattr(fp, "_sq_cy", 0.5) hs = getattr(fp, "_sq_half", 0.2) ang = getattr(fp, "_sq_angle", 0.0) corners = np.array([ [-hs, -hs], [hs, -hs], [hs, hs], [-hs, hs], [-hs, -hs] ]) if ang != 0: c, s = np.cos(ang), np.sin(ang) corners = corners @ np.array([[c, -s], [s, c]]).T corners[:, 0] += sq[0] corners[:, 1] += sq[1] ax.plot(corners[:, 0], corners[:, 1], color="cyan", linewidth=1.5, linestyle="--") elif geometry == "multi_circle": circles = getattr(fp, "_circles", []) for c in circles: theta = np.linspace(0, 2 * np.pi, 128) ax.plot(c["cx"] + c["radius"] * np.cos(theta), c["cy"] + c["radius"] * np.sin(theta), color="cyan", linewidth=1.5, linestyle="--") elif geometry == "circle": cx = getattr(fp, "_cx", 0.5) cy = getattr(fp, "_cy", 0.5) r = getattr(fp, "_radius", 0.2) theta = np.linspace(0, 2 * np.pi, 128) ax.plot(cx + r * np.cos(theta), cy + r * np.sin(theta), color="cyan", linewidth=1.5, linestyle="--") def _save_pngs(steps, stem, checkpoint_path, k, geometry, env, ref_grid): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import matplotlib.tri as tri per_step_dir = f"{stem}_steps" os.makedirs(os.path.dirname(stem) or ".", exist_ok=True) os.makedirs(per_step_dir, exist_ok=True) # ── Overview grid ── n = len(steps) ncols = min(n, 4) nrows = (n + ncols - 1) // ncols fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 3.5 * nrows)) axes_flat = np.array([axes]) if nrows * ncols == 1 else np.array(axes).flatten() for i, step_data in enumerate(steps): mesh, scalar, err_val, n_elem = step_data[:4] pts = mesh.p.T tg = tri.Triangulation(pts[:, 0], pts[:, 1], mesh.t.T) s = np.abs(scalar) if np.iscomplexobj(scalar) else scalar vmin, vmax = s.min(), s.max() if vmax - vmin < 1e-12: vmin, vmax = vmin - 0.5, vmax + 0.5 tcf = _render_field(axes_flat[i], tg, s, f"Step {i}: {n_elem} elem, err={err_val:.4f}", vmin, vmax) fig.colorbar(tcf, ax=axes_flat[i], fraction=0.046, pad=0.04) _draw_scatterer(axes_flat[i], geometry, env) for j in range(n, len(axes_flat)): axes_flat[j].set_visible(False) fig.subplots_adjust(left=0.04, right=0.90, top=0.90, bottom=0.06, wspace=0.15, hspace=0.30) geo_label = {"square": "Square", "multi_circle": "Multi-Circle", "circle": "Circle"}.get(geometry, geometry) fig.suptitle( f"Helmholtz |E_scat| [{geo_label}] — {os.path.basename(checkpoint_path)}\n" f"k={k:.1f} eps_r info in scatterer overlay", fontsize=12, ) fig.savefig(f"{stem}.png", dpi=200, bbox_inches="tight") plt.close(fig) print(f"[Viz] Overview → {stem}.png") # ── Per-step panels (FEM + Reference + Error) ── for i, step_data in enumerate(steps): mesh, scalar, err_val, n_elem = step_data[:4] u_ref_at_verts = step_data[4] if len(step_data) > 4 else None pts = mesh.p.T tg = tri.Triangulation(pts[:, 0], pts[:, 1], mesh.t.T) coarse_val = np.abs(scalar) if np.iscomplexobj(scalar) else scalar fig2, axes2 = plt.subplots(1, 3, figsize=(18, 6)) axes2 = list(np.atleast_1d(axes2)) # Panel 1: FEM cvmin, cvmax = coarse_val.min(), coarse_val.max() if cvmax - cvmin < 1e-12: cvmin, cvmax = cvmin - 0.5, cvmax + 0.5 tcf1 = _render_field(axes2[0], tg, coarse_val, f"Step {i}: FEM |E_scat| ({n_elem} elem)", cvmin, cvmax) _draw_scatterer(axes2[0], geometry, env) fig2.colorbar(tcf1, ax=axes2[0], fraction=0.046, pad=0.04) # Panel 2: Fine FEM reference on grid if ref_grid is not None: g = ref_grid gm = np.abs(g["E_scat"]) mvmin, mvmax = gm.min(), gm.max() if mvmax - mvmin < 1e-12: mvmin, mvmax = mvmin - 0.5, mvmax + 0.5 im2 = axes2[1].pcolormesh(g["X"], g["Y"], gm, shading="gouraud", cmap="jet", vmin=mvmin, vmax=mvmax) axes2[1].set_title("Fine FEM Ref |E_scat|", fontsize=9) axes2[1].set_aspect("equal") axes2[1].set_xticks([]) axes2[1].set_yticks([]) _draw_scatterer(axes2[1], geometry, env) fig2.colorbar(im2, ax=axes2[1], fraction=0.046, pad=0.04) # Panel 3: Pointwise error if u_ref_at_verts is not None: u_fem_abs = np.abs(scalar) u_ref_abs = np.abs(u_ref_at_verts) error_abs = np.abs(u_fem_abs - u_ref_abs) evmin, evmax = 0.0, error_abs.max() or 1.0 if evmax - evmin < 1e-12: evmax = evmin + 1.0 tcf3 = _render_field(axes2[2], tg, error_abs, f"||FEM|−|Ref|| L2={err_val:.4f}", evmin, evmax) _draw_scatterer(axes2[2], geometry, env) fig2.colorbar(tcf3, ax=axes2[2], fraction=0.046, pad=0.04) fig2.tight_layout() fig2.savefig(f"{per_step_dir}/step{i:02d}.png", dpi=150, bbox_inches="tight") plt.close(fig2) print(f"[Viz] Per-step PNGs → {per_step_dir}/ ({n} files)") # ═══════════════════════════════════════════════════════════════════════ # Scatterer config injection # ═══════════════════════════════════════════════════════════════════════ def _inject_scatterer_config(base_config: dict, geometry: str, sc_cfg: dict, k_test: float): """Inject scatterer params from test config into the base config's helmholtz section. Returns (config, factory) where factory is the geometry-specific create function. """ hc = (base_config.setdefault("environment", {}) .setdefault("mesh_refinement", {}) .setdefault("fem", {}) .setdefault("helmholtz", {})) sc = hc.setdefault("scatterer", {}) sc["mode"] = "fixed" sc["eps_r"] = float(sc_cfg.get("eps_r", 3.0)) if geometry == "square": sc["square"] = { "cx": float(sc_cfg.get("cx", 0.5)), "cy": float(sc_cfg.get("cy", 0.5)), "half_side": float(sc_cfg.get("half_side", 0.15)), "angle": float(sc_cfg.get("angle", 0.0)), } elif geometry == "multi_circle": circles_raw = sc_cfg.get("circles", []) circles = [] for c in circles_raw: circles.append({ "cx": float(c["cx"]), "cy": float(c["cy"]), "radius": float(c["radius"]), "eps_r": float(c.get("eps_r", sc_cfg.get("eps_r", 3.0))), }) sc["circles"] = circles elif geometry == "circle": sc["cx"] = float(sc_cfg.get("cx", 0.5)) sc["cy"] = float(sc_cfg.get("cy", 0.5)) sc["radius"] = float(sc_cfg.get("radius", 0.2)) hc["wave_number_mode"] = "fixed" hc["wave_number"] = float(k_test) factory = _GEOMETRY_FACTORIES.get(geometry) return base_config, factory # ═══════════════════════════════════════════════════════════════════════ # Main test function # ═══════════════════════════════════════════════════════════════════════ def test_alt_media( base_config: dict, test_cfg: dict, cli_overrides: Optional[dict] = None, ): """Run AFEM inference with config-driven parameters. Args: base_config: loaded from config.yaml (model/network/algo) test_cfg: loaded from test_config.yaml (test-specific params) cli_overrides: optional CLI arg overrides dict """ ov = cli_overrides or {} # ── Resolve parameters: test_cfg < CLI override ── tc = test_cfg.get("test", {}) ref_cfg = test_cfg.get("reference", {}) sc_cfg = test_cfg.get("scatterer", {}) geometry = ov.get("geometry") or tc.get("geometry", "circle") checkpoint_path = ov.get("checkpoint") or tc.get("checkpoint", "checkpoints/model_final.pt") output_path = ov.get("output") or tc.get("output", "result/test_media.png") seed = ov.get("seed") or tc.get("seed", 99) k_test = ov.get("k_test") or test_cfg.get("k_test", 8.0) n_refine_vertex = ov.get("n_refine_vertex") or ref_cfg.get("n_refine_vertex", 2) n_refine_grid = ov.get("n_refine_grid") or ref_cfg.get("n_refine_grid", 3) grid_resolution = ov.get("grid_resolution") or ref_cfg.get("grid_resolution", 200) # Allow CLI override of scatterer params for key in ("cx", "cy", "radius", "eps_r", "half_side", "angle"): if ov.get(key) is not None: sc_cfg[key] = ov[key] if ov.get("circles") is not None: sc_cfg["circles"] = ov["circles"] algo = base_config.get("algorithm", {}) # ── 1. Inject scatterer config ── config, factory = _inject_scatterer_config( copy.deepcopy(base_config), geometry, sc_cfg, k_test) # ── 2. Create env with alt factory ── import environment.fem_problem as fem_problem_module _orig_create = None if factory is not None: _orig_create = fem_problem_module.create_helmholtz_problem fem_problem_module.create_helmholtz_problem = factory from environment.mesh_refinement import MeshRefinement env = MeshRefinement( environment_config=config.get("environment", {}).get("mesh_refinement", {}), seed=seed, ) # ── 3. Load model ── model = create_model(env, config.get("network", {}), algo.get("ppo", {})) load_checkpoint(model, checkpoint_path) model.eval() dev = next(model.parameters()).device print(f"[Device] {dev}") model = model.to(dev) # ── 4. Reset env ── print(f"[Test] Geometry: {geometry} k={k_test:.3f}") obs = env.reset() # ── 5. Patch epsilon_r_elements (after reset) ── _patch_epsilon_r(env) # Restore original factory if _orig_create is not None: fem_problem_module.create_helmholtz_problem = _orig_create # ── 6. Print scatterer info ── fp = env.fem_problem.fem_problem if geometry == "square": print(f"[Test] Square: center=({getattr(fp, '_sq_cx', 0.5):.3f}, " f"{getattr(fp, '_sq_cy', 0.5):.3f}) half_side={getattr(fp, '_sq_half', 0.2):.3f}") elif geometry == "multi_circle": circles_attr = getattr(fp, "_circles", []) for i, c in enumerate(circles_attr): print(f"[Test] Circle {i}: center=({c['cx']:.3f}, {c['cy']:.3f}) " f"r={c['radius']:.3f} eps_r={c['eps_r']:.1f}") elif geometry == "circle": print(f"[Test] Circle: center=({getattr(fp, '_cx', 0.5):.3f}, " f"{getattr(fp, '_cy', 0.5):.3f}) r={getattr(fp, '_radius', 0.2):.3f}") # ── 7. Compute fine-FEM reference ONCE on initial mesh ── n_init = env.mesh.t.shape[1] print(f"[Test] Initial mesh: {n_init} elements") print(f"[Test] Computing fine-FEM reference (n_refine_vertex={n_refine_vertex}, " f"n_refine_grid={n_refine_grid}, grid={grid_resolution})...") t0 = time.time() u_ref_initial, ref_mesh, ref_sol = _compute_fine_fem_reference(env, n_refine=n_refine_vertex) ref_grid = _compute_ref_grid(env, n_refine=n_refine_grid, resolution=grid_resolution) print(f"[Test] Reference ready ({time.time() - t0:.1f}s, grid {ref_grid['X'].shape})") # ── 8. Run inference ── stem = output_path.rsplit(".", 1)[0] if "." in output_path else output_path init_mesh = env.mesh init_sol = env.scalar_solution init_err = _compute_step_error(init_sol, u_ref_initial) steps = [(init_mesh, init_sol, init_err, env.num_agents, u_ref_initial)] n_elem_init = env.num_agents print(f" Step 0: reward=--- err={init_err:.4f} elements={n_elem_init}") done = False step_idx = 0 total_reward = 0.0 while not done: obs_g = obs.to(dev) with torch.no_grad(): actions, _, _ = model(Batch.from_data_list([obs_g]), deterministic=True) obs, reward, done, info = env.step(actions.cpu().numpy()) step_r = float(np.sum(reward)) total_reward += step_r step_idx += 1 # Interpolate cached reference to current mesh vertices (no re-solve) u_ref_current = _interpolate_ref_to_mesh(env.mesh.p.T, ref_mesh, ref_sol) step_err = _compute_step_error(env.scalar_solution, u_ref_current) steps.append((env.mesh, env.scalar_solution, step_err, env.num_agents, u_ref_current)) print(f" Step {step_idx:2d}: reward={step_r:+.4f} err={step_err:.4f} " f"elements={info.get('num_elements', '?')} " f"sel={info.get('selected_count', 0)} " f"done={done}") print(f"\n[Test] total_reward={total_reward:.4f} final_err={steps[-1][2]:.4f} " f"final_elements={steps[-1][3]}") # ── 9. Visualize ── _save_pngs(steps, stem, checkpoint_path, k_test, geometry, env, ref_grid) print(f"[Viz] Done → {output_path}") # ═══════════════════════════════════════════════════════════════════════ # CLI # ═══════════════════════════════════════════════════════════════════════ def _load_yaml(path: str) -> dict: """Load a YAML file, resolving relative paths against project root.""" import yaml if not os.path.isabs(path): path = os.path.join(_project_root, path) with open(path, "r") as f: return yaml.safe_load(f) def main(): parser = argparse.ArgumentParser( description="Test AFEM trained model on alternative scatterer geometries") # Config parser.add_argument("--config", default="src/test_config.yaml", help="Test config YAML (default: src/test_config.yaml)") # Test scenario overrides parser.add_argument("--geometry", choices=["square", "multi_circle", "circle"], help="Scatterer geometry (overrides config)") parser.add_argument("--checkpoint", help="Model checkpoint path (overrides config)") parser.add_argument("--output", help="Output image path (overrides config)") parser.add_argument("--seed", type=int, help="Random seed (overrides config)") parser.add_argument("--k-test", type=float, help="Wave number (overrides config)") # Scatterer overrides parser.add_argument("--cx", type=float, help="Scatterer center x") parser.add_argument("--cy", type=float, help="Scatterer center y") parser.add_argument("--radius", type=float, help="Scatterer radius (circle)") parser.add_argument("--eps-r", type=float, help="Dielectric constant eps_r") parser.add_argument("--half-side", type=float, help="Half side length (square)") parser.add_argument("--angle", type=float, help="Rotation angle in radians (square)") parser.add_argument("--circles", nargs="*", default=None, help="Circle specs: 'cx,cy,radius[,eps_r]' (multi_circle)") # Reference computation overrides parser.add_argument("--n-refine-vertex", type=int, help="Uniform refinement levels for vertex error reference") parser.add_argument("--n-refine-grid", type=int, help="Uniform refinement levels for grid heatmap reference") parser.add_argument("--grid-resolution", type=int, help="Grid resolution N for heatmap (N x N)") args = parser.parse_args() # ── Load test config ── test_cfg = _load_yaml(args.config) # ── Load base config ── base_config_path = test_cfg.get("base_config", "src/config.yaml") base_config = _load_yaml(base_config_path) # ── Build CLI overrides dict (only non-None values) ── cli_overrides = {} for key in ("geometry", "checkpoint", "output", "seed", "k_test", "cx", "cy", "radius", "eps_r", "half_side", "angle", "n_refine_vertex", "n_refine_grid", "grid_resolution"): val = getattr(args, key.replace("-", "_"), None) if val is not None: cli_overrides[key] = val # Parse --circles if provided if args.circles is not None: circles = [] for spec in args.circles: parts = [float(x.strip()) for x in spec.split(",")] circles.append({ "cx": parts[0], "cy": parts[1], "radius": parts[2], "eps_r": parts[3] if len(parts) > 3 else 3.0, }) cli_overrides["circles"] = circles # ── Set seeds ── seed = cli_overrides.get("seed", test_cfg.get("test", {}).get("seed", 99)) torch.manual_seed(seed) np.random.seed(seed) test_alt_media( base_config=base_config, test_cfg=test_cfg, cli_overrides=cli_overrides, ) if __name__ == "__main__": main()