first commit
This commit is contained in:
commit
377b645008
|
|
@ -0,0 +1,396 @@
|
|||
# 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.5，helmholtz.pre_asymptotic_N 可配）。约 1.5 点/波长，刚好跨过渐近区门槛，赋予初始网格基本相位解析能力但不过度消耗物理预算，为 RL agent 留出充分的选择性细化空间
|
||||
- **后验误差**: 残差型 indicator（Ainsworth & 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（嵌入）
|
||||
├─ MessagePassingStack（2 层消息传递，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.0（std≈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_i，x 越小 ⇒ 优先级越高（纯排序，不设正负门槛）
|
||||
- `initial_log_std=-2.0`（std≈0.135），clamp 在 [-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_i：x 越小 ⇒ 优先级越高（纯排序，不设正负门槛）
|
||||
- 不再使用 `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` λ·(n−1), λ=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_i−1)` λ=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 过早收敛
|
||||
|
|
@ -0,0 +1,154 @@
|
|||
---
|
||||
ASMR++ 奖励计算完整分析
|
||||
|
||||
默认配置使用 reward_type: spatial_max + error_metric: maximum。整个奖励信号链分以下步骤：
|
||||
|
||||
Step 1: 误差估计 — 精细网格参考解
|
||||
|
||||
参考网格 (初始网格细化6次)
|
||||
↓ FEM求解
|
||||
参考解 u_ref (视为"真值")
|
||||
↓
|
||||
粗网格解 u_coarse 在每个积分点(参考网格元素中点)与 u_ref 比较
|
||||
↓
|
||||
绝对误差 |u_ref - u_coarse| per 积分点
|
||||
↓ scatter_max per 粗元素
|
||||
error_per_element: 每个粗网格元素内的最大误差 (num_elements, solution_dim)
|
||||
|
||||
用精细网格做数值积分 (error_integrator.py:86-169)，支持三种积分方式：mean(积分平均值)、squared(积分平方误差)、maximum(元素内最大误差)。默认Poisson 是标量 PDE，solution_dim=1。
|
||||
|
||||
Step 2: spatial_max 奖励计算
|
||||
|
||||
核心代码在 mesh_refinement.py:657-714，以下是逐步推导：
|
||||
|
||||
奖励基准 (reward_per_agent_and_dim)
|
||||
= previous_error_per_element ← 细分前该元素的误差
|
||||
|
||||
┌─────┬─────┐
|
||||
│ │ │
|
||||
├─────┼─────┤
|
||||
│ │ │
|
||||
└─────┴─────┘
|
||||
|
||||
父元素 i (error=0.8) 子元素: j1(0.3), j2(0.5), j3(0.6), j4(0.1)
|
||||
↓ scatter_max per agent_mapping
|
||||
max_mapped_error[i] = max(0.3, 0.5, 0.6, 0.1) = 0.6
|
||||
↓
|
||||
reward_raw[i] = 0.8 - 0.6 = +0.2 ✅ 误差最大的子元素也比父元素好
|
||||
|
||||
关键：spatial_max 只奖励"所有子元素误差都下降"的情况。如果有任一子元素误差仍等于原父元素误差，reward=0。
|
||||
|
||||
父元素 j (error=0.5) 子元素: k1(0.5), k2(0.1), k3(0.2), k4(0.05)
|
||||
↓
|
||||
max_mapped_error[j] = max(0.5, 0.1, 0.2, 0.05) = 0.5
|
||||
↓
|
||||
reward_raw[j] = 0.5 - 0.5 = 0 ❌ 有一个子元素仍未改善
|
||||
|
||||
对比 spatial (非 max) 模式：
|
||||
|
||||
reward_raw[i] = previous_error[i] - Σ_j φ_ij * error[j]
|
||||
= 标量加法 (np.add.at) 把所有子元素误差从父元素误差中减去
|
||||
这种模式下即使部分子元素没有改善，整体仍有正奖励。
|
||||
|
||||
Step 3: 归一化 + 降维到标量
|
||||
|
||||
# 除以初始网格的误差 → 把误差改善量归一化到 [0, ~1] 区间
|
||||
reward_per_agent_and_dim = reward_per_agent_and_dim / initial_approximation_error
|
||||
|
||||
# 多维 PDE 降维: dot product with solution_dimension_weights
|
||||
# Poisson 是标量PDE, weights=[1.0], 即恒等变换
|
||||
reward_per_agent = project_to_scalar(reward_per_agent_and_dim)
|
||||
# = np.dot(reward_per_agent_and_dim, [1.0]) = reward_per_agent_and_dim
|
||||
|
||||
Step 4: 元素惩罚 (Element Penalty)
|
||||
|
||||
# 统计每个父元素产生了多少子元素
|
||||
element_counts = unique(agent_mapping, return_counts=True)[1] # 每个父元素→子元素的数量
|
||||
element_counts = element_counts - 1 # 减1因为是"新增的"子元素数
|
||||
|
||||
# 默认 λ ~ 0.01 (loguniform 采样于 [1e-3, 1e-1])
|
||||
element_penalty = λ * element_counts
|
||||
|
||||
┌──────────────────────────┬────────────────┬──────────────────┐
|
||||
│ 场景 │ element_counts │ penalty (λ=0.01) │
|
||||
├──────────────────────────┼────────────────┼──────────────────┤
|
||||
│ 未细分元素 │ 0 │ 0 │
|
||||
├──────────────────────────┼────────────────┼──────────────────┤
|
||||
│ 分裂为 4 个子三角 │ 3 │ 0.03 │
|
||||
├──────────────────────────┼────────────────┼──────────────────┤
|
||||
│ 被波及细分 (Rivara 平滑) │ 1-3 │ 0.01-0.03 │
|
||||
└──────────────────────────┴────────────────┴──────────────────┘
|
||||
|
||||
作用: 惩罚是正则化项，防止策略无节制细分所有元素。只在"误差改善 > 细分代价"时细分才有利。
|
||||
|
||||
Step 5: 元素上限惩罚 (Element Limit Penalty)
|
||||
|
||||
if num_elements > maximum_elements (20000):
|
||||
element_limit_penalty = 1000 / previous_num_elements # ≈ 0.05~0.5 per agent
|
||||
else:
|
||||
element_limit_penalty = 0
|
||||
|
||||
Step 6: 最终每 Agent 奖励
|
||||
|
||||
r_i = error_improvement_i / initial_error
|
||||
- λ * new_elements_created_by_i
|
||||
- limit_penalty
|
||||
|
||||
形状为 (num_agents_t,) — 每个 agent（父元素）一个标量奖励。
|
||||
|
||||
Step 7: 奖励到 TD 误差 — 与论文公式 (3) 的对应
|
||||
|
||||
Buffer 存储:
|
||||
r_i(s_t, a_t) ← 父元素 i 的奖励 (num_agents_t,)
|
||||
V_i(s_t) ← 父元素 i 的价值 (num_agents_t,)
|
||||
φ_ij = agent_mapping ← 子元素j → 父元素i 的映射
|
||||
V_j(s_{t+1}) ← 子元素的价值 (num_agents_{t+1},)
|
||||
|
||||
GAE Delta 计算:
|
||||
projected_V = scatter_sum(V_j(s_{t+1}), index=φ_ij) ← Σ_j φ_ij·V_j(s_{t+1})
|
||||
δ_i = r_i + γ * projected_V_i - V_i(s_t)
|
||||
|
||||
对应论文 (3): δ_i^t = r(s^t, a^t)_i + γ·Σ_j φ_ij^t·V_j(s^{t+1}) - V_i(s^t)
|
||||
|
||||
Step 8: 混合奖励 (Mixed Return, global_weight=0.5)
|
||||
|
||||
在 MixedOnPolicyBuffer 中额外计算：
|
||||
|
||||
# 全局奖励 (均值)
|
||||
r_global = mean(r_i) # 所有agent的平均奖励
|
||||
|
||||
# 全局价值 (均值)
|
||||
V_global = mean(V_i) # 所有agent的平均价值
|
||||
|
||||
# 全局 GAE
|
||||
δ_global = r_global + γ·V_global' - V_global
|
||||
|
||||
# 局部 GAE
|
||||
δ_local_i = 上述 per-agent GAE
|
||||
|
||||
# 混合 Advantage
|
||||
A_i = (1 - 0.5) * A_local_i + 0.5 * A_global
|
||||
|
||||
完整奖励流总结
|
||||
|
||||
FEM求解 → 逐元素误差估计 (±积分 vs 参考网格)
|
||||
↓
|
||||
spatial_max: error_before - max_error_of_children
|
||||
↓
|
||||
归一化 (/ initial_error)
|
||||
↓
|
||||
- λ * new_elements + limit_penalty
|
||||
↓
|
||||
r_i (per agent) ────────────→ 局部 GAE → A_local_i
|
||||
│ ↓
|
||||
└→ r_global = mean(r_i) → 全局 GAE → A_global
|
||||
↓
|
||||
A_i = 0.5·A_local_i + 0.5·A_global
|
||||
↓
|
||||
送入 PPO policy_loss
|
||||
|
||||
设计精巧之处：
|
||||
1. 空间奖励 + agent_mapping：每个元素独立计算误差改善，通过 agent_mapping φ_ij 追踪父→子关系
|
||||
2. spatial_max 语义：reward 表示"最差子元素的误差下降量"——驱动策略优先细分误差最大的区域
|
||||
3. 元素惩罚：防止盲目细分，精确到每个 agent 独立计算代价
|
||||
4. 混合奖励：局部信号指导细粒度决策 + 全局信号稳定整体训练
|
||||
|
|
@ -0,0 +1,255 @@
|
|||
# ASMR++ 网络架构与数据流 (默认配置)
|
||||
|
||||
> 基于 `configs/asmr_pp/asmr_default.yaml` — `value_function_aggr: spatial`, `projection_type: sum`
|
||||
|
||||
## 架构总览
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
subgraph ENV["♻️ 环境: MeshRefinement"]
|
||||
A1["FEMProblemCircularQueue<br/>随机采样 PDE 问题"]
|
||||
A2["生成初始粗网格<br/>(meshpy, 2D 三角剖分)"]
|
||||
A3["FEM 求解器<br/>计算 PDE 解和逐单元误差"]
|
||||
A4["构建观测图<br/>(节点=单元, 边=邻接关系)"]
|
||||
A1 --> A2 --> A3 --> A4
|
||||
end
|
||||
|
||||
subgraph GRAPH["📊 观测图 (torch_geometric Data)"]
|
||||
B1["<b>节点特征 (x)</b><br/>━━━━━━━━━━━━━━━━<br/>solution_mean / solution_std<br/>volume / timestep<br/>element_penalty<br/>source_term (PDE 特征)<br/>共 ~10-15 维"]
|
||||
B2["<b>边特征 (edge_attr)</b><br/>━━━━━━━━━━━━━━━━<br/>euclidean_distance<br/>共 1 维"]
|
||||
B3["<b>边索引 (edge_index)</b><br/>━━━━━━━━━━━━━━━━<br/>双向邻接 + 自环"]
|
||||
end
|
||||
|
||||
subgraph NORM["📏 观测归一化器"]
|
||||
C1["node.x: running mean/std"]
|
||||
C2["edge_attr: running mean/std"]
|
||||
end
|
||||
|
||||
subgraph HMPN["🧠 HMPN 基础网络 (HomogeneousMessagePassingBase)"]
|
||||
subgraph EMBED["输入嵌入"]
|
||||
D1["节点嵌入: Linear(in→64)"]
|
||||
D2["边嵌入: Linear(in→64)"]
|
||||
end
|
||||
subgraph STACK["消息传递堆栈 (num_steps=2, residual=inner, layernorm=inner)"]
|
||||
subgraph STEP1["Step 1/2"]
|
||||
E1["<b>边更新</b> HomogeneousEdgeModule<br/>concat[src(64), dst(64), edge(64)]<br/>→ LatentMLP(192→64, 2层, LeakyReLU)<br/>→ LayerNorm → +inner residual"]
|
||||
E2["<b>节点更新</b> HomogeneousMessagePassingNodeModule<br/>scatter_mean(edge→dest) → concat[node(64), agg(64)]<br/>→ LatentMLP(128→64, 2层, LeakyReLU)<br/>→ LayerNorm → +inner residual"]
|
||||
E1 --> E2
|
||||
end
|
||||
subgraph STEP2["Step 2/2"]
|
||||
F1["边更新 (同上)"]
|
||||
F2["节点更新 (同上)"]
|
||||
F1 --> F2
|
||||
end
|
||||
STEP1 --> STEP2
|
||||
end
|
||||
D1 --> STEP1
|
||||
D2 --> STEP1
|
||||
STEP2 --> G["输出: 节点潜在特征 (num_nodes, 64)"]
|
||||
end
|
||||
|
||||
subgraph HEADS["🎯 策略与价值头 (share_base=False, 各自独立 GNN)"]
|
||||
subgraph ACTOR["Actor 头"]
|
||||
H1["Policy MLP<br/>2层, Tanh<br/>64→64→64"]
|
||||
H2["Linear(64→action_dim)"]
|
||||
H3["log_std (可学习)"]
|
||||
H4["DiagGaussian(μ, σ)<br/>每节点输出独立动作"]
|
||||
H1 --> H2 --> H4
|
||||
H3 --> H4
|
||||
end
|
||||
subgraph CRITIC["Critic 头 — 逐节点价值，不做 scatter 聚合"]
|
||||
I1["Value MLP<br/>2层, Tanh<br/>64→64→1"]
|
||||
I2["输出形状: (num_agents, 1)<br/>每个 agent 独立 V_i(s)<br/><b>value_function_aggr=spatial<br/>→ 不聚合，保持逐节点</b>"]
|
||||
I1 --> I2
|
||||
end
|
||||
G --> H1
|
||||
G --> I1
|
||||
end
|
||||
|
||||
subgraph BUFFER["🗃️ MixedOnPolicyBuffer (global_weight=0.5)"]
|
||||
J1["<b>局部 GAE (逐节点)</b><br/>δ_i = r_i + γ·Σ_j φ_ij·V_j(s') - V_i(s)<br/>projection_type='sum': Σ 通过 agent_mapping 反投影"]
|
||||
J2["<b>全局 GAE (图级别)</b><br/>δ_global = r_global + γ·V_mean(s') - V_mean(s)"]
|
||||
J3["<b>混合 Advantage</b><br/>A_i = (1-0.5)·A_i_local + 0.5·A_global"]
|
||||
J1 --> J3
|
||||
J2 --> J3
|
||||
end
|
||||
|
||||
subgraph PPO["🔄 PPO 训练"]
|
||||
K1["256 步 Rollout"]
|
||||
K2["5 Epochs, batch_size=32"]
|
||||
K3["policy_loss + 0.5·value_loss<br/>clip_range=0.2"]
|
||||
K4["梯度裁剪 0.5, Adam lr=3e-4"]
|
||||
K1 --> K2 --> K3 --> K4
|
||||
end
|
||||
|
||||
ENV --> GRAPH --> NORM --> HMPN
|
||||
HMPN --> HEADS
|
||||
ACTOR -->|动作| ENV
|
||||
CRITIC -->|"V_i(s) 逐节点"| BUFFER
|
||||
ENV -->|"r_i, agent_mapping φ"| BUFFER
|
||||
BUFFER --> PPO
|
||||
PPO -->|更新参数| HMPN
|
||||
PPO -->|更新参数| HEADS
|
||||
```
|
||||
|
||||
## 核心纠正: projection_type 的真实作用
|
||||
|
||||
**之前的错误理解**:
|
||||
- ~~Critic 输出 scatter_sum → 图级别价值~~ ❌
|
||||
|
||||
**正确理解**:
|
||||
- `value_function_aggr: "spatial"` → Critic **不做任何聚合**，输出 `(num_agents, 1)` 逐节点价值 ✅
|
||||
- `projection_type: "sum"` → 在 **Buffer** 中通过 `agent_mapping` 反投影下一步价值时使用 ✅
|
||||
|
||||
两个参数作用于完全不同的位置:
|
||||
|
||||
| 参数 | 作用位置 | 作用 |
|
||||
|------|----------|------|
|
||||
| `value_function_aggr: "spatial"` | `SwarmPPOActorCritic._get_values_and_distribution()` | 控制 Critic 输出是否聚合: `"spatial"` → 保持逐节点 |
|
||||
| `projection_type: "sum"` | `SpatialOnPolicyBuffer._project_to_previous_step()` | 控制 agent_mapping 反投影方式: sum→子元素价值求和回父元素 |
|
||||
|
||||
## 详细数据流 (序列图)
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
actor Trainer
|
||||
participant Env as MeshRefinement
|
||||
participant Norm as Normalizer
|
||||
participant GNN as HMPN Base
|
||||
participant Actor as Policy Head
|
||||
participant Critic as Value Head
|
||||
participant Buffer as MixedOnPolicyBuffer
|
||||
|
||||
Note over Trainer,Buffer: === Rollout (256 步) ===
|
||||
|
||||
Trainer->>Env: reset()
|
||||
Env->>Env: 随机 Poisson PDE + 随机域 + GMM 负载
|
||||
Env->>Env: 初始粗网格 → FEM 求解 → 构建观测图
|
||||
|
||||
loop 256 步
|
||||
Env-->>Norm: 观测图 (原始 node.x, edge_attr)
|
||||
Norm-->>GNN: 归一化后图
|
||||
GNN->>GNN: Edge Dropout (0.1, 仅训练)
|
||||
GNN->>GNN: 嵌入 → MP Step1 → MP Step2
|
||||
GNN-->>Actor: node_features (num_nodes, 64)
|
||||
GNN-->>Critic: node_features (num_nodes, 64)
|
||||
|
||||
Actor->>Actor: MLP → μ, σ → 采样动作
|
||||
Critic->>Critic: MLP(64→1) → <b>V_i(s): (num_agents, 1) 逐节点</b>
|
||||
|
||||
Actor-->>Env: actions (num_agents, 1)
|
||||
Env->>Env: 元素选择 → 网格细分
|
||||
Env->>Env: FEM 求解 → 计算空间奖励 r_i
|
||||
Env-->>Buffer: (obs, a, r_i, V_i, log_prob, agent_mapping φ)
|
||||
end
|
||||
|
||||
Note over Buffer: === GAE 计算 (逐节点 + 混合奖励) ===
|
||||
|
||||
Buffer->>Buffer: <b>局部 δ_i(t) = r_i + γ·Σ_j φ_ij·V_j(t+1) - V_i(t)</b>
|
||||
Buffer->>Buffer: projection_type='sum': Σ_j 通过 agent_mapping 反投影
|
||||
Buffer->>Buffer: 局部 GAE → A_local_i (逐节点)
|
||||
Buffer->>Buffer: 全局 GAE → A_global (图级, 用 mean(V_i) 算)
|
||||
Buffer->>Buffer: <b>A_i = 0.5·A_local_i + 0.5·A_global</b>
|
||||
Buffer->>Buffer: R_i = A_i + V_i(s)
|
||||
|
||||
Note over Trainer,Buffer: === 训练 (5 Epochs × batch 32) ===
|
||||
|
||||
loop 5 Epochs
|
||||
Buffer-->>Trainer: (obs, a, old_log_prob, old_V_i, A_i, R_i)
|
||||
Trainer->>GNN: 重新前向传播
|
||||
GNN-->>Actor: node_features
|
||||
GNN-->>Critic: node_features
|
||||
Actor->>Actor: 新 log_prob
|
||||
Critic->>Critic: 新 V_i (逐节点)
|
||||
Trainer->>Trainer: ratio = exp(log_prob_new - log_prob_old)
|
||||
Trainer->>Trainer: policy_loss = -min(ratio·A_i, clip(ratio,0.8,1.2)·A_i)
|
||||
Trainer->>Trainer: value_loss = 0.5·clip(V_new, V_old±0.2) vs R_i
|
||||
Trainer->>Trainer: backward() + grad_clip(0.5) + Adam.step()
|
||||
end
|
||||
```
|
||||
|
||||
## 论文公式 (3) 与代码对应
|
||||
|
||||
论文中的 TD 误差公式:
|
||||
|
||||
$$\delta^t_i = r(s^t, a^t)_i + \gamma \sum_j \phi_{ij}^t V_j(s^{t+1}) - V_i(s^t)$$
|
||||
|
||||
在代码中的实现路径 (`spatial_on_policy_buffer.py:174-178`):
|
||||
|
||||
```python
|
||||
# _get_agent_wise_advantages_and_returns()
|
||||
for step in range(self.buffer_size):
|
||||
if self.dones[step]:
|
||||
delta = self.rewards[step] - self.values[step] # r_i - V_i(s)
|
||||
else:
|
||||
delta = self.rewards[step] \
|
||||
+ self.discount_factor * projected_next_values[step] \ # + γ·Σ_j φ_ij·V_j(s')
|
||||
- self.values[step] # - V_i(s)
|
||||
```
|
||||
|
||||
其中 `projected_next_values[step]` 由 `_project_to_previous_step()` 产生:
|
||||
|
||||
```python
|
||||
# projection_type='sum'
|
||||
projected_value = scatter_sum(values[step], index=agent_mappings[step], dim=0)
|
||||
# ^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^^
|
||||
# V_j(s_{t+1}) φ_ij: 新agent j → 旧agent i
|
||||
```
|
||||
|
||||
## 关键默认参数
|
||||
|
||||
| 参数 | 值 | 代码位置 |
|
||||
|------|-----|----------|
|
||||
| **算法** | PPO | `config["algorithm"]["name"]` |
|
||||
| **网络骨架** | Homogeneous MPN | `config["network"]["type_of_base"]` |
|
||||
| **GNN 架构** | mpn (message passing) | `config["network"]["base"]["architecture"]` |
|
||||
| **潜在维度** | 64 | `config["network"]["latent_dimension"]` |
|
||||
| **MP 步数** | 2 | `config["network"]["base"]["stack"]["num_steps"]` |
|
||||
| **残差** | inner | `config["network"]["base"]["stack"]["residual_connections"]` |
|
||||
| **层归一化** | inner | `config["network"]["base"]["stack"]["layer_norm"]` |
|
||||
| **边→节点聚合** | mean | `config["network"]["base"]["scatter_reduce"]` |
|
||||
| **Base MLP** | 2层, LeakyReLU | `config["network"]["base"]["stack"]["mlp"]` |
|
||||
| **Actor MLP** | 2层, Tanh | `config["network"]["actor"]["mlp"]` |
|
||||
| **Critic MLP** | 2层, Tanh | `config["network"]["critic"]["mlp"]` |
|
||||
| **价值函数范围** | **spatial** (逐节点, 不聚合) | `config["algorithm"]["ppo"]["value_function_aggr"]` |
|
||||
| **价值投影方式** | **sum** (agent_mapping 反投影用) | `config["algorithm"]["ppo"]["projection_type"]` |
|
||||
| **混合奖励权重** | 0.5 | `config["algorithm"]["mixed_return"]["global_weight"]` |
|
||||
| **共享 Base** | False (Actor/Critic 各自独立 GNN) | `config["network"]["share_base"]` |
|
||||
| **动作分布** | DiagGaussian (连续) | 动作空间为 `gym.spaces.Box` |
|
||||
| **Rollout 步数** | 256 | `config["algorithm"]["ppo"]["num_rollout_steps"]` |
|
||||
| **训练轮次** | 5 | `config["algorithm"]["ppo"]["epochs_per_iteration"]` |
|
||||
| **Batch 大小** | 32 | `config["algorithm"]["batch_size"]` |
|
||||
| **GAE λ** | 0.95 | `config["algorithm"]["ppo"]["gae_lambda"]` |
|
||||
| **折现 γ** | 1.0 | `config["algorithm"]["discount_factor"]` |
|
||||
| **PPO clip** | 0.2 | `config["algorithm"]["ppo"]["clip_range"]` |
|
||||
| **梯度裁剪** | 0.5 | `config["algorithm"]["ppo"]["max_grad_norm"]` |
|
||||
| **学习率** | 3e-4 | `config["network"]["training"]["learning_rate"]` |
|
||||
| **边 Dropout** | 0.1 | `config["network"]["base"]["edge_dropout"]` |
|
||||
| **Episode 步数** | 6 | `config["environment"]["mesh_refinement"]["num_timesteps"]` |
|
||||
| **PDE** | Poisson (GMM 负载, zero Dirichlet) | `config["environment"]["mesh_refinement"]["fem"]["pde_type"]` |
|
||||
|
||||
## projection_type 的两种职责
|
||||
|
||||
`projection_type` 在 Buffer 中有**两处**使用，都是通过 `agent_mapping` 做跨时间步的 agent 反投影:
|
||||
|
||||
### 1. 价值反投影 — 公式 (3) 的 Σ 项
|
||||
```python
|
||||
# _project_to_previous_step() — spatial_on_policy_buffer.py:33
|
||||
projected_value = scatter_sum(values[step], index=agent_mappings[step], dim=0)
|
||||
# 下一步的 V_j(s_{t+1}) 按 agent_mapping φ_ij 求和回当前步的 agent i
|
||||
```
|
||||
|
||||
### 2. GAE 时间差分反投影 — 动态规划递推
|
||||
```python
|
||||
# _get_agent_wise_advantages_and_returns() — spatial_on_policy_buffer.py:169
|
||||
projected_last_gae = scatter_sum(last_gae, index=self._agent_mappings[step], dim=0)
|
||||
# 上一步累积的 GAE 按 agent_mapping 反投影
|
||||
```
|
||||
|
||||
## 核心创新点
|
||||
|
||||
1. **Swarm 视角 + 变长 Agent**: 每个网格元素是一个 agent，元素分裂后 agent 数量动态增长
|
||||
2. **空间奖励 + agent_mapping**: 通过 `agent_mapping φ_ij` 追踪父→子关系，支持逐节点的 TD 误差计算（公式 3）
|
||||
3. **混合奖励学习**: 局部逐节点 Advantage + 全局图级 Advantage 加权混合 (0.5:0.5)
|
||||
4. **MPN 通信**: 边更新 + 节点更新的消息传递，元素通过共享三角形边交换 PDE 解信息
|
||||
5. **自适应细化**: 连续动作 → 概率性元素选择 → 非均匀网格，资源集中在误差大的区域
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,112 @@
|
|||
"""Compare iter100 vs iter150 checkpoints: action_mean diff and refine_mask equality."""
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch_geometric.data import Batch
|
||||
|
||||
from src.network import create_model
|
||||
from src.utils import load_checkpoint, setup_helmholtz_config
|
||||
|
||||
|
||||
def load_config():
|
||||
from src.utils import load_config as _lc
|
||||
from pathlib import Path
|
||||
cfg_path = Path(__file__).resolve().parent / "src" / "config.yaml"
|
||||
return _lc(str(cfg_path))
|
||||
|
||||
|
||||
def compare_checkpoints(ckpt_a, ckpt_b, label_a="iter100", label_b="iter150"):
|
||||
config = load_config()
|
||||
setup_helmholtz_config(config)
|
||||
algo = config.get("algorithm", {})
|
||||
|
||||
from environment.mesh_refinement import MeshRefinement
|
||||
|
||||
env = MeshRefinement(
|
||||
environment_config=config.get("environment", {}).get("mesh_refinement", {}),
|
||||
seed=99,
|
||||
)
|
||||
|
||||
# ── Load both models ──
|
||||
model_a = create_model(env, config.get("network", {}), algo.get("ppo", {}))
|
||||
load_checkpoint(model_a, ckpt_a)
|
||||
model_a.eval()
|
||||
|
||||
model_b = create_model(env, config.get("network", {}), algo.get("ppo", {}))
|
||||
load_checkpoint(model_b, ckpt_b)
|
||||
model_b.eval()
|
||||
|
||||
# ── Get same initial observation ──
|
||||
env.reset()
|
||||
obs = env.reset() # second reset ensures same state
|
||||
|
||||
with torch.no_grad():
|
||||
batch = Batch.from_data_list([obs])
|
||||
|
||||
# Model A
|
||||
shared_a, batch_a = model_a._encode(batch)
|
||||
latent_pi_a = model_a.policy_mlp(shared_a)
|
||||
action_mean_a = model_a.action_out(latent_pi_a).cpu().numpy().flatten()
|
||||
dist_a = model_a._make_distribution(latent_pi_a)
|
||||
actions_a = dist_a.get_actions(deterministic=True).cpu().numpy().flatten()
|
||||
|
||||
# Model B
|
||||
shared_b, batch_b = model_b._encode(batch)
|
||||
latent_pi_b = model_b.policy_mlp(shared_b)
|
||||
action_mean_b = model_b.action_out(latent_pi_b).cpu().numpy().flatten()
|
||||
dist_b = model_b._make_distribution(latent_pi_b)
|
||||
actions_b = dist_b.get_actions(deterministic=True).cpu().numpy().flatten()
|
||||
|
||||
# ── Compare action_mean ──
|
||||
diff = action_mean_a - action_mean_b
|
||||
print(f"\n{'='*60}")
|
||||
print(f" 1. action_mean comparison")
|
||||
print(f"{'='*60}")
|
||||
print(f" {label_a} action_mean: min={action_mean_a.min():.6f} max={action_mean_a.max():.6f} mean={action_mean_a.mean():.6f} std={action_mean_a.std():.6f}")
|
||||
print(f" {label_b} action_mean: min={action_mean_b.min():.6f} max={action_mean_b.max():.6f} mean={action_mean_b.mean():.6f} std={action_mean_b.std():.6f}")
|
||||
print(f" ---")
|
||||
print(f" |diff|: min={np.abs(diff).min():.8f} max={np.abs(diff).max():.8f} mean={np.abs(diff).mean():.8f}")
|
||||
print(f" diff = 0 exactly: {int(np.sum(diff == 0))} / {len(diff)} ({100 * np.sum(diff == 0) / len(diff):.2f}%)")
|
||||
print(f" |diff| < 1e-6: {int(np.sum(np.abs(diff) < 1e-6))} / {len(diff)}")
|
||||
print(f" |diff| < 1e-4: {int(np.sum(np.abs(diff) < 1e-4))} / {len(diff)}")
|
||||
print(f" cos similarity: {np.dot(action_mean_a, action_mean_b) / (np.linalg.norm(action_mean_a) * np.linalg.norm(action_mean_b) + 1e-12):.8f}")
|
||||
|
||||
# ── Compare refine_mask (action > 0) ──
|
||||
mask_a = actions_a > 0.0
|
||||
mask_b = actions_b > 0.0
|
||||
mask_equal = np.array_equal(mask_a, mask_b)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print(f" 2. refine_mask comparison")
|
||||
print(f"{'='*60}")
|
||||
print(f" {label_a} refine_mask: sum={mask_a.sum()} / {len(mask_a)} ({100 * mask_a.sum() / len(mask_a):.1f}%)")
|
||||
print(f" {label_b} refine_mask: sum={mask_b.sum()} / {len(mask_b)} ({100 * mask_b.sum() / len(mask_b):.1f}%)")
|
||||
print(f" refine_mask exactly equal: {mask_equal}")
|
||||
print(f" mask XOR sum: {(mask_a ^ mask_b).sum()} / {len(mask_a)}")
|
||||
|
||||
if not mask_equal:
|
||||
diff_idx = np.where(mask_a != mask_b)[0]
|
||||
print(f" First 20 differing indices: {diff_idx[:20].tolist()}")
|
||||
print(f" At those indices, {label_a} action_mean: {action_mean_a[diff_idx[:10]]}")
|
||||
print(f" At those indices, {label_b} action_mean: {action_mean_b[diff_idx[:10]]}")
|
||||
|
||||
# ── 3. Parameter-level diff ──
|
||||
print(f"\n{'='*60}")
|
||||
print(f" 3. Model parameter weight diff (L2 norm)")
|
||||
print(f"{'='*60}")
|
||||
sd_a = torch.load(ckpt_a, map_location="cpu")["model_state_dict"]
|
||||
sd_b = torch.load(ckpt_b, map_location="cpu")["model_state_dict"]
|
||||
for k in sorted(sd_a.keys()):
|
||||
w_a = sd_a[k].float()
|
||||
w_b = sd_b[k].float()
|
||||
l2 = torch.norm(w_a - w_b).item()
|
||||
rel = l2 / (torch.norm(w_a).item() + 1e-12)
|
||||
print(f" {k:55s} |Δ|₂={l2:.6e} rel={rel:.6e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
d1 = sys.argv[1] if len(sys.argv) > 1 else "checkpoints/model_iter0100.pt"
|
||||
d2 = sys.argv[2] if len(sys.argv) > 2 else "checkpoints/model_iter0150.pt"
|
||||
l1 = sys.argv[3] if len(sys.argv) > 3 else "iter100"
|
||||
l2 = sys.argv[4] if len(sys.argv) > 4 else "iter150"
|
||||
compare_checkpoints(d1, d2, l1, l2)
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,70 @@
|
|||
import copy
|
||||
from typing import Any, Dict, Union
|
||||
|
||||
import numpy as np
|
||||
from skfem import MeshTri1
|
||||
|
||||
|
||||
class Domain:
|
||||
"""Square domain [0,1]x[0,1] with initial coarse mesh and fine integration mesh."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
domain_config: Dict[Union[str, int], Any],
|
||||
random_state: np.random.RandomState,
|
||||
):
|
||||
xmin, ymin, xmax, ymax = domain_config.get("boundary", [0.0, 0.0, 1.0, 1.0])
|
||||
self._boundary = np.array([xmin, ymin, xmax, ymax])
|
||||
self._random_state = random_state
|
||||
|
||||
num_elements = domain_config.get("initial_num_elements", None)
|
||||
if num_elements is not None:
|
||||
domain_area = (xmax - xmin) * (ymax - ymin)
|
||||
self._max_volume = 2.0 * domain_area / float(num_elements)
|
||||
else:
|
||||
self._max_volume = domain_config.get("max_initial_element_volume", 0.05)
|
||||
|
||||
self._initial_mesh = self._create_initial_mesh()
|
||||
|
||||
@property
|
||||
def initial_mesh(self) -> MeshTri1:
|
||||
return copy.deepcopy(self._initial_mesh)
|
||||
|
||||
def replace_initial_mesh(self, mesh: MeshTri1) -> None:
|
||||
"""Replace the stored initial mesh (e.g. after Nyquist enforcement)."""
|
||||
self._initial_mesh = mesh
|
||||
|
||||
def get_integration_mesh(self) -> MeshTri1:
|
||||
return self._initial_mesh.refined(4)
|
||||
|
||||
@property
|
||||
def boundary_line_segments(self) -> np.ndarray:
|
||||
boundary_edges = self._initial_mesh.boundary_facets()
|
||||
boundary_node_indices = self._initial_mesh.facets[:, boundary_edges]
|
||||
return self._initial_mesh.p[:, boundary_node_indices].T.reshape(-1, 4)
|
||||
|
||||
def _create_initial_mesh(self) -> MeshTri1:
|
||||
return self._meshpy_square()
|
||||
|
||||
def _meshpy_square(self) -> MeshTri1:
|
||||
import meshpy.triangle as triangle
|
||||
|
||||
xmin, ymin, xmax, ymax = self._boundary
|
||||
points = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
|
||||
facets = [(0, 1), (1, 2), (2, 3), (3, 0)]
|
||||
|
||||
info = triangle.MeshInfo()
|
||||
info.set_points(points)
|
||||
info.set_facets(facets)
|
||||
|
||||
mesh = triangle.build(info, max_volume=self._max_volume)
|
||||
vertices = np.array(mesh.points).T
|
||||
triangles = np.array(mesh.elements).T
|
||||
return MeshTri1(vertices, triangles)
|
||||
|
||||
|
||||
def create_domain(
|
||||
*, domain_config: Dict[Union[str, int], Any], random_state: np.random.RandomState
|
||||
) -> Domain:
|
||||
return Domain(domain_config=domain_config, random_state=random_state)
|
||||
|
|
@ -0,0 +1,197 @@
|
|||
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
|
||||
|
||||
# ---- 额外的 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"])
|
||||
|
|
@ -0,0 +1,54 @@
|
|||
import numpy as np
|
||||
from skfem import Mesh
|
||||
|
||||
|
||||
def get_element_midpoints(mesh: Mesh, transpose: bool = True) -> np.ndarray:
|
||||
midpoints = np.mean(mesh.p[:, mesh.t], axis=1)
|
||||
return midpoints.T if transpose else midpoints
|
||||
|
||||
# 算三个顶点的mean/std/...
|
||||
def get_aggregation_per_element(
|
||||
solution: np.ndarray,
|
||||
element_indices: np.ndarray,
|
||||
aggregation_function_str: str = "mean",
|
||||
) -> np.ndarray:
|
||||
vals = solution[element_indices]
|
||||
if aggregation_function_str == "mean":
|
||||
return vals.mean(axis=1)
|
||||
elif aggregation_function_str == "std":
|
||||
return vals.std(axis=1)
|
||||
elif aggregation_function_str == "min":
|
||||
return vals.min(axis=1)
|
||||
elif aggregation_function_str == "max":
|
||||
return vals.max(axis=1)
|
||||
elif aggregation_function_str == "median":
|
||||
return np.median(vals, axis=1)
|
||||
raise ValueError(f"Unknown aggregation function: {aggregation_function_str}")
|
||||
|
||||
|
||||
# 计算三角形面积
|
||||
def get_triangle_areas_from_indices(
|
||||
positions: np.ndarray, triangle_indices: np.ndarray
|
||||
) -> np.ndarray:
|
||||
i0, i1, i2 = triangle_indices[:, 0], triangle_indices[:, 1], triangle_indices[:, 2]
|
||||
return np.abs(0.5 * (
|
||||
(positions[i1, 0] - positions[i0, 0]) * (positions[i2, 1] - positions[i0, 1])
|
||||
- (positions[i2, 0] - positions[i0, 0]) * (positions[i1, 1] - positions[i0, 1])
|
||||
))
|
||||
|
||||
|
||||
# penalty:\alpha的采样方式
|
||||
def sample_in_range(max_value: float, min_value: float, sampling_type: str) -> float:
|
||||
if sampling_type == "uniform":
|
||||
return np.random.uniform(min_value, max_value)
|
||||
elif sampling_type == "loguniform":
|
||||
return np.exp(np.random.uniform(np.log(min_value), np.log(max_value)))
|
||||
raise ValueError(f"Unknown sampling type: {sampling_type}")
|
||||
|
||||
|
||||
def construct_sizing_field_1d(x: np.ndarray, eps: float = 1e-4) -> np.ndarray:
|
||||
"""Softplus 激活 → 目标网格面积 (numpy 版)。"""
|
||||
def _softplus(x):
|
||||
return np.log1p(np.exp(np.clip(x, -50, 50)))
|
||||
x = np.atleast_1d(np.asarray(x, dtype=np.float64))
|
||||
return _softplus(x) + eps
|
||||
|
|
@ -0,0 +1,619 @@
|
|||
import copy
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
from skfem import Basis, ElementTriP1, Mesh, asm, solve
|
||||
from skfem.assembly import BilinearForm, FacetBasis, LinearForm
|
||||
from skfem.helpers import dot, grad
|
||||
|
||||
from .domain import create_domain
|
||||
from .fem_util import get_aggregation_per_element, get_element_midpoints
|
||||
|
||||
|
||||
class HelmholtzProblem:
|
||||
"""2D Helmholtz scattering FEM solver with Sommerfeld BC."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
fem_config: Dict[Union[str, int], Any],
|
||||
random_state: np.random.RandomState = np.random.RandomState(),
|
||||
):
|
||||
helmholtz_config = fem_config.get("helmholtz", {})
|
||||
|
||||
# ── 1. 波数 k ──
|
||||
wave_number_mode = helmholtz_config.get("wave_number_mode", "fixed")
|
||||
if wave_number_mode == "random_uniform":
|
||||
k_min = helmholtz_config.get("wave_number_min", 2.0)
|
||||
k_max = helmholtz_config.get("wave_number_max", 8.0)
|
||||
self._k = float(random_state.uniform(k_min, k_max))
|
||||
else:
|
||||
self._k = float(helmholtz_config.get("wave_number", 10.0))
|
||||
|
||||
# ── 2. 介质散射体参数 ──
|
||||
sc = helmholtz_config.get("scatterer", {})
|
||||
scatterer_mode = sc.get("mode", "fixed")
|
||||
|
||||
if scatterer_mode == "random_uniform":
|
||||
self._cx = float(
|
||||
random_state.uniform(sc.get("cx_min", 0.3), sc.get("cx_max", 0.7))
|
||||
)
|
||||
self._cy = float(
|
||||
random_state.uniform(sc.get("cy_min", 0.3), sc.get("cy_max", 0.7))
|
||||
)
|
||||
self._radius = float(
|
||||
random_state.uniform(
|
||||
sc.get("radius_min", 0.1), sc.get("radius_max", 0.25)
|
||||
)
|
||||
)
|
||||
self._eps_r = float(
|
||||
random_state.uniform(
|
||||
sc.get("eps_r_min", 2.0), sc.get("eps_r_max", 7.0)
|
||||
)
|
||||
)
|
||||
else:
|
||||
self._cx = float(sc.get("cx", 0.5))
|
||||
self._cy = float(sc.get("cy", 0.5))
|
||||
self._radius = float(sc.get("radius", 0.2))
|
||||
self._eps_r = float(sc.get("eps_r", 2.0))
|
||||
|
||||
# ── 3. 组装 FEM 双线性和线性形式 ──
|
||||
self._bilin_form = self._make_bilinear_form()
|
||||
self._lin_form_real = self._make_linear_form_real()
|
||||
self._lin_form_imag = self._make_linear_form_imag()
|
||||
|
||||
# ── 4. 初始化域（k^exponent 自适应网格密度 × domain area 线性缩放）──
|
||||
# exponent 和 k_ref 均可通过 helmholtz config 配置
|
||||
# exponent=2: P1 Helmholtz 理论最优 (污染误差 ∝ (kh)^2, N ∝ k^2)
|
||||
# exponent=1.5: 工程折中，避免高 k 初始过密
|
||||
# domain area 缩放: 保证不同域尺寸下每单位面积单元数一致 → h 不变
|
||||
domain_cfg = copy.deepcopy(fem_config.get("domain"))
|
||||
boundary = domain_cfg.get("boundary", [0, 0, 1, 1])
|
||||
domain_area = (boundary[2] - boundary[0]) * (boundary[3] - boundary[1])
|
||||
k_ref = helmholtz_config.get("k_ref", 6.0)
|
||||
k_exponent = helmholtz_config.get("k_exponent", 1.5)
|
||||
base_elements = domain_cfg.get("initial_num_elements", 400)
|
||||
scaled_elements = int(base_elements * (self._k / k_ref) ** k_exponent * domain_area)
|
||||
domain_cfg["initial_num_elements"] = max(scaled_elements, int(base_elements * domain_area))
|
||||
self._domain = create_domain(
|
||||
domain_config=domain_cfg,
|
||||
random_state=copy.deepcopy(random_state),
|
||||
)
|
||||
|
||||
# ── 4.5. 介质区前渐近区边缘约束 ──
|
||||
# 放宽 Nyquist (N=4) → 前渐近区边缘 (N=1~1.5)，赋予介质内初始网格基本相位解析能力
|
||||
# 约束: h_init ≤ λ_local / N，λ_local = 2π/(k√ε_r)
|
||||
# N=1.5 对应约 1.5 点/波长，刚好跨过渐近区门槛，不撑爆物理预算
|
||||
pre_asymptotic_N = helmholtz_config.get("pre_asymptotic_N", 1.5)
|
||||
pre_asymptotic_mesh = self._enforce_nyquist_in_dielectric(
|
||||
self._domain.initial_mesh, N=pre_asymptotic_N
|
||||
)
|
||||
self._domain.replace_initial_mesh(pre_asymptotic_mesh)
|
||||
|
||||
# ── 5. PDE 特征名称 ──
|
||||
pde_config = fem_config.get(fem_config.get("pde_type", "helmholtz"), {})
|
||||
self._element_feature_names = [
|
||||
name
|
||||
for name, include in pde_config.get("element_features", {}).items()
|
||||
if include
|
||||
]
|
||||
|
||||
# ── Public interface ─────────────────────────────────────
|
||||
|
||||
def mesh_to_basis(self, mesh: Mesh) -> Basis:
|
||||
return Basis(mesh, ElementTriP1())
|
||||
|
||||
def calculate_solution(self, basis: Basis, cache: bool = False) -> np.ndarray:
|
||||
K = asm(self._bilin_form, basis)
|
||||
f = asm(self._lin_form_real, basis) + 1j * asm(self._lin_form_imag, basis)
|
||||
|
||||
boundary_facets = basis.mesh.boundary_facets()
|
||||
facet_basis = FacetBasis(basis.mesh, basis.elem, facets=boundary_facets)
|
||||
|
||||
@BilinearForm
|
||||
def boundary_mass(u, v, w):
|
||||
return u * v
|
||||
|
||||
M_boundary = asm(boundary_mass, facet_basis)
|
||||
K_total = K.astype(np.complex128) - 1j * self._k * M_boundary
|
||||
u_scat = solve(K_total, f)
|
||||
|
||||
return u_scat
|
||||
|
||||
def get_error_estimate_per_element(
|
||||
self, basis: Basis, solution: np.ndarray
|
||||
) -> Dict[str, np.ndarray]:
|
||||
eps_r_arr = _compute_eps_r_at_midpoints(basis.mesh, self._cx, self._cy, self._radius, self._eps_r)
|
||||
return {"indicator": _compute_residual_indicator(basis.mesh, solution, k=self._k, eps_r=eps_r_arr)}
|
||||
|
||||
def element_features(self, mesh: Mesh, element_feature_names: List[str]) -> Optional[np.ndarray]:
|
||||
features_list = []
|
||||
if "epsilon_r" in element_feature_names:
|
||||
features_list.append(
|
||||
_compute_eps_r_at_midpoints(mesh, self._cx, self._cy, self._radius, self._eps_r)[:, None]
|
||||
)
|
||||
return np.concatenate(features_list, axis=1) if features_list else None
|
||||
|
||||
def _enforce_nyquist_in_dielectric(self, mesh: Mesh, N: float = 1.5, max_iter: int = 10) -> Mesh:
|
||||
"""Iteratively refine elements inside the dielectric until h_K ≤ λ_d/N.
|
||||
|
||||
λ_d = 2π/(k√ε_r) is the wavelength inside the dielectric.
|
||||
N=1.5 corresponds to the edge of the pre-asymptotic regime (~1.5 points
|
||||
per wavelength) — just enough for the wave field to exhibit basic phase
|
||||
resolution without exhausting the physical element budget. This relaxes
|
||||
the old Nyquist N=4 constraint, leaving headroom for the RL agent to
|
||||
selectively refine where residual indicators demand it.
|
||||
"""
|
||||
lambda_d = 2.0 * np.pi / (self._k * np.sqrt(self._eps_r))
|
||||
h_max = lambda_d / N
|
||||
|
||||
for _ in range(max_iter):
|
||||
i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
|
||||
x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
|
||||
x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
|
||||
x2, y2 = mesh.p[0, i2], mesh.p[1, i2]
|
||||
|
||||
e01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
|
||||
e12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
|
||||
e20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
|
||||
h_K = np.maximum(np.maximum(e01, e12), e20)
|
||||
|
||||
midpoints = np.mean(mesh.p[:, mesh.t], axis=1).T
|
||||
in_dielectric = (
|
||||
(midpoints[:, 0] - self._cx) ** 2
|
||||
+ (midpoints[:, 1] - self._cy) ** 2
|
||||
<= self._radius**2
|
||||
)
|
||||
|
||||
to_refine = np.where(in_dielectric & (h_K > h_max))[0]
|
||||
if len(to_refine) == 0:
|
||||
break
|
||||
mesh = mesh.refined(to_refine)
|
||||
|
||||
return mesh
|
||||
|
||||
# ── Properties ───────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def initial_mesh(self) -> Mesh:
|
||||
return self._domain.initial_mesh
|
||||
|
||||
@property
|
||||
def boundary_line_segments(self) -> np.ndarray:
|
||||
return self._domain.boundary_line_segments
|
||||
|
||||
|
||||
@staticmethod
|
||||
def project_to_scalar(values: np.ndarray) -> np.ndarray:
|
||||
return values
|
||||
|
||||
def additional_plots_from_mesh(self, mesh: Mesh) -> Dict:
|
||||
return {}
|
||||
|
||||
# ── FEM form assembly ────────────────────────────────────
|
||||
|
||||
def _eps_r_at_quad_points(self, x, y):
|
||||
in_cyl = (x - self._cx) ** 2 + (y - self._cy) ** 2 <= self._radius**2
|
||||
return np.where(in_cyl, self._eps_r, 1.0)
|
||||
|
||||
def _make_bilinear_form(self):
|
||||
k = self._k
|
||||
|
||||
@BilinearForm
|
||||
def bilin(u, v, w):
|
||||
eps_r = self._eps_r_at_quad_points(w.x[0], w.x[1])
|
||||
return dot(grad(u), grad(v)) - k**2 * eps_r * u * v
|
||||
|
||||
return bilin
|
||||
|
||||
def _make_linear_form_real(self):
|
||||
k = self._k
|
||||
|
||||
@LinearForm
|
||||
def lin(v, w):
|
||||
eps_r = self._eps_r_at_quad_points(w.x[0], w.x[1])
|
||||
return k**2 * (eps_r - 1.0) * np.cos(k * w.x[0]) * v
|
||||
|
||||
return lin
|
||||
|
||||
def _make_linear_form_imag(self):
|
||||
k = self._k
|
||||
|
||||
@LinearForm
|
||||
def lin(v, w):
|
||||
eps_r = self._eps_r_at_quad_points(w.x[0], w.x[1])
|
||||
return k**2 * (eps_r - 1.0) * np.sin(k * w.x[0]) * v
|
||||
|
||||
return lin
|
||||
|
||||
|
||||
# ── 辅助函数 ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _compute_eps_r_at_midpoints(
|
||||
mesh: Mesh,
|
||||
cx: float = 0.5,
|
||||
cy: float = 0.5,
|
||||
radius: float = 0.2,
|
||||
eps_r_in: float = 2.0,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
计算每个单元中点处的相对介电常数 ε_r。
|
||||
|
||||
判断单元中点是否落在介质圆柱内:
|
||||
- 在圆柱内 → ε_r = eps_r_in (如 2.0)
|
||||
- 在圆柱外 → ε_r = 1.0 (真空)
|
||||
|
||||
Returns:
|
||||
eps_r: shape (num_elements,)
|
||||
"""
|
||||
midpoints = get_element_midpoints(mesh)
|
||||
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_in, 1.0)
|
||||
|
||||
|
||||
def _compute_residual_indicator(
|
||||
mesh: Mesh,
|
||||
u_h: np.ndarray,
|
||||
k: float = 10.0,
|
||||
eps_r: Union[float, np.ndarray] = 1.0,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
基于残差的逐单元后验误差估计 — 无量纲化版本。
|
||||
|
||||
引入局部波数 k_local = k√ε_r 消除纯几何尺度 h 带来的特征偏差，
|
||||
使误差指示子反映"相位分辨率残差"而非"网格粗疏程度"。
|
||||
|
||||
P1 单元三项:
|
||||
1. r_int = (h_K/k_local)·√V_K · |k²ε_r·u_h + k²(ε_r-1)·u_inc|
|
||||
2. r_jump = √(½ Σ_{e∈∂K} (h_e/k_local)·|[[∇u_h·n]]|²)
|
||||
3. r_sbc = (h_bnd/k_local)·|∂u/∂n - i·k_local·u|
|
||||
|
||||
Returns:
|
||||
eta_elements: shape (num_elements,) 的逐单元误差指标
|
||||
"""
|
||||
n_elements = mesh.t.shape[1]
|
||||
eps_r = np.asarray(eps_r)
|
||||
k_local = k * np.sqrt(np.maximum(eps_r, 1.0))
|
||||
|
||||
# ── 1. 单元几何量 ──
|
||||
i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
|
||||
x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
|
||||
x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
|
||||
x2, y2 = mesh.p[0, i2], mesh.p[1, i2]
|
||||
|
||||
det_J = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
|
||||
element_areas = np.abs(det_J) / 2.0
|
||||
|
||||
edge_len_01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
|
||||
edge_len_12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
|
||||
edge_len_20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
|
||||
h_K = np.maximum(np.maximum(edge_len_01, edge_len_12), edge_len_20)
|
||||
|
||||
# ── 2. 梯度（常数，因为是 P1 单元）──
|
||||
u0, u1, u2 = u_h[i0], u_h[i1], u_h[i2]
|
||||
inv_det = np.divide(1.0, det_J, where=det_J != 0, out=np.zeros_like(det_J))
|
||||
du10, du20 = u1 - u0, u2 - u0
|
||||
|
||||
grad_x = ((y2 - y0) * du10 - (y1 - y0) * du20) * inv_det
|
||||
grad_y = (-(x2 - x0) * du10 + (x1 - x0) * du20) * inv_det
|
||||
grad_per_element = np.column_stack([grad_x, grad_y])
|
||||
|
||||
# ── 3. 单元内部残差 ──
|
||||
u_mid = (u0 + u1 + u2) / 3.0
|
||||
x_mid = (x0 + x1 + x2) / 3.0
|
||||
u_inc_mid = np.exp(1j * k * x_mid)
|
||||
f_mid = (k**2) * (eps_r - 1.0) * u_inc_mid
|
||||
r_mid = f_mid + (k**2) * eps_r * u_mid
|
||||
|
||||
cell_residual_sq = (h_K**2) * element_areas * np.abs(r_mid) ** 2 / (k_local ** 2)
|
||||
cell_residual_sq[element_areas < 1e-15] = 0.0
|
||||
|
||||
# ── 4. 内部边梯度跳变 ──
|
||||
interior_facets_idx = np.where(mesh.f2t[1] != -1)[0]
|
||||
elem_left = mesh.f2t[0, interior_facets_idx]
|
||||
elem_right = mesh.f2t[1, interior_facets_idx]
|
||||
|
||||
edges_p1 = mesh.p[:, mesh.facets[0, interior_facets_idx]].T
|
||||
edges_p2 = mesh.p[:, mesh.facets[1, interior_facets_idx]].T
|
||||
edge_vectors = edges_p2 - edges_p1
|
||||
h_e = np.linalg.norm(edge_vectors, axis=1)
|
||||
n_e = np.c_[edge_vectors[:, 1], -edge_vectors[:, 0]] / (h_e[:, None] + 1e-15)
|
||||
|
||||
grad_left = grad_per_element[elem_left]
|
||||
grad_right = grad_per_element[elem_right]
|
||||
jump_val = np.abs(np.sum((grad_left - grad_right) * n_e, axis=1))
|
||||
jump_val_sq = jump_val ** 2
|
||||
|
||||
jump_residual_sq = np.zeros(n_elements)
|
||||
np.add.at(jump_residual_sq, elem_left, 0.5 * h_e * jump_val_sq / k_local[elem_left])
|
||||
np.add.at(jump_residual_sq, elem_right, 0.5 * h_e * jump_val_sq / k_local[elem_right])
|
||||
|
||||
# ── 5. 合并 ──
|
||||
eta_sq = cell_residual_sq + jump_residual_sq
|
||||
|
||||
# ── 6. SBC 边界残差 ──
|
||||
boundary_facets_idx = np.where(mesh.f2t[1] == -1)[0]
|
||||
if len(boundary_facets_idx) > 0:
|
||||
bnd_elem = mesh.f2t[0, boundary_facets_idx]
|
||||
bnd_p1 = mesh.p[:, mesh.facets[0, boundary_facets_idx]].T
|
||||
bnd_p2 = mesh.p[:, mesh.facets[1, boundary_facets_idx]].T
|
||||
bnd_vectors = bnd_p2 - bnd_p1
|
||||
h_bnd = np.linalg.norm(bnd_vectors, axis=1)
|
||||
n_bnd = np.c_[bnd_vectors[:, 1], -bnd_vectors[:, 0]] / (h_bnd[:, None] + 1e-15)
|
||||
|
||||
grad_bnd = grad_per_element[bnd_elem]
|
||||
du_dn = np.sum(grad_bnd * n_bnd, axis=1)
|
||||
|
||||
if eps_r.ndim == 1:
|
||||
k_local = k * np.sqrt(np.maximum(eps_r[bnd_elem], 1.0))
|
||||
else:
|
||||
k_local = k
|
||||
|
||||
u_edge_mean = (
|
||||
u_h[mesh.facets[0, boundary_facets_idx]]
|
||||
+ u_h[mesh.facets[1, boundary_facets_idx]]
|
||||
) / 2.0
|
||||
sbc_residual = du_dn - 1j * k_local * u_edge_mean
|
||||
sbc_residual_sq = (h_bnd ** 2) * np.abs(sbc_residual) ** 2 / (k_local ** 2)
|
||||
np.add.at(eta_sq, bnd_elem, sbc_residual_sq)
|
||||
|
||||
eta_sq = np.maximum(eta_sq, 0.0)
|
||||
return np.sqrt(eta_sq)
|
||||
|
||||
|
||||
def _compute_residual_components(
|
||||
mesh: Mesh,
|
||||
u_h: np.ndarray,
|
||||
k: float = 10.0,
|
||||
eps_r: Union[float, np.ndarray] = 1.0,
|
||||
apply_log: bool = True,
|
||||
) -> Dict[str, np.ndarray]:
|
||||
"""
|
||||
计算逐单元的三项 PDE 物理残差（分离版，无量纲化）。
|
||||
|
||||
引入 k_local = k√ε_r 消除几何尺度偏差，使 GNN 跨介质公平感知"相位分辨率残差"。
|
||||
保留源项信息（k²(ε_r-1)·u_inc），确保极粗网格下介质内部巨大物理激励仍可被网络捕捉。
|
||||
|
||||
P1 单元返回:
|
||||
internal_residual: (h_K/k_local)·√V_i · |k²ε_r·u + k²(ε_r-1)·u_inc|
|
||||
gradient_jump: √(½ Σ_{e∈∂K_i} (h_e/k_local)·|[[∇u·n]]|²)
|
||||
sbc_residual: (h_bnd/k_local)·|∂u/∂n - i·k_local·u|
|
||||
element_areas: 单元面积
|
||||
is_sbc_boundary: 该单元是否与 SBC 边界相邻 (0/1)
|
||||
|
||||
Args:
|
||||
apply_log: True → log10 压缩（喂 GNN）；False → 原始值（喂 reward）
|
||||
"""
|
||||
n_elements = mesh.t.shape[1]
|
||||
eps_r = np.asarray(eps_r)
|
||||
k_local = k * np.sqrt(np.maximum(eps_r, 1.0))
|
||||
|
||||
# ── 1. 单元几何量 ──
|
||||
i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
|
||||
x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
|
||||
x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
|
||||
x2, y2 = mesh.p[0, i2], mesh.p[1, i2]
|
||||
|
||||
det_J = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
|
||||
element_areas = np.abs(det_J) / 2.0
|
||||
|
||||
edge_len_01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
|
||||
edge_len_12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
|
||||
edge_len_20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
|
||||
h_K = np.maximum(np.maximum(edge_len_01, edge_len_12), edge_len_20)
|
||||
|
||||
# ── 2. 梯度（常数，因为是 P1 单元）──
|
||||
u0, u1, u2 = u_h[i0], u_h[i1], u_h[i2]
|
||||
inv_det = np.divide(1.0, det_J, where=det_J != 0, out=np.zeros_like(det_J))
|
||||
du10, du20 = u1 - u0, u2 - u0
|
||||
|
||||
grad_x = ((y2 - y0) * du10 - (y1 - y0) * du20) * inv_det
|
||||
grad_y = (-(x2 - x0) * du10 + (x1 - x0) * du20) * inv_det
|
||||
grad_per_element = np.column_stack([grad_x, grad_y])
|
||||
|
||||
# P1 单元内部残差: ∇²u_h = 0（线性元二阶导为零），故仅含反应项
|
||||
# 完整强形式: |∇²u + k²·ε_r·u + k²·(ε_r-1)·u_inc|
|
||||
# 对 P1: ∇²u_h ≡ 0 → 残差 = |k²·ε_r·u + k²·(ε_r-1)·u_inc|
|
||||
u_mid = (u0 + u1 + u2) / 3.0
|
||||
x_mid = (x0 + x1 + x2) / 3.0
|
||||
u_inc_mid = np.exp(1j * k * x_mid)
|
||||
f_mid = (k**2) * (eps_r - 1.0) * u_inc_mid
|
||||
r_mid = f_mid + (k**2) * eps_r * u_mid
|
||||
internal_residual = (h_K / k_local) * np.sqrt(element_areas) * np.abs(r_mid)
|
||||
internal_residual[element_areas < 1e-15] = 0.0
|
||||
|
||||
# ── 4. 内部边梯度跳变 (逐单元) ──
|
||||
interior_facets_idx = np.where(mesh.f2t[1] != -1)[0]
|
||||
elem_left = mesh.f2t[0, interior_facets_idx]
|
||||
elem_right = mesh.f2t[1, interior_facets_idx]
|
||||
|
||||
edges_p1 = mesh.p[:, mesh.facets[0, interior_facets_idx]].T
|
||||
edges_p2 = mesh.p[:, mesh.facets[1, interior_facets_idx]].T
|
||||
edge_vectors = edges_p2 - edges_p1
|
||||
h_e = np.linalg.norm(edge_vectors, axis=1)
|
||||
n_e = np.c_[edge_vectors[:, 1], -edge_vectors[:, 0]] / (h_e[:, None] + 1e-15)
|
||||
|
||||
grad_left = grad_per_element[elem_left]
|
||||
grad_right = grad_per_element[elem_right]
|
||||
jump_val = np.abs(np.sum((grad_left - grad_right) * n_e, axis=1))
|
||||
|
||||
gradient_jump = np.zeros(n_elements, dtype=np.float64)
|
||||
jump_sq_per_edge = jump_val ** 2
|
||||
np.add.at(gradient_jump, elem_left, 0.5 * h_e * jump_sq_per_edge / k_local[elem_left])
|
||||
np.add.at(gradient_jump, elem_right, 0.5 * h_e * jump_sq_per_edge / k_local[elem_right])
|
||||
gradient_jump = np.sqrt(gradient_jump)
|
||||
|
||||
# ── 5. SBC 边界残差 + 边界标记 ──
|
||||
sbc_residual = np.zeros(n_elements, dtype=np.float64)
|
||||
is_sbc_boundary = np.zeros(n_elements, dtype=np.float32)
|
||||
boundary_facets_idx = np.where(mesh.f2t[1] == -1)[0]
|
||||
if len(boundary_facets_idx) > 0:
|
||||
bnd_elem = mesh.f2t[0, boundary_facets_idx]
|
||||
bnd_p1 = mesh.p[:, mesh.facets[0, boundary_facets_idx]].T
|
||||
bnd_p2 = mesh.p[:, mesh.facets[1, boundary_facets_idx]].T
|
||||
bnd_vectors = bnd_p2 - bnd_p1
|
||||
h_bnd = np.linalg.norm(bnd_vectors, axis=1)
|
||||
n_bnd = np.c_[bnd_vectors[:, 1], -bnd_vectors[:, 0]] / (h_bnd[:, None] + 1e-15)
|
||||
|
||||
grad_bnd = grad_per_element[bnd_elem]
|
||||
du_dn = np.sum(grad_bnd * n_bnd, axis=1)
|
||||
|
||||
if eps_r.ndim == 1:
|
||||
k_local = k * np.sqrt(np.maximum(eps_r[bnd_elem], 1.0))
|
||||
else:
|
||||
k_local = k
|
||||
|
||||
u_edge_mean = (
|
||||
u_h[mesh.facets[0, boundary_facets_idx]]
|
||||
+ u_h[mesh.facets[1, boundary_facets_idx]]
|
||||
) / 2.0
|
||||
sbc_val = np.abs(du_dn - 1j * k_local * u_edge_mean)
|
||||
np.add.at(sbc_residual, bnd_elem, (h_bnd / k_local) * sbc_val)
|
||||
is_sbc_boundary[bnd_elem] = 1.0
|
||||
|
||||
# ── 对数预处理：压缩跨数量级动态范围（仅 GNN 特征需要）──
|
||||
if apply_log:
|
||||
_log_eps = 1e-8
|
||||
internal_residual = np.log10(np.maximum(internal_residual, _log_eps))
|
||||
gradient_jump = np.log10(np.maximum(gradient_jump, _log_eps))
|
||||
sbc_residual = np.log10(np.maximum(sbc_residual, _log_eps))
|
||||
|
||||
return {
|
||||
"internal_residual": internal_residual.astype(np.float32),
|
||||
"gradient_jump": gradient_jump.astype(np.float32),
|
||||
"sbc_residual": sbc_residual.astype(np.float32),
|
||||
"element_areas": element_areas.astype(np.float32),
|
||||
"is_sbc_boundary": is_sbc_boundary,
|
||||
}
|
||||
|
||||
|
||||
def _compute_residual_density(
|
||||
mesh: Mesh,
|
||||
u_h: np.ndarray,
|
||||
k: float = 10.0,
|
||||
eps_r: Union[float, np.ndarray] = 1.0,
|
||||
) -> Dict[str, np.ndarray]:
|
||||
"""
|
||||
Compute intensive (h-free) residual density components for reward.
|
||||
|
||||
Unlike _compute_residual_components which includes h-scaling
|
||||
(h_K·√V, h_e·|jump|, h_bnd·|sbc|), this returns the raw PDE residuals
|
||||
that are independent of element size — true "error densities".
|
||||
|
||||
Returns:
|
||||
rho_int: |k²·ε_r·u + k²·(ε_r-1)·u_inc| per element
|
||||
rho_jump: √(mean_{e∈∂K_int} |[[∇u·n]]|²) per element
|
||||
rho_sbc: √(mean_{e∈∂K∩Γ_sbc} |∂u/∂n - i·k·u|²) per element
|
||||
"""
|
||||
n_elements = mesh.t.shape[1]
|
||||
eps_r = np.asarray(eps_r)
|
||||
|
||||
# ── 1. element geometry ──
|
||||
i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
|
||||
x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
|
||||
x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
|
||||
x2, y2 = mesh.p[0, i2], mesh.p[1, i2]
|
||||
|
||||
det_J = (x1 - x0) * (y2 - y0) - (x2 - x0) * (y1 - y0)
|
||||
|
||||
# ── 2. gradient (constant per P1 element) ──
|
||||
u0, u1, u2 = u_h[i0], u_h[i1], u_h[i2]
|
||||
inv_det = np.divide(1.0, det_J, where=det_J != 0, out=np.zeros_like(det_J))
|
||||
du10, du20 = u1 - u0, u2 - u0
|
||||
|
||||
grad_x = ((y2 - y0) * du10 - (y1 - y0) * du20) * inv_det
|
||||
grad_y = (-(x2 - x0) * du10 + (x1 - x0) * du20) * inv_det
|
||||
grad_per_element = np.column_stack([grad_x, grad_y])
|
||||
|
||||
# ── 3. interior residual density: |k²·ε_r·u_mid + k²·(ε_r-1)·u_inc_mid| ──
|
||||
u_mid = (u0 + u1 + u2) / 3.0
|
||||
x_mid = (x0 + x1 + x2) / 3.0
|
||||
u_inc_mid = np.exp(1j * k * x_mid)
|
||||
r_mid = (k**2) * eps_r * u_mid + (k**2) * (eps_r - 1.0) * u_inc_mid
|
||||
rho_int = np.abs(r_mid)
|
||||
|
||||
# ── 4. gradient jump density: √(mean |[[∇u·n]]|²) per element ──
|
||||
interior_facets_idx = np.where(mesh.f2t[1] != -1)[0]
|
||||
elem_left = mesh.f2t[0, interior_facets_idx]
|
||||
elem_right = mesh.f2t[1, interior_facets_idx]
|
||||
|
||||
edges_p1 = mesh.p[:, mesh.facets[0, interior_facets_idx]].T
|
||||
edges_p2 = mesh.p[:, mesh.facets[1, interior_facets_idx]].T
|
||||
edge_vectors = edges_p2 - edges_p1
|
||||
h_e = np.linalg.norm(edge_vectors, axis=1)
|
||||
n_e = np.c_[edge_vectors[:, 1], -edge_vectors[:, 0]] / (h_e[:, None] + 1e-15)
|
||||
|
||||
grad_left = grad_per_element[elem_left]
|
||||
grad_right = grad_per_element[elem_right]
|
||||
jump_val_sq = np.abs(np.sum((grad_left - grad_right) * n_e, axis=1)) ** 2
|
||||
|
||||
jump_sq_sum = np.zeros(n_elements, dtype=np.float64)
|
||||
jump_count = np.zeros(n_elements, dtype=np.float64)
|
||||
np.add.at(jump_sq_sum, elem_left, jump_val_sq)
|
||||
np.add.at(jump_sq_sum, elem_right, jump_val_sq)
|
||||
np.add.at(jump_count, elem_left, 1)
|
||||
np.add.at(jump_count, elem_right, 1)
|
||||
|
||||
rho_jump = np.zeros(n_elements, dtype=np.float64)
|
||||
mask_jump = jump_count > 0
|
||||
rho_jump[mask_jump] = np.sqrt(jump_sq_sum[mask_jump] / jump_count[mask_jump])
|
||||
|
||||
# ── 5. SBC boundary density: √(mean |∂u/∂n - i·k·u|²) per element ──
|
||||
rho_sbc = np.zeros(n_elements, dtype=np.float64)
|
||||
boundary_facets_idx = np.where(mesh.f2t[1] == -1)[0]
|
||||
if len(boundary_facets_idx) > 0:
|
||||
bnd_elem = mesh.f2t[0, boundary_facets_idx]
|
||||
bnd_p1 = mesh.p[:, mesh.facets[0, boundary_facets_idx]].T
|
||||
bnd_p2 = mesh.p[:, mesh.facets[1, boundary_facets_idx]].T
|
||||
bnd_vectors = bnd_p2 - bnd_p1
|
||||
h_bnd = np.linalg.norm(bnd_vectors, axis=1)
|
||||
n_bnd = np.c_[bnd_vectors[:, 1], -bnd_vectors[:, 0]] / (h_bnd[:, None] + 1e-15)
|
||||
|
||||
grad_bnd = grad_per_element[bnd_elem]
|
||||
du_dn = np.sum(grad_bnd * n_bnd, axis=1)
|
||||
|
||||
if eps_r.ndim == 1:
|
||||
k_local = k * np.sqrt(np.maximum(eps_r[bnd_elem], 1.0))
|
||||
else:
|
||||
k_local = k
|
||||
|
||||
u_edge_mean = (
|
||||
u_h[mesh.facets[0, boundary_facets_idx]]
|
||||
+ u_h[mesh.facets[1, boundary_facets_idx]]
|
||||
) / 2.0
|
||||
sbc_val_sq = np.abs(du_dn - 1j * k_local * u_edge_mean) ** 2
|
||||
|
||||
sbc_sq_sum = np.zeros(n_elements, dtype=np.float64)
|
||||
sbc_count = np.zeros(n_elements, dtype=np.float64)
|
||||
np.add.at(sbc_sq_sum, bnd_elem, sbc_val_sq)
|
||||
np.add.at(sbc_count, bnd_elem, 1)
|
||||
|
||||
mask_sbc = sbc_count > 0
|
||||
rho_sbc[mask_sbc] = np.sqrt(sbc_sq_sum[mask_sbc] / sbc_count[mask_sbc])
|
||||
|
||||
return {
|
||||
"rho_int": rho_int.astype(np.float64),
|
||||
"rho_jump": rho_jump.astype(np.float64),
|
||||
"rho_sbc": rho_sbc.astype(np.float64),
|
||||
}
|
||||
|
||||
|
||||
# ── 工厂函数 ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def create_helmholtz_problem(
|
||||
*, fem_config: Dict[Union[str, int], Any], random_state: np.random.RandomState
|
||||
) -> HelmholtzProblem:
|
||||
"""
|
||||
创建 Helmholtz 问题实例。
|
||||
|
||||
Args:
|
||||
fem_config: FEM 配置字典
|
||||
random_state: 随机状态
|
||||
|
||||
Returns:
|
||||
HelmholtzProblem 实例
|
||||
"""
|
||||
return HelmholtzProblem(fem_config=fem_config, random_state=random_state)
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,202 @@
|
|||
"""2D Mie scattering analytical solution for a dielectric cylinder (TM polarization).
|
||||
|
||||
Computes the exact scattered and total fields for a circular dielectric cylinder
|
||||
under plane-wave illumination u_inc = exp(i·k0·x).
|
||||
|
||||
Line-by-line translation of the validated MATLAB reference (result/mie.py).
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from scipy.special import jv, hankel1
|
||||
from typing import Optional, Tuple
|
||||
|
||||
|
||||
def mie_scattered_field(
|
||||
points: np.ndarray,
|
||||
k0: float,
|
||||
eps_r: float,
|
||||
radius: float,
|
||||
cx: float = 0.5,
|
||||
cy: float = 0.5,
|
||||
) -> np.ndarray:
|
||||
"""Compute the scattered E_z field at arbitrary query points.
|
||||
|
||||
The scattered field is u_scat = u_total − u_inc, valid both inside and
|
||||
outside the cylinder. This matches the FEM scattered-field formulation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
points : (N, 2) np.ndarray — (x, y) coordinates
|
||||
k0 : float — vacuum wavenumber
|
||||
eps_r : float — relative permittivity
|
||||
radius : float — cylinder radius
|
||||
cx, cy : float — cylinder centre
|
||||
|
||||
Returns
|
||||
-------
|
||||
E_scat : (N,) np.complex128
|
||||
"""
|
||||
m = np.sqrt(eps_r)
|
||||
k1 = k0 * m # wavenumber inside cylinder
|
||||
x_size = k0 * radius # size parameter
|
||||
|
||||
# ── polar coordinates relative to cylinder centre ──
|
||||
dx = points[:, 0] - cx
|
||||
dy = points[:, 1] - cy
|
||||
R = np.sqrt(dx * dx + dy * dy)
|
||||
Phi = np.arctan2(dy, dx) # [-π, π], matches MATLAB cart2pol
|
||||
|
||||
# ── Wiscombe truncation (matches MATLAB round(…)) ──
|
||||
N_trunc = int(np.round(x_size + 4.05 * x_size ** (1.0 / 3.0) + 2))
|
||||
N_trunc = max(N_trunc, 3)
|
||||
|
||||
E_scat = np.zeros(len(points), dtype=np.complex128)
|
||||
E_int = np.zeros(len(points), dtype=np.complex128)
|
||||
|
||||
for n in range(-N_trunc, N_trunc + 1):
|
||||
# boundary values — matches MATLAB besselj / besselh(…, 1, …)
|
||||
J_nx = jv(n, x_size)
|
||||
J_nmx = jv(n, k1 * radius)
|
||||
H_nx = hankel1(n, x_size)
|
||||
|
||||
# derivatives via recurrence Z'_n = ½ (Z_{n-1} − Z_{n+1})
|
||||
J_nx_p = 0.5 * (jv(n - 1, x_size) - jv(n + 1, x_size))
|
||||
J_nmx_p = 0.5 * (jv(n - 1, k1 * radius) - jv(n + 1, k1 * radius))
|
||||
H_nx_p = 0.5 * (hankel1(n - 1, x_size) - hankel1(n + 1, x_size))
|
||||
|
||||
# TM scattering coefficient a_n
|
||||
num_a = m * J_nx * J_nmx_p - J_nx_p * J_nmx
|
||||
den_a = J_nmx * H_nx_p - m * J_nmx_p * H_nx
|
||||
a_n = num_a / den_a
|
||||
|
||||
# internal coefficient c_n
|
||||
num_c = J_nx * H_nx_p - J_nx_p * H_nx # Wronskian (2i/(π x) from theory)
|
||||
c_n = num_c / den_a
|
||||
|
||||
# phase factor iⁿ · exp(i·n·φ)
|
||||
phase = (1j) ** n * np.exp(1j * n * Phi)
|
||||
|
||||
# scattered field (valid outside the cylinder)
|
||||
out = R >= radius
|
||||
if out.any():
|
||||
E_scat[out] += a_n * hankel1(n, k0 * R[out]) * phase[out]
|
||||
|
||||
# internal total field (valid inside the cylinder)
|
||||
inside = R < radius
|
||||
if inside.any():
|
||||
E_int[inside] += c_n * jv(n, k1 * R[inside]) * phase[inside]
|
||||
|
||||
# phase reference at cylinder centre (matches MATLAB phase_shift)
|
||||
phase_shift = np.exp(1j * k0 * cx)
|
||||
E_scat *= phase_shift
|
||||
E_int *= phase_shift
|
||||
|
||||
# ── scattered field inside cylinder = internal total − incident ──
|
||||
E_inc = np.exp(1j * k0 * points[:, 0])
|
||||
inside = R < radius
|
||||
if inside.any():
|
||||
E_scat[inside] = E_int[inside] - E_inc[inside]
|
||||
|
||||
return E_scat
|
||||
|
||||
|
||||
def mie_total_field(
|
||||
points: np.ndarray,
|
||||
k0: float,
|
||||
eps_r: float,
|
||||
radius: float,
|
||||
cx: float = 0.5,
|
||||
cy: float = 0.5,
|
||||
) -> np.ndarray:
|
||||
"""Compute the total E_z field.
|
||||
|
||||
Outside: u_inc + u_scat
|
||||
Inside: internal field (refracted wave)
|
||||
"""
|
||||
m = np.sqrt(eps_r)
|
||||
k1 = k0 * m
|
||||
x_size = k0 * radius
|
||||
|
||||
dx = points[:, 0] - cx
|
||||
dy = points[:, 1] - cy
|
||||
R = np.sqrt(dx * dx + dy * dy)
|
||||
Phi = np.arctan2(dy, dx)
|
||||
|
||||
N_trunc = int(np.round(x_size + 4.05 * x_size ** (1.0 / 3.0) + 2))
|
||||
N_trunc = max(N_trunc, 3)
|
||||
|
||||
E_scat = np.zeros(len(points), dtype=np.complex128)
|
||||
E_int = np.zeros(len(points), dtype=np.complex128)
|
||||
|
||||
for n in range(-N_trunc, N_trunc + 1):
|
||||
J_nx = jv(n, x_size)
|
||||
J_nmx = jv(n, k1 * radius)
|
||||
H_nx = hankel1(n, x_size)
|
||||
|
||||
J_nx_p = 0.5 * (jv(n - 1, x_size) - jv(n + 1, x_size))
|
||||
J_nmx_p = 0.5 * (jv(n - 1, k1 * radius) - jv(n + 1, k1 * radius))
|
||||
H_nx_p = 0.5 * (hankel1(n - 1, x_size) - hankel1(n + 1, x_size))
|
||||
|
||||
num_a = m * J_nx * J_nmx_p - J_nx_p * J_nmx
|
||||
den_a = J_nmx * H_nx_p - m * J_nmx_p * H_nx
|
||||
a_n = num_a / den_a
|
||||
|
||||
num_c = J_nx * H_nx_p - J_nx_p * H_nx
|
||||
c_n = num_c / den_a
|
||||
|
||||
phase = (1j) ** n * np.exp(1j * n * Phi)
|
||||
|
||||
out = R >= radius
|
||||
if out.any():
|
||||
E_scat[out] += a_n * hankel1(n, k0 * R[out]) * phase[out]
|
||||
|
||||
inside = R < radius
|
||||
if inside.any():
|
||||
E_int[inside] += c_n * jv(n, k1 * R[inside]) * phase[inside]
|
||||
|
||||
phase_shift = np.exp(1j * k0 * cx)
|
||||
E_scat *= phase_shift
|
||||
E_int *= phase_shift
|
||||
|
||||
E_inc = np.exp(1j * k0 * points[:, 0])
|
||||
|
||||
E_total = np.zeros(len(points), dtype=np.complex128)
|
||||
E_total[R >= radius] = E_inc[R >= radius] + E_scat[R >= radius]
|
||||
E_total[R < radius] = E_int[R < radius]
|
||||
|
||||
return E_total
|
||||
|
||||
|
||||
def mie_grid_solution(
|
||||
k0: float,
|
||||
eps_r: float,
|
||||
radius: float,
|
||||
cx: float = 0.5,
|
||||
cy: float = 0.5,
|
||||
x_range: Tuple[float, float] = (0.0, 1.0),
|
||||
y_range: Tuple[float, float] = (0.0, 1.0),
|
||||
Nx: int = 400,
|
||||
Ny: int = 400,
|
||||
) -> dict:
|
||||
"""Compute Mie solution on a regular grid (for plotting / visual checks).
|
||||
|
||||
Returns a dict with keys: X, Y, R, Phi, E_inc, E_scat, E_total.
|
||||
"""
|
||||
x_vec = np.linspace(x_range[0], x_range[1], Nx)
|
||||
y_vec = np.linspace(y_range[0], y_range[1], Ny)
|
||||
X, Y = np.meshgrid(x_vec, y_vec)
|
||||
|
||||
points = np.column_stack([X.ravel(), Y.ravel()])
|
||||
dx = points[:, 0] - cx
|
||||
dy = points[:, 1] - cy
|
||||
R = np.sqrt(dx * dx + dy * dy).reshape(Ny, Nx)
|
||||
Phi = np.arctan2(dy, dx).reshape(Ny, Nx)
|
||||
|
||||
E_inc = np.exp(1j * k0 * X)
|
||||
E_scat = mie_scattered_field(points, k0, eps_r, radius, cx, cy).reshape(Ny, Nx)
|
||||
E_total = mie_total_field(points, k0, eps_r, radius, cx, cy).reshape(Ny, Nx)
|
||||
|
||||
return {
|
||||
"X": X, "Y": Y, "R": R, "Phi": Phi,
|
||||
"E_inc": E_inc, "E_scat": E_scat, "E_total": E_total,
|
||||
}
|
||||
|
|
@ -0,0 +1,92 @@
|
|||
"""
|
||||
环境层通用工具
|
||||
=============
|
||||
提供数组拼接、索引采样、tensor→numpy 转换等辅助功能。
|
||||
"""
|
||||
|
||||
from typing import Dict, Iterable, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
from numpy import ndarray
|
||||
from torch import Tensor
|
||||
from torch_geometric.data.data import BaseData
|
||||
|
||||
|
||||
def save_concatenate(
|
||||
arrays: Iterable[np.ndarray], *args, **kwargs
|
||||
) -> Optional[np.ndarray]:
|
||||
"""
|
||||
安全拼接多个数组。自动过滤 None 值，空列表返回 None。
|
||||
|
||||
Args:
|
||||
arrays: 要拼接的数组列表（可能包含 None）
|
||||
|
||||
Returns:
|
||||
拼接后的数组；若全为 None 则返回 None
|
||||
|
||||
Example:
|
||||
>>> result = save_concatenate([arr1, None, arr2], axis=1)
|
||||
"""
|
||||
arrays = [array for array in arrays if array is not None]
|
||||
if len(arrays) == 0:
|
||||
return None
|
||||
return np.concatenate(arrays, *args, **kwargs)
|
||||
|
||||
|
||||
class IndexSampler:
|
||||
"""
|
||||
随机索引采样器 — 用于循环缓冲区中随机抽取 PDE 实例。
|
||||
|
||||
内部维护一个随机排列的索引数组，每次调用 next() 返回一个索引。
|
||||
遍历完所有索引后自动重新洗牌。
|
||||
|
||||
Example:
|
||||
>>> sampler = IndexSampler(100, np.random.RandomState(42))
|
||||
>>> idx = sampler.next() # 随机抽取一个索引
|
||||
"""
|
||||
|
||||
def __init__(self, size: int, random_state: np.random.RandomState):
|
||||
self._size = size
|
||||
self._indices = np.arange(size)
|
||||
self._random_state = random_state
|
||||
self._reset()
|
||||
|
||||
def next(self) -> int:
|
||||
"""返回下一个随机索引，到底后自动洗牌重排。"""
|
||||
if self._position == self._size:
|
||||
self._reset()
|
||||
index = self._indices[self._position]
|
||||
self._position += 1
|
||||
return index
|
||||
|
||||
def _reset(self):
|
||||
self._position = 0
|
||||
self._random_state.shuffle(self._indices)
|
||||
|
||||
def __len__(self):
|
||||
return self._size
|
||||
|
||||
|
||||
def detach(
|
||||
tensor: Union[Tensor, Dict[str, Tensor], List[Tensor]],
|
||||
) -> Union[ndarray, Dict[str, ndarray], List[ndarray], BaseData]:
|
||||
"""
|
||||
将 PyTorch tensor 安全转换为 numpy 数组（自动处理 GPU→CPU）。
|
||||
|
||||
Args:
|
||||
tensor: PyTorch tensor、tensor 字典或 tensor 列表
|
||||
|
||||
Returns:
|
||||
对应的 numpy 数组
|
||||
|
||||
Example:
|
||||
>>> action_np = detach(actions_tensor) # → np.ndarray
|
||||
"""
|
||||
if isinstance(tensor, dict):
|
||||
return {key: detach(value) for key, value in tensor.items()}
|
||||
elif isinstance(tensor, list):
|
||||
return [detach(value) for value in tensor]
|
||||
if tensor.is_cuda:
|
||||
return tensor.cpu().detach().numpy()
|
||||
else:
|
||||
return tensor.detach().numpy()
|
||||
|
|
@ -0,0 +1,69 @@
|
|||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import plotly.graph_objects as go
|
||||
from plotly.basedatatypes import BaseTraceType
|
||||
from skfem import Mesh
|
||||
|
||||
|
||||
# 将网格与标量场转为 plotly 三角形 traces + 布局，供 RL 环境实时渲染
|
||||
def get_plotly_mesh_traces_and_layout(
|
||||
mesh: Mesh,
|
||||
scalars: np.ndarray,
|
||||
title: str = "Mesh",
|
||||
mesh_dimension: int = 2,
|
||||
boundary: Optional[np.ndarray] = None,
|
||||
) -> Tuple[List[BaseTraceType], Dict[str, Any]]:
|
||||
vertices = mesh.p
|
||||
triangles = mesh.t
|
||||
n_elements = triangles.shape[1]
|
||||
s = np.asarray(scalars, dtype=np.float64).flatten()
|
||||
|
||||
x_tri = vertices[0, triangles].T
|
||||
y_tri = vertices[1, triangles].T
|
||||
intensity_tri = s[triangles].T
|
||||
|
||||
vmin, vmax = s.min(), s.max()
|
||||
|
||||
traces = []
|
||||
for elem_idx in range(n_elements):
|
||||
x_e, y_e, s_e = x_tri[elem_idx], y_tri[elem_idx], intensity_tri[elem_idx]
|
||||
traces.append(go.Scatter(
|
||||
x=x_e.tolist() + [x_e[0]],
|
||||
y=y_e.tolist() + [y_e[0]],
|
||||
mode="lines",
|
||||
fill="toself",
|
||||
fillcolor=_get_color(float(np.mean(s_e)), vmin, vmax),
|
||||
line=dict(color="black", width=0.5),
|
||||
showlegend=False,
|
||||
hoverinfo="skip",
|
||||
))
|
||||
|
||||
if traces:
|
||||
traces[0].marker = dict(
|
||||
color=s.min(), colorscale="RdBu_r", showscale=True,
|
||||
colorbar=dict(title="Solution"),
|
||||
)
|
||||
|
||||
layout = {
|
||||
"title": title,
|
||||
"xaxis": {"title": "x", "scaleanchor": "y"},
|
||||
"yaxis": {"title": "y"},
|
||||
"showlegend": False,
|
||||
}
|
||||
if boundary is not None:
|
||||
layout["xaxis"]["range"] = [boundary[0], boundary[2]]
|
||||
layout["yaxis"]["range"] = [boundary[1], boundary[3]]
|
||||
|
||||
return traces, layout
|
||||
|
||||
|
||||
# 标量值 → matplotlib RdBu_r 色表映射的 RGBA 字符串
|
||||
def _get_color(value: float, vmin: float, vmax: float) -> str:
|
||||
import matplotlib.cm as cm
|
||||
import matplotlib.colors as mcolors
|
||||
|
||||
norm = mcolors.Normalize(vmin=vmin, vmax=vmax)
|
||||
rgba = cm.RdBu_r(norm(value))
|
||||
r, g, b, a = rgba
|
||||
return f"rgba({int(r * 255)},{int(g * 255)},{int(b * 255)},{a:.2f})"
|
||||
|
|
@ -0,0 +1,94 @@
|
|||
clc; clear; close all;
|
||||
|
||||
% ================= 1. 物理参数定义 =================
|
||||
r = 0.1; % 圆柱半径
|
||||
eps_r = 5.0; % 相对介电常数
|
||||
m = sqrt(eps_r); % 相对折射率 m = ~1.414
|
||||
k0 = 6; % 背景真空中波数 (k=6)
|
||||
k1 = k0 * m; % 圆柱内部波数
|
||||
x_size = k0 * r; % 尺寸参数 x = k0*a
|
||||
|
||||
% ================= 2. 计算域网格设置 =================
|
||||
x_range = 1;
|
||||
y_range = 1;
|
||||
Nx = 500;
|
||||
Ny = 500;
|
||||
x_vec = linspace(0, x_range, Nx);
|
||||
y_vec = linspace(0, y_range, Ny);
|
||||
[X, Y] = meshgrid(x_vec, y_vec);
|
||||
|
||||
xc = 0.5; yc = 0.5;
|
||||
[Phi, R] = cart2pol(X - xc, Y - yc); % 转换为极坐标
|
||||
|
||||
% ================= 3. 场初始化 =================
|
||||
E_scat = zeros(size(X)); % 散射场
|
||||
E_int = zeros(size(X)); % 内部场
|
||||
|
||||
% Wiscombe 截断准则（决定级数展开需要算到第几阶）
|
||||
N_trunc = round(x_size + 4.05 * x_size^(1/3) + 2);
|
||||
|
||||
% ================= 4. 2D Mie 级数展开计算 =================
|
||||
% 2D 圆柱级数从 -N 到 +N
|
||||
for n = -N_trunc : N_trunc
|
||||
|
||||
% 边界处的贝塞尔函数值
|
||||
J_nx = besselj(n, x_size);
|
||||
J_nmx = besselj(n, k1 * r);
|
||||
H_nx = besselh(n, 1, x_size);
|
||||
|
||||
% 边界处的导数值 (利用递推公式 Z_n' = 0.5 * (Z_{n-1} - Z_{n+1}))
|
||||
J_nx_p = 0.5 * (besselj(n-1, x_size) - besselj(n+1, x_size));
|
||||
J_nmx_p = 0.5 * (besselj(n-1, k1*r) - besselj(n+1, k1*r));
|
||||
H_nx_p = 0.5 * (besselh(n-1, 1, x_size) - besselh(n+1, 1, x_size));
|
||||
|
||||
% 计算 TM 偏振下的散射系数 a_n (对应 E_z)
|
||||
num_a = m .* J_nx .* J_nmx_p - J_nx_p .* J_nmx;
|
||||
den_a = J_nmx .* H_nx_p - m .* J_nmx_p .* H_nx;
|
||||
a_n = num_a ./ den_a;
|
||||
|
||||
% 计算内部透射系数 c_n
|
||||
num_c = J_nx .* H_nx_p - J_nx_p .* H_nx; % 这其实是 Wronskian
|
||||
c_n = num_c ./ den_a;
|
||||
|
||||
% 空间相位因子: i^n * exp(i*n*phi)
|
||||
phase = (1i)^n * exp(1i * n * Phi);
|
||||
|
||||
% 累加外部散射场 (仅在 R >= r 区域有效)
|
||||
out_idx = R >= r;
|
||||
E_scat(out_idx) = E_scat(out_idx) + a_n .* besselh(n, 1, k0 * R(out_idx)) .* phase(out_idx);
|
||||
|
||||
% 累加内部总场 (仅在 R < r 区域有效)
|
||||
in_idx = R < r;
|
||||
E_int(in_idx) = E_int(in_idx) + c_n .* besselj(n, k1 * R(in_idx)) .* phase(in_idx);
|
||||
end
|
||||
|
||||
% ================= 5. 组装全场并绘图 =================
|
||||
% 入射平面波: u_inc = exp(i*k0*x)
|
||||
phase_shift = exp(1i * k0 * xc);
|
||||
E_scat = E_scat .* phase_shift;
|
||||
E_int = E_int .* phase_shift;
|
||||
|
||||
E_inc = exp(1i * k0 * X);
|
||||
|
||||
% 总场 = 外部(入射 + 散射) + 内部场
|
||||
% 组装总场
|
||||
E_total = zeros(size(X));
|
||||
E_total(R >= r) = E_inc(R >= r) + E_scat(R >= r);
|
||||
E_total(R < r) = E_int(R < r);
|
||||
|
||||
|
||||
% 绘图
|
||||
figure('Color','w');
|
||||
pcolor(X, Y, real(E_total-E_inc));
|
||||
max_E_real = max(max(real(E_total-E_inc)));
|
||||
shading interp;
|
||||
axis equal tight;
|
||||
colorbar;
|
||||
colormap jet;
|
||||
title(sprintf('2D Cylinder Mie Scattering |E_{scatter}| (Max = %.4f)', max_E_real));
|
||||
|
||||
% 绘制圆柱边界
|
||||
hold on;
|
||||
theta_circle = linspace(0, 2*pi, 100);
|
||||
plot(xc + r * cos(theta_circle), yc + r * sin(theta_circle), 'k--', 'LineWidth', 1.5);
|
||||
hold off;
|
||||
|
|
@ -0,0 +1,14 @@
|
|||
一、 引入因果律：对偶加权残差法（Dual-Weighted Residual, DWR）与其让 GNN 在空间中盲目摸索残差的传播规律，不如直接利用偏微分方程的伴随算子（Adjoint Operator）显式求解误差的传播路径。在 DWR 理论中，我们定义一个关心的目标泛函 $J(e)$（例如远场总场的误差）。为了找到局部残差 $R(u_h)$ 是如何影响 $J(e)$ 的，我们需要求解原方程的对偶（伴随）问题：$$\mathcal{L}^* z = J'(\cdot)$$由于亥姆霍兹方程是自伴随或复对称的，对偶解 $z$ 本质上就是一个以目标区域为源的反向传播波（Green's function 的叠加）。严格的误差表示定理（Error Representation Theorem）给出：$$J(e) = \sum_{K \in \Omega_h} \left( \langle r_{\text{int}}, z - z_h \rangle_K + \langle r_{\text{jump}}, z - z_h \rangle_{\partial K} \right)$$第一性原理 AI 方案：物理先验特征：在 FEM 求解器中，顺手在极粗网格上解一次对偶问题得到 $z_h$（计算代价极小）。将权重项 $\omega_K = |z - z_h|_K$（或者启发式地使用 $|z_h|_K$ 的梯度）作为 GNN 的节点输入特征。自然适配：网络会立刻“看”到，虽然介质外部的 $r_{\text{jump}}$ 很大，但那里的对偶权重 $\omega_K$ 极小；而介质内部的对偶权重巨大。网络在不加任何人为截断的情况下，自然顺着物理因果律将算力投向介质内部。
|
||||
|
||||
COMSOL 的自适应往往隐式或显式地结合了对偶加权残差（DWR），能够识别“远场误差是由哪里传播过来的”。
|
||||
|
||||
二、 相空间与动量解耦：Wigner 分布与相空间光学残差在含有横向动量（如余弦载波项）和复杂色散介质的全场计算中，空间域的标量残差 $\eta_K$ 掩盖了误差的物理本质。污染效应的核心在于波矢（动量 $\mathbf{k}$）方向的失配。从相空间光学的角度来看，可以用维格纳分布函数（Wigner Distribution Function, WDF） 将标量场映射到位置-动量相空间 $W(\mathbf{x}, \mathbf{p})$。在渐近区，波场满足相空间的射线输运方程。数值解 $u_h$ 与真实解的差异，在空间域表现为弥散的干涉条纹，但在相空间中，却能清晰地表现为动量谱的分叉与频移。第一性原理 AI 方案：抛弃纯空间域的 $L_2$ 残差聚合。在误差提取步骤，对全场残差提取局部波矢谱（类似于短时傅里叶变换或 WDF 近似提取）。将动量偏差（Momentum Mismatch）作为核心 Reward。当且仅当一个细化动作能够将数值波阵面的 $\mathbf{k}$ 矢量方向拉回到正确的理论物理色散面上时，才给予正向激励。这样，网络优化的不再是单纯的数值差异，而是逼近真实的物理色散关系。
|
||||
|
||||
三、 算子层面的修正：变分稳定化（GLS / Trefftz 方法）目前的强化学习框架试图用网格细化（$h$-refinement）去填补 P1 单元固有的色散缺陷。从底层物理看，这是在用极高的计算成本为糟糕的基函数买单。如果从变分形式（Weak Form）出发，标准的 Galerkin 方法在亥姆霍兹算子下会失去最佳逼近性（Céa 引理中的稳定性常数随波数爆炸）。我们需要在算子层面进行修正。第一性原理 AI 方案：Galerkin Least-Squares (GLS) 稳定化：在标准的变分方程中，加入与残差相关的稳定项：$$B_{GLS}(u_h, v_h) = B_{Gal}(u_h, v_h) + \sum_K \tau_K \langle \mathcal{L}u_h - f, \mathcal{L}v_h \rangle_K$$通过精心设计稳定化参数 $\tau_K$，可以直接在 FEM 矩阵组装层面抵消 P1 单元的色散误差。此时，外部的虚假污染误差会在物理求解阶段被直接压制，GNN 面对的将是一个干净、局域化的残差场。物理信息的基函数（Trefftz / Plane Wave Basis）：放弃多项式基函数。对于散射总场问题，介质内部和外部的物理场本质上是局部平面波或柱面波的叠加。如果在单元内部使用满足 $\nabla^2 \phi + k^2\varepsilon_r \phi = 0$ 的平面波作为基函数（即 Trefftz 方法或平面波非连续 Galerkin 方法 PWDG），网格内部残差 $r_{\text{int}}$ 将恒等于零。此时，所有的物理误差将以第一性原理的方式，极其干净地全部集中在介质与空气交界面的梯度/通量跳变 $r_{\text{jump}}$ 上。网络只需要专注于处理界面处的阻抗匹配即可，彻底根除了污染效应。
|
||||
|
||||
|
||||
1. 优先推进：对偶加权残差法 (DWR) 的 AI 赋能这是目前投资回报率（ROI）最高、最能快速落地的方案。可行性 (极高)： 你现有的 ASMR++ 框架已经极其完善（GNN 观测 + 连续尺寸场 + PPO）。引入 DWR 不需要重构底层的 FEM 求解核心。你只需要在粗网格计算时，额外配置一个右端项（目标泛函的导数）求解一次伴随方程，将其作为额外的 GNN 节点特征。代码改动量最小，且能够迅速验证效果。创新性 (中高)： DWR 本身是传统自适应有限元（AFEM）的经典理论，但在传统计算中，求解伴随方程的开销往往被认为过大。通过 RL 与 GNN，让智能体“学习” DWR 提供的因果律，从而在极少步骤内预测出最优的网格尺寸场，这是一个极其 solid 的 AI4S 创新点。物理信息嵌入 (高)： 完美解决了“污染效应”中的非局部性问题。智能体的图神经网络不再是盲目地卷积局部几何残差，而是顺着伴随场（反向传播的波）的指引，直接“看”到了误差的因果律。发文章角度： 非常适合投往计算力学或物理机器学习的顶级期刊（如 JCP, CMAME）。故事主线明确：“通过强化学习结合 DWR，打破高频 Helmholtz 方程自适应网格细化中的污染效应陷阱”。
|
||||
|
||||
2. 旗舰目标：相空间动量解耦 (Wigner 分布)这是上限最高、最颠覆性的方案，也是构建科研护城河的终极武器。可行性 (较高挑战)： 计算二维波场的 Wigner 分布函数 (WDF) 会带来维度爆炸（2D 空间 $\rightarrow$ 4D 相空间），将其放入 RL 的每个 Reward 计算 loop 中会导致严重的效率瓶颈。你需要设计一种轻量级的局部波矢提取算法。创新性 (极高)： 目前 AI4S 领域的 PDE 求解和网格优化几乎全部停留在空间域（$L_2$ 或 $H^1$ 范数）。将相空间光学的概念引入有限元误差估计，是从根本上切换了视角。物理信息嵌入 (极高)： 若要真正挑战跨越不同介质的零样本泛化 (Zero-shot generalization)，单纯的空间域残差是极其脆弱的。因为不同 $\varepsilon_r$ 对应的空间波长和残差量级完全不同。但在 WDF 描述的相空间中，不同介质的波传播都遵循统一的射线哈密顿力学。以动量失配（Momentum Mismatch）作为 Reward，智能体优化的不再是表象的干涉条纹，而是底层的色散流形。发文章角度： 冲击综合性或交叉学科顶刊（如 Nature Computational Science, Light: Science & Applications, 或 PRL）。结合在相位恢复和 WDF 重构上已有的技术积累，这可以包装成一个完全超越传统 FEM 思维的“相空间 AI 自适应物理引擎”。
|
||||
|
||||
3. 基础支撑：算子层面的变分稳定化 (GLS / Trefftz)这是一个偏传统计算力学但极其硬核的方案。可行性 (中等)： 需要深入修改你的 helmholtz.py，改变弱形式（Weak Form）的矩阵组装过程。特别是 Trefftz 方法或平面波不连续伽辽金 (PWDG)，其积分规则和界面通量定义与标准 P1 连续元完全不同。创新性 (高)： Trefftz 方法本身就自带极强的物理先验（基函数严格满足局部齐次方程）。用 RL 智能体去动态配置界面处的阻抗匹配和通量惩罚，是一个极具技术深度的方向。物理信息嵌入 (最高)： 它是唯一从算子理论层面彻底消灭 P1 单元色散误差（Dispersion Error）的方案。网格内部毫无误差，所有优化预算全部分配在界面跳变上。发文章角度： 属于极其硬核的数值分析与 AI 结合工作，更受传统数学和力学审稿人的青睐。
|
||||
File diff suppressed because it is too large
Load Diff
Binary file not shown.
|
|
@ -0,0 +1,50 @@
|
|||
# QA Report: AFEM 组会汇报 PPTX
|
||||
|
||||
## 构建状态
|
||||
- **状态**: OK
|
||||
- **文件**: `output/final_presentation_cn.pptx`
|
||||
- **大小**: 70.4 KB
|
||||
- **页数**: 15
|
||||
- **格式**: 16:9 宽屏 (13.3 x 7.5 inches)
|
||||
- **语言**: 中文（全中文标题与正文，英文保留技术术语）
|
||||
|
||||
## 验证结果
|
||||
- python-pptx 重新打开: OK
|
||||
- 全部 15 页均有中文文本内容
|
||||
- 幻灯片结构符合大纲设计
|
||||
|
||||
## 15 页结构
|
||||
|
||||
| # | 标题 | 类型 |
|
||||
|---|------|------|
|
||||
| 1 | AFEM：基于 GNN + PPO 强化学习的自适应网格细化方法 | 标题页 |
|
||||
| 2 | 研究背景：为什么自适应网格细化很重要 | 背景 |
|
||||
| 3 | 知识缺口与技术瓶颈 | 缺口/动机 |
|
||||
| 4 | 系统架构：RL 自适应网格细化闭环管线 | 技术路线 |
|
||||
| 5 | 创新 [1]：无量纲化残差误差估计 -- 消除几何尺度偏差 | 创新 |
|
||||
| 6 | 创新 [2]：12 维增强输入特征 -- 赋予 GNN 几何与物理感知 | 创新 |
|
||||
| 7 | 创新 [3]：Score-based 连续尺寸场 + 物理预算约束 + 动作掩码 | 创新 |
|
||||
| 8 | 创新 [4]：L2 聚合奖励设计 -- 保证非负，永不惩罚细化 | 创新 |
|
||||
| 9 | 奖励标度校准：随机策略下各分量量级实测 | 证据 |
|
||||
| 10 | 创新 [5]：尺度不变性架构 -- 从 1x1 到 2x2 的泛化 | 创新 |
|
||||
| 11 | 双 GNN 架构与 PPO 训练细节 | 架构 |
|
||||
| 12 | 训练观察与诊断：奖励稀疏性与大波数泛化 | 诊断 |
|
||||
| 13 | 创新点汇总与可复用价值 | 综合 |
|
||||
| 14 | 局限性与未解决问题 | 局限 |
|
||||
| 15 | 总结 | 总结 |
|
||||
|
||||
## 图片/资源
|
||||
- 未提取外部图片（纯 python-pptx 绘制）
|
||||
- 所有视觉效果为原生 PPTX 图形和文本框
|
||||
- `output/assets/figures/` 目录已创建（空）
|
||||
|
||||
## 已知局限
|
||||
1. **无渲染预览** -- 环境中无可用的无头渲染器 (LibreOffice)，未做逐页视觉 QA
|
||||
2. **无外部图片** -- 建议后续将 `result/visualization*.png` 的网格截图添加到 slides 5-8 的关键证据页
|
||||
3. **字体依赖** -- 使用 'Microsoft YaHei'，在 macOS/Linux 上可能回退到系统默认无衬线字体
|
||||
4. **技术词汇混用** -- 关键术语 (eta_K, k_local, GNN, PPO, GAE 等) 保留英文，其余为中文
|
||||
|
||||
## 建议手动补充
|
||||
1. 将 `result/visualization*.png` 中的网格对比截图添加到对应的创新页
|
||||
2. 在汇报机器上验证字体渲染效果
|
||||
3. 如有需要，为关键证据页添加口头讲稿备注
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
{
|
||||
"venvPath": ".",
|
||||
"venv": ".venv",
|
||||
"typeCheckingMode": "off",
|
||||
"reportPrivateImportUsage": false,
|
||||
"reportMissingImports": true
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.5 MiB |
|
|
@ -0,0 +1,97 @@
|
|||
clc; clear; close all;
|
||||
|
||||
% ================= 1. 物理参数定义 =================
|
||||
r = 0.1; % 圆柱半径
|
||||
eps_r = 5.0; % 相对介电常数
|
||||
m = sqrt(eps_r); % 相对折射率 m = ~1.414
|
||||
k0 = 50; % 背景真空中波数 (k=6)
|
||||
k1 = k0 * m; % 圆柱内部波数
|
||||
x_size = k0 * r; % 尺寸参数 x = k0*a
|
||||
|
||||
% ================= 2. 计算域网格设置 =================
|
||||
x_range = 1;
|
||||
y_range = 1;
|
||||
Nx = 500;
|
||||
Ny = 500;
|
||||
x_vec = linspace(0, x_range, Nx);
|
||||
y_vec = linspace(0, y_range, Ny);
|
||||
[X, Y] = meshgrid(x_vec, y_vec);
|
||||
|
||||
xc = 0.5; yc = 0.5;
|
||||
[Phi, R] = cart2pol(X - xc, Y - yc); % 转换为极坐标
|
||||
|
||||
% ================= 3. 场初始化 =================
|
||||
E_scat = zeros(size(X)); % 散射场
|
||||
E_int = zeros(size(X)); % 内部场
|
||||
|
||||
% Wiscombe 截断准则（决定级数展开需要算到第几阶）
|
||||
N_trunc = round(x_size + 4.05 * x_size^(1/3) + 2);
|
||||
|
||||
% ================= 4. 2D Mie 级数展开计算 =================
|
||||
% 2D 圆柱级数从 -N 到 +N
|
||||
for n = -N_trunc : N_trunc
|
||||
|
||||
% 边界处的贝塞尔函数值
|
||||
J_nx = besselj(n, x_size);
|
||||
J_nmx = besselj(n, k1 * r);
|
||||
H_nx = besselh(n, 1, x_size);
|
||||
|
||||
% 边界处的导数值 (利用递推公式 Z_n' = 0.5 * (Z_{n-1} - Z_{n+1}))
|
||||
J_nx_p = 0.5 * (besselj(n-1, x_size) - besselj(n+1, x_size));
|
||||
J_nmx_p = 0.5 * (besselj(n-1, k1*r) - besselj(n+1, k1*r));
|
||||
H_nx_p = 0.5 * (besselh(n-1, 1, x_size) - besselh(n+1, 1, x_size));
|
||||
|
||||
% 计算 TM 偏振下的散射系数 a_n (对应 E_z)
|
||||
num_a = m .* J_nx .* J_nmx_p - J_nx_p .* J_nmx;
|
||||
den_a = J_nmx .* H_nx_p - m .* J_nmx_p .* H_nx;
|
||||
a_n = num_a ./ den_a;
|
||||
|
||||
% 计算内部透射系数 c_n
|
||||
num_c = J_nx .* H_nx_p - J_nx_p .* H_nx; % 这其实是 Wronskian
|
||||
c_n = num_c ./ den_a;
|
||||
|
||||
% 空间相位因子: i^n * exp(i*n*phi)
|
||||
phase = (1i)^n * exp(1i * n * Phi);
|
||||
|
||||
% 累加外部散射场 (仅在 R >= r 区域有效)
|
||||
out_idx = R >= r;
|
||||
E_scat(out_idx) = E_scat(out_idx) + a_n .* besselh(n, 1, k0 * R(out_idx)) .* phase(out_idx);
|
||||
|
||||
% 累加内部总场 (仅在 R < r 区域有效)
|
||||
in_idx = R < r;
|
||||
E_int(in_idx) = E_int(in_idx) + c_n .* besselj(n, k1 * R(in_idx)) .* phase(in_idx);
|
||||
end
|
||||
|
||||
% ================= 5. 组装全场并绘图 =================
|
||||
% 入射平面波: u_inc = exp(i*k0*x)
|
||||
phase_shift = exp(1i * k0 * xc);
|
||||
E_scat = E_scat .* phase_shift;
|
||||
E_int = E_int .* phase_shift;
|
||||
|
||||
E_inc = exp(1i * k0 * X);
|
||||
|
||||
% 总场 = 外部(入射 + 散射) + 内部场
|
||||
% 组装总场
|
||||
E_total = zeros(size(X));
|
||||
E_total(R >= r) = E_inc(R >= r) + E_scat(R >= r);
|
||||
E_total(R < r) = E_int(R < r);
|
||||
%
|
||||
% % 提取最大场强做对比
|
||||
% max_E_val = max(abs(E_total(:)));
|
||||
% fprintf('2D 理论解析解中心区域最大场强 (max |E_total|): %.4f\n', max_E_val);
|
||||
|
||||
% 绘图
|
||||
figure('Color','w');
|
||||
pcolor(X, Y, abs(E_total-E_inc));
|
||||
max_E_real = max(max(abs(E_total-E_inc)));
|
||||
shading interp;
|
||||
axis equal tight;
|
||||
colorbar;
|
||||
colormap jet;
|
||||
title(sprintf('2D Cylinder Mie Scattering |E_{scatter}| (Max = %.4f)', max_E_real));
|
||||
|
||||
% 绘制圆柱边界
|
||||
hold on;
|
||||
theta_circle = linspace(0, 2*pi, 100);
|
||||
plot(xc + r * cos(theta_circle), yc + r * sin(theta_circle), 'k--', 'LineWidth', 1.5);
|
||||
hold off;
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 2.0 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.6 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.7 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.8 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.8 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.7 MiB |
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,108 @@
|
|||
#############################
|
||||
# 训练：
|
||||
# CUDA_VISIBLE_DEVICES=7 python src/main.py --mode train --config src/config.yaml
|
||||
# 测试：
|
||||
# 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
|
||||
|
||||
# 可视化：
|
||||
# python src/main.py --mode viz --checkpoint checkpoints/model_iter0400.pt
|
||||
# python src/main.py --mode viz --checkpoint checkpoints/model_iter0100.pt --k-test 8.0 --center 0.6,0.5 --radius 0.1
|
||||
###########################
|
||||
|
||||
algorithm:
|
||||
batch_size: 32
|
||||
discount_factor: 1.0
|
||||
ppo:
|
||||
clip_range: 0.2
|
||||
entropy_coefficient: 0.001
|
||||
epochs_per_iteration: 5 # 每轮迭代对同一批 rollout 数据重复训练几个 epoch
|
||||
gae_lambda: 0.95
|
||||
initial_log_std: -2.0 # 初始动作 log 标准差，exp(-2)≈0.135
|
||||
max_grad_norm: 0.5
|
||||
num_rollout_steps: 256
|
||||
value_function_coefficient: 0.5
|
||||
use_gpu: true
|
||||
environment:
|
||||
mesh_refinement:
|
||||
edge_features:
|
||||
euclidean_distance: true
|
||||
element_features:
|
||||
element_penalty: true
|
||||
is_sbc_boundary: true
|
||||
k_local_sqrt_vol: true
|
||||
solution_std: true
|
||||
timestep: true
|
||||
volume: true
|
||||
wave_number: true
|
||||
x_position: false
|
||||
y_position: false
|
||||
dist_to_interface: true
|
||||
element_limit_penalty: 10000
|
||||
element_penalty:
|
||||
sample_penalty: false
|
||||
value: 0.06
|
||||
fem:
|
||||
domain:
|
||||
boundary:
|
||||
- 0
|
||||
- 0
|
||||
- 3
|
||||
- 3
|
||||
initial_num_elements: 75
|
||||
helmholtz:
|
||||
k_ref: 6.0
|
||||
k_exponent: 2.0
|
||||
scatterer:
|
||||
cx: 1.5
|
||||
cx_max: 0.8
|
||||
cx_min: 0.2
|
||||
cy: 1.5
|
||||
cy_max: 0.8
|
||||
cy_min: 0.2
|
||||
eps_r: 5.0
|
||||
eps_r_max: 8.0
|
||||
eps_r_min: 2.0
|
||||
mode: random_uniform
|
||||
radius: 0.2
|
||||
radius_max: 0.2
|
||||
radius_min: 0.05
|
||||
wave_number: 30.0
|
||||
wave_number_max: 3.0
|
||||
wave_number_min: 15.0
|
||||
wave_number_mode: random_uniform
|
||||
num_pdes: 100
|
||||
pde_type: helmholtz
|
||||
pre_asymptotic_N: 1.5
|
||||
maximum_elements: 50000
|
||||
num_timesteps: 4
|
||||
refinement_strategy: continuous_sizing_field
|
||||
reward_type: spatial
|
||||
global_reward_alpha: 0.5 # 全局奖励权重
|
||||
# rho_weights:
|
||||
# w_int: 0.0 # ρ_int 权重 (代码自动除以 k²)
|
||||
# w_jump: 1.0 # ρ_jump 权重 (代码自动除以 k)
|
||||
# w_sbc: 20.0 # ρ_sbc 权重 (代码自动除以 k)
|
||||
iterations: 401
|
||||
network:
|
||||
actor:
|
||||
mlp:
|
||||
activation_function: tanh
|
||||
num_layers: 2
|
||||
base:
|
||||
edge_dropout: 0.1
|
||||
scatter_reduce: mean
|
||||
stack:
|
||||
mlp:
|
||||
activation_function: leakyrelu
|
||||
num_layers: 2
|
||||
num_steps: 2
|
||||
critic:
|
||||
mlp:
|
||||
activation_function: tanh
|
||||
num_layers: 2
|
||||
latent_dimension: 64
|
||||
training:
|
||||
learning_rate: 0.0003
|
||||
lr_decay: 0.995
|
||||
optimizer: adam
|
||||
|
|
@ -0,0 +1,159 @@
|
|||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch_geometric.data import Batch
|
||||
|
||||
logging.getLogger("skfem").setLevel(logging.ERROR)
|
||||
|
||||
_project_root = Path(__file__).resolve().parent.parent
|
||||
if str(_project_root) not in sys.path:
|
||||
sys.path.insert(0, str(_project_root))
|
||||
|
||||
from src.network import create_model
|
||||
from src.ppo import PPOTrainer
|
||||
from src.utils import load_checkpoint, load_config, parse_center, save_checkpoint, setup_helmholtz_config
|
||||
from src.visualize import visualize
|
||||
|
||||
|
||||
def train(config: dict, iterations: int, checkpoint_dir: str = "checkpoints", save_freq: int = 50):
|
||||
t0 = time.time()
|
||||
algo = config.get("algorithm", {})
|
||||
dev = torch.device("cuda" if torch.cuda.is_available() and algo.get("use_gpu") else "cpu")
|
||||
print(f"[Device] {dev}")
|
||||
|
||||
from environment.mesh_refinement import MeshRefinement
|
||||
|
||||
env = MeshRefinement(
|
||||
environment_config=config.get("environment", {}).get("mesh_refinement", {}),
|
||||
seed=42,
|
||||
)
|
||||
print(f"[Env] node_feats={env.num_node_features} edge_feats={env.num_edge_features} act_dim={env.action_dimension}")
|
||||
|
||||
model = create_model(env, config.get("network", {}), algo.get("ppo", {}), device=dev)
|
||||
print(f"[Model] params={sum(p.numel() for p in model.parameters()):,}")
|
||||
|
||||
trainer = PPOTrainer(model, env, algo, device=dev)
|
||||
os.makedirs(checkpoint_dir, exist_ok=True)
|
||||
|
||||
for it in range(1, iterations + 1):
|
||||
t1 = time.time()
|
||||
metrics = trainer.fit_iteration()
|
||||
print(
|
||||
f" {it:4d}/{iterations} | loss={metrics['loss']:.4f} ev={metrics['explained_variance']:.3f} "
|
||||
f"agents={metrics['num_agents']:.0f} avg_r={metrics['avg_reward']:.4f} sum_r={metrics['sum_reward']:.2f} "
|
||||
f"x<0={metrics.get('neg_action_ratio', 0):.2f} "
|
||||
f"elig={metrics.get('eligible_ratio', 0):.2f} "
|
||||
f"sel={metrics.get('selected_count', 0):.0f} "
|
||||
f"{time.time() - t1:.1f}s"
|
||||
)
|
||||
if it % save_freq == 0 or it == iterations:
|
||||
save_checkpoint(model, model.optimizer, it, os.path.join(checkpoint_dir, f"model_iter{it:04d}.pt"))
|
||||
|
||||
save_checkpoint(model, model.optimizer, iterations, os.path.join(checkpoint_dir, "model_final.pt"))
|
||||
print(f"[Train] done, total time {time.time() - t0:.1f}s")
|
||||
|
||||
|
||||
def _eval_mie_error_test(env) -> float:
|
||||
"""Compute relative L2 error of FEM vs Mie analytical solution."""
|
||||
fp = getattr(env.fem_problem, "fem_problem", None)
|
||||
if fp is None:
|
||||
return float("nan")
|
||||
_eps_r = getattr(fp, "_eps_r", None)
|
||||
_radius = getattr(fp, "_radius", None)
|
||||
_cx = getattr(fp, "_cx", None)
|
||||
_cy = getattr(fp, "_cy", None)
|
||||
_k = getattr(fp, "_k", None)
|
||||
if any(v is None for v in [_eps_r, _radius, _cx, _cy, _k]):
|
||||
return float("nan")
|
||||
|
||||
from environment.mie_solution import mie_scattered_field
|
||||
pts = env.mesh.p.T
|
||||
u_mie = mie_scattered_field(pts, k0=_k, eps_r=_eps_r, radius=_radius, cx=_cx, cy=_cy)
|
||||
u_fem = env.scalar_solution
|
||||
diff = np.abs(u_fem - u_mie)
|
||||
denom = np.linalg.norm(np.abs(u_mie))
|
||||
if denom < 1e-12:
|
||||
denom = 1.0
|
||||
return float(np.linalg.norm(diff) / denom)
|
||||
|
||||
|
||||
def test(config: dict, checkpoint_path: str, k_test=None, center=None, radius=None, eps_test=None):
|
||||
setup_helmholtz_config(config, k_test=k_test, center=center, radius=radius, eps_test=eps_test)
|
||||
algo = config.get("algorithm", {})
|
||||
|
||||
from environment.mesh_refinement import MeshRefinement
|
||||
|
||||
env = MeshRefinement(
|
||||
environment_config=config.get("environment", {}).get("mesh_refinement", {}),
|
||||
seed=99,
|
||||
)
|
||||
model = create_model(env, config.get("network", {}), algo.get("ppo", {}))
|
||||
load_checkpoint(model, checkpoint_path)
|
||||
model.eval()
|
||||
|
||||
obs = env.reset()
|
||||
done = False
|
||||
step = 0
|
||||
n_elem_init = getattr(env, "_num_elements", env.num_agents)
|
||||
mie_err_0 = _eval_mie_error_test(env)
|
||||
print(f" Step {step:2d}: reward=--- mie_err={mie_err_0:.4f} elements={n_elem_init}"
|
||||
f" budget={getattr(env, '_n_budget', '?')}")
|
||||
|
||||
total_reward = 0.0
|
||||
while not done:
|
||||
with torch.no_grad():
|
||||
actions, _, _ = model(Batch.from_data_list([obs]), deterministic=True)
|
||||
obs, reward, done, info = env.step(actions.cpu().numpy())
|
||||
step_r = float(np.sum(reward))
|
||||
total_reward += step_r
|
||||
step += 1
|
||||
mie_err = _eval_mie_error_test(env)
|
||||
print(f" Step {step:2d}: reward={step_r:+.4f} mie_err={mie_err:.4f}"
|
||||
f" elements={info.get('num_elements', '?')} "
|
||||
f"x<0={info.get('neg_action_ratio', 0):.2f} sel={info.get('selected_count', 0)}")
|
||||
|
||||
print(f"\n[Test] total_reward={total_reward:.4f} final_mie_error={mie_err:.4f}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="AFEM — Adaptive FEM with PPO RL")
|
||||
parser.add_argument("--mode", required=True, choices=["train", "test", "viz"])
|
||||
parser.add_argument("--config", default="src/config.yaml")
|
||||
parser.add_argument("--iterations", type=int, default=None)
|
||||
parser.add_argument("--checkpoint", default="checkpoints/model_final.pt")
|
||||
parser.add_argument("--checkpoint-dir", default="checkpoints")
|
||||
parser.add_argument("--save-freq", type=int, default=50)
|
||||
parser.add_argument("--output", default="result/visualization.png")
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--k-test", type=float, default=None)
|
||||
parser.add_argument("--center", type=str, default=None)
|
||||
parser.add_argument("--radius", type=float, default=None)
|
||||
parser.add_argument("--eps-test", type=float, default=None)
|
||||
|
||||
args = parser.parse_args()
|
||||
torch.manual_seed(args.seed)
|
||||
np.random.seed(args.seed)
|
||||
|
||||
cfg_path = args.config if os.path.isabs(args.config) else os.path.join(_project_root, args.config)
|
||||
config = load_config(cfg_path)
|
||||
if args.iterations is not None:
|
||||
config["iterations"] = args.iterations
|
||||
|
||||
center = parse_center(args.center)
|
||||
|
||||
if args.mode == "train":
|
||||
train(config, config.get("iterations", 100), args.checkpoint_dir, args.save_freq)
|
||||
elif args.mode == "test":
|
||||
test(config, args.checkpoint, k_test=args.k_test, center=center, radius=args.radius, eps_test=args.eps_test)
|
||||
elif args.mode == "viz":
|
||||
visualize(config, args.checkpoint, output_path=args.output, k_test=args.k_test, center=center, radius=args.radius, eps_test=args.eps_test)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,419 @@
|
|||
import copy
|
||||
|
||||
import gym
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torch_geometric.data import Data
|
||||
from torch_geometric.utils import dropout_edge
|
||||
from torch_scatter import scatter_mean
|
||||
|
||||
|
||||
def get_scatter_reduce(name: str):
|
||||
name = name.lower()
|
||||
if name == "mean":
|
||||
from torch_scatter import scatter_mean
|
||||
return scatter_mean
|
||||
if name == "sum":
|
||||
from torch_scatter import scatter_add
|
||||
return scatter_add
|
||||
if name == "max":
|
||||
from torch_scatter import scatter_max
|
||||
return lambda *a, **kw: scatter_max(*a, **kw)[0]
|
||||
if name == "min":
|
||||
from torch_scatter import scatter_min
|
||||
return lambda *a, **kw: scatter_min(*a, **kw)[0]
|
||||
if name == "std":
|
||||
from torch_scatter import scatter_std
|
||||
return scatter_std
|
||||
raise ValueError(f"Unknown scatter reduce '{name}'")
|
||||
|
||||
|
||||
# ──
|
||||
# 1. LatentMLP — GNN 内部使用的 MLP（保持隐层维度不变）
|
||||
# ──
|
||||
class LatentMLP(nn.Module):
|
||||
"""
|
||||
MLP that operates entirely in latent space (dim in == dim out == latent_dim).
|
||||
Used inside EdgeModule and NodeModule.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features: int, latent_dim: int, config: dict):
|
||||
super().__init__()
|
||||
num_layers = config.get("num_layers", 2)
|
||||
activation = config.get("activation_function", "leakyrelu").lower()
|
||||
add_output = config.get("add_output_layer", False)
|
||||
|
||||
layers = []
|
||||
prev_dim = in_features
|
||||
for i in range(num_layers):
|
||||
layers.append(nn.Linear(prev_dim, latent_dim))
|
||||
layers.append(_get_activation(activation))
|
||||
prev_dim = latent_dim
|
||||
|
||||
if add_output:
|
||||
layers.append(nn.Linear(prev_dim, latent_dim))
|
||||
|
||||
self.mlp = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.mlp(x)
|
||||
|
||||
|
||||
def _get_activation(name: str) -> nn.Module:
|
||||
name = name.lower()
|
||||
if name == "relu":
|
||||
return nn.ReLU()
|
||||
elif name == "leakyrelu":
|
||||
return nn.LeakyReLU()
|
||||
elif name == "elu":
|
||||
return nn.ELU()
|
||||
elif name in ("swish", "silu"):
|
||||
return nn.SiLU()
|
||||
elif name == "mish":
|
||||
return nn.Mish()
|
||||
elif name == "gelu":
|
||||
return nn.GELU()
|
||||
elif name == "tanh":
|
||||
return nn.Tanh()
|
||||
raise ValueError(f"Unknown activation '{name}'")
|
||||
|
||||
|
||||
# ──
|
||||
# 2. EdgeModule — 边更新：MLP([src_node | dst_node | edge_attr])
|
||||
# ──
|
||||
class EdgeModule(nn.Module):
|
||||
"""Update edge features from sender node, receiver node, and existing edge features."""
|
||||
|
||||
def __init__(self, latent_dim: int, mlp_config: dict):
|
||||
super().__init__()
|
||||
in_features = 3 * latent_dim # [src_node, dst_node, edge_attr]
|
||||
self.mlp = LatentMLP(in_features, latent_dim, mlp_config)
|
||||
|
||||
def forward(self, graph: Data):
|
||||
src, dst = graph.edge_index
|
||||
agg = torch.cat([graph.x[src], graph.x[dst], graph.edge_attr], dim=-1)
|
||||
graph.edge_attr = self.mlp(agg)
|
||||
|
||||
|
||||
# ──
|
||||
# 4. NodeModule — 节点更新：MLP([node | scatter(入边)])
|
||||
# ──
|
||||
class NodeModule(nn.Module):
|
||||
"""Update node features from own features and aggregated incoming edge features."""
|
||||
|
||||
def __init__(self, latent_dim: int, mlp_config: dict, scatter_reducer):
|
||||
super().__init__()
|
||||
in_features = 2 * latent_dim # [node, aggregated_edges]
|
||||
self.mlp = LatentMLP(in_features, latent_dim, mlp_config)
|
||||
self.scatter = scatter_reducer
|
||||
|
||||
def forward(self, graph: Data):
|
||||
_, dst = graph.edge_index
|
||||
agg_edges = self.scatter(
|
||||
graph.edge_attr, dst, dim=0, dim_size=graph.x.shape[0]
|
||||
)
|
||||
agg = torch.cat([graph.x, agg_edges], dim=-1)
|
||||
graph.x = self.mlp(agg)
|
||||
|
||||
|
||||
# ──
|
||||
# 5. MessagePassingStep — 单步消息传递
|
||||
# ──
|
||||
class MessagePassingStep(nn.Module):
|
||||
"""
|
||||
One full message-passing step:
|
||||
1. Edge update
|
||||
2. Edge inner residual + LayerNorm
|
||||
3. Node update
|
||||
4. Node inner residual + LayerNorm
|
||||
"""
|
||||
|
||||
def __init__(self, latent_dim: int, stack_config: dict, scatter_reducer):
|
||||
super().__init__()
|
||||
mlp_config = stack_config["mlp"]
|
||||
|
||||
self.edge_module = EdgeModule(latent_dim, mlp_config)
|
||||
self.node_module = NodeModule(latent_dim, mlp_config, scatter_reducer)
|
||||
|
||||
self.node_ln = nn.LayerNorm(latent_dim)
|
||||
self.edge_ln = nn.LayerNorm(latent_dim)
|
||||
|
||||
def forward(self, graph: Data):
|
||||
old_x = graph.x
|
||||
old_edge = graph.edge_attr
|
||||
|
||||
# Edge update
|
||||
self.edge_module(graph)
|
||||
graph.edge_attr = self.edge_ln(graph.edge_attr + old_edge)
|
||||
|
||||
# Node update
|
||||
self.node_module(graph)
|
||||
graph.x = self.node_ln(graph.x + old_x)
|
||||
|
||||
|
||||
# ──
|
||||
# 6. MessagePassingStack — 堆叠 N 个 Step
|
||||
# ──
|
||||
class MessagePassingStack(nn.Module):
|
||||
"""Stack of multiple MessagePassingSteps with optional step repeats."""
|
||||
|
||||
def __init__(self, latent_dim: int, stack_config: dict, scatter_reducer):
|
||||
super().__init__()
|
||||
num_steps = stack_config.get("num_steps", 2)
|
||||
self.num_step_repeats = stack_config.get("num_step_repeats", 1)
|
||||
self.steps = nn.ModuleList(
|
||||
[
|
||||
MessagePassingStep(latent_dim, stack_config, scatter_reducer)
|
||||
for _ in range(num_steps)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, graph: Data):
|
||||
for step in self.steps:
|
||||
for _ in range(self.num_step_repeats):
|
||||
step(graph)
|
||||
|
||||
|
||||
# ──
|
||||
# 7. MessagePassingBase — GNN 基座
|
||||
# ──
|
||||
class MessagePassingBase(nn.Module):
|
||||
"""
|
||||
Full GNN base: Linear → Stack → unpacked output.
|
||||
Returns (node_features_dict, edge_features, None, batch_dict)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_node_features: int,
|
||||
in_edge_features: int,
|
||||
latent_dim: int,
|
||||
base_config: dict,
|
||||
device=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.edge_dropout = base_config.get("edge_dropout", 0.0)
|
||||
self.create_copy = base_config.get("create_graph_copy", True)
|
||||
|
||||
scatter_name = base_config.get("scatter_reduce", "mean")
|
||||
self.scatter_reducer = get_scatter_reduce(scatter_name)
|
||||
|
||||
self.node_embedding = nn.Linear(in_node_features, latent_dim)
|
||||
self.edge_embedding = nn.Linear(in_edge_features, latent_dim)
|
||||
|
||||
# Stack
|
||||
stack_config = base_config.get("stack", {})
|
||||
self.stack = MessagePassingStack(latent_dim, stack_config, self.scatter_reducer)
|
||||
|
||||
if device is not None:
|
||||
self.to(device)
|
||||
|
||||
def forward(self, graph: Data):
|
||||
if self.create_copy:
|
||||
graph = copy.deepcopy(graph)
|
||||
|
||||
# Edge dropout (training only)
|
||||
if self.edge_dropout > 0 and self.training:
|
||||
graph.edge_index, mask = dropout_edge(
|
||||
graph.edge_index, p=self.edge_dropout, training=True
|
||||
)
|
||||
graph.edge_attr = graph.edge_attr[mask]
|
||||
|
||||
# Embed
|
||||
graph.x = self.node_embedding(graph.x)
|
||||
graph.edge_attr = self.edge_embedding(graph.edge_attr)
|
||||
|
||||
# Message passing
|
||||
self.stack(graph)
|
||||
|
||||
# Unpack
|
||||
node_name = "element" # homogeneous graph node type for mesh refinement
|
||||
batch = (
|
||||
graph.batch
|
||||
if hasattr(graph, "batch") and graph.batch is not None
|
||||
else torch.zeros(graph.x.shape[0], dtype=torch.long, device=graph.x.device)
|
||||
)
|
||||
|
||||
edge_key = f"{node_name}2{node_name}"
|
||||
return (
|
||||
{node_name: graph.x},
|
||||
{
|
||||
edge_key: {
|
||||
"edge_index": graph.edge_index.long(),
|
||||
"edge_attr": graph.edge_attr,
|
||||
}
|
||||
},
|
||||
None,
|
||||
{node_name: batch},
|
||||
)
|
||||
|
||||
|
||||
# ──
|
||||
# 8. MLP — Actor/Critic 头使用的 MLP
|
||||
# ──
|
||||
class MLP(nn.Module):
|
||||
"""Feedforward MLP for actor/critic heads."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
config: dict,
|
||||
latent_dim: int = None,
|
||||
out_features: int = None,
|
||||
device=None,
|
||||
):
|
||||
super().__init__()
|
||||
activation = config.get("activation_function", "tanh").lower()
|
||||
num_layers = config.get("num_layers", 2)
|
||||
|
||||
layers = []
|
||||
prev = in_features
|
||||
dim = latent_dim or 64
|
||||
for _ in range(num_layers):
|
||||
layers.append(nn.Linear(prev, dim))
|
||||
layers.append(_get_activation(activation))
|
||||
prev = dim
|
||||
|
||||
if out_features is not None:
|
||||
layers.append(nn.Linear(prev, out_features))
|
||||
self._out_features = out_features
|
||||
else:
|
||||
self._out_features = prev
|
||||
|
||||
self.net = nn.Sequential(*layers)
|
||||
if device is not None:
|
||||
self.to(device)
|
||||
|
||||
@property
|
||||
def out_features(self) -> int:
|
||||
return self._out_features
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.net(x)
|
||||
|
||||
|
||||
# ──
|
||||
# 9. ActorCritic — PPO Actor-Critic 网络
|
||||
# ──
|
||||
def create_model(env, network_config: dict, ppo_config: dict, device=None):
|
||||
"""Factory function: create Actor-Critic model from environment and configs."""
|
||||
return ActorCritic(
|
||||
environment=env,
|
||||
network_config=network_config,
|
||||
ppo_config=ppo_config,
|
||||
device=device,
|
||||
)
|
||||
|
||||
|
||||
class ActorCritic(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self, environment, network_config: dict, ppo_config: dict, device=None
|
||||
):
|
||||
super().__init__()
|
||||
latent_dim = network_config.get("latent_dimension", 64)
|
||||
base_config = network_config.get("base", {})
|
||||
train_config = network_config.get("training", {})
|
||||
actor_cfg = network_config.get("actor", {}).get("mlp", {})
|
||||
critic_cfg = network_config.get("critic", {}).get("mlp", {})
|
||||
|
||||
self.value_function_aggr = ppo_config.get("value_function_aggr", "spatial")
|
||||
self.agent_node_type = "element"
|
||||
|
||||
self.base = MessagePassingBase(
|
||||
in_node_features=environment.num_node_features,
|
||||
in_edge_features=environment.num_edge_features,
|
||||
latent_dim=latent_dim,
|
||||
base_config=base_config,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self.policy_mlp = MLP(
|
||||
latent_dim, actor_cfg, latent_dim=latent_dim, device=device
|
||||
)
|
||||
action_dim = environment.action_dimension
|
||||
if isinstance(environment._action_space, gym.spaces.Box):
|
||||
from stable_baselines3.common.distributions import DiagGaussianDistribution
|
||||
|
||||
self.action_dist = DiagGaussianDistribution(action_dim)
|
||||
self.action_out, self.log_std = self.action_dist.proba_distribution_net(
|
||||
latent_dim=self.policy_mlp.out_features,
|
||||
log_std_init=ppo_config.get("initial_log_std", 0.0),
|
||||
)
|
||||
else:
|
||||
from stable_baselines3.common.distributions import CategoricalDistribution
|
||||
|
||||
self.action_dist = CategoricalDistribution(action_dim)
|
||||
self.action_out = self.action_dist.proba_distribution_net(
|
||||
latent_dim=self.policy_mlp.out_features
|
||||
)
|
||||
self.log_std = None
|
||||
|
||||
self.value_mlp = MLP(
|
||||
latent_dim, critic_cfg, latent_dim=latent_dim, out_features=1, device=device
|
||||
)
|
||||
|
||||
self._setup_optimizer(train_config)
|
||||
|
||||
if device is not None:
|
||||
self.to(device)
|
||||
|
||||
def _setup_optimizer(self, train_config: dict):
|
||||
lr = train_config.get("learning_rate", 3e-4)
|
||||
wd = train_config.get("l2_norm", 0)
|
||||
params = list(self.parameters())
|
||||
if self.log_std is not None and not any(p is self.log_std for p in params):
|
||||
params.append(self.log_std)
|
||||
self.optimizer = optim.Adam(params, lr=lr, weight_decay=wd)
|
||||
sched_rate = train_config.get("lr_decay", train_config.get("lr_scheduling_rate", 1))
|
||||
self.lr_scheduler = (
|
||||
optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=sched_rate)
|
||||
if sched_rate is not None and sched_rate < 1
|
||||
else None
|
||||
)
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return next(self.parameters()).device
|
||||
|
||||
def _encode(self, observations):
|
||||
"""Run shared GNN backbone once, return (shared_features, batch_indices)."""
|
||||
observations = observations.to(self.device)
|
||||
node_feats, _, _, batches = self.base(observations)
|
||||
batch = batches[self.agent_node_type]
|
||||
feats = node_feats[self.agent_node_type]
|
||||
return feats, batch
|
||||
|
||||
def _make_distribution(self, latent_pi):
|
||||
mean_actions = self.action_out(latent_pi)
|
||||
if self.log_std is not None:
|
||||
return self.action_dist.proba_distribution(mean_actions, self.log_std)
|
||||
return self.action_dist.proba_distribution(mean_actions)
|
||||
|
||||
def _aggregate_values(self, values, batch):
|
||||
if self.value_function_aggr == "mean":
|
||||
return scatter_mean(values, batch, dim=0)
|
||||
elif self.value_function_aggr == "sum":
|
||||
from torch_scatter import scatter_add
|
||||
return scatter_add(values, batch, dim=0)
|
||||
elif self.value_function_aggr == "max":
|
||||
from torch_scatter import scatter_max
|
||||
return scatter_max(values, batch, dim=0)[0]
|
||||
return values
|
||||
|
||||
def forward(self, observations, deterministic: bool = False):
|
||||
shared_feats, batch = self._encode(observations)
|
||||
dist = self._make_distribution(self.policy_mlp(shared_feats))
|
||||
actions = dist.get_actions(deterministic=deterministic)
|
||||
log_probs = dist.log_prob(actions)
|
||||
values = self._aggregate_values(self.value_mlp(shared_feats).squeeze(-1), batch)
|
||||
return actions, values, log_probs
|
||||
|
||||
def evaluate_actions(self, observations, actions):
|
||||
actions = actions.to(self.device)
|
||||
shared_feats, batch = self._encode(observations)
|
||||
dist = self._make_distribution(self.policy_mlp(shared_feats))
|
||||
values = self._aggregate_values(self.value_mlp(shared_feats).squeeze(-1), batch)
|
||||
return values, dist.log_prob(actions), dist.entropy()
|
||||
|
|
@ -0,0 +1,274 @@
|
|||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch_geometric.data import Batch
|
||||
from torch_scatter import scatter_add
|
||||
|
||||
|
||||
class RolloutBuffer:
|
||||
|
||||
def __init__(self, buffer_size: int,
|
||||
gae_lambda: float,
|
||||
discount_factor: float,
|
||||
device=None,
|
||||
):
|
||||
self.buffer_size = buffer_size
|
||||
self.gae_lambda = gae_lambda
|
||||
self.discount_factor = discount_factor
|
||||
self.device = device
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.observations = []
|
||||
self.actions = []
|
||||
self.log_probs = []
|
||||
self.rewards = [] # per-agent rewards (list of tensors, varying shapes)
|
||||
self.values = [] # per-agent values (list of tensors, varying shapes)
|
||||
self.dones = []
|
||||
self.agent_mappings = [] # mapping from new → old agent indices per step
|
||||
self.pos = 0
|
||||
|
||||
def add(
|
||||
self, observation, actions, reward, done, value, log_probs,
|
||||
agent_mapping=None,
|
||||
):
|
||||
dev = self.device
|
||||
self.observations.append(observation.to(dev))
|
||||
self.actions.append(actions.to(dev))
|
||||
self.log_probs.append(log_probs.to(dev))
|
||||
self.rewards.append(torch.as_tensor(reward, dtype=torch.float32, device=dev).flatten())
|
||||
self.values.append(value.flatten().to(dev))
|
||||
self.dones.append(float(done))
|
||||
self.agent_mappings.append(
|
||||
torch.as_tensor(agent_mapping, dtype=torch.long, device=dev).flatten()
|
||||
)
|
||||
self.pos += 1
|
||||
|
||||
def compute_returns_and_advantage(self, last_value):
|
||||
"""Single-path GAE: potential-shaped per-agent reward with scatter_add for mesh refinement."""
|
||||
last_value = last_value.to(self.device).flatten()
|
||||
n = self.buffer_size
|
||||
|
||||
dones = torch.as_tensor(self.dones, device=self.device)
|
||||
|
||||
# ---- 0. Normalize rewards to unit scale ----
|
||||
all_rews = torch.cat([r.flatten() for r in self.rewards])
|
||||
rew_mean = all_rews.mean()
|
||||
rew_std = all_rews.std()
|
||||
if rew_std > 1e-8:
|
||||
self.rewards = [(r - rew_mean) / rew_std for r in self.rewards]
|
||||
|
||||
# ---- 1. Per-agent GAE (scatter_add for mesh refinement) ----
|
||||
advantages = [None] * n
|
||||
deltas = []
|
||||
next_values = self.values[1:] + [last_value]
|
||||
|
||||
for step in range(n):
|
||||
if dones[step]:
|
||||
next_val = self.values[step]
|
||||
else:
|
||||
next_val = scatter_add(next_values[step], self.agent_mappings[step], dim=0)
|
||||
delta = self.rewards[step] + (0 if dones[step] else self.discount_factor * next_val) - self.values[step]
|
||||
deltas.append(delta)
|
||||
|
||||
last_gae = torch.zeros_like(self.agent_mappings[-1], dtype=torch.float32, device=self.device)
|
||||
for step in reversed(range(n)):
|
||||
if dones[step]:
|
||||
last_gae = deltas[step]
|
||||
else:
|
||||
last_gae = deltas[step] + self.discount_factor * self.gae_lambda * scatter_add(last_gae, self.agent_mappings[step], dim=0)
|
||||
advantages[step] = last_gae
|
||||
|
||||
self.returns = [adv + val for adv, val in zip(advantages, self.values)]
|
||||
|
||||
# ---- 2. Normalize advantages (per-batch, zero-mean unit-std) ----
|
||||
all_advs = torch.cat([a.flatten() for a in advantages])
|
||||
adv_mean = all_advs.mean()
|
||||
adv_std = all_advs.std()
|
||||
if adv_std > 1e-8:
|
||||
advantages = [(a - adv_mean) / adv_std for a in advantages]
|
||||
# NOTE: returns and values keep their original scale — no unit-scale normalization,
|
||||
# so the value network sees a stable regression target across iterations.
|
||||
|
||||
self.advantages = [ret - val for ret, val in zip(self.returns, self.values)]
|
||||
|
||||
def get(self, batch_size: int):
|
||||
"""Yield random minibatches from the buffer."""
|
||||
indices = np.random.permutation(self.buffer_size)
|
||||
start = 0
|
||||
while start < self.buffer_size:
|
||||
batch_idx = indices[start : start + batch_size]
|
||||
start += batch_size
|
||||
|
||||
obs_batch = Batch.from_data_list([self.observations[i] for i in batch_idx])
|
||||
acts = torch.cat([self.actions[i] for i in batch_idx], dim=0)
|
||||
lps = torch.cat([self.log_probs[i].flatten() for i in batch_idx], dim=0)
|
||||
vals = torch.cat([self.values[i].flatten() for i in batch_idx], dim=0)
|
||||
advs = torch.cat([self.advantages[i].flatten() for i in batch_idx], dim=0)
|
||||
rets = torch.cat([self.returns[i].flatten() for i in batch_idx], dim=0)
|
||||
|
||||
obs_batch, acts, lps, vals, advs, rets = (
|
||||
x.to(self.device) for x in (obs_batch, acts, lps, vals, advs, rets)
|
||||
)
|
||||
yield obs_batch, acts, lps, vals, advs, rets
|
||||
|
||||
@property
|
||||
def full(self):
|
||||
return self.pos >= self.buffer_size
|
||||
|
||||
@property
|
||||
def explained_variance(self):
|
||||
all_vals = torch.cat([v.flatten() for v in self.values])
|
||||
all_rets = torch.cat([r.flatten() for r in self.returns])
|
||||
var_ret = torch.var(all_rets)
|
||||
if var_ret < 1e-12:
|
||||
return 0.0
|
||||
return float(1.0 - torch.var(all_rets - all_vals) / var_ret)
|
||||
|
||||
|
||||
# ── PPO losses ────────────────────────────────────────────
|
||||
def policy_loss(advantages: torch.Tensor, ratio: torch.Tensor, clip_range: float) -> torch.Tensor:
|
||||
"""Clipped PPO policy loss."""
|
||||
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
|
||||
loss1 = advantages * ratio
|
||||
loss2 = advantages * torch.clamp(ratio, 1.0 - clip_range, 1.0 + clip_range)
|
||||
return -torch.min(loss1, loss2).mean()
|
||||
|
||||
|
||||
def value_loss(
|
||||
returns: torch.Tensor, values: torch.Tensor,
|
||||
old_values: torch.Tensor, clip_range: float,
|
||||
) -> torch.Tensor:
|
||||
"""Clipped value function loss."""
|
||||
vf_loss = F.mse_loss(returns, values)
|
||||
if clip_range > 0:
|
||||
v_clipped = old_values + (values - old_values).clamp(-clip_range, clip_range)
|
||||
vf_loss = torch.max(vf_loss, F.mse_loss(returns, v_clipped))
|
||||
return vf_loss
|
||||
|
||||
|
||||
def entropy_loss(entropy) -> torch.Tensor:
|
||||
"""Entropy bonus for exploration."""
|
||||
return -torch.mean(entropy)
|
||||
|
||||
|
||||
class PPOTrainer:
|
||||
|
||||
def __init__(self, actor_critic, environment, config: dict, device=None):
|
||||
self.policy = actor_critic
|
||||
self.env = environment
|
||||
self.device = device
|
||||
|
||||
ppo_cfg = config.get("ppo", {})
|
||||
self.num_rollout_steps = ppo_cfg.get("num_rollout_steps", 256)
|
||||
self.epochs_per_iteration = ppo_cfg.get("epochs_per_iteration", 5)
|
||||
self.batch_size = config.get("batch_size", 32)
|
||||
self.clip_range = ppo_cfg.get("clip_range", 0.2)
|
||||
self.max_grad_norm = ppo_cfg.get("max_grad_norm", 0.5)
|
||||
self.entropy_coef = ppo_cfg.get("entropy_coefficient", 0.0)
|
||||
self.vf_coef = ppo_cfg.get("value_function_coefficient", 0.5)
|
||||
self.vf_clip_range = ppo_cfg.get("value_function_clip_range", 0.2)
|
||||
self.gae_lambda = ppo_cfg.get("gae_lambda", 0.95)
|
||||
self.discount_factor = config.get("discount_factor", 1.0)
|
||||
|
||||
self.buffer = RolloutBuffer(
|
||||
buffer_size=self.num_rollout_steps,
|
||||
gae_lambda=self.gae_lambda,
|
||||
discount_factor=self.discount_factor,
|
||||
device=device,
|
||||
)
|
||||
|
||||
def collect_rollouts(self):
|
||||
self.policy.eval()
|
||||
self.buffer.reset()
|
||||
obs = self.env.reset()
|
||||
step_rewards, step_num_agents = [], []
|
||||
_rho_keys = ("rho_int_mean", "rho_jump_mean", "rho_sbc_mean",
|
||||
"w_rho_int", "w_rho_jump", "w_rho_sbc")
|
||||
rho_accum = {k: 0.0 for k in _rho_keys}
|
||||
diag_keys = ("neg_action_ratio", "eligible_ratio", "selected_count")
|
||||
diag_accum = {k: 0.0 for k in diag_keys}
|
||||
diag_steps = 0
|
||||
|
||||
for _ in range(self.num_rollout_steps):
|
||||
with torch.no_grad():
|
||||
actions, values, log_probs = self.policy(
|
||||
Batch.from_data_list([obs]), deterministic=False
|
||||
)
|
||||
values = values.flatten()
|
||||
next_obs, reward, done, info = self.env.step(actions.cpu().numpy())
|
||||
step_rewards.append(float(np.sum(reward)))
|
||||
step_num_agents.append(int(len(reward)))
|
||||
for k in _rho_keys:
|
||||
if k in info:
|
||||
rho_accum[k] += float(info[k])
|
||||
for k in diag_keys:
|
||||
if k in info:
|
||||
diag_accum[k] += float(info[k])
|
||||
diag_steps += 1
|
||||
|
||||
self.buffer.add(
|
||||
observation=obs, actions=actions, reward=reward,
|
||||
done=float(done), value=values, log_probs=log_probs,
|
||||
agent_mapping=self.env.agent_mapping,
|
||||
)
|
||||
obs = self.env.reset() if done else next_obs
|
||||
|
||||
with torch.no_grad():
|
||||
_, last_value, _ = self.policy(Batch.from_data_list([obs]), deterministic=True)
|
||||
last_value = last_value.squeeze(-1).flatten()
|
||||
self.buffer.compute_returns_and_advantage(last_value)
|
||||
|
||||
n = max(1, self.num_rollout_steps)
|
||||
metrics = {
|
||||
"num_agents": step_num_agents[-1], "reward": step_rewards[-1],
|
||||
"avg_agents": np.mean(step_num_agents),
|
||||
"avg_reward": np.mean(step_rewards),
|
||||
"min_reward": np.min(step_rewards),
|
||||
"max_reward": np.max(step_rewards),
|
||||
"sum_reward": np.sum(step_rewards),
|
||||
}
|
||||
# rho diagnostics for weight calibration (averaged over rollout)
|
||||
for k in _rho_keys:
|
||||
metrics[k] = rho_accum[k] / n
|
||||
# score-based refinement diagnostics
|
||||
n_diag = max(1, diag_steps)
|
||||
for k in diag_keys:
|
||||
metrics[k] = diag_accum[k] / n_diag
|
||||
return metrics
|
||||
|
||||
def train_step(self):
|
||||
self.policy.train()
|
||||
total_losses = []
|
||||
for _ in range(self.epochs_per_iteration):
|
||||
for obs_batch, acts, old_lp, old_vals, advs, rets in self.buffer.get(self.batch_size):
|
||||
values, log_probs, entropy = self.policy.evaluate_actions(obs_batch, acts)
|
||||
values = values.squeeze(-1)
|
||||
|
||||
ratio = torch.exp(log_probs - old_lp)
|
||||
|
||||
pl = policy_loss(advs, ratio, self.clip_range)
|
||||
vl = self.vf_coef * value_loss(rets, values, old_vals, self.vf_clip_range)
|
||||
el = self.entropy_coef * entropy_loss(entropy)
|
||||
loss = pl + vl + el
|
||||
|
||||
self.policy.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
|
||||
self.policy.optimizer.step()
|
||||
if self.policy.log_std is not None:
|
||||
self.policy.log_std.data.clamp_(-4.0, -1.0)
|
||||
total_losses.append(loss.item())
|
||||
|
||||
if self.policy.lr_scheduler is not None:
|
||||
self.policy.lr_scheduler.step()
|
||||
|
||||
return {
|
||||
"loss": np.mean(total_losses) if total_losses else 0.0,
|
||||
"explained_variance": self.buffer.explained_variance,
|
||||
}
|
||||
|
||||
def fit_iteration(self):
|
||||
metrics = self.collect_rollouts()
|
||||
metrics.update(self.train_step())
|
||||
return metrics
|
||||
|
|
@ -0,0 +1,63 @@
|
|||
import os
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
|
||||
def load_config(path: str) -> dict:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
return yaml.safe_load(f)
|
||||
|
||||
|
||||
def save_checkpoint(model, optimizer: torch.optim.Optimizer, iteration: int, path: str):
|
||||
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
||||
torch.save(
|
||||
{
|
||||
"iteration": iteration,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"optimizer_state_dict": optimizer.state_dict(),
|
||||
},
|
||||
path,
|
||||
)
|
||||
print(f"[Checkpoint] saved → {path}")
|
||||
|
||||
|
||||
def load_checkpoint(model, path: str, device=None) -> int:
|
||||
ckpt = torch.load(path, map_location=device or "cpu")
|
||||
model.load_state_dict(ckpt["model_state_dict"], strict=False)
|
||||
if "optimizer_state_dict" in ckpt and hasattr(model, "optimizer"):
|
||||
try:
|
||||
model.optimizer.load_state_dict(ckpt["optimizer_state_dict"])
|
||||
except Exception:
|
||||
pass
|
||||
it = ckpt.get("iteration", 0)
|
||||
print(f"[Checkpoint] loaded ← {path} (iter {it})")
|
||||
return it
|
||||
|
||||
|
||||
def setup_helmholtz_config(config: dict, k_test=None, center=None, radius=None, eps_test=None) -> float:
|
||||
"""Lock scatterer/helmholtz config for test/viz. Returns wave number k."""
|
||||
hc = config.setdefault("environment", {}).setdefault("mesh_refinement", {}).setdefault("fem", {}).setdefault("helmholtz", {})
|
||||
sc = hc.setdefault("scatterer", {})
|
||||
sc["mode"] = "fixed"
|
||||
if center is not None:
|
||||
sc["cx"], sc["cy"] = center[0], center[1]
|
||||
if radius is not None:
|
||||
sc["radius"] = radius
|
||||
if eps_test is not None:
|
||||
sc["eps_r"] = eps_test
|
||||
if k_test is not None:
|
||||
hc["wave_number_mode"] = "fixed"
|
||||
hc["wave_number"] = k_test
|
||||
return hc.get("wave_number", 6.0)
|
||||
|
||||
|
||||
def parse_center(center_str: Optional[str]) -> Optional[Tuple[float, float]]:
|
||||
if center_str is None:
|
||||
return None
|
||||
parts = center_str.split(",")
|
||||
if len(parts) != 2:
|
||||
raise ValueError(f"Invalid --center format (expected 'cx,cy'): {center_str}")
|
||||
return (float(parts[0].strip()), float(parts[1].strip()))
|
||||
|
|
@ -0,0 +1,293 @@
|
|||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch_geometric.data import Batch
|
||||
|
||||
|
||||
# ── 高分辨率 FEM 参考解（保留作为回退） ──────────────────────────
|
||||
def _compute_fem_reference(env):
|
||||
from skfem import Basis, ElementTriP1
|
||||
|
||||
fp = env.fem_problem.fem_problem
|
||||
ref_mesh = fp._domain.get_integration_mesh()
|
||||
ref_basis = Basis(ref_mesh, ElementTriP1())
|
||||
ref_sol = fp.calculate_solution(ref_basis, cache=False)
|
||||
return ref_mesh, ref_sol
|
||||
|
||||
|
||||
# ── Mie 解析参考解 ──────────────────────────────────────────────
|
||||
def _compute_mie_reference(env):
|
||||
"""Return Mie scattered field sampled at FEM mesh vertices.
|
||||
|
||||
Falls back to FEM reference if scatterer is non-circular.
|
||||
"""
|
||||
from environment.mie_solution import mie_scattered_field
|
||||
|
||||
fp = getattr(env.fem_problem, "fem_problem", None)
|
||||
if fp is None:
|
||||
return _compute_fem_reference(env), None
|
||||
|
||||
_eps_r = getattr(fp, "_eps_r", None)
|
||||
_radius = getattr(fp, "_radius", None)
|
||||
_cx = getattr(fp, "_cx", None)
|
||||
_cy = getattr(fp, "_cy", None)
|
||||
_k = getattr(fp, "_k", None)
|
||||
|
||||
if any(v is None for v in [_eps_r, _radius, _cx, _cy, _k]):
|
||||
return _compute_fem_reference(env), None
|
||||
|
||||
pts = env.mesh.p.T
|
||||
u_mie = mie_scattered_field(pts, k0=_k, eps_r=_eps_r, radius=_radius, cx=_cx, cy=_cy)
|
||||
|
||||
from environment.mie_solution import mie_grid_solution
|
||||
import matplotlib.tri as tri
|
||||
|
||||
xlim = (pts[:, 0].min(), pts[:, 0].max())
|
||||
ylim = (pts[:, 1].min(), pts[:, 1].max())
|
||||
grid = mie_grid_solution(_k, _eps_r, _radius, _cx, _cy,
|
||||
x_range=xlim, y_range=ylim, Nx=500, Ny=500)
|
||||
|
||||
mie_info = {
|
||||
"grid": grid,
|
||||
"eps_r": _eps_r, "radius": _radius,
|
||||
"cx": _cx, "cy": _cy, "k": _k,
|
||||
}
|
||||
return u_mie, mie_info
|
||||
|
||||
|
||||
# ── 渲染辅助 ─────────────────────────────────────────────────────
|
||||
def _render_field(ax, x, y, triang, values, title, vmin, vmax, show_mesh=True, cmap="jet"):
|
||||
tcf = ax.tripcolor(triang, values, shading="gouraud", cmap=cmap, vmin=vmin, vmax=vmax)
|
||||
if show_mesh and triang is not None:
|
||||
n = triang.triangles.shape[0]
|
||||
ax.triplot(triang, lw=(0.5 if n < 500 else 0.3), color="black",
|
||||
alpha=(0.7 if n < 2000 else 0.5))
|
||||
ax.set_xlim(x.min(), x.max())
|
||||
ax.set_ylim(y.min(), y.max())
|
||||
ax.set_aspect("equal")
|
||||
ax.set_title(title, fontsize=9)
|
||||
ax.set_xticks([])
|
||||
ax.set_yticks([])
|
||||
return tcf
|
||||
|
||||
|
||||
# ── 保存 PNG ─────────────────────────────────────────────────────
|
||||
def _save_png(steps, stem, checkpoint_path, k, cx=0.5, cy=0.5, radius=0.2, eps_r=2.0,
|
||||
mie_info=None):
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.tri as tri
|
||||
|
||||
per_step_dir = f"{stem}_steps"
|
||||
os.makedirs(os.path.dirname(stem) or ".", exist_ok=True)
|
||||
os.makedirs(per_step_dir, exist_ok=True)
|
||||
|
||||
n = len(steps)
|
||||
ncols = min(n, 4)
|
||||
nrows = (n + ncols - 1) // ncols
|
||||
fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 3.5 * nrows))
|
||||
if nrows * ncols == 1:
|
||||
axes = np.array([axes])
|
||||
else:
|
||||
axes = np.array(axes).flatten()
|
||||
|
||||
for i, step_data in enumerate(steps):
|
||||
mesh, scalar, err_val, n_elem = step_data[:4]
|
||||
pts = mesh.p.T
|
||||
tg = tri.Triangulation(pts[:, 0], pts[:, 1], mesh.t.T)
|
||||
s = np.abs(scalar) if np.iscomplexobj(scalar) else scalar
|
||||
lmin, lmax = s.min(), s.max()
|
||||
if lmax - lmin < 1e-12:
|
||||
lmin, lmax = lmin - 0.5, lmax + 0.5
|
||||
tcf = _render_field(axes[i], pts[:, 0], pts[:, 1], tg, s,
|
||||
f"Step {i}: {n_elem} elem, err={err_val:.4f}",
|
||||
lmin, lmax, cmap="jet")
|
||||
fig.colorbar(tcf, ax=axes[i], fraction=0.046, pad=0.04)
|
||||
axes[i].add_patch(plt.Circle((cx, cy), radius, fill=False,
|
||||
edgecolor="cyan", linewidth=1.5, linestyle="--"))
|
||||
|
||||
for j in range(n, len(axes)):
|
||||
axes[j].set_visible(False)
|
||||
|
||||
fig.subplots_adjust(left=0.04, right=0.90, top=0.90, bottom=0.06, wspace=0.15, hspace=0.30)
|
||||
k_str = f"k={k:.1f}" if k is not None else "k=?"
|
||||
ref_tag = " [Mie ref]" if mie_info is not None else ""
|
||||
fig.suptitle(
|
||||
f"Helmholtz |E_scat|{ref_tag} — {checkpoint_path}\n"
|
||||
f"{k_str}, eps_r={eps_r:.1f} at ({cx:.2f},{cy:.2f}) r={radius:.2f}",
|
||||
fontsize=12,
|
||||
)
|
||||
fig.savefig(f"{stem}.png", dpi=200, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f"[Viz] Overview → {stem}.png")
|
||||
|
||||
for i, step_data in enumerate(steps):
|
||||
mesh, scalar, err_val, n_elem = step_data[:4]
|
||||
u_mie_at_verts = step_data[4] if len(step_data) > 4 else None
|
||||
|
||||
pts = mesh.p.T
|
||||
tg_coarse = tri.Triangulation(pts[:, 0], pts[:, 1], mesh.t.T)
|
||||
coarse_val = np.abs(scalar) if np.iscomplexobj(scalar) else scalar
|
||||
|
||||
has_mie = u_mie_at_verts is not None
|
||||
ncols = 3 if has_mie else 1
|
||||
fig2, axes2 = plt.subplots(1, ncols, figsize=(6 * ncols, 6))
|
||||
axes2 = [axes2] if ncols == 1 else list(np.atleast_1d(axes2))
|
||||
|
||||
# ── Panel 1: FEM scattered field ──
|
||||
cvmin, cvmax = coarse_val.min(), coarse_val.max()
|
||||
if cvmax - cvmin < 1e-12:
|
||||
cvmin, cvmax = cvmin - 0.5, cvmax + 0.5
|
||||
tcf1 = _render_field(axes2[0], pts[:, 0], pts[:, 1], tg_coarse, coarse_val,
|
||||
f"Step {i}: FEM |E_scat| ({n_elem} elem) max={cvmax:.4f}",
|
||||
cvmin, cvmax, cmap="jet")
|
||||
axes2[0].add_patch(plt.Circle((cx, cy), radius, fill=False,
|
||||
edgecolor="cyan", linewidth=1.5, linestyle="--"))
|
||||
fig2.colorbar(tcf1, ax=axes2[0], fraction=0.046, pad=0.04)
|
||||
|
||||
im2 = None
|
||||
if has_mie:
|
||||
# ── Panel 2: Mie scattered field (smooth grid, not FEM vertices) ──
|
||||
if mie_info is not None and "grid" in mie_info:
|
||||
g = mie_info["grid"]
|
||||
gm = np.abs(g["E_scat"])
|
||||
mvmin, mvmax = gm.min(), gm.max()
|
||||
if mvmax - mvmin < 1e-12:
|
||||
mvmin, mvmax = mvmin - 0.5, mvmax + 0.5
|
||||
im2 = axes2[1].pcolormesh(g["X"], g["Y"], gm,
|
||||
shading="gouraud", cmap="jet",
|
||||
vmin=mvmin, vmax=mvmax)
|
||||
axes2[1].set_title(f"Mie |E_scat| max={mvmax:.4f}", fontsize=9)
|
||||
else:
|
||||
mie_abs = np.abs(u_mie_at_verts)
|
||||
mvmin, mvmax = mie_abs.min(), mie_abs.max()
|
||||
if mvmax - mvmin < 1e-12:
|
||||
mvmin, mvmax = mvmin - 0.5, mvmax + 0.5
|
||||
im2 = _render_field(axes2[1], pts[:, 0], pts[:, 1], tg_coarse, mie_abs,
|
||||
f"Mie |E_scat| max={mvmax:.4f}",
|
||||
mvmin, mvmax, show_mesh=False, cmap="jet")
|
||||
axes2[1].set_aspect("equal")
|
||||
axes2[1].set_xticks([])
|
||||
axes2[1].set_yticks([])
|
||||
axes2[1].add_patch(plt.Circle((cx, cy), radius, fill=False,
|
||||
edgecolor="cyan", linewidth=1.5, linestyle="--"))
|
||||
if im2 is not None:
|
||||
fig2.colorbar(im2, ax=axes2[1], fraction=0.046, pad=0.04)
|
||||
|
||||
# ── Panel 3: ||FEM| - |Mie|| error ──
|
||||
mie_abs = np.abs(u_mie_at_verts)
|
||||
error_abs = np.abs(coarse_val - mie_abs)
|
||||
evmin, evmax = 0.0, error_abs.max() or 1.0
|
||||
if evmax - evmin < 1e-12:
|
||||
evmax = evmin + 1.0
|
||||
tcf3 = _render_field(axes2[2], pts[:, 0], pts[:, 1], tg_coarse, error_abs,
|
||||
f"||FEM|-|Mie|| L2={err_val:.4f} max={error_abs.max():.4f}",
|
||||
evmin, evmax, show_mesh=True, cmap="hot")
|
||||
axes2[2].add_patch(plt.Circle((cx, cy), radius, fill=False,
|
||||
edgecolor="cyan", linewidth=1.5, linestyle="--"))
|
||||
fig2.colorbar(tcf3, ax=axes2[2], fraction=0.046, pad=0.04)
|
||||
|
||||
fig2.tight_layout()
|
||||
fig2.savefig(f"{per_step_dir}/step{i:02d}.png", dpi=150, bbox_inches="tight")
|
||||
plt.close(fig2)
|
||||
|
||||
print(f"[Viz] Per-step PNGs → {per_step_dir}/ ({n} files)")
|
||||
|
||||
|
||||
# ── Viz 模式入口 ──────────────────────────────────────────────────
|
||||
def visualize(config: dict, checkpoint_path: str, output_path: str = "result/visualization.png",
|
||||
k_test=None, center=None, radius=None, eps_test=None):
|
||||
from src.network import create_model
|
||||
from src.utils import load_checkpoint, setup_helmholtz_config
|
||||
|
||||
k = setup_helmholtz_config(config, k_test=k_test, center=center, radius=radius,
|
||||
eps_test=eps_test)
|
||||
algo = config.get("algorithm", {})
|
||||
|
||||
from environment.mesh_refinement import MeshRefinement
|
||||
|
||||
env = MeshRefinement(
|
||||
environment_config=config.get("environment", {}).get("mesh_refinement", {}),
|
||||
seed=99,
|
||||
)
|
||||
model = create_model(env, config.get("network", {}), algo.get("ppo", {}))
|
||||
load_checkpoint(model, checkpoint_path)
|
||||
model.eval()
|
||||
|
||||
stem = output_path.rsplit(".", 1)[0] if "." in output_path else output_path
|
||||
|
||||
print(f"\n[Viz] Initializing...")
|
||||
obs = env.reset()
|
||||
|
||||
_fp = getattr(env.fem_problem, "fem_problem", None)
|
||||
_cx = getattr(_fp, "_cx", 0.5) if _fp is not None else 0.5
|
||||
_cy = getattr(_fp, "_cy", 0.5) if _fp is not None else 0.5
|
||||
_radius = getattr(_fp, "_radius", 0.2) if _fp is not None else 0.2
|
||||
_eps_r = getattr(_fp, "_eps_r", 2.0) if _fp is not None else 2.0
|
||||
|
||||
print(f"[Viz] Helmholtz params: k={k:.3f} eps_r={_eps_r:.2f} "
|
||||
f"center=({_cx:.3f}, {_cy:.3f}) radius={_radius:.3f}")
|
||||
|
||||
# ── Mie analytical reference ──
|
||||
print(f"[Viz] Computing Mie reference solution...")
|
||||
u_mie_ref, mie_info = _compute_mie_reference(env)
|
||||
if mie_info is not None:
|
||||
print(f"[Viz] Mie reference ready (analytical, no domain truncation error)")
|
||||
|
||||
# ── Initial step ──
|
||||
init_mesh = env.mesh
|
||||
init_sol = env.scalar_solution
|
||||
init_err = _compute_step_error(env, u_mie_ref)
|
||||
steps = [(init_mesh, init_sol, init_err, env.num_agents, u_mie_ref)]
|
||||
|
||||
print(f"[Viz] Running inference...")
|
||||
done = False
|
||||
step_idx = 0
|
||||
while not done:
|
||||
with torch.no_grad():
|
||||
actions, _, _ = model(Batch.from_data_list([obs]), deterministic=True)
|
||||
obs, _, done, _ = env.step(actions.cpu().numpy())
|
||||
step_idx += 1
|
||||
sol = env.scalar_solution
|
||||
n_elem = env.num_agents
|
||||
u_mie_current = _eval_mie_on_mesh(env, mie_info)
|
||||
step_err = _compute_step_error(env, u_mie_current)
|
||||
|
||||
diag_n_sel = getattr(env, "_diag_selected_count", -1)
|
||||
diag_n_elig = int(getattr(env, "_diag_eligible_ratio", 0) * env.num_agents)
|
||||
diag_n_mask = int(getattr(env, "_diag_masked_ratio", 0) * env.num_agents)
|
||||
remaining = getattr(env, "_n_budget", 0) - env.num_agents
|
||||
print(f" Step {step_idx}: verts={env.mesh.p.shape[1]} elem={n_elem} "
|
||||
f"mie_err={step_err:.4f} "
|
||||
f"sel={diag_n_sel} elig={diag_n_elig} masked={diag_n_mask} "
|
||||
f"remaining={remaining} done={done}")
|
||||
|
||||
steps.append((env.mesh, sol, step_err, n_elem, u_mie_current))
|
||||
|
||||
_save_png(steps, stem, checkpoint_path, k, cx=_cx, cy=_cy, radius=_radius,
|
||||
eps_r=_eps_r, mie_info=mie_info)
|
||||
print(f"[Viz] Done → {output_path}")
|
||||
|
||||
|
||||
def _compute_step_error(env, u_mie_ref) -> float:
|
||||
"""相对 L₂ 误差: ||u_fem − u_mie||₂ / ||u_mie||₂ (复数，含幅值+相位)。"""
|
||||
if u_mie_ref is None:
|
||||
return float("nan")
|
||||
u_fem = env.scalar_solution # complex scattered field
|
||||
diff = np.abs(u_fem - u_mie_ref) # pointwise |complex difference|
|
||||
denom = np.linalg.norm(np.abs(u_mie_ref))
|
||||
if denom < 1e-12:
|
||||
denom = 1.0
|
||||
return float(np.linalg.norm(diff) / denom)
|
||||
|
||||
|
||||
def _eval_mie_on_mesh(env, mie_info):
|
||||
"""Re-evaluate Mie scattered field on current FEM mesh vertices."""
|
||||
if mie_info is None:
|
||||
return None
|
||||
from environment.mie_solution import mie_scattered_field
|
||||
pts = env.mesh.p.T
|
||||
return mie_scattered_field(pts, k0=mie_info["k"], eps_r=mie_info["eps_r"],
|
||||
radius=mie_info["radius"], cx=mie_info["cx"], cy=mie_info["cy"])
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
# ================= 配置区 =================
|
||||
$ServerA_User = "dxw"
|
||||
$ServerA_IP = "222.20.97.222"
|
||||
$RemotePath = "/public/home/dxw/Codes/afem" # 服务器A上项目的绝对路径
|
||||
$LocalPath = "F:\ASMRplusplus-main" # 本地项目路径
|
||||
# ==========================================
|
||||
|
||||
Write-Host ">>> Step 1: Downloading code from Server A..." -ForegroundColor Cyan
|
||||
scp -r "${ServerA_User}@${ServerA_IP}:${RemotePath}/*" $LocalPath
|
||||
|
||||
Write-Host ">>> Step 2: Preparing to commit to Git..." -ForegroundColor Cyan
|
||||
Set-Location $LocalPath
|
||||
git add .
|
||||
|
||||
$date = Get-Date -Format "yyyy-MM-dd HH:mm:ss"
|
||||
git commit -m "Auto-sync from Server A at $date"
|
||||
|
||||
Write-Host ">>> Step 3: Pushing to Git Server B..." -ForegroundColor Cyan
|
||||
git push origin main
|
||||
|
||||
Write-Host "`n[Success] All operations completed!" -ForegroundColor Green
|
||||
Pause
|
||||
Loading…
Reference in New Issue