afem/README.md

395 lines
19 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=2.0, k_ref=6.0),保证不同域尺寸下每单位面积单元数一致
- 可配 exponent^2 = P1 Helmholtz 理论最优 (污染误差 ∝ k²)。建议 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 残差 + 振幅 + 相位方向 + 物理参数,节点 13 维 + 边 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 层消息传递 + GVN 全局广播inner 残差 + LayerNorm
│ ├─ MessagePassingStep × N
│ │ ├─ EdgeModule: MLP([src | dst | edge_attr])
│ │ └─ NodeModule: MLP([node | scatter(入边)])
│ └─ GlobalVirtualNode (GVN): η_K 加权注意力池化 → 注意力门控广播
│ h_V = Σ(η_v/Ση)·h_vα_v = σ(W_att[h_v || h_V])h_v ← h_v + α_v ⊙ W_V·h_V
└─ 输出: 节点隐向量
```
| 超参数 | 值 |
|--------|-----|
| 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 过早收敛到下限
---
## 输入特征
### 节点特征13 维)
| 维度 | 来源 | 名称 | 说明 |
|------|------|------|------|
| 1 | cfg | `volume` | 无量纲单元面积volume / λ² |
| 3 | cfg | `internal_residual` / `gradient_jump` / `sbc_residual` | PDE 残差三分量(真空波数 k 归一化,经 log₁₀ 压缩):<br>`(h_K/k)·√V·|r|` / `√(½Σ h_e·\|jump\|²/k)` / `(h_bnd/k)·\|SBC\|` |
| 1 | cfg | `element_penalty` | 单元惩罚系数 λ |
| 1 | cfg | `timestep` | 当前 rollout 步数 |
| 1 | cfg | `k_local_sqrt_vol` | k × √(ε_r) × √(V)(局域波数 × 特征长度) |
| 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` | 散射场振幅 \|u_scat\|per-element 均值) |
| 1 | fix | `cos_phase` | Re(u) / (\|u\| + 1e-8),相位方向余弦,∈ [1, 1],无分支切割 |
| 1 | fix | `sin_phase` | Im(u) / (\|u\| + 1e-8),相位方向正弦,与 cos 联合编码相位 |
> - **cfg**: 由 `element_features` 配置控制
> - **fix**: 始终启用Helmholtz 振幅 + 相位方向,硬编码)
>
> GNN 输入用 `_compute_residual_components`(真空波数 k 归一化log₁₀ 压缩。Reward 用逐单元 η_K`_eta_indicator`),与 GNN 特征公式一致但不经 log 压缩。SBC 边界条件保留 `k_local`。
### 边特征1 维)
| 维度 | 名称 | 说明 |
|------|------|------|
| 1 | `phase_distance` | 相邻单元中点相位距离 = d × √(k_local_src·k_local_dst) / 2π — 介质内短波长自然放大,赋予 GNN k 不变性 |
---
## 调用逻辑
```
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` | 被 Reverse Dörfler 剔除的噪声尾部占比(累积能量 <1% 总误差的底部单元 | 因场景而异非固定比例 |
| `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} \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} \cdot \left| [[\nabla u \cdot n]] \right|^2_e} \tag{2}$$
$$r_{\text{sbc}} = \frac{h_{bnd}}{k} \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 \sim [L]^{-1}$$h_e \sim [L]$$|\text{jump}|^2 \sim [L]^{-2}$
三项均严格无量纲$h_e/k \cdot |\text{jump}|^2 \sim [L]^2 \cdot [L]^{-2} = 1$。
SBC 边界条件仍用 $k_{local}$物理正确仅归一化因子改用 $k$。
介质内残差不再被 $\sqrt{\varepsilon_r}$ 压低Agent 获得正确的介质内/外优先级信号
`η_K` 的计算`_compute_residual_indicator` GNN 输入特征`_compute_residual_components`公式完全一致特征仅多一层 log₁₀ 压缩关键验证点
- 内部残差P1 ∇²u_h 0仅含反应项 `k²ε_r·u + k²(ε_r-1)·u_inc`真空波数 k 归一化
- 梯度跳变`(h_e/k)·|jump|²`,½ 分配给相邻左右单元$h_e$ 保留边积分路径细化后自然衰减
- SBC 项归一化用 k物理条件保留 k_local`(h_bnd²/k²)·|∂u/∂n i·k_local·u|²`
### 连续尺寸场策略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 ∈ Reverse Dörfler 保留集 // 数值底线 + 能量尾部淘汰 (ε_noise=0.01, ≥20% floor)
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 个单元细化预算耗尽或 Reverse Dörfler 屏蔽所有候选 episode 自动结束不再浪费 FEM 求解
- **k_exponent 可配**初始网格缩放指数可通过 `helmholtz.k_exponent` 配置默认 2.0),² P1 Helmholtz 理论最优 k=30 $N_{init}$ k=6 25×
- **动作掩码 (Reverse Dörfler)** η_K 升序排列剔除累积平方误差贡献 < ε_noise·Ση² 的底部单元数值噪声/已收敛区)。基于能量分布而非密度分位数在重尾和均匀误差分布下均自适应保留率不低于 20% 确保 Agent 始终有充分的选择空间
### 奖励计算
---
#### 变量
| 符号 | 含义 |
|------|------|
| `η_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)
// 数值底线 + Reverse Dörfler 能量尾部淘汰
V_min_safeguard = 1e-10 × domain_area // 纯数值安全底线,防止 FEM 退化
η_sq = η_old²; total_energy = Σ η_sq
k_dorfler = searchsorted(cumsum(sort_asc(η_sq)), ε_noise·total_energy) // ε_noise=0.01
k = min(k_dorfler, N max(1, N//5)) // ≥20% floor
eligible = {i | V_old[i] > V_min_safeguard AND i ∈ sort_asc_idx[k:] }
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 = max(0.01 × η_P95, 1e-6)
// Refined parents: r_local + zero-sum bonus penalty
if i refined_parents:
r_i = log(η_old + ε) log(√(Σ η_child²) + ε) // r_local 0 (L 聚合)
+ 0.3 × (η_old / μ 1.0) // zero-sum bonus (Σ = 0)
0.06 // action penalty
// Unrefined parents: causal isolation
else:
r_i = 0
> **零和奖金**:α·(η/μ1) 全场求和为零。细化高于均值的单元得正奖金,低于均值的倒扣。
> 这是 Dörfler 准则的 RL 对偶Agent 必须选出误差超过全均水平的单元。
> **因果隔离**:未细化单元 r ≡ 0。零和奖金本身足够强介质内 +0.51)、
> 不再需要忽视惩罚的推力,排序机制自动淘汰不划算的单元。
> **L₂ 聚合**:√(Σ η_child²) ≤ η_parent 天然成立r_local ≥ 0 永不惩罚细化。
**Step 4 — 全局误差(仅诊断)**
global_bonus = α·[log(E_old) log(E_new)]α = 0.5
不注入 Actor rewardHelmholtz 污染误差可使 E_new > E_old 在正确细化后发生,
注入 global_bonus 导致因果断裂。Actor 仅优化 Step 3 的 per-element reward。
---
#### 奖励标度校准旧尺寸场下测量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仅 refined parents | L₂ 保证 r_local ≥ 0int 主导 +0.69 |
| 零和奖金 `0.3×(η/μ1)` | 仅 refined parents | Σ=0高于 μ 得正奖,低于 μ 倒扣 (Dörfler 准则的 RL 对偶) |
| 动作惩罚 `λ=0.06` | per-refined-parent | 轻微抑制网格膨胀1-to-4 扣 0.06 |
| 因果隔离 `r=0` | unrefined parents | 零和奖金足够强,不需额外推力 |
| 全局项 `α·ΔlogE` α=0.5 | 仅诊断 | 不注入 Actor避免污染误差因果断裂 |
---
## PPO 关键细节
- **单路 GAE**: r_local 自身已闭合因果(细化单元的局部误差改善),不需要势函数塑形。用 `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 `[-3.0, -1.0]`σ [0.05, 0.37]<br>
初始化 `-2.0` (σ≈0.135),放宽下限防止策略过早确定化
- **熵正则**: `entropy_coefficient=0.005`,施加有意义的探索压力防止 x<0 崩塌
- **epochs_per_iteration**: 3减少对同一 rollout 的过拟合