import copy import os from typing import Any, Dict, List, Optional, Union import numpy as np from skfem import Basis, Mesh from .fem_util import get_element_midpoints from .helmholtz import HelmholtzProblem, create_helmholtz_problem from .utils import IndexSampler class FEMProblemWrapper: """Wraps a HelmholtzProblem, managing mesh, solution cache, and refinement history.""" def __init__( self, *, fem_config: Dict[Union[str, int], Any], fem_problem: HelmholtzProblem, pde_features: Dict[str, List[str]], ): self._fem_config = fem_config self.fem_problem = fem_problem self._pde_element_feature_names = pde_features["element_features"] self._mesh: Optional[Mesh] = None self._previous_mesh: Optional[Mesh] = None self._solution: Optional[np.ndarray] = None self._nodal_solution: Optional[np.ndarray] = None self._refinements_per_element: Optional[np.ndarray] = None self._plot_boundary = np.array(fem_config.get("domain", {}).get("boundary", [0, 0, 1, 1])) def reset(self): self._mesh = self.fem_problem.initial_mesh self._previous_mesh = copy.deepcopy(self._mesh) self._refinements_per_element = np.zeros(self.num_elements, dtype=np.int32) def calculate_solution_and_get_error(self) -> Dict[str, np.ndarray]: self.calculate_solution() return self.get_error_estimate_per_element() def calculate_solution(self) -> None: self._solution = self.fem_problem.calculate_solution(basis=self._basis, cache=True) self._nodal_solution = self._solution def get_error_estimate_per_element(self) -> Dict[str, np.ndarray]: return self.fem_problem.get_error_estimate_per_element( basis=self._basis, solution=self._solution ) def refine_mesh(self, elements_to_refine: np.ndarray) -> np.ndarray: if len(elements_to_refine) > 0: refined_mesh = self._mesh.refined(elements_to_refine) new_midpoints = refined_mesh.p[:, refined_mesh.t].mean(axis=1) element_finder = self._mesh.element_finder() corresponding_elements = element_finder(*new_midpoints) element_indices, inverse_indices, counts = np.unique( corresponding_elements, return_counts=True, return_inverse=True ) self._refinements_per_element[element_indices] += counts - 1 self._refinements_per_element = self._refinements_per_element[inverse_indices] else: refined_mesh = self._mesh inverse_indices = np.arange(self._mesh.t.shape[1]).astype(np.int64) self.mesh = refined_mesh return inverse_indices # ---- PDE 相关的单元特征（source_term 等）---- def element_features(self) -> np.ndarray: return self.fem_problem.element_features( mesh=self._mesh, element_feature_names=self._pde_element_feature_names ) # ---- 将多分量值归约为标量（Helmholtz 取实部）---- def project_to_scalar(self, values: np.ndarray) -> np.ndarray: return self.fem_problem.project_to_scalar(values=values) # ---- 当前 FEM 网格 ---- @property def mesh(self) -> Optional[Mesh]: return self._mesh @mesh.setter def mesh(self, mesh: Mesh) -> None: self._previous_mesh = copy.deepcopy(self._mesh) self._mesh = mesh # ---- P1 线性基函数 ---- @property def _basis(self) -> Basis: return self.fem_problem.mesh_to_basis(self._mesh) # ---- 细化前的网格（奖励计算中回溯用）---- @property def previous_mesh(self) -> Mesh: return self._previous_mesh # ---- 当前网格单元总数 ---- @property def num_elements(self) -> int: return self._mesh.t.shape[1] # ---- 每个单元被细化的次数 ---- @property def refinements_per_element(self) -> np.ndarray: return self._refinements_per_element # ---- 顶点上的 FEM 解 ---- @property def nodal_solution(self) -> np.ndarray: assert self._nodal_solution is not None, "Solution not computed yet" return self._nodal_solution # ---- 单元中点坐标 (num_elements, 2) ---- @property def element_midpoints(self) -> np.ndarray: return get_element_midpoints(self._mesh) # ---- 单元顶点索引 (num_elements, 3) ---- @property def element_indices(self) -> np.ndarray: return self._mesh.t.T # ---- 顶点坐标 (num_vertices, 2) ---- @property def vertex_positions(self) -> np.ndarray: return self._mesh.p.T # ---- 网格边（相邻顶点对索引）---- @property def mesh_edges(self) -> np.ndarray: return self._mesh.facets # ---- 每个单元的相邻单元（排除边界）---- @property def element_neighbors(self) -> np.ndarray: return self._mesh.f2t[:, self._mesh.f2t[1] != -1] # ---- 可视化用的计算域边界框 ---- @property def plot_boundary(self): return self._plot_boundary # ---- 额外的 plotly 渲染图层 ---- def additional_plots(self) -> Dict: return self.fem_problem.additional_plots_from_mesh(self._mesh) class FEMProblemCircularQueue: """Circular buffer of Helmholtz instances for training generalization.""" def __init__( self, *, fem_config: Dict[Union[str, int], Any], random_state: np.random.RandomState = np.random.RandomState(), ): self._fem_config = fem_config self._random_state = random_state num_pdes = fem_config.get("num_pdes", 100) self._use_buffer = num_pdes is not None and num_pdes > 0 num_pdes = num_pdes if self._use_buffer else 1 self._index_sampler = IndexSampler(num_pdes, random_state=self._random_state) self._fem_problems: List[Optional[FEMProblemWrapper]] = [None for _ in range(num_pdes)] pde_config = fem_config.get(fem_config.get("pde_type", "helmholtz"), {}) self._pde_features = { "element_features": [ name for name, include in pde_config.get("element_features", {}).items() if include ], } def next(self) -> FEMProblemWrapper: return self._next_from_idx(pde_idx=self._index_sampler.next()) def _next_from_idx(self, pde_idx: int) -> FEMProblemWrapper: if (not self._use_buffer) or self._fem_problems[pde_idx] is None: new_seed = self._random_state.randint(0, 2**31) new_problem = create_helmholtz_problem( fem_config=self._fem_config, random_state=np.random.RandomState(seed=new_seed), ) self._fem_problems[pde_idx] = FEMProblemWrapper( fem_config=self._fem_config, fem_problem=new_problem, pde_features=self._pde_features, ) self._fem_problems[pde_idx].reset() return self._fem_problems[pde_idx] # PDE 提供的单元特征个数 @property def num_pde_element_features(self) -> int: return len(self._pde_features["element_features"])