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