from typing import Any, Callable, Dict, List, Optional, Tuple, Union import gym import numpy as np import plotly.graph_objects as go import torch from plotly.basedatatypes import BaseTraceType from skfem import Mesh from torch_geometric.data import Data from .fem_problem import FEMProblemCircularQueue, FEMProblemWrapper from .fem_util import ( construct_sizing_field_1d, get_aggregation_per_element, get_triangle_areas_from_indices, sample_in_range, ) from .utils import save_concatenate from .visualization import get_plotly_mesh_traces_and_layout class MeshRefinement(gym.Env): """Graph-based 2D mesh refinement RL environment using scikit-FEM backend.""" def __init__( self, environment_config: Dict[Union[str, int], Any], seed: Optional[int] = None ): """ Args: environment_config: Config for the environment. Details can be found in the configs/references/mesh_refinement_reference.yaml example file seed: Optional seed for the random number generator. """ self._environment_config = environment_config self._random_state: np.random.RandomState = np.random.RandomState(seed=seed) self._num_node_features: int = 0 self._num_edge_features: int = 0 self.fem_problem_queue = FEMProblemCircularQueue( fem_config=environment_config.get("fem"), random_state=np.random.RandomState(seed=seed), ) self.fem_problem: Optional[FEMProblemWrapper] = None ################################################ # general environment parameters # ################################################ self._refinement_strategy: str = environment_config.get("refinement_strategy") self._max_timesteps = environment_config.get("num_timesteps") self._element_limit_penalty = environment_config.get("element_limit_penalty") self._maximum_elements = environment_config.get("maximum_elements") self._element_penalty_config = self._environment_config.get("element_penalty") self._sample_penalty = self._element_penalty_config.get("sample_penalty") ################################################ # graph connectivity, feature and action space # ################################################ self._reward_type = environment_config.get("reward_type") self._global_reward_alpha = float(environment_config.get("global_reward_alpha", 0.2)) _rho_w = environment_config.get("rho_weights", {}) self._w_rho_int = float(_rho_w.get("w_int", 1.0)) self._w_rho_jump = float(_rho_w.get("w_jump", 1.0)) self._w_rho_sbc = float(_rho_w.get("w_sbc", 1.0)) self._include_vertices = environment_config.get("include_vertices") self._set_graph_sizes() ################################################ # internal state and cache # ################################################ self._timestep: int = 0 self._element_penalty_lambda = None # 0 # set default value self._initial_approximation_errors: Optional[Dict[str, float]] = None self._reward = None self._cumulative_return: np.array = 0 # return of the environment # dictionary containing the error estimation for the current solution for different error evaluation metrics self._error_estimation_dict: Optional[Dict[str, np.array]] = None self._initial_error_norm = None self.current_error = None # scalar total error, updated each step self.initial_error = ( None # scalar initial total error, fixed as normalization baseline ) # last-step history for delta-based rewards and plotting self._previous_error_per_element: Optional[np.array] = None self._previous_num_elements: Optional[int] = None self._previous_agent_mapping = None self._previous_element_volumes = None self._previous_std_per_element = None self._previous_eta_components: Optional[Dict[str, np.ndarray]] = None self._previous_rho_components: Optional[Dict[str, np.ndarray]] = None # fields/internal variables for spatial mesh refinement, especially a spatial reward self._agent_mapping = None # mapping List[old_element_indices] of size new_element_indices that maps self._reward_per_agent: Optional[np.array] = ( 0 # cumulative return of the environment per agent ) self._cumulative_reward_per_agent: Optional[np.array] = ( 0 # cumulative reward of the environment per agent ) # additional policy information that is not passed through the graph self._include_additional_policy_information = environment_config.get( "include_additional_policy_information" ) self._manual_normalization = environment_config.get( "manual_normalization", None ) # manually normalize the error ################################################ # recording and plotting # ################################################ self._initial_num_elements = None def _set_graph_sizes(self): """ Internally sets the * action dimension * number of node types and node features for each type * number of edge types and edge features for each type depending on the configuration. Uses the same edge features for all edge types. Returns: """ edge_feature_config = self._environment_config.get("edge_features") self._edge_features = [ feature_name for feature_name, include_feature in edge_feature_config.items() if include_feature ] # set number of edge features num_edge_features = 0 if "euclidean_distance" in self._edge_features: num_edge_features += 1 self._element_feature_functions = self._register_element_features() self._num_node_features = len(self._element_feature_functions) self._num_node_features += self.fem_problem_queue.num_pde_element_features self._num_edge_features = num_edge_features def _register_element_features(self) -> Dict[str, Callable[[], np.array]]: cfg = self._environment_config.get("element_features") names = [n for n, inc in cfg.items() if inc] feats = {} if "x_position" in names: feats["x_position"] = lambda: self._element_midpoints[:, 0] if "y_position" in names: feats["y_position"] = lambda: self._element_midpoints[:, 1] if "volume" in names: feats["volume"] = lambda: self._volume_normalized if "solution_std" in names: feats["internal_residual"] = lambda: self._residual_components["internal_residual"] feats["gradient_jump"] = lambda: self._residual_components["gradient_jump"] feats["sbc_residual"] = lambda: self._residual_components["sbc_residual"] if "element_penalty" in names: feats["element_penalty"] = lambda: np.repeat(self._element_penalty_lambda, self._num_elements) if "timestep" in names: feats["timestep"] = lambda: np.repeat(self._timestep, self._num_elements) if "wave_number" in names: feats["wave_number"] = lambda: np.repeat(self._wave_number, self._num_elements) if "k_local_sqrt_vol" in names: feats["k_local_sqrt_vol"] = lambda: self._k_local_sqrt_vol if "is_sbc_boundary" in names: feats["is_sbc_boundary"] = lambda: self._residual_components["is_sbc_boundary"] if "dist_to_interface" in names: feats["dist_to_interface"] = lambda: self._dist_to_interface # Complex field decomposition (always present for Helmholtz) feats["epsilon_r"] = lambda: self._epsilon_r_elements feats["total_solution_magnitude"] = lambda: np.abs(self._complex_solution_mean) return feats def reset(self) -> Data: """ Resets the environment and returns an (initial) observation of the next rollout according to the reset environment state Returns: The observation of the initial state. """ # get the next fem problem. This samples a new domain and new load function, resets the mesh and the solution. self.fem_problem = self.fem_problem_queue.next() # calculate the solution of the finite element problem for the initial mesh and retrieve an error per element self._error_estimation_dict = ( self.fem_problem.calculate_solution_and_get_error() ) # reset the internal state of the environment. This includes the current timestep, the current element penalty # and some values for calculating the reward and plotting the env self._reset_internal_state() observation = self.last_observation return observation def _reset_internal_state(self): """ Resets the internal state of the environment Returns: """ self._agent_mapping = np.arange(self._num_elements).astype( np.int64 ) # map to identity at first step self._previous_agent_mapping = np.arange(self._num_elements).astype( np.int64 ) # map to identity at first step self._previous_element_volumes = self.element_volumes self._previous_eta_indicator = self._eta_indicator self._previous_eta_components = self._eta_components_raw self._previous_rho_components = self._rho_components self._previous_solution_l2_norm = self._compute_solution_l2_norm() self._reward_per_agent = np.zeros(self.num_agents) self._cumulative_reward_per_agent = np.zeros(self._num_elements) # reset timestep and rewards self._timestep = 0 self._reward = 0 self._cumulative_return = 0 self._diag_selected_count = -1 # 防止跨 episode 残留触发 is_terminal # reset internal state that tracks statistics over the episode self._previous_error_per_element = self.error_per_element # collect a dictionary of initial errors to normalize them when calculating metrics during evaluation self._initial_approximation_errors = ( self._calculate_initial_approximation_errors() ) self._previous_num_elements = self._num_elements self._initial_num_elements = self._num_elements self._initial_median_area = float(np.median(self.element_volumes)) k = self._wave_number eps_r_elem = self._epsilon_r_elements lambda_local = 2.0 * np.pi / (k * np.sqrt(np.maximum(eps_r_elem, 1.0))) A_budget = 0.5 * (lambda_local / 6.0) ** 2 self._element_budget_area = A_budget N_phys = int(np.ceil(np.sum(self.element_volumes / A_budget))) rho_min = 5.0 self._n_budget = max(N_phys, int(np.ceil(rho_min * self._num_elements))) if self.error_per_element is not None: self._initial_error_norm = np.linalg.norm(self.error_per_element, axis=0) # Record initial total error as normalization baseline for reward calculation self.current_error = self._compute_total_error() self.initial_error = self.current_error # Protection against near-zero initial error (prevents division by zero) if self.initial_error < 1e-8: self.initial_error = 1.0 # reset the element penalty, necessary if it is sampled if self._sample_penalty: sampling_type = self._element_penalty_config.get( "sampling_type", "loguniform" ) min_value = self._element_penalty_config.get("min") max_value = self._element_penalty_config.get("max") element_penalty_lambda = sample_in_range( max_value, min_value, sampling_type ) self._element_penalty_lambda = element_penalty_lambda else: # element penalty is a scalar value self._element_penalty_lambda = self._element_penalty_config.get("value") def _calculate_initial_approximation_errors(self): if self._manual_normalization: return { error_name: self._manual_normalization for error_name in self.error_estimation_dict } else: result = {} for error_name, errors in self.error_estimation_dict.items(): errors = np.atleast_1d(np.asarray(errors, dtype=np.float64)) val = np.sqrt(np.sum(errors ** 2)) result[error_name] = float(val) + 1e-12 return result def step(self, action: np.ndarray) -> Tuple[Data, np.array, bool, Dict[str, Any]]: """ Performs a step of the Mesh Refinement task. Wrapped in try-except to prevent program crashes from ill-conditioned FEM solves caused by degenerate meshes (especially in early training when the continuous sizing field produces extreme element shapes). On FEM failure: returns done=True with an extreme penalty reward (-10000) to implicitly teach the agent to avoid generating invalid meshes. Args: action: the action the agents will take in this step. Has shape (num_agents, action_dimension) Given as an array of shape (num_agents, action_dimension) Returns: A 4-tuple (observations, reward, done, info), where * observations is a graph of the agents and their positions, in this case of the refined mesh * reward is a single scalar shared between all agents, i.e., per **graph** * done is a boolean flag that says whether the current rollout is finished or not * info is a dictionary containing additional information """ assert not self.is_terminal, ( f"Tried to perform a step on a terminated environment. Currently on step " f"{self._timestep:} of {self._max_timesteps:} " f"with {self._num_elements}/{self._maximum_elements} elements." ) self._timestep += 1 # ================================================================ # 核心逻辑: try-except 物理防崩盘机制 # 捕获 FEM 求解器因畸形网格抛出的任何异常 # ================================================================ try: self._set_previous_step() # refine mesh and store which element has become which set of new elements self._agent_mapping = self._refine_mesh(action=action) # solve equation and calculate error per element/element self._previous_error_per_element = self.error_per_element self._error_estimation_dict = ( self.fem_problem.calculate_solution_and_get_error() ) # query returns observation = self.last_observation reward_dict = self._get_reward_dict() metric_dict = self._get_metric_dict() action_dict = self._get_action_dict(action=action) # done after a given number of steps or if the mesh becomes too large done = self.is_terminal info = ( reward_dict | metric_dict | action_dict | { "is_truncated": self.is_truncated, "return": self._cumulative_return, "neg_action_ratio": getattr(self, "_diag_neg_ratio", 0.0), "eligible_ratio": getattr(self, "_diag_eligible_ratio", 0.0), "masked_ratio": getattr(self, "_diag_masked_ratio", 0.0), "selected_count": getattr(self, "_diag_selected_count", 0), "n_budget": self._n_budget, } ) return observation, self._reward, done, info # except (np.linalg.LinAlgError, ValueError, RuntimeError, Exception) as e: except (np.linalg.LinAlgError, ValueError, RuntimeError) as e: # ============================================================ # FEM 物理崩溃捕获 # 可能原因: # 1. 畸形网格导致刚度矩阵奇异 (LinAlgError) # 2. 连续动作产生了退化元素 (ValueError) # 3. scikit-fem 内部网格操作异常 (RuntimeError) # # 策略: 立即终止本回合，给予极端惩罚，迫使智能体学习 # 避免产生会导致 FEM 崩溃的网格。 # ============================================================ import sys if not hasattr(self, "_crash_print_count"): self._crash_print_count = 0 if self._crash_print_count < 5: print( f"[FEM Crash] step={self._timestep}, " f"elements_before={self._previous_num_elements if self._previous_num_elements is not None else '?'}, " f"type={type(e).__name__}: {str(e)[:300]}", file=sys.stderr, flush=True, ) self._crash_print_count += 1 elif self._crash_print_count == 5: print( f"[FEM Crash] ... suppressing further crash prints ...", file=sys.stderr, flush=True, ) self._crash_print_count += 1 crash_penalty = -10000.0 # 使用细化前的元素数，确保 reward 尺寸与 policy 输出的 values 一致 # self._previous_num_elements 已在 _set_previous_step() 中保存 num_agents = ( self._previous_num_elements if self._previous_num_elements is not None else (self.num_agents if self.num_agents > 0 else 1) ) self._reward = np.full(num_agents, crash_penalty, dtype=np.float32) self._cumulative_return = self._cumulative_return + np.sum(self._reward) # 确保 agent_mapping 与 reward/values 维度一致 self._agent_mapping = np.arange(num_agents, dtype=np.int64) # _num_elements is a property, cannot be set directly # 返回当前观测 (如果可用) 或空图 try: observation = self.last_observation except Exception: # 创建一个最小空图作为 fallback observation = Data( x=torch.zeros( (num_agents, self.num_node_features), dtype=torch.float32 ), edge_index=torch.zeros((2, 0), dtype=torch.long), edge_attr=torch.zeros( (0, self.num_edge_features), dtype=torch.float32 ), ) info = { "is_truncated": False, "return": float(np.sum(self._reward)), "weighted_remaining_error": float("inf"), "num_elements": self._num_elements if self.fem_problem is not None else 0, "num_agents": num_agents, "fem_crash": True, "crash_reason": str(e)[:200], # 截断以防日志过长 } self._timestep = self._max_timesteps # 强制终止 return observation, self._reward, True, info def inference_step( self, action: np.ndarray ) -> Tuple[Data, float, bool, Dict[str, Any]]: """ Performs a step of the Mesh Refinement task *without* calculating the reward or difference to the fine-grained reference. This is used for inference Args: action: the action the agents will take in this step. Has shape (num_agents, action_dimension) Given as an array of shape (num_agents, action_dimension) Returns: A 4-tuple (observations, reward, done, info), where * observations is a graph of the agents and their positions, in this case of the refined mesh * reward is a single scalar shared between all agents, i.e., per **graph** * done is a boolean flag that says whether the current rollout is finished or not * info is a dictionary containing additional information """ assert not self.is_terminal, ( f"Tried to perform a step on a terminated environment. Currently on step " f"{self._timestep:} of {self._max_timesteps:} " f"with {self._num_elements}/{self._maximum_elements} elements." ) self._timestep += 1 self._agent_mapping = self._refine_mesh(action=action) # solve equation self.fem_problem.calculate_solution() observation = self.last_observation done = self.is_terminal info = {} return observation, self._reward, done, info def _set_previous_step(self): """ Sets variables for the previous timestep. These are used for the reward function, as well as for different kinds of plots and metrics """ self._previous_num_elements = self._num_elements self._previous_agent_mapping = self._agent_mapping self._previous_element_volumes = self.element_volumes self._previous_eta_indicator = self._eta_indicator self._previous_eta_components = self._eta_components_raw self._previous_rho_components = self._rho_components self._previous_solution_l2_norm = self._compute_solution_l2_norm() def _compute_solution_l2_norm(self) -> float: """Approximate ||u_h||_{L2(Ω)} via element centroids: sqrt(Σ_K |ū_K|² · area_K).""" u_scat = self.fem_problem.nodal_solution # complex (n_vertices,) elem_idx = self._element_indices # (n_elements, 3) vols = self.element_volumes # (n_elements,) u_elem = u_scat[elem_idx] # (n_elements, 3) u_elem_mean = np.mean(u_elem, axis=1) # (n_elements,) complex mean u_mag = np.abs(u_elem_mean) return float(np.sqrt(np.sum(u_mag ** 2 * vols))) def _refine_mesh(self, action: np.array) -> np.array: """ Refines fem_problem.mesh by splitting all faces/elements for which the average of agent activation surpasses a threshold. If this refinement exceeds the maximum number of nodes allowed in the environment, we return a boolean flag that indicates so and stops the environment Optionally smoothens the newly created mesh as a post-processing step Args: action: An action/activation per element. - continuous_sizing_field: shape (num_agents, 1) or (num_agents,) → 目标网格面积 - absolute/absolute_discrete: shape (num_agents,) or (num_agents, 1) → scalar threshold Returns: An array of mapped element indices """ # 标量动作统一 flatten 到 1D action = action.flatten() elements_to_refine = self._get_elements_to_refine(action) # updates self.fem_problem.mesh element_mapping = self.fem_problem.refine_mesh(elements_to_refine) return element_mapping def _get_elements_to_refine(self, action: np.array) -> np.array: """ Calculate which elements to refine based on the action, refinement strategy and the maximum number of elements allowed in the environment Args: action: An action/activation per agent, i.e., per element. 1D array of shape (num_agents,). - continuous_sizing_field: 每个 agent 输出 1 个标量 → Softplus → 期望最大单元面积 - absolute/absolute_discrete: scalar threshold Returns: An array of ids corresponding to elements_to_refine """ # select elements to refine based on the average actions of its surrounding agents/nodes if self._refinement_strategy == "continuous_sizing_field": # ================================================================ # Score-based 细化选择（由 actor 直接排序，物理预算约束） # # Actor 输出标量 x_i: x_i < 0 → 希望细化; x_i > 0 → 不希望细化 # 排序依据 score = -x_i，在预算和上限内选 top-k # # 物理预算 N_budget: Σ area_K / A_budget，其中 # A_budget = ½(λ_local/6)²，对应每局部波长方向 ~6 个尺度点 # # 动作掩码 (Dörfler-P95): η_K < 0.05·η_P95 的单元移出候选池， # P95 锚定物理误差尺度，免疫远场噪声稀释，强制预算投入误差主导区 # ================================================================ x = action.flatten() # ── 训练监控指标（在所有 early return 之前计算）── self._diag_neg_ratio = float(np.mean(x < 0.0)) remaining = self._n_budget - self._num_elements max_parents_by_budget = max(0, remaining // 3) if max_parents_by_budget <= 0: self._diag_eligible_ratio = 0.0 self._diag_selected_count = 0 return np.array([], dtype=np.int64) # 动态计算每单元预算面积（仅用于 N_budget 全局资源上限） eps_r_elem = self._epsilon_r_elements k = self._wave_number lambda_local = 2.0 * np.pi / (k * np.sqrt(np.maximum(eps_r_elem, 1.0))) A_budget = 0.5 * (lambda_local / 6.0) ** 2 # 纯数值安全底线：仅防止 scikit-fem 因浮点精度导致的退化/奇异。 # 不再用 0.25*A_budget —— RL 应自主学会"多细才够"， # 而非被人为启发式 (12 点/波长) 限制。 domain_area = float(np.prod(self.fem_problem.plot_boundary[2:] - self.fem_problem.plot_boundary[:2])) V_min_safeguard = 1e-10 * domain_area # Filter 1: numerical safeguard only — no physics heuristic area_eligible = np.where(self.element_volumes > V_min_safeguard)[0] # Filter 2: Dörfler-style action mask — exclude elements below 5% of η_P95 # P95 anchors the threshold to physically meaningful error scale, # immune to far-field noise dilution (unlike median or mean). # η_K < 0.05·η_P95 → not worth the refinement budget. eta_current = self._eta_indicator eta_p95 = np.percentile(eta_current, 95) error_eligible = np.where(eta_current >= 0.05 * eta_p95)[0] eligible = np.intersect1d(area_eligible, error_eligible) self._diag_eligible_ratio = float(len(eligible)) / max(self._num_elements, 1) self._diag_masked_ratio = ( 1.0 - float(len(eligible)) / max(len(area_eligible), 1) if len(area_eligible) > 0 else 0.0 ) num = min( len(eligible), max(1, self._num_elements // 4), max_parents_by_budget, ) if num <= 0: self._diag_selected_count = 0 return np.array([], dtype=np.int64) # x 越小 ⇒ 优先级越高（纯排序，不设正负门槛） score = -x selected = eligible[np.argsort(score[eligible])[-num:]] self._diag_selected_count = len(selected) elements_to_refine = selected elif self._refinement_strategy in ["absolute", "absolute_discrete"]: elements_to_refine = np.argwhere(action > 0.0).flatten() else: raise ValueError( f"Unknown refinement strategy '{self._refinement_strategy}" ) return elements_to_refine def render(self, mode: str = "human", render_intermediate_steps: bool = False, *args, **kwargs): if not (render_intermediate_steps or self.is_terminal): return [], {} remaining_error = self._get_remaining_error(return_dimensions=False) title = ( f"Solution. Element Penalty: {self._element_penalty_lambda:.1e} " f"Reward: {np.sum(self._reward):.3f} Return: {np.sum(self._cumulative_return):.3f} " f"Agents: {self.num_agents} Remaining Error: {remaining_error:.3f}" ) traces, layout = get_plotly_mesh_traces_and_layout( mesh=self.mesh, scalars=np.real(self.scalar_solution), mesh_dimension=2, title=title, boundary=self.fem_problem.plot_boundary, ) _fp = self.fem_problem.fem_problem cx = getattr(_fp, "_cx", 0.5) cy = getattr(_fp, "_cy", 0.5) r = getattr(_fp, "_radius", 0.2) traces.append(go.Scatter3d( x=cx + r * np.cos(np.linspace(0, 2 * np.pi, 128)), y=cy + r * np.sin(np.linspace(0, 2 * np.pi, 128)), z=np.zeros(128), mode="lines", line=dict(color="cyan", width=2, dash="dash"), name="Dielectric", showlegend=True, )) return traces, {"layout": layout} def _get_remaining_error( self, return_dimensions: bool = False ) -> Union[np.array, Tuple]: """ Get the remaining error by aggregating over all elements and taking the convex sum of all solution dimensions """ remaining_error_per_dimension = np.sqrt( np.sum(self.error_per_element**2, axis=0) ) # Collapse per-element/per-dim initial error to scalar if needed norm = np.atleast_1d(np.asarray(self.initial_approximation_error, dtype=np.float64)) if remaining_error_per_dimension.ndim < norm.ndim or ( remaining_error_per_dimension.ndim == norm.ndim and remaining_error_per_dimension.shape != norm.shape ): norm = np.sqrt(np.sum(norm**2)) remaining_error_per_dimension = ( remaining_error_per_dimension / norm ) # normalize remaining_error = self.project_to_scalar(remaining_error_per_dimension) # Ensure scalar output (defensive against (1,) or (Ne,) arrays from 1D PDEs) remaining_error = float(np.asarray(remaining_error).ravel()[0]) if return_dimensions: return remaining_error, remaining_error_per_dimension else: return remaining_error def _compute_total_error(self, error_per_element: np.ndarray = None) -> float: """ Compute a scalar total error from a per-element error array. Uses the same aggregation (sum or max) as _get_remaining_error, but without normalization by the initial approximation error. Args: error_per_element: Per-element error array of shape (num_elements, solution_dimension). If None, uses the current error_per_element. Returns: A scalar total error. """ if error_per_element is None: error_per_element = self.error_per_element error_per_dim = np.sqrt(np.sum(error_per_element**2, axis=0)) return float(self.project_to_scalar(error_per_dim)) @property def last_observation(self) -> Data: """ Retrieve an observation graph for the current state of the environment. We use an additional self.last_observation wrapper to make sure that classes that inherit from this one have access to node and edge features outside the Data() structure Returns: A Data() object of the graph that describes the current state of this environment """ graph_dict = {} graph_dict = graph_dict | self._get_graph_nodes() graph_dict = graph_dict | self._get_graph_edges() observation_graph = Data(**graph_dict) return observation_graph def _get_graph_nodes(self) -> Dict[str, Dict[str, torch.Tensor]]: """ Returns a dictionary of node features that are used to describe the current state of this environment. Returns: A dictionary of node features. This dictionary has the format {"x": element_features} where element and node features depend on the context, but include things like the evaluation of the target function, the degree of the node, etc. """ # Builds feature matrix of shape (#elements, #features) # by iterating over the functions in self._element_feature_functions. element_features = np.array( [fn() for key, fn in self._element_feature_functions.items()] ).T element_features = save_concatenate( [element_features, self.fem_problem.element_features()], axis=1 ) element_features = torch.tensor(element_features, dtype=torch.float32) node_dict = {"x": element_features} return node_dict def _get_graph_edges( self, ) -> Dict[Union[str, Tuple[str, str, str]], Dict[str, torch.Tensor]]: """ Returns a dictionary of edge features that are used to describe the current state of this environment. Note that we always use symmetric graphs and self edges. Returns: A dictionary of edge features. This dictionary has the format { "edge_index": indices, "edge_attr": features } """ edge_index, edge_attr = self._element2element_features(self._num_edge_features) edge_dict = {"edge_index": edge_index, "edge_attr": edge_attr} return edge_dict def _element2element_features(self, num_edge_features: int): # concatenate incoming, outgoing and self edges of each node to get an undirected graph src_nodes = np.concatenate( ( self._element_neighbors[0], self._element_neighbors[1], np.arange(self._num_elements), ), axis=0, ) dest_nodes = np.concatenate( ( self._element_neighbors[1], self._element_neighbors[0], np.arange(self._num_elements), ), axis=0, ) num_edges = self._element_neighbors.shape[1] * 2 + self._num_elements edge_features = np.empty(shape=(num_edges, num_edge_features)) edge_feature_position = 0 if "euclidean_distance" in self._edge_features: euclidean_distances = np.linalg.norm( self._element_midpoints[dest_nodes] - self._element_midpoints[src_nodes], axis=1, ) lam = 2.0 * np.pi / self._wave_number edge_features[:, edge_feature_position] = euclidean_distances / lam edge_feature_position += 1 edge_index = torch.tensor(np.vstack((src_nodes, dest_nodes))).long() edge_attr = torch.tensor(edge_features, dtype=torch.float32) return edge_index, edge_attr def _get_reward_dict(self) -> Dict[str, np.float32]: """ Calculate the reward for the current timestep depending on the environment states and the action the agents took. Args: Returns: Dictionary that must contain "reward" as well as partial reward data """ reward, reward_dict = self._get_reward_by_type() self._reward = reward self._cumulative_return = self._cumulative_return + np.sum(self._reward) return reward_dict def _get_metric_dict(self) -> Dict[str, Any]: remaining_error, remaining_error_per_dimension = self._get_remaining_error( return_dimensions=True ) metric_dict = { "weighted_remaining_error": remaining_error, "error_times_agents": remaining_error * self.num_agents, "delta_elements": self._num_elements - self._previous_num_elements, "avg_total_refinements": np.log( self._num_elements / self._initial_num_elements ) / np.log(4), "avg_step_refinements": np.log( self._num_elements / self._previous_num_elements ) / np.log(4), "num_elements": self._num_elements, "num_agents": self.num_agents, "reached_element_limits": self.reached_element_limits, "refinement_std": self._refinements_per_element.std(), } for error_metric, element_errors in self.error_estimation_dict.items(): if element_errors.ndim >= 1 and element_errors.shape[0] == self._num_elements: error_per_dimension = np.max(element_errors, axis=0) else: error_per_dimension = element_errors error_per_dimension = ( error_per_dimension / self._initial_approximation_errors[error_metric] ) remaining_error = self.project_to_scalar(error_per_dimension) metric_dict[f"{error_metric}_error"] = remaining_error return metric_dict def _get_action_dict(self, action: np.ndarray) -> Dict[str, Any]: """ Returns a dictionary of information about the action that was taken in the current timestep Args: action: The action that was taken in the current timestep Returns: A dictionary of information about the action that was taken in the current timestep """ action_dict = {} if self._refinement_strategy in ["absolute", "absolute_discrete"]: action_dict["action_mean"] = np.mean(action) action_dict["action_std"] = np.std(action) return action_dict def _get_reward_by_type(self) -> Tuple[np.array, Dict]: """ Potential-based reward shaping on η indicator. spatial_max — Per-agent reward (parent i → children C(i)): r_local_i = log(η_old_i + ε_dynamic) − log(max_{j∈C(i)} η_new_j + ε_dynamic) − λ·(|C(i)| − 1) spatial — Per-agent reward (parent i → children C(i)): r_local_i = log(η_old_i + ε_dynamic) − log(√(Σ_{j∈C(i)} η_new_j²) + ε_dynamic) − λ·(|C(i)| − 1) ε_dynamic = max(0.01 × η_P95, 1e-6) — anchored to P95 of residual, immune to far-field dilution; prevents reward hacking on near-zero-η elements. Potential function: Φ(s) = −log(E_global) E_global = √(Σ_K η_K²) / ||u_h||_{L2(Ω)} (dimensionless) Shaped reward: r_i = r_local_i + α · (log E_old − log E_new) """ reward_dict = {} # Dynamic epsilon anchored to P95 of η — immune to far-field dilution # that plagues mean-based approaches. P95 is driven by physically # meaningful error in the dielectric, not background noise. # ε_dynamic = max(0.01 × η_P95, 1e-6) eta_current_raw = self._eta_indicator eta_p95 = float(np.percentile(eta_current_raw, 95)) eps = max(0.01 * eta_p95, 1e-6) old_eta = self._previous_eta_indicator + eps new_eta = eta_current_raw + eps if self._reward_type == "spatial_max": from torch_scatter import scatter_max agent_mapping = torch.tensor(self.agent_mapping) child_eta = torch.tensor(new_eta) max_child_eta, _ = scatter_max( src=child_eta, index=agent_mapping, dim=0, dim_size=old_eta.shape[0], ) max_child_eta = max_child_eta.numpy() reward_per_agent_and_dim = np.log(old_eta) - np.log(max_child_eta) elif self._reward_type == "spatial": from torch_scatter import scatter_add agent_mapping = torch.tensor(self.agent_mapping) # L₂ aggregation: √(Σ η_child²) — never punishes refinement child_eta = torch.tensor(new_eta) sum_sq_child_eta = scatter_add( src=child_eta * child_eta, index=agent_mapping, dim=0, dim_size=old_eta.shape[0], ) l2_child_eta = np.sqrt(sum_sq_child_eta.numpy()) + eps reward_per_agent_and_dim = np.log(old_eta) - np.log(l2_child_eta) else: raise ValueError(f"Unknown reward type {self._reward_type}") reward_per_agent = self.project_to_scalar(reward_per_agent_and_dim) # apply action/element penalty unique_old, counts = np.unique(self.agent_mapping, return_counts=True) element_penalty = np.zeros(len(reward_per_agent), dtype=reward_per_agent.dtype) element_penalty[unique_old] = self._element_penalty_lambda * (counts - 1) element_limit_penalty = ( (self._element_limit_penalty / self._previous_num_elements) if self.reached_element_limits else 0 ) reward_per_agent = ( reward_per_agent - element_penalty - element_limit_penalty ) # ── Potential-based shaping: only refined parents get the global bonus ── l2_old = self._previous_solution_l2_norm l2_new = self._compute_solution_l2_norm() eta_l2_old = float(np.sqrt(np.sum(old_eta ** 2))) eta_l2_new = float(np.sqrt(np.sum(new_eta ** 2))) eps_l2 = 1e-12 E_old = eta_l2_old / max(l2_old, eps_l2) E_new = eta_l2_new / max(l2_new, eps_l2) global_bonus = self._global_reward_alpha * float(np.log(E_old + eps_l2) - np.log(E_new + eps_l2)) refined_parents = unique_old[counts > 1] reward_per_agent[refined_parents] += global_bonus self._reward_per_agent = reward_per_agent self._cumulative_reward_per_agent = ( self._cumulative_reward_per_agent[self._previous_agent_mapping] + reward_per_agent ) reward = reward_per_agent reward_dict["reward"] = reward reward_dict["potential_bonus"] = global_bonus reward_dict["penalty"] = -reward reward_dict["element_limit_penalty"] = element_limit_penalty reward_dict["element_penalty"] = element_penalty reward_dict["element_penalty_lambda"] = self._element_penalty_lambda return reward, reward_dict @property def mesh(self) -> Mesh: """ Returns the current mesh. """ return self.fem_problem.mesh @property def agent_node_type(self) -> str: return "element" @property def _vertex_positions(self) -> np.array: """ Returns the positions of all vertices/nodes of the mesh. Returns: np.array of shape (num_vertices, 2) """ return self.fem_problem.vertex_positions @property def _element_indices(self) -> np.array: return self.fem_problem.element_indices @property def _element_midpoints(self) -> np.array: """ Returns the midpoints of all elements/faces. Returns: np.array of shape (num_elements, 2) """ return self.fem_problem.element_midpoints @property def _mesh_edges(self) -> np.array: """ Returns: the edges of all vertices/nodes of the mesh. Shape (2, num_edges) """ return self.fem_problem.mesh_edges @property def _element_neighbors(self) -> np.array: """ Find neighbors of each element. Shape (2, num_neighbors) Returns: """ # f2t are element/face neighborhoods, which are set to -1 for boundaries return self.fem_problem.element_neighbors @property def _num_elements(self) -> int: return len(self._element_indices) @property def _num_vertices(self) -> int: return len(self._vertex_positions) @property def element_volumes(self) -> np.array: return get_triangle_areas_from_indices( positions=self._vertex_positions, triangle_indices=self._element_indices ) @property def num_node_features(self) -> int: return self._num_node_features @property def num_edge_features(self) -> int: return self._num_edge_features @property def action_dimension(self) -> int: """ Returns: The dimensionality of the action space. - continuous_sizing_field: 1D continuous output → 目标网格面积 (Softplus 激活) - absolute_discrete: 2 discrete actions (refine / don't refine) - others: single continuous scalar """ if self._refinement_strategy == "continuous_sizing_field": return 1 # 1D 连续标量 → Softplus → 目标面积 S_i elif self._refinement_strategy == "absolute_discrete": return 2 else: # single continuous value return 1 @property def num_agents(self) -> int: if self.fem_problem is not None and self.fem_problem.mesh is not None: return self._num_elements else: return 1 # placeholder @property def _action_space(self) -> gym.Space: """ Returns: The **current** action space of the environment. Bound to change, since the number of agents changes """ if self._refinement_strategy in ["absolute_discrete", "argmax", "single_agent"]: return gym.spaces.MultiDiscrete([self.action_dimension] * self.num_agents) elif self._refinement_strategy == "continuous_sizing_field": # 连续 1D 输出: 每个 agent 输出 1 个标量 → Softplus → 目标网格面积 # 无界连续空间，PPO Gaussian policy 负责探索 return gym.spaces.Box( low=-1e5, high=1e5, shape=(self.num_agents, self.action_dimension), dtype=np.float32, ) else: return gym.spaces.Box( low=-1e5, high=1e5, shape=( self.num_agents, self.action_dimension, ), dtype=np.float32, ) @property def agent_mapping(self) -> np.array: assert self._agent_mapping is not None, "Element mapping not initialized" return self._agent_mapping @property def previous_agent_mapping(self) -> np.array: assert self._previous_agent_mapping is not None, ( "Previous element mapping not initialized" ) return self._previous_agent_mapping @property def reached_element_limits(self) -> bool: """ True if the number of elements/faces in the mesh is above the maximum allowed value. Returns: """ return self._num_elements > self._maximum_elements @property def is_truncated(self) -> bool: return self._timestep >= self._max_timesteps @property def is_terminal(self) -> bool: # Agent selected nothing to refine — budget exhausted or # Doerfler mask filtered everything. Episode converged naturally. # -1 = not yet evaluated (reset state), 0 = nothing selected this step. sc = getattr(self, "_diag_selected_count", -1) if sc == 0: return True return self.reached_element_limits or self.is_truncated @property def solution(self) -> np.array: """ Returns: solution vector per *vertex* of the mesh. An array (num_vertices, solution_dimension), where every entry corresponds to the solution of the parameterized fem_problem equation at the position of the respective node/vertex. """ return self.fem_problem.nodal_solution def project_to_scalar(self, values: np.array) -> np.array: """ Projects a value per node or graph and solution dimension to a scalar value per node. Args: values: A vector of shape ([num_vertices/nodes,] solution_dimension) Returns: A scalar value per vertex """ return self.fem_problem.project_to_scalar(values) @property def scalar_solution(self): return self.project_to_scalar(self.solution) @property def error_per_element(self) -> np.array: """ Returns: error per element of the mesh. np.array of shape (num_elements, solution_dimension) """ return self._error_estimation_dict.get("indicator") @property def initial_approximation_error(self) -> np.array: """ Returns: error per element of the mesh. np.array of shape (num_elements, solution_dimension) """ return self._initial_approximation_errors.get("indicator") @property def error_estimation_dict(self) -> Dict[str, np.array]: """ Returns a dictionary of all error estimation methods and their respective errors. These errors may be per element/face, or per integration point, depending on the metric. Returns: """ return self._error_estimation_dict @property def _refinements_per_element(self) -> np.array: return self.fem_problem.refinements_per_element @property def _solution_std_per_element(self) -> np.array: """ Computes the standard deviation of the solution per element. Returns: np.array of shape (num_elements, solution_dimension) Note: 此属性仅用于 backward compatibility； 新代码使用 _residual_components 替代。 """ return get_aggregation_per_element( self.solution, self._element_indices, aggregation_function_str="std" ) # ========================================================================= # PDE 物理残差特征 (替代 solution_std) # ========================================================================= @property def _residual_components(self) -> Dict[str, np.ndarray]: """逐单元的三项 PDE 残差 + 边界标记。""" from .helmholtz import _compute_residual_components fp = self.fem_problem.fem_problem k = getattr(fp, "_k", 10.0) u_scat = self.fem_problem.nodal_solution eps_r = self._epsilon_r_elements return _compute_residual_components( self.fem_problem.mesh, u_scat, k=k, eps_r=eps_r ) @property def _k_local_sqrt_vol(self) -> np.ndarray: """每个单元的 k_local × sqrt(volume)。""" k = self._wave_number eps_r = self._epsilon_r_elements k_local = k * np.sqrt(np.maximum(eps_r, 1.0)) return (k_local * np.sqrt(self.element_volumes)).astype(np.float32) @property def _volume_normalized(self) -> np.ndarray: """无量纲单元面积: volume / lambda^2。""" lam = 2.0 * np.pi / self._wave_number return (self.element_volumes / (lam * lam)).astype(np.float32) @property def _eta_indicator(self) -> np.ndarray: """ 标准 FEM 残差误差指示器，用于 reward 计算。 η_i = √(R_int_i² + J_grad_i² + R_sbc_i²) 其中: R_int_i = h_K · √V_i · |k²ε_r u + k²(ε_r-1)u_inc| J_grad_i = √(½ Σ_{e∈∂K_i} h_e² · |[[∇u·n]]|²) R_sbc_i = √h_bnd · |∂u/∂n - i·k_local·u| 与 _compute_residual_indicator 的公式完全一致。 Returns: shape (num_elements,) float64 """ from .helmholtz import _compute_residual_components fp = self.fem_problem.fem_problem k = getattr(fp, "_k", 10.0) u_scat = self.fem_problem.nodal_solution comps = _compute_residual_components( self.fem_problem.mesh, u_scat, k=k, eps_r=self._epsilon_r_elements, apply_log=False, ) self._cached_eta_components_raw = comps return np.sqrt( comps["internal_residual"] ** 2 + comps["gradient_jump"] ** 2 + comps["sbc_residual"] ** 2 ) @property def _eta_components_raw(self) -> Dict[str, np.ndarray]: """返回逐单元的三项原始残差分量（apply_log=False），由 _eta_indicator 缓存。""" if not hasattr(self, "_cached_eta_components_raw") or self._cached_eta_components_raw is None: _ = self._eta_indicator # triggers caching return self._cached_eta_components_raw @property def _rho_components(self) -> Dict[str, np.ndarray]: """返回逐单元的残差密度三分量（不含 h-缩放），用于 reward 计算。 Returns: rho_int: |k²·ε_r·u + k²·(ε_r-1)·u_inc| rho_jump: √(mean |[[∇u·n]]|²) per element rho_sbc: √(mean |∂u/∂n - i·k·u|²) per element """ from .helmholtz import _compute_residual_density fp = self.fem_problem.fem_problem k = getattr(fp, "_k", 10.0) u_scat = self.fem_problem.nodal_solution return _compute_residual_density( self.fem_problem.mesh, u_scat, k=k, eps_r=self._epsilon_r_elements, ) # ========================================================================= # SBC 状态空间辅助属性：介电常数 + 复数场均值 # ========================================================================= @property def _wave_number(self) -> float: """Helmholtz 波数 k，从当前 FEM 问题实例读取（支持随机采样）。""" fp = self.fem_problem.fem_problem return getattr(fp, '_k', 10.0) @property def _epsilon_r_elements(self) -> np.ndarray: """ 每个单元的相对介电常数 εr。 从 FEM 问题实例读取介质几何参数，按单元中点判断是否在介质内。 Returns: shape (num_elements,) float64 array """ fp = self.fem_problem.fem_problem cx = getattr(fp, "_cx", 0.5) cy = getattr(fp, "_cy", 0.5) radius = getattr(fp, "_radius", 0.2) eps_r = getattr(fp, "_eps_r", 2.0) midpoints = self._element_midpoints 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, 1.0) @property def _dist_to_interface(self) -> np.ndarray: """每个单元中点到介质圆柱边界的带符号距离（内部为负，外部为正）。 用真空波长 lambda = 2*pi/k 做无量纲归一化，再经 sign(d)·ln(1+|d|) 压缩。 ln 压缩保留近场分辨力（小 |d| 时近似线性），远场自然对数压缩， 与残差特征的 log₁₀ 压缩风格一致。无硬截断，处处可导。 """ fp = self.fem_problem.fem_problem cx = getattr(fp, "_cx", 0.5) cy = getattr(fp, "_cy", 0.5) radius = getattr(fp, "_radius", 0.2) midpoints = self._element_midpoints dist = np.sqrt((midpoints[:, 0] - cx) ** 2 + (midpoints[:, 1] - cy) ** 2) lam = 2.0 * np.pi / self._wave_number d = (dist - radius) / lam return (np.sign(d) * np.log1p(np.abs(d))).astype(np.float32) @property def _eps_r_global(self) -> float: """散射体材料的相对介电常数（全局标量）。""" fp = self.fem_problem.fem_problem return getattr(fp, "_eps_r", 2.0) @property def _complex_solution_mean(self) -> np.ndarray: """ 每个单元内复数 FEM 解的均值 (complex128)。 SBC 边界条件下解为复数值；内部残差和边界残差均基于复数场。 使用 P1 节点值的三点平均作为单元代表值。 Returns: shape (num_elements,) complex128 array """ return get_aggregation_per_element( self.solution, self._element_indices, aggregation_function_str="mean", ) @property def sample_penalty(self) -> bool: return self._sample_penalty @property def refinement_strategy(self) -> str: return self._refinement_strategy @property def has_homogeneous_graph(self) -> bool: return not self._include_vertices @property def mesh_dimension(self) -> int: return 2 def set_element_penalty_lambda( self, position_or_value: float, from_position: bool = True ): """ Sets the element penalty lambda from the provided position. Args: position_or_value: A float between 0 and 1 that determines the element penalty lambda if from_position. Otherwise, the value of the element penalty lambda. from_position: If True, the element penalty lambda is taken log-uniformly from the provided position, regardless of how the value is usually sampled. Returns: None Note: Sets self._element_penalty_lambda """ element_penalty_config = self._environment_config.get("element_penalty") if element_penalty_config.get("sample_penalty"): if from_position: # sample element penalty loguniformly for comparison between methods log_min = np.log(element_penalty_config.get("min")) log_max = np.log(element_penalty_config.get("max")) self._element_penalty_lambda = np.exp( position_or_value * log_min + (1 - position_or_value) * log_max ) else: # fixed element penalty self._element_penalty_lambda = position_or_value else: # element penalty is fixed self._element_penalty_lambda = element_penalty_config.get("value") #################### # additional plots # #################### def _plot_value_per_element( self, value_per_element: np.array, title: str, normalize_by_element_volume: bool = False, mesh: Optional[Mesh] = None, ) -> go.Figure: """ only return traces if asked or at the last step to avoid overlay of multiple steps Args: value_per_element: A numpy array of shape (num_elements,). title: The title of the plot. normalize_by_element_volume: If True, the values are normalized by the element volume as value /= element_volume. mesh: The mesh to plot the values on. If None, the mesh of the current state is used. Returns: A plotly figure with an outline of the mesh and value per element in the element midpoints. """ if mesh is None: assert len(value_per_element) == self.num_agents, ( f"Need to provide a value per agent, given " f"'{value_per_element.shape}' and '{self.num_agents}'" ) mesh = self.fem_problem.mesh mesh_dimension = 2 else: mesh_dimension = mesh.dim() if normalize_by_element_volume: value_per_element = value_per_element / self.element_volumes boundary = self.fem_problem.plot_boundary traces, layout = get_plotly_mesh_traces_and_layout( mesh=mesh, scalars=value_per_element, mesh_dimension=mesh_dimension, title=title, boundary=boundary, ) value_per_element_plot = go.Figure(data=traces, layout=layout) return value_per_element_plot def _plot_error_per_element( self, normalize_by_element_volume: bool = True ) -> go.Figure: weighted_remaining_error = self._get_remaining_error(return_dimensions=False) return self._plot_value_per_element( value_per_element=self.project_to_scalar(self.error_per_element), normalize_by_element_volume=normalize_by_element_volume, title=f"Element Errors. Remaining total error: {weighted_remaining_error:.4f}", ) def additional_plots( self, iteration: int, policy_step_function: Optional[callable] = None ) -> Dict[str, go.Figure]: """ Function that takes an algorithm iteration as input and returns a number of additional plots about the current environment as output. Some plots may be always selected, some only on e.g., iteration 0. Args: iteration: The current iteration of the algorithm. policy_step_function: (Optional) A function that takes a graph as input and returns the action(s) and (q)-value(s) for each agent. """ _, remaining_error_per_solution_dimension = self._get_remaining_error( return_dimensions=True ) additional_plots = { "refinements_per_element": self._plot_value_per_element( value_per_element=self._refinements_per_element, title="Refinements per element", ), "scalar_solution_std_per_element": self._plot_value_per_element( value_per_element=self.project_to_scalar( self._solution_std_per_element ), title=f"Element Std of Solution Norm", ), "scalar_solution_error_per_element": self._plot_error_per_element( normalize_by_element_volume=False ), } if policy_step_function is not None: from .utils import detach actions, values = policy_step_function(observation=self.last_observation) if len(actions) == self._num_elements: additional_plots["final_actor_evaluation"] = ( self._plot_value_per_element( detach(actions), title=f"Action per Agent at step {self._timestep}", ) ) if len(values) == self._num_elements: additional_plots["final_critic_evaluation"] = ( self._plot_value_per_element( detach(values), title=f"Critic Evaluation at step {self._timestep}", ) ) if self._reward_type in ["spatial", "spatial_max", "spatial_volume"]: additional_plots["cumulative_reward"] = self._plot_value_per_element( value_per_element=self._cumulative_reward_per_agent, title="Cumulative Reward", mesh=self.fem_problem.previous_mesh, ) additional_plots["reward_per_agent"] = self._plot_value_per_element( value_per_element=self._reward_per_agent, title="Final Reward", mesh=self.fem_problem.previous_mesh, ) additional_plots |= self.fem_problem.additional_plots() return additional_plots def __deepcopy__(self, memo): """ Overwrite deepcopy to reinitialize stateless (lambda-) functions it is sufficient to call the register functions, as only new objects for the stateless lambda functions have to be created Args: memo: Returns: """ from copy import deepcopy cls = self.__class__ result = cls.__new__(cls) memo[id(self)] = result for k, v in self.__dict__.items(): setattr(result, k, deepcopy(v, memo)) setattr( result, "_element_feature_functions", result._register_element_features() ) return result