1578 lines
66 KiB
Python
1578 lines
66 KiB
Python
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)
|
||
# amplitude + phase direction (cos/sin ∈ [−1,1]), ε=1e-8 at |u|→0 nodes
|
||
feats["epsilon_r"] = lambda: self._epsilon_r_elements
|
||
feats["total_solution_magnitude"] = lambda: np.abs(self._complex_solution_mean)
|
||
feats["cos_phase"] = lambda: np.real(self._complex_solution_mean) / (np.abs(self._complex_solution_mean) + 1e-8)
|
||
feats["sin_phase"] = lambda: np.imag(self._complex_solution_mean) / (np.abs(self._complex_solution_mean) + 1e-8)
|
||
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
|
||
self._diag_dorfler_tail_ratio = 0.0
|
||
self._diag_dorfler_floor_active = False
|
||
|
||
# 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),
|
||
"dorfler_tail_ratio": getattr(self, "_diag_dorfler_tail_ratio", 0.0),
|
||
"dorfler_floor_active": float(getattr(self, "_diag_dorfler_floor_active", False)),
|
||
"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 个尺度点
|
||
#
|
||
# 动作掩码 (Reverse Dörfler): 按 η_K 升序排列，剔除累积平方误差
|
||
# 贡献 < ε_noise·Ση² 的底部单元（数值噪声/已收敛区），保留 ≥20%
|
||
# 的单元确保 Agent 始终有充分的选择空间
|
||
# ================================================================
|
||
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
|
||
self._diag_dorfler_tail_ratio = 0.0
|
||
self._diag_dorfler_floor_active = False
|
||
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: Reverse Dörfler — eliminate the noise tail, not select the elite.
|
||
# Sort η_K ascending; remove the smallest elements whose cumulative η²
|
||
# contributes < ε_noise of total error energy. These are numerically
|
||
# converged or noise — not worth the agent's attention.
|
||
# A 20% floor on the eligible ratio guarantees the agent meaningful
|
||
# choices even in heavy-tailed distributions where energy is concentrated.
|
||
eta_current = self._eta_indicator
|
||
eta_sq = eta_current ** 2
|
||
total_energy = np.sum(eta_sq)
|
||
|
||
if total_energy > 0:
|
||
idx_asc = np.argsort(eta_current) # ascending
|
||
cumsum_asc = np.cumsum(eta_sq[idx_asc])
|
||
eps_noise = 0.01 # bottom 1% of energy = noise tail
|
||
k_dorfler = int(np.searchsorted(cumsum_asc, eps_noise * total_energy))
|
||
self._diag_dorfler_tail_ratio = float(k_dorfler) / max(self._num_elements, 1)
|
||
# floor: keep at least 20% of elements for RL agent choice
|
||
min_keep = max(1, self._num_elements // 5)
|
||
k = min(k_dorfler, self._num_elements - min_keep)
|
||
self._diag_dorfler_floor_active = k < k_dorfler
|
||
error_eligible = idx_asc[k:]
|
||
else:
|
||
self._diag_dorfler_tail_ratio = 0.0
|
||
self._diag_dorfler_floor_active = False
|
||
error_eligible = np.arange(self._num_elements)
|
||
|
||
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)
|
||
observation_graph.eta = torch.tensor(self._eta_indicator, dtype=torch.float32)
|
||
|
||
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,
|
||
)
|
||
# Phase distance: physical edge length in local wavelengths.
|
||
# k_local = k·√ε_r adapts to the medium — two elements are "farther
|
||
# apart" in phase inside high-ε regions where the wave oscillates
|
||
# faster. This gives the GNN a k-invariant metric for generalisation.
|
||
k_local_src = self._wave_number * np.sqrt(np.maximum(
|
||
self._epsilon_r_elements[src_nodes], 1.0))
|
||
k_local_dst = self._wave_number * np.sqrt(np.maximum(
|
||
self._epsilon_r_elements[dest_nodes], 1.0))
|
||
k_edge = np.sqrt(k_local_src * k_local_dst) # geometric mean
|
||
edge_features[:, edge_feature_position] = euclidean_distances * k_edge / (2.0 * np.pi)
|
||
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)
|
||
|
||
# ── Causal isolation + bounded signals ──
|
||
# r_local: clipped to [−1, +1] — prevents pollution-error inversions
|
||
# (±4.6) from hijacking the Critic's value estimate.
|
||
# r_bonus: 0.5·tanh(η/μ − 1) — linear near μ (preserves Dörfler),
|
||
# saturates at ±0.5 for extreme η, bounded and safe.
|
||
# Unrefined parents: r = 0 (causal isolation).
|
||
unique_old, counts = np.unique(self.agent_mapping, return_counts=True)
|
||
refined_mask = np.zeros(len(reward_per_agent), dtype=bool)
|
||
refined_mask[unique_old[counts > 1]] = True
|
||
|
||
# Clip r_local to prevent outlier-driven value collapse
|
||
reward_per_agent = np.clip(reward_per_agent, -1.0, 1.0)
|
||
|
||
# Bounded state bonus: tanh preserves Dörfler near μ, caps at extreme η
|
||
eta_raw = self._previous_eta_indicator
|
||
mu_eta = float(np.mean(eta_raw))
|
||
reward_per_agent[refined_mask] += 0.5 * np.tanh(
|
||
eta_raw[refined_mask] / (mu_eta + 1e-8) - 1.0
|
||
)
|
||
|
||
# Unrefined: clean zero (causal isolation)
|
||
reward_per_agent[~refined_mask] = 0.0
|
||
|
||
# apply action/element penalty (refined parents only)
|
||
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
|
||
)
|
||
|
||
# ── Global error change (diagnostic only, NOT injected into Actor reward) ──
|
||
# Removing global_bonus from per-element reward eliminates the broken causal
|
||
# chain: Helmholtz pollution error can make E_new > E_old even when the
|
||
# selected elements were the right choice, punishing agents for physics
|
||
# they didn't cause. Actor optimises r_local only; Critic captures global
|
||
# effects through value estimation.
|
||
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))
|
||
# global_bonus intentionally NOT added to reward_per_agent — see above.
|
||
|
||
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
|
||
# Reverse Dörfler 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
|