635 lines
25 KiB
Python
635 lines
25 KiB
Python
import copy
|
||
import time
|
||
from typing import Any, Dict, List, Optional, Union
|
||
|
||
import numpy as np
|
||
from skfem import Basis, ElementTriP1, Mesh, asm, solve
|
||
from skfem.assembly import BilinearForm, FacetBasis, LinearForm
|
||
from skfem.helpers import dot, grad
|
||
|
||
from .domain import create_domain
|
||
from .fem_util import get_aggregation_per_element, get_element_midpoints
|
||
|
||
|
||
class HelmholtzProblem:
|
||
"""2D Helmholtz scattering FEM solver with Sommerfeld BC."""
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
fem_config: Dict[Union[str, int], Any],
|
||
random_state: np.random.RandomState = np.random.RandomState(),
|
||
):
|
||
helmholtz_config = fem_config.get("helmholtz", {})
|
||
|
||
# ── 1. 波数 k ──
|
||
wave_number_mode = helmholtz_config.get("wave_number_mode", "fixed")
|
||
if wave_number_mode == "random_uniform":
|
||
k_min = helmholtz_config.get("wave_number_min", 2.0)
|
||
k_max = helmholtz_config.get("wave_number_max", 8.0)
|
||
self._k = float(random_state.uniform(k_min, k_max))
|
||
else:
|
||
self._k = float(helmholtz_config.get("wave_number", 10.0))
|
||
|
||
# ── 2. 介质散射体参数 ──
|
||
sc = helmholtz_config.get("scatterer", {})
|
||
scatterer_mode = sc.get("mode", "fixed")
|
||
|
||
if scatterer_mode == "random_uniform":
|
||
self._cx = float(
|
||
random_state.uniform(sc.get("cx_min", 0.3), sc.get("cx_max", 0.7))
|
||
)
|
||
self._cy = float(
|
||
random_state.uniform(sc.get("cy_min", 0.3), sc.get("cy_max", 0.7))
|
||
)
|
||
self._radius = float(
|
||
random_state.uniform(
|
||
sc.get("radius_min", 0.1), sc.get("radius_max", 0.25)
|
||
)
|
||
)
|
||
self._eps_r = float(
|
||
random_state.uniform(
|
||
sc.get("eps_r_min", 2.0), sc.get("eps_r_max", 7.0)
|
||
)
|
||
)
|
||
else:
|
||
self._cx = float(sc.get("cx", 0.5))
|
||
self._cy = float(sc.get("cy", 0.5))
|
||
self._radius = float(sc.get("radius", 0.2))
|
||
self._eps_r = float(sc.get("eps_r", 2.0))
|
||
|
||
# ── 3. 组装 FEM 双线性和线性形式 ──
|
||
self._bilin_form = self._make_bilinear_form()
|
||
self._lin_form_real = self._make_linear_form_real()
|
||
self._lin_form_imag = self._make_linear_form_imag()
|
||
|
||
# ── 4. 初始化域（k^exponent 自适应网格密度 × domain area 线性缩放）──
|
||
# exponent 和 k_ref 均可通过 helmholtz config 配置
|
||
# exponent=2: P1 Helmholtz 理论最优 (污染误差 ∝ (kh)^2, N ∝ k^2)
|
||
# exponent=1.5: 工程折中，避免高 k 初始过密
|
||
# domain area 缩放: 保证不同域尺寸下每单位面积单元数一致 → h 不变
|
||
domain_cfg = copy.deepcopy(fem_config.get("domain"))
|
||
boundary = domain_cfg.get("boundary", [0, 0, 1, 1])
|
||
domain_area = (boundary[2] - boundary[0]) * (boundary[3] - boundary[1])
|
||
k_ref = helmholtz_config.get("k_ref", 6.0)
|
||
k_exponent = helmholtz_config.get("k_exponent", 2.0)
|
||
base_elements = domain_cfg.get("initial_num_elements", 400)
|
||
scaled_elements = int(base_elements * (self._k / k_ref) ** k_exponent * domain_area)
|
||
domain_cfg["initial_num_elements"] = max(scaled_elements, int(base_elements * domain_area))
|
||
self._domain = create_domain(
|
||
domain_config=domain_cfg,
|
||
random_state=copy.deepcopy(random_state),
|
||
)
|
||
|
||
# ── 4.5. 介质区前渐近区边缘约束 ──
|
||
# 放宽 Nyquist (N=4) → 前渐近区边缘 (N=1~1.5)，赋予介质内初始网格基本相位解析能力
|
||
# 约束: h_init ≤ λ_local / N，λ_local = 2π/(k√ε_r)
|
||
# N=1.5 对应约 1.5 点/波长，刚好跨过渐近区门槛，不撑爆物理预算
|
||
pre_asymptotic_N = helmholtz_config.get("pre_asymptotic_N", 1.5)
|
||
pre_asymptotic_mesh = self._enforce_nyquist_in_dielectric(
|
||
self._domain.initial_mesh, N=pre_asymptotic_N
|
||
)
|
||
self._domain.replace_initial_mesh(pre_asymptotic_mesh)
|
||
|
||
# ── 5. PDE 特征名称 ──
|
||
pde_config = fem_config.get(fem_config.get("pde_type", "helmholtz"), {})
|
||
self._element_feature_names = [
|
||
name
|
||
for name, include in pde_config.get("element_features", {}).items()
|
||
if include
|
||
]
|
||
|
||
# ── Public interface ─────────────────────────────────────
|
||
|
||
def mesh_to_basis(self, mesh: Mesh) -> Basis:
|
||
return Basis(mesh, ElementTriP1())
|
||
|
||
def calculate_solution(self, basis: Basis, cache: bool = False) -> np.ndarray:
|
||
_t = {}
|
||
|
||
_t0 = time.perf_counter()
|
||
K = asm(self._bilin_form, basis)
|
||
_t1 = time.perf_counter()
|
||
f = asm(self._lin_form_real, basis) + 1j * asm(self._lin_form_imag, basis)
|
||
_t2 = time.perf_counter()
|
||
|
||
boundary_facets = basis.mesh.boundary_facets()
|
||
facet_basis = FacetBasis(basis.mesh, basis.elem, facets=boundary_facets)
|
||
|
||
@BilinearForm
|
||
def boundary_mass(u, v, w):
|
||
return u * v
|
||
|
||
M_boundary = asm(boundary_mass, facet_basis)
|
||
_t3 = time.perf_counter()
|
||
K_total = K.astype(np.complex128) - 1j * self._k * M_boundary
|
||
u_scat = solve(K_total, f)
|
||
_t4 = time.perf_counter()
|
||
|
||
_t["assemble_K"] = _t1 - _t0
|
||
_t["assemble_f"] = _t2 - _t1
|
||
_t["assemble_boundary"] = _t3 - _t2
|
||
_t["solve"] = _t4 - _t3
|
||
_t["total"] = _t4 - _t0
|
||
_t["n_dof"] = int(basis.mesh.p.shape[1])
|
||
self._last_solve_timing = _t
|
||
|
||
return u_scat
|
||
|
||
def get_error_estimate_per_element(
|
||
self, basis: Basis, solution: np.ndarray
|
||
) -> Dict[str, np.ndarray]:
|
||
eps_r_arr = _compute_eps_r_at_midpoints(basis.mesh, self._cx, self._cy, self._radius, self._eps_r)
|
||
return {"indicator": _compute_residual_indicator(basis.mesh, solution, k=self._k, eps_r=eps_r_arr)}
|
||
|
||
def element_features(self, mesh: Mesh, element_feature_names: List[str]) -> Optional[np.ndarray]:
|
||
features_list = []
|
||
if "epsilon_r" in element_feature_names:
|
||
features_list.append(
|
||
_compute_eps_r_at_midpoints(mesh, self._cx, self._cy, self._radius, self._eps_r)[:, None]
|
||
)
|
||
return np.concatenate(features_list, axis=1) if features_list else None
|
||
|
||
def _enforce_nyquist_in_dielectric(self, mesh: Mesh, N: float = 1.5, max_iter: int = 10) -> Mesh:
|
||
"""Iteratively refine elements inside the dielectric until h_K ≤ λ_d/N.
|
||
|
||
λ_d = 2π/(k√ε_r) is the wavelength inside the dielectric.
|
||
N=1.5 corresponds to the edge of the pre-asymptotic regime (~1.5 points
|
||
per wavelength) — just enough for the wave field to exhibit basic phase
|
||
resolution without exhausting the physical element budget. This relaxes
|
||
the old Nyquist N=4 constraint, leaving headroom for the RL agent to
|
||
selectively refine where residual indicators demand it.
|
||
"""
|
||
lambda_d = 2.0 * np.pi / (self._k * np.sqrt(self._eps_r))
|
||
h_max = lambda_d / N
|
||
|
||
for _ in range(max_iter):
|
||
i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
|
||
x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
|
||
x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
|
||
x2, y2 = mesh.p[0, i2], mesh.p[1, i2]
|
||
|
||
e01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
|
||
e12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
|
||
e20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
|
||
h_K = np.maximum(np.maximum(e01, e12), e20)
|
||
|
||
midpoints = np.mean(mesh.p[:, mesh.t], axis=1).T
|
||
in_dielectric = (
|
||
(midpoints[:, 0] - self._cx) ** 2
|
||
+ (midpoints[:, 1] - self._cy) ** 2
|
||
<= self._radius**2
|
||
)
|
||
|
||
to_refine = np.where(in_dielectric & (h_K > h_max))[0]
|
||
if len(to_refine) == 0:
|
||
break
|
||
mesh = mesh.refined(to_refine)
|
||
|
||
return mesh
|
||
|
||
# ── Properties ───────────────────────────────────────────
|
||
|
||
@property
|
||
def initial_mesh(self) -> Mesh:
|
||
return self._domain.initial_mesh
|
||
|
||
@property
|
||
def boundary_line_segments(self) -> np.ndarray:
|
||
return self._domain.boundary_line_segments
|
||
|
||
|
||
@staticmethod
|
||
def project_to_scalar(values: np.ndarray) -> np.ndarray:
|
||
return values
|
||
|
||
def additional_plots_from_mesh(self, mesh: Mesh) -> Dict:
|
||
return {}
|
||
|
||
# ── FEM form assembly ────────────────────────────────────
|
||
|
||
def _eps_r_at_quad_points(self, x, y):
|
||
in_cyl = (x - self._cx) ** 2 + (y - self._cy) ** 2 <= self._radius**2
|
||
return np.where(in_cyl, self._eps_r, 1.0)
|
||
|
||
def _make_bilinear_form(self):
|
||
k = self._k
|
||
|
||
@BilinearForm
|
||
def bilin(u, v, w):
|
||
eps_r = self._eps_r_at_quad_points(w.x[0], w.x[1])
|
||
return dot(grad(u), grad(v)) - k**2 * eps_r * u * v
|
||
|
||
return bilin
|
||
|
||
def _make_linear_form_real(self):
|
||
k = self._k
|
||
|
||
@LinearForm
|
||
def lin(v, w):
|
||
eps_r = self._eps_r_at_quad_points(w.x[0], w.x[1])
|
||
return k**2 * (eps_r - 1.0) * np.cos(k * w.x[0]) * v
|
||
|
||
return lin
|
||
|
||
def _make_linear_form_imag(self):
|
||
k = self._k
|
||
|
||
@LinearForm
|
||
def lin(v, w):
|
||
eps_r = self._eps_r_at_quad_points(w.x[0], w.x[1])
|
||
return k**2 * (eps_r - 1.0) * np.sin(k * w.x[0]) * v
|
||
|
||
return lin
|
||
|
||
|
||
# ── 辅助函数 ──────────────────────────────────────────────────
|
||
|
||
|
||
def _compute_eps_r_at_midpoints(
|
||
mesh: Mesh,
|
||
cx: float = 0.5,
|
||
cy: float = 0.5,
|
||
radius: float = 0.2,
|
||
eps_r_in: float = 2.0,
|
||
) -> np.ndarray:
|
||
"""
|
||
计算每个单元中点处的相对介电常数 ε_r。
|
||
|
||
判断单元中点是否落在介质圆柱内:
|
||
- 在圆柱内 → ε_r = eps_r_in (如 2.0)
|
||
- 在圆柱外 → ε_r = 1.0 (真空)
|
||
|
||
Returns:
|
||
eps_r: shape (num_elements,)
|
||
"""
|
||
midpoints = get_element_midpoints(mesh)
|
||
x_mid, y_mid = midpoints[:, 0], midpoints[:, 1]
|
||
in_cylinder = (x_mid - cx) ** 2 + (y_mid - cy) ** 2 <= radius**2
|
||
return np.where(in_cylinder, eps_r_in, 1.0)
|
||
|
||
|
||
def _compute_residual_indicator(
|
||
mesh: Mesh,
|
||
u_h: np.ndarray,
|
||
k: float = 10.0,
|
||
eps_r: Union[float, np.ndarray] = 1.0,
|
||
) -> np.ndarray:
|
||
"""
|
||
基于残差的逐单元后验误差估计 — 无量纲化版本。
|
||
|
||
使用真空波数 k₀ 归一化（非 k_local），使误差指示子反映"绝对物理误差"
|
||
而非"相对局部波长的分辨率"。介质内短波（ε_r>1）的残差在 k_local 下被
|
||
压低 √ε_r 倍，改用 k₀ 后介质内 η 自然放大，Agent 获得正确优先级。
|
||
|
||
P1 单元三项:
|
||
1. r_int = (h_K/k)·√V_K · |k²ε_r·u_h + k²(ε_r-1)·u_inc|
|
||
2. r_jump = √(½ Σ_{e∈∂K} (h_e/k)·|[[∇u_h·n]]|²)
|
||
3. r_sbc = (h_bnd/k)·|∂u/∂n - i·k_local·u| (SBC 条件仍用 k_local)
|
||
|
||
Returns:
|
||
eta_elements: shape (num_elements,) 的逐单元误差指标
|
||
"""
|
||
n_elements = mesh.t.shape[1]
|
||
eps_r = np.asarray(eps_r)
|
||
|
||
# ── 1. 单元几何量 ──
|
||
i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
|
||
x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
|
||
x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
|
||
x2, y2 = mesh.p[0, i2], mesh.p[1, i2]
|
||
|
||
det_J = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
|
||
element_areas = np.abs(det_J) / 2.0
|
||
|
||
edge_len_01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
|
||
edge_len_12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
|
||
edge_len_20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
|
||
h_K = np.maximum(np.maximum(edge_len_01, edge_len_12), edge_len_20)
|
||
|
||
# ── 2. 梯度（常数，因为是 P1 单元）──
|
||
u0, u1, u2 = u_h[i0], u_h[i1], u_h[i2]
|
||
inv_det = np.divide(1.0, det_J, where=det_J != 0, out=np.zeros_like(det_J))
|
||
du10, du20 = u1 - u0, u2 - u0
|
||
|
||
grad_x = ((y2 - y0) * du10 - (y1 - y0) * du20) * inv_det
|
||
grad_y = (-(x2 - x0) * du10 + (x1 - x0) * du20) * inv_det
|
||
grad_per_element = np.column_stack([grad_x, grad_y])
|
||
|
||
# ── 3. 单元内部残差 ──
|
||
u_mid = (u0 + u1 + u2) / 3.0
|
||
x_mid = (x0 + x1 + x2) / 3.0
|
||
u_inc_mid = np.exp(1j * k * x_mid)
|
||
f_mid = (k**2) * (eps_r - 1.0) * u_inc_mid
|
||
r_mid = f_mid + (k**2) * eps_r * u_mid
|
||
|
||
cell_residual_sq = (h_K**2) * element_areas * np.abs(r_mid) ** 2 / (k ** 2)
|
||
cell_residual_sq[element_areas < 1e-15] = 0.0
|
||
|
||
# ── 4. 内部边梯度跳变 ──
|
||
interior_facets_idx = np.where(mesh.f2t[1] != -1)[0]
|
||
elem_left = mesh.f2t[0, interior_facets_idx]
|
||
elem_right = mesh.f2t[1, interior_facets_idx]
|
||
|
||
edges_p1 = mesh.p[:, mesh.facets[0, interior_facets_idx]].T
|
||
edges_p2 = mesh.p[:, mesh.facets[1, interior_facets_idx]].T
|
||
edge_vectors = edges_p2 - edges_p1
|
||
h_e = np.linalg.norm(edge_vectors, axis=1)
|
||
n_e = np.c_[edge_vectors[:, 1], -edge_vectors[:, 0]] / (h_e[:, None] + 1e-15)
|
||
|
||
grad_left = grad_per_element[elem_left]
|
||
grad_right = grad_per_element[elem_right]
|
||
jump_val = np.abs(np.sum((grad_left - grad_right) * n_e, axis=1))
|
||
jump_val_sq = jump_val ** 2
|
||
|
||
jump_residual_sq = np.zeros(n_elements)
|
||
np.add.at(jump_residual_sq, elem_left, 0.5 * h_e * jump_val_sq / k)
|
||
np.add.at(jump_residual_sq, elem_right, 0.5 * h_e * jump_val_sq / k)
|
||
|
||
# ── 5. 合并 ──
|
||
eta_sq = cell_residual_sq + jump_residual_sq
|
||
|
||
# ── 6. SBC 边界残差 ──
|
||
boundary_facets_idx = np.where(mesh.f2t[1] == -1)[0]
|
||
if len(boundary_facets_idx) > 0:
|
||
bnd_elem = mesh.f2t[0, boundary_facets_idx]
|
||
bnd_p1 = mesh.p[:, mesh.facets[0, boundary_facets_idx]].T
|
||
bnd_p2 = mesh.p[:, mesh.facets[1, boundary_facets_idx]].T
|
||
bnd_vectors = bnd_p2 - bnd_p1
|
||
h_bnd = np.linalg.norm(bnd_vectors, axis=1)
|
||
n_bnd = np.c_[bnd_vectors[:, 1], -bnd_vectors[:, 0]] / (h_bnd[:, None] + 1e-15)
|
||
|
||
grad_bnd = grad_per_element[bnd_elem]
|
||
du_dn = np.sum(grad_bnd * n_bnd, axis=1)
|
||
|
||
if eps_r.ndim == 1:
|
||
k_local = k * np.sqrt(np.maximum(eps_r[bnd_elem], 1.0))
|
||
else:
|
||
k_local = k
|
||
|
||
u_edge_mean = (
|
||
u_h[mesh.facets[0, boundary_facets_idx]]
|
||
+ u_h[mesh.facets[1, boundary_facets_idx]]
|
||
) / 2.0
|
||
sbc_residual = du_dn - 1j * k_local * u_edge_mean
|
||
sbc_residual_sq = (h_bnd ** 2) * np.abs(sbc_residual) ** 2 / (k ** 2)
|
||
np.add.at(eta_sq, bnd_elem, sbc_residual_sq)
|
||
|
||
eta_sq = np.maximum(eta_sq, 0.0)
|
||
return np.sqrt(eta_sq)
|
||
|
||
|
||
def _compute_residual_components(
|
||
mesh: Mesh,
|
||
u_h: np.ndarray,
|
||
k: float = 10.0,
|
||
eps_r: Union[float, np.ndarray] = 1.0,
|
||
apply_log: bool = True,
|
||
) -> Dict[str, np.ndarray]:
|
||
"""
|
||
计算逐单元的三项 PDE 物理残差（分离版，无量纲化）。
|
||
|
||
使用真空波数 k₀ 归一化 — 介质内短波残差不再被 k_local 压低，GNN 获得
|
||
正确的介质内/外优先级信号。
|
||
|
||
P1 单元返回:
|
||
internal_residual: (h_K/k)·√V_i · |k²ε_r·u + k²(ε_r-1)·u_inc|
|
||
gradient_jump: √(½ Σ_{e∈∂K_i} (h_e/k)·|[[∇u·n]]|²)
|
||
sbc_residual: (h_bnd/k)·|∂u/∂n - i·k_local·u| (SBC 条件仍用 k_local)
|
||
element_areas: 单元面积
|
||
is_sbc_boundary: 该单元是否与 SBC 边界相邻 (0/1)
|
||
|
||
Args:
|
||
apply_log: True → log10 压缩（喂 GNN）；False → 原始值（喂 reward）
|
||
"""
|
||
n_elements = mesh.t.shape[1]
|
||
eps_r = np.asarray(eps_r)
|
||
|
||
# ── 1. 单元几何量 ──
|
||
i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
|
||
x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
|
||
x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
|
||
x2, y2 = mesh.p[0, i2], mesh.p[1, i2]
|
||
|
||
det_J = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
|
||
element_areas = np.abs(det_J) / 2.0
|
||
|
||
edge_len_01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
|
||
edge_len_12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
|
||
edge_len_20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
|
||
h_K = np.maximum(np.maximum(edge_len_01, edge_len_12), edge_len_20)
|
||
|
||
# ── 2. 梯度（常数，因为是 P1 单元）──
|
||
u0, u1, u2 = u_h[i0], u_h[i1], u_h[i2]
|
||
inv_det = np.divide(1.0, det_J, where=det_J != 0, out=np.zeros_like(det_J))
|
||
du10, du20 = u1 - u0, u2 - u0
|
||
|
||
grad_x = ((y2 - y0) * du10 - (y1 - y0) * du20) * inv_det
|
||
grad_y = (-(x2 - x0) * du10 + (x1 - x0) * du20) * inv_det
|
||
grad_per_element = np.column_stack([grad_x, grad_y])
|
||
|
||
# P1 单元内部残差: ∇²u_h = 0（线性元二阶导为零），故仅含反应项
|
||
# 完整强形式: |∇²u + k²·ε_r·u + k²·(ε_r-1)·u_inc|
|
||
# 对 P1: ∇²u_h ≡ 0 → 残差 = |k²·ε_r·u + k²·(ε_r-1)·u_inc|
|
||
u_mid = (u0 + u1 + u2) / 3.0
|
||
x_mid = (x0 + x1 + x2) / 3.0
|
||
u_inc_mid = np.exp(1j * k * x_mid)
|
||
f_mid = (k**2) * (eps_r - 1.0) * u_inc_mid
|
||
r_mid = f_mid + (k**2) * eps_r * u_mid
|
||
internal_residual = (h_K / k) * np.sqrt(element_areas) * np.abs(r_mid)
|
||
internal_residual[element_areas < 1e-15] = 0.0
|
||
|
||
# ── 4. 内部边梯度跳变 (逐单元) ──
|
||
interior_facets_idx = np.where(mesh.f2t[1] != -1)[0]
|
||
elem_left = mesh.f2t[0, interior_facets_idx]
|
||
elem_right = mesh.f2t[1, interior_facets_idx]
|
||
|
||
edges_p1 = mesh.p[:, mesh.facets[0, interior_facets_idx]].T
|
||
edges_p2 = mesh.p[:, mesh.facets[1, interior_facets_idx]].T
|
||
edge_vectors = edges_p2 - edges_p1
|
||
h_e = np.linalg.norm(edge_vectors, axis=1)
|
||
n_e = np.c_[edge_vectors[:, 1], -edge_vectors[:, 0]] / (h_e[:, None] + 1e-15)
|
||
|
||
grad_left = grad_per_element[elem_left]
|
||
grad_right = grad_per_element[elem_right]
|
||
jump_val = np.abs(np.sum((grad_left - grad_right) * n_e, axis=1))
|
||
|
||
gradient_jump = np.zeros(n_elements, dtype=np.float64)
|
||
jump_sq_per_edge = jump_val ** 2
|
||
np.add.at(gradient_jump, elem_left, 0.5 * h_e * jump_sq_per_edge / k)
|
||
np.add.at(gradient_jump, elem_right, 0.5 * h_e * jump_sq_per_edge / k)
|
||
gradient_jump = np.sqrt(gradient_jump)
|
||
|
||
# ── 5. SBC 边界残差 + 边界标记 ──
|
||
sbc_residual = np.zeros(n_elements, dtype=np.float64)
|
||
is_sbc_boundary = np.zeros(n_elements, dtype=np.float32)
|
||
boundary_facets_idx = np.where(mesh.f2t[1] == -1)[0]
|
||
if len(boundary_facets_idx) > 0:
|
||
bnd_elem = mesh.f2t[0, boundary_facets_idx]
|
||
bnd_p1 = mesh.p[:, mesh.facets[0, boundary_facets_idx]].T
|
||
bnd_p2 = mesh.p[:, mesh.facets[1, boundary_facets_idx]].T
|
||
bnd_vectors = bnd_p2 - bnd_p1
|
||
h_bnd = np.linalg.norm(bnd_vectors, axis=1)
|
||
n_bnd = np.c_[bnd_vectors[:, 1], -bnd_vectors[:, 0]] / (h_bnd[:, None] + 1e-15)
|
||
|
||
grad_bnd = grad_per_element[bnd_elem]
|
||
du_dn = np.sum(grad_bnd * n_bnd, axis=1)
|
||
|
||
if eps_r.ndim == 1:
|
||
k_local = k * np.sqrt(np.maximum(eps_r[bnd_elem], 1.0))
|
||
else:
|
||
k_local = k
|
||
|
||
u_edge_mean = (
|
||
u_h[mesh.facets[0, boundary_facets_idx]]
|
||
+ u_h[mesh.facets[1, boundary_facets_idx]]
|
||
) / 2.0
|
||
sbc_val = np.abs(du_dn - 1j * k_local * u_edge_mean)
|
||
np.add.at(sbc_residual, bnd_elem, (h_bnd / k) * sbc_val)
|
||
is_sbc_boundary[bnd_elem] = 1.0
|
||
|
||
# ── 对数预处理：压缩跨数量级动态范围（仅 GNN 特征需要）──
|
||
if apply_log:
|
||
_log_eps = 1e-8
|
||
internal_residual = np.log10(np.maximum(internal_residual, _log_eps))
|
||
gradient_jump = np.log10(np.maximum(gradient_jump, _log_eps))
|
||
sbc_residual = np.log10(np.maximum(sbc_residual, _log_eps))
|
||
|
||
return {
|
||
"internal_residual": internal_residual.astype(np.float32),
|
||
"gradient_jump": gradient_jump.astype(np.float32),
|
||
"sbc_residual": sbc_residual.astype(np.float32),
|
||
"element_areas": element_areas.astype(np.float32),
|
||
"is_sbc_boundary": is_sbc_boundary,
|
||
}
|
||
|
||
|
||
def _compute_residual_density(
|
||
mesh: Mesh,
|
||
u_h: np.ndarray,
|
||
k: float = 10.0,
|
||
eps_r: Union[float, np.ndarray] = 1.0,
|
||
) -> Dict[str, np.ndarray]:
|
||
"""
|
||
Compute intensive (h-free) residual density components for reward.
|
||
|
||
Unlike _compute_residual_components which includes h-scaling
|
||
(h_K·√V, h_e·|jump|, h_bnd·|sbc|), this returns the raw PDE residuals
|
||
that are independent of element size — true "error densities".
|
||
|
||
Returns:
|
||
rho_int: |k²·ε_r·u + k²·(ε_r-1)·u_inc| per element
|
||
rho_jump: √(mean_{e∈∂K_int} |[[∇u·n]]|²) per element
|
||
rho_sbc: √(mean_{e∈∂K∩Γ_sbc} |∂u/∂n - i·k·u|²) per element
|
||
"""
|
||
n_elements = mesh.t.shape[1]
|
||
eps_r = np.asarray(eps_r)
|
||
|
||
# ── 1. element geometry ──
|
||
i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
|
||
x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
|
||
x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
|
||
x2, y2 = mesh.p[0, i2], mesh.p[1, i2]
|
||
|
||
det_J = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
|
||
|
||
# ── 2. gradient (constant per P1 element) ──
|
||
u0, u1, u2 = u_h[i0], u_h[i1], u_h[i2]
|
||
inv_det = np.divide(1.0, det_J, where=det_J != 0, out=np.zeros_like(det_J))
|
||
du10, du20 = u1 - u0, u2 - u0
|
||
|
||
grad_x = ((y2 - y0) * du10 - (y1 - y0) * du20) * inv_det
|
||
grad_y = (-(x2 - x0) * du10 + (x1 - x0) * du20) * inv_det
|
||
grad_per_element = np.column_stack([grad_x, grad_y])
|
||
|
||
# ── 3. interior residual density: |k²·ε_r·u_mid + k²·(ε_r-1)·u_inc_mid| ──
|
||
u_mid = (u0 + u1 + u2) / 3.0
|
||
x_mid = (x0 + x1 + x2) / 3.0
|
||
u_inc_mid = np.exp(1j * k * x_mid)
|
||
r_mid = (k**2) * eps_r * u_mid + (k**2) * (eps_r - 1.0) * u_inc_mid
|
||
rho_int = np.abs(r_mid)
|
||
|
||
# ── 4. gradient jump density: √(mean |[[∇u·n]]|²) per element ──
|
||
interior_facets_idx = np.where(mesh.f2t[1] != -1)[0]
|
||
elem_left = mesh.f2t[0, interior_facets_idx]
|
||
elem_right = mesh.f2t[1, interior_facets_idx]
|
||
|
||
edges_p1 = mesh.p[:, mesh.facets[0, interior_facets_idx]].T
|
||
edges_p2 = mesh.p[:, mesh.facets[1, interior_facets_idx]].T
|
||
edge_vectors = edges_p2 - edges_p1
|
||
h_e = np.linalg.norm(edge_vectors, axis=1)
|
||
n_e = np.c_[edge_vectors[:, 1], -edge_vectors[:, 0]] / (h_e[:, None] + 1e-15)
|
||
|
||
grad_left = grad_per_element[elem_left]
|
||
grad_right = grad_per_element[elem_right]
|
||
jump_val_sq = np.abs(np.sum((grad_left - grad_right) * n_e, axis=1)) ** 2
|
||
|
||
jump_sq_sum = np.zeros(n_elements, dtype=np.float64)
|
||
jump_count = np.zeros(n_elements, dtype=np.float64)
|
||
np.add.at(jump_sq_sum, elem_left, jump_val_sq)
|
||
np.add.at(jump_sq_sum, elem_right, jump_val_sq)
|
||
np.add.at(jump_count, elem_left, 1)
|
||
np.add.at(jump_count, elem_right, 1)
|
||
|
||
rho_jump = np.zeros(n_elements, dtype=np.float64)
|
||
mask_jump = jump_count > 0
|
||
rho_jump[mask_jump] = np.sqrt(jump_sq_sum[mask_jump] / jump_count[mask_jump])
|
||
|
||
# ── 5. SBC boundary density: √(mean |∂u/∂n - i·k·u|²) per element ──
|
||
rho_sbc = np.zeros(n_elements, dtype=np.float64)
|
||
boundary_facets_idx = np.where(mesh.f2t[1] == -1)[0]
|
||
if len(boundary_facets_idx) > 0:
|
||
bnd_elem = mesh.f2t[0, boundary_facets_idx]
|
||
bnd_p1 = mesh.p[:, mesh.facets[0, boundary_facets_idx]].T
|
||
bnd_p2 = mesh.p[:, mesh.facets[1, boundary_facets_idx]].T
|
||
bnd_vectors = bnd_p2 - bnd_p1
|
||
h_bnd = np.linalg.norm(bnd_vectors, axis=1)
|
||
n_bnd = np.c_[bnd_vectors[:, 1], -bnd_vectors[:, 0]] / (h_bnd[:, None] + 1e-15)
|
||
|
||
grad_bnd = grad_per_element[bnd_elem]
|
||
du_dn = np.sum(grad_bnd * n_bnd, axis=1)
|
||
|
||
if eps_r.ndim == 1:
|
||
k_local = k * np.sqrt(np.maximum(eps_r[bnd_elem], 1.0))
|
||
else:
|
||
k_local = k
|
||
|
||
u_edge_mean = (
|
||
u_h[mesh.facets[0, boundary_facets_idx]]
|
||
+ u_h[mesh.facets[1, boundary_facets_idx]]
|
||
) / 2.0
|
||
sbc_val_sq = np.abs(du_dn - 1j * k_local * u_edge_mean) ** 2
|
||
|
||
sbc_sq_sum = np.zeros(n_elements, dtype=np.float64)
|
||
sbc_count = np.zeros(n_elements, dtype=np.float64)
|
||
np.add.at(sbc_sq_sum, bnd_elem, sbc_val_sq)
|
||
np.add.at(sbc_count, bnd_elem, 1)
|
||
|
||
mask_sbc = sbc_count > 0
|
||
rho_sbc[mask_sbc] = np.sqrt(sbc_sq_sum[mask_sbc] / sbc_count[mask_sbc])
|
||
|
||
return {
|
||
"rho_int": rho_int.astype(np.float64),
|
||
"rho_jump": rho_jump.astype(np.float64),
|
||
"rho_sbc": rho_sbc.astype(np.float64),
|
||
}
|
||
|
||
|
||
# ── 工厂函数 ──────────────────────────────────────────────────
|
||
|
||
|
||
def create_helmholtz_problem(
|
||
*, fem_config: Dict[Union[str, int], Any], random_state: np.random.RandomState
|
||
) -> HelmholtzProblem:
|
||
"""
|
||
创建 Helmholtz 问题实例。
|
||
|
||
Args:
|
||
fem_config: FEM 配置字典
|
||
random_state: 随机状态
|
||
|
||
Returns:
|
||
HelmholtzProblem 实例
|
||
"""
|
||
return HelmholtzProblem(fem_config=fem_config, random_state=random_state)
|