610 lines
26 KiB
Python
610 lines
26 KiB
Python
#!/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()
|