afem/README.md

397 lines
18 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# AFEM — 自适应网格细化的 GNN + PPO 强化学习
## 项目架构
```
afem/
├── src/ # 应用层
│ ├── config.yaml # 配置文件
│ ├── main.py # 入口:解析命令行 → train / test / viz
│ ├── network.py # GNN + Actor-Critic 完整网络定义
│ ├── ppo.py # RolloutBuffer + PPOTrainer
│ ├── utils.py # 读配置、保存/加载 checkpoint
│ └── visualize.py # viz 模式:加载模型 → 推理 → 存 PNG
├── environment/ # 仿真环境层
│ ├── mesh_refinement.py # ★ 核心:网格细化 RL 环境
│ │ # - GNN 图观测构建(节点 + 边特征)
│ │ # - continuous_sizing_field (score-based + budget) 细化策略
│ │ # - spatial 奖励
│ ├── helmholtz.py # Helmholtz FEM 求解器 + 残差误差估计
│ ├── fem_problem.py # FEM 问题封装 + PDE 循环缓冲区
│ ├── fem_util.py # 三角形面积、中点、随机采样、尺寸场函数
│ ├── domain.py # 计算域meshpy 三角剖分
│ ├── utils.py # 数组拼接、随机索引采样
│ └── visualization.py # plotly 网格渲染RL 环境用)
├── checkpoints/ # 模型保存
├── result/ # 可视化输出
└── README.md
```
---
## 项目简介
### 物理场景
二维 Helmholtz 电磁散射:
```
∇²u_scat + k²·ε_r·u_scat = k²·(1-ε_r)·u_inc
```
- **入射波**: 沿 -x 方向的平面波 `u_inc = exp(i·k·x)`
- **散射体**: 圆形介质柱ε_r 随机采样),位置和半径可配
- **边界条件**: SBC (Sommerfeld) `∂u/∂n = i·k·u`
- **域**: 可配矩形域,初始网格密度自适应 + domain area 线性缩放:`N_init = N_base × (k/k_ref)^k_exponent × domain_area`。k_ref 和 k_exponent 均可通过 helmholtz config 配置(默认 k_exponent=1.5, k_ref=6.0),保证不同域尺寸下每单位面积单元数一致
- 可配 exponent^2 = P1 Helmholtz 理论最优 (污染误差 ∝ k²)^1.5 = 工程折中。建议 N_base 配合 exponent 调整,使 N_init 约为 COMSOL 目标 (λ/10√ε_r) 的 30-50%,为 RL agent 留出充分细化空间
- **介质区前渐近区边缘约束**: 介质内 λ_d = 2π/(k√ε_r) 更短,强制迭代细化至 h ≤ λ_d/N默认 N=1.5helmholtz.pre_asymptotic_N 可配)。约 1.5 点/波长,刚好跨过渐近区门槛,赋予初始网格基本相位解析能力但不过度消耗物理预算,为 RL agent 留出充分的选择性细化空间
- **后验误差**: 残差型 indicatorAinsworth & Oden 风格),含单元内部残差 + 梯度跳变 + SBC 边界残差
### 强化学习建模
| 概念 | 对应实体 |
|------|---------|
| **智能体** | 每个三角形网格单元 |
| **状态** | GNN 节点特征(几何 + PDE 残差 + 复数场分解 + 物理参数,节点 12 维 + 边 1 维) |
| **动作** | 1 维连续标量 x_i → score = -x_i 排序,在物理预算内 top-k 选细化单元x 越小优先级越高) |
| **奖励** | 局部子单元 η 的 log-ratio 改善spatial: sum 聚合 / spatial_max: max 聚合)+ α 衰减全局 η log-ratio shaping |
| **终止** | 达到最大步数或超过最大单元数 |
---
## 网络架构
双 GNN 架构policy / value 各自独立基座):
```
图观测 → MessagePassingBase → MLP → 动作分布 / value 标量
├─ nn.Linear嵌入
├─ MessagePassingStack2 层消息传递inner 残差 + LayerNorm
│ └─ MessagePassingStep × N
│ ├─ EdgeModule: MLP([src | dst | edge_attr])
│ └─ NodeModule: MLP([node | scatter(入边)])
└─ 输出: 节点隐向量
```
| 超参数 | 值 |
|--------|-----|
| latent_dim | 64 |
| 消息传递层数 | 2 |
| 残差连接 | inner |
| 归一化 | inner LayerNorm |
| 边 dropout | 0.1 |
| Actor MLP | 2 层 tanh |
| Critic MLP | 2 层 tanh |
| Optimizer | Adam, lr=3e-4, lr_decay=0.995 |
| **动作分布** | `DiagGaussianDistribution`(连续 Box 动作空间),`log_std` 可学习clamp 在 [-4.0, -1.0] |
| **log_std 策略** | 初始化 -2.0std≈0.135),每步 optimizer.step() 后 clamp 到 [-4.0, -1.0]std ∈ [0.018, 0.368]),熵系数 0.001 |
### 动作分布策略说明
环境定义的是 `_action_space`(下划线前缀),网络初始化时必须用 `environment._action_space` 而非 `environment.action_space`(后者默认为 None会错误回退到 `CategoricalDistribution(1)`,导致 policy gradient 恒为零)。
`continuous_sizing_field`score-based的动作有效范围约 [-3, 3]
- score = -x_ix 越小 ⇒ 优先级越高(纯排序,不设正负门槛)
- `initial_log_std=-2.0`std≈0.135clamp 在 [-4.0, -1.0]std ∈ [0.018, 0.368]
-`entropy_coefficient=0.001` 提供微弱探索压力,避免 log_std 过早收敛到下限
---
## 输入特征
### 节点特征12 维)
| 维度 | 来源 | 名称 | 说明 |
|------|------|------|------|
| 1 | cfg | `volume` | 无量纲单元面积volume / λ² |
| 3 | cfg | `internal_residual` / `gradient_jump` / `sbc_residual` | PDE 残差三分量(无量纲化,经 log₁₀ 压缩):<br>`(h_K/k_local)·√V·|r|` / `√(½Σ h_e·\|jump\|²/k_local)` / `(h_bnd/k_local)·\|SBC\|` |
| 1 | cfg | `element_penalty` | 单元惩罚系数 λ |
| 1 | cfg | `timestep` | 当前 rollout 步数 |
| 1 | cfg | `wave_number` | Helmholtz 波数 k |
| 1 | cfg | `k_local_sqrt_vol` | k × √体积(局域波数 × 特征长度) |
| 1 | cfg | `is_sbc_boundary` | 是否与 SBC 吸收边界相邻 (0/1) |
| 1 | cfg | `dist_to_interface` | 到介质圆柱边界的带符号距离,无量纲化后经 sign·ln(1+|d|) 压缩:`sign(d)·ln(1+|(dist-radius)/λ|)` — 近场近似线性保留分辨力,远场对数压缩避免 OOD与残差 log₁₀ 风格一致 |
| 1 | fix | `epsilon_r` | 单元中点相对介电常数(圆柱内 = εᵣ,外 = 1.0 |
| 1 | fix | `total_solution_magnitude` | 散射场复数解的振幅 |
> - **cfg**: 由 `element_features` 配置控制
> - **fix**: 始终启用Helmholtz 复数场分解,硬编码)
>
> GNN 输入用 `_compute_residual_components`k_local 无量纲化log₁₀ 压缩。Reward 用逐单元 η_K`_eta_indicator`),与 GNN 特征公式一致但不经 log 压缩。
### 边特征1 维)
| 维度 | 名称 | 说明 |
|------|------|------|
| 1 | `euclidean_distance` | 相邻单元中点欧几里得距离 / λ(无量纲边特征) |
---
## 调用逻辑
```
main.py --mode train/test/viz
├─→ utils.load_config() # 读 YAML
├─→ environment.MeshRefinement # 创建 RL 环境
│ └─→ FEMProblemCircularQueue # 管理 N 个随机 PDE 实例
│ └─→ HelmholtzProblem # FEM 求解 + 残差误差
├─→ network.create_model() # 创建 ActorCritic
└─ [train] → ppo.PPOTrainer.fit_iteration() 循环
├─ collect_rollouts() # 256 步 rollout
│ └─ buffer.compute_returns_and_advantage()
│ └─ 单路 GAE # 逐 agent 时序差分scatter_add 处理网格细化),奖励含势函数塑形项
│ └─ Return / value 归一化
└─ train_step() # 多 epoch PPO 更新
├─ policy_loss() # Clipped PPO
├─ value_loss() # Clipped value loss
└─ entropy_loss() # 熵正则
```
### 环境内部调用
```
MeshRefinement.reset()
└─→ FEMProblemWrapper.reset()
└─→ initial_mesh (meshpy → 介质内 前渐近区边缘迭代细化)
MeshRefinement.step(action)
├─→ score = -x 排序 + 物理预算约束 → top-k 细化单元
├─→ FEMProblemWrapper.refine_mesh() # scikit-fem refine
├─→ calculate_solution_and_get_error()
│ ├─→ HelmholtzProblem.calculate_solution() # FEM 求解
│ └─→ _compute_residual_indicator() # 残差误差
├─→ _get_reward_by_type() # spatial 奖励
└─→ last_observation # 构建 Data(x, edge_index, edge_attr)
```
### 训练
```bash
CUDA_VISIBLE_DEVICES=7 python src/main.py --mode train --config src/config.yaml
```
首次迭代需收集 256 步 rollout含 FEM 求解),后续打印:
```
it | loss ev agents reward x<0 elig sel time
```
| 字段 | 含义 | 健康范围 |
|------|------|---------|
| `x<0` | `mean(x_i < 0)`,负值动作比例(纯诊断) | 越负的单元优先级越高 |
| `elig` | 通过双过滤器的候选占比 | 排除数值退化 + 低误差的单元 |
| `mask` | 被 Dörfler-P95 掩码 (η<0.05·η_P95) 滤掉的占比 | 因场景而异非固定比例 |
| `sel` | 实际选中的细化单元数 | 每步最多 N_current // 4 |
| `n_budget` | 全局物理预算 episode 固定 | k=30 ~1800 |
### 测试
```bash
python src/main.py --mode test --checkpoint checkpoints/model_final.pt --k-test 6.0
python src/main.py --mode test --checkpoint checkpoints/model_final.pt \
--k-test 6.0 --center 0.3,0.6 --radius 0.15
```
输出
```
Step 0: reward=--- error=1.0000 elements=174 budget=1885
Step 1: reward=+12.345 error=0.7160 elements=618 x<0=0.45 sel=87
...
```
每步打印 `reward error elements x<0 sel` 0 步额外显示 `N_budget`
### 可视化
```bash
python src/main.py --mode viz --checkpoint checkpoints/model_final.pt --k-test 30.0
```
输出: `result/visualization.png`总览+ `result/visualization_steps/step*.png`逐步对比)。
---
## 后验误差估计
### 残差 indicator 公式(无量纲化)
引入局部波数 $k_{local} = k\sqrt{\max(\varepsilon_r, 1.0)}$消除纯几何尺度 $h$ 带来的特征偏差
使误差指示子反映"相位分辨率残差"而非"网格粗疏程度"。
P1 三角单元 K三项残差分量为
$$r_{\text{int}} = \frac{h_K}{k_{local}} \sqrt{V_K} \cdot \left| k^2\varepsilon_r u + k^2(\varepsilon_r-1)u_{inc} \right|_K \tag{1}$$
$$r_{\text{jump}} = \sqrt{\frac{1}{2}\sum_{e\in\partial K} \frac{h_e}{k_{local}} \cdot \left| [[\nabla u \cdot n]] \right|^2_e} \tag{2}$$
$$r_{\text{sbc}} = \frac{h_{bnd}}{k_{local}} \cdot \left| \frac{\partial u}{\partial n} - ik_{local}u \right| \tag{3}$$
**逐单元误差指示子**
$$\eta_K = \sqrt{r_{\text{int}}^2 + r_{\text{jump}}^2 + r_{\text{sbc}}^2}$$
量纲分析$k_{local} \sim [L]^{-1}$$h_e \sim [L]$$|\text{jump}|^2 \sim [L]^{-2}$
三项均严格无量纲$h_e/k_{local} \cdot |\text{jump}|^2 \sim [L]^2 \cdot [L]^{-2} = 1$。
细化后 $h_e$ 缩小直接降低跳变项 RL agent 提供可感知的正向 reward 信号
`η_K` 的计算`_compute_residual_indicator` GNN 输入特征`_compute_residual_components`公式完全一致特征仅多一层 log₁₀ 压缩关键验证点
- 内部残差P1 ∇²u_h 0仅含反应项 `k²ε_r·u + k²(ε_r-1)·u_inc`除以 `k_local` 后跨介质公平可比
- 梯度跳变`(h_e/k_local)·|jump|²`,½ 分配给相邻左右单元$h_e$ 保留边积分路径细化后自然衰减
- SBC 项在 η_K² 中为 `(h_bnd²/k_local²)·|B|²`分量 `r_sbc = (h_bnd/k_local)·|B|`
### 连续尺寸场策略score-based + 物理预算约束 + 动作掩码)
Actor 输出标量 x_i score = -x_i 直接排序在预算和上限内选 top-k
```
A_budget_i = ½(λ_local_i / 6)² // 每局部波长方向 ~6 尺度点(仅用于 N_budget 计算)
λ_local_i = 2π / (k · √ε_r_i)
N_budget = max(N_phys, ⌈5·N_init⌉) // rho_min=5.0,至少 5× 初始单元数,保证 RL 多步细化空间
N_phys = ⌈ Σ |K_i| / A_budget_i ⌉ // 全局物理预算k=30 真空 ~1800
remaining = N_budget N_current
V_min_safeguard = 1e-10 × domain_area // 纯数值底线(防止 FEM 求解器退化)
eligible: area > V_min_safeguard AND η_K ≥ 0.05·η_P95 // 数值底线 + Dörfler-P95
num = min(|eligible|, N_current//4, remaining//3)
selected = top-k by score = -x_i → 1-to-4 切分
```
- score = -x_ix 越小 优先级越高纯排序不设正负门槛
- 不再使用 `0.25·A_budget` 启发式面积地板RL 应自主学会"细化到多细"而非被人类经验 (12 /波长) 限制仅保留数值底线 V_min_safeguard = 1e-10 × domain_area 防止浮点精度问题
- per-step cap 从固定 200 改为自适应 `N_current // 4`随网格规模缩放但增速更缓避免大网格时单步消耗过多预算rho_min 3.0 提升到 5.0赋予更多预算余量
- **sel=0 提前终止** agent 选中 0 个单元细化预算耗尽或 Dörfler 屏蔽所有候选 episode 自动结束不再浪费 FEM 求解
- **k_exponent 可配**初始网格缩放指数可通过 `helmholtz.k_exponent` 配置默认 1.5),² P1 Helmholtz 理论最优
- **动作掩码 (Dörfler-P95)**η_K < 0.05·η_P95 的单元移出候选池P95 锚定物理误差尺度免疫远场噪声稀释 median/mean 不同确保只有误差达标的区域消耗细化预算
### 奖励计算
---
#### 变量
| 符号 | 含义 |
|------|------|
| `η_K = √(r_int² + r_jump² + r_sbc²)` | 逐单元误差指示子`r_*` 定义见式 (1)(3) |
| `C(i)` | 父单元 i 1-to-4 切分产生的子单元集合 |
| `M_new[j]` | 子单元 j 对应的父单元索引 |
| `n_i = |C(i)|` | 父单元 i 的子单元数1 表示未切分 |
| `E_global = √(Σ η_K²) / \|\|u_h\|\|_{L₂(Ω)}` | 全局无量纲误差 |
---
#### 算法
**Step 0 — 保存旧状态** (`_set_previous_step`)
```
η_old ← 旧逐单元 η_K
||u_h_old|| ← 旧解 L₂ 范数 (≈ √(Σ |ū_K|² · area_K))
```
**Step 1 — 网格细化** (`_refine_mesh`)
```
x = action.flatten()
score = -x // x 越小 ⇒ 优先级越高
remaining = N_budget N_old
max_by_budget = max(0, remaining // 3)
// 数值底线 + Dörfler-P95 掩码
V_min_safeguard = 1e-10 × domain_area // 纯数值安全底线,防止 FEM 退化
η_p95 = percentile(η_old, 95)
eligible = {i | V_old[i] > V_min_safeguard AND η_old_i ≥ 0.05·η_p95}
num = min(|eligible|, N_old//3, max_by_budget)
elements_to_refine = top-k of eligible by score
M_new[j] ∈ {0,…,N_old-1} // 子→父映射
```
**Step 2 — FEM 求解 + 误差估计**
```
η_new ← 新逐单元 η_K
||u_h_new|| ← 新解 L₂ 范数
```
**Step 3 — 局部奖励**动态截断 ε_dynamic
ε_dynamic = max(0.01 × η_P95, 1e-6) // P95 锚定免疫远场噪声稀释
ε_dynamic = max(0.05 × mean(η_new), 1e-6) // 自适应钳制切断远场低 η reward hacking
spatial: r_local_i = log(η_old_i + ε_dynamic) log( √(Σ_{j: M_new[j]=i} η_new_j²) + ε_dynamic )
spatial_max: r_local_i = log(η_old_i + ε_dynamic) log( max_{j: M_new[j]=i} η_new_j + ε_dynamic )
```
> **L₂ 聚合保证 r_local ≥ 0**: 对 1-to-4 切分:
> ```
> Σ η_child² = int²/4 + jump² + sbc² ≤ η_parent² = int² + jump² + sbc²
> → r_local = ½[log(η_parent²) log(Σ η_child²)] ≥ 0
> ```
> - 纯 int 主导: r_local = log(2) ≈ 0.69(强正奖励)
> - 纯 jump/sbc 主导: r_local = 0中性不惩罚不奖励
> - **永远不会惩罚细化**——与 L₁ sum 不同L₂ 天然避免了对 jump/sbc 主导区的结构性负偏置。
**Step 4 — 动作惩罚**
```
penalty_i = λ · (n_i 1) // λ = 0.06
+ (λ_limit / N_old) · 𝟙[达到最大单元数上限] // λ_limit = 10000
r_local_i r_local_i penalty_i
```
**Step 5 — 全局势函数塑形**(仅发给被细化的父单元)
```
E_global = √(Σ_K η_K²) / ||u_h||_{L₂(Ω)}
global_bonus = α · [ log(E_global_old) log(E_global_new) ] // α = 0.2
r_i = r_local_i penalty_i + global_bonus · 𝟙[i 被细化] // 未细化的单元 reward 0
```
> 全局改进信号只分配给实际参与细化的单元,避免被未细化单元稀释。
---
#### 奖励标度校准旧尺寸场下测量score-based 后需重新标定)
在随机策略下实测各分量量级1321 refined-parent 样本
| 分量 | 均值 | r_local 比例 |
|------|------|:---:|
| `r_local` ( refined parents) | +0.364 | |
| `penalty` λ·(n1), λ=0.02 | +0.045 | 1/8 |
| `α·ΔlogE` α=0.2 | +0.069 | 1/5 |
| **net** | **+0.387** | |
满足 `r_local ≫ penalty` `α·ΔlogE ≈ r_local / 5`局部 credit assignment 不被全局信号淹没
---
#### 设计要点
| 组件 | 聚合 | 作用 |
|------|------|------|
| 局部项 `log(η_old / √(Σ η_child²))` | scatter_add父求平方和再开方 | L 聚合保证 r_local 0不惩罚任何细化int 主导区获强正奖励 (≈+0.69) jump/sbc 区中性 |
| 动作惩罚 `λ(n_i1)` λ=0.02 | per-parent | 轻微抑制网格膨胀1-to-4 切分扣 0.06仅占 r_local ~16% |
| 元素上限惩罚 | 达到 20000 上限时触发 | 极端情况兜底λ_limit / N_old 0.05~0.5 per agent |
| 全局项 `α·ΔlogE` α=0.2 | 仅细化父单元 | L 无量纲全局误差下降趋势只发给实际参与细化的单元避免被未细化单元稀释 |
---
## PPO 关键细节
- **单路 GAE**: 势函数塑形后的奖励已包含全局改进信号 `scatter_add` 将细化后的子单元值聚合回父单元单路 GAE 即可
- **奖励归一化**: rollout reward z-score 归一化std < 1e-8 则跳过
- **Value clipping**: 默认 clip_range=0.2
- **梯度裁剪**: max_grad_norm=0.5
- **log_std clamp**: 每步 `optimizer.step()` 后将 `log_std` clamp `[-4.0, -1.0]`std [0.018, 0.368]<br>
初始化 `-2.0` (std≈0.135),避免 `continuous_sizing_field` 有效范围 [-3, 3] 内噪声过大
- **熵正则**: `entropy_coefficient=0.001`,防止 log_std 过早收敛