# Outlook: GNN-Guided Adaptive Mesh Refinement for 2D Helmholtz Scattering

## 1. 问题定义

求解 2D 介质圆柱的电磁散射（散射场公式）：

```
∇²u + k²·ε_r·u = −k²·(ε_r − 1)·u_inc
∂u/∂n − i·k_local·u = 0                    (Sommerfeld 辐射 BC)
```

- 入射波：`u_inc = exp(i·k·x)`，参考解：Mie 解析解
- 参数空间：k∈[3,15]，eps_r∈[2,8]，cx/cy∈[0.2,0.8]，radius∈[0.05,0.25]
- 核心目标：在参数化 Helmholtz 散射问题中，学习一个无在线求解的预算约束网格预测器，用于近似 residual-AMR 的最终加密分布，并在低预算下优于简单物理启发式网格。

---

## 2. 算法全流程

```
Step 1. 数据生成 (gen.py)
    残差驱动 AMR → 每步保存 cell 状态 + 标签
            │
            ▼
Step 2. 训练 (train_correction.py)
    features(15-dim) + physics_score → GNN → sigmoid → 二分类 teacher_mark
            │
            ▼
Step 3. 评估
    3a. 离线指标 (test_correction.py) — top-k overlap / AUC
    3b. Rollout 评估 (eval_correction.py) — 迭代加密 → FEM → aw_rel
            │
            ▼
Step 4. 可视化 (viz_correction.py)
    amr 模式: GNN 驱动完整 AMR → 网格 + 场 + 误差
    step 模式: 单步 GNN vs teacher vs physics 标记对比
```

---

## 3. 数据生成 (`gen.py`)

### 3.1 初始网格

物理自适应初始网格（`build_physics_safe_initial_mesh`）：

- 均匀基底网格 + 迭代局部加密
- 介质外：h ≤ λ₀/q，介质内：h ≤ λ_eff/q（q=2，每波长 2 个单元）
- 加密准则：`score = max_edge / h_target > 1` 的单元被加密
- 二分搜索控制每批加密数量，不超预算

### 3.2 AMR 循环（每步保存）

```
初始 mesh → FEM solve → 残差估计器 η → teacher_mark (η top-k)
                                    ├── physics_score → physics_mark (top-k)
                                    ├── correction_label = teacher_mark − physics_mark
                                    └── 保存 .npz → 按 η 选单元加密 → 下一步
```

**残差估计器 η**（`environment/helmholtz.py:_compute_residual_indicator`）：
- 内部残差：`h_K/k · √V_K · |k²ε_r·u_h + k²(ε_r−1)·u_inc|`
- 梯度跳变：`√(½ Σ h_e/k · |[[∇u_h·n]]|²)`
- SBC 边界残差：`h_bnd/k · |∂u/∂n − i·k_local·u|`

**标记策略**：
- teacher_mark：按 η 取 top `mark_fraction`（默认 3%）
- physics_mark：按 `physics_score` 取 top `mark_fraction`
- correction_label = +1（teacher 独有）、0（一致）、−1（physics 独有）

**安全过滤**：
1. 面积过滤：排除面积 ≤ 1e-10 的退化单元
2. 反向 Dörfler：排除误差贡献最低 1% 的单元

### 3.3 输出格式

```
outlook/data_correction/
├── params_list.npz              # (n_samples, 5) PDE 参数
├── sample0000_step000.npz       # 逐样本逐步数据
└── summary.json
```

每个 step .npz：

| 字段 | 形状 | 说明 |
|------|------|------|
| `features` | `(n_elem, 15)` | 几何/物理特征 |
| `edge_index` | `(2, n_edges)` | 网格图结构 |
| `physics_score` | `(n_elem,)` | h/λ_eff |
| `teacher_eta` | `(n_elem,)` | 残差估计器 |
| `teacher_mark` | `(n_elem,)` | η top-k 标记 (0/1) |
| `physics_mark` | `(n_elem,)` | physics top-k 标记 (0/1) |
| `correction_label` | `(n_elem,)` | teacher − physics (−1/0/+1) |

---

## 4. 特征工程

### 4.1 15 维基础特征

| 维 | 特征 | 说明 |
|----|------|------|
| 0,1 | x, y | 单元中点坐标 |
| 2 | area | 单元面积 |
| 3 | dist_to_center | 到圆柱中心距离 |
| 4 | signed_dist | dist − radius（负=介质内） |
| 5 | inside | 是否在圆柱内 (0/1) |
| 6 | k | 波数 |
| 7 | eps_r | 介电常数 |
| 8 | radius | 半径 |
| 9,10 | cx, cy | 圆柱中心 |
| 11 | k_h | k × √area |
| 12 | k_eps_h | k × √eps_r × √area |
| 13,14 | sin(k·x), cos(k·x) | 入射波相位 |

### 4.2 Physics Score（第 16 维输入）

```python
lambda_eff = 2π / (k · √eps_r)   # 介质内
          或 2π / k                # 介质外
physics_score = max_edge / lambda_eff   # >1 = 分辨率不足
```

### 4.3 归一化

训练集 z-score 归一化，推理复用同一统计量：
```python
x_norm = (concat(features, physics_score) - mean) / scale
```

---

## 5. GNN 架构

### 5.1 CorrectionGNN

基于 DensityGNN 骨干，替换密度回归头为二分类 logit 头：

```
Input: (n_cells, 16) node features
  │
  ├── node_embedding: Linear(16 → latent_dim)
  ├── edge_embedding: Linear(16 → latent_dim)
  │
  ├── N × MessagePassingStep:
  │   ├── EdgeModule: MLP([src | dst | edge_attr]) → latent_dim
  │   ├── NodeModule: MLP([node | mean(incoming_edges)]) → latent_dim
  │   └── LayerNorm + Residual
  │
  ├── GlobalVirtualNode: mean_pool → attention_gate → broadcast
  │
  └── head: Linear(latent → hidden) → ReLU → Linear(hidden → 1 logit)
```

### 5.2 关键设计

- 边特征：`edge_attr = |x[src] - x[dst]|`（节点特征差绝对值）
- 边丢弃：训练 0.1，推理 0.0
- 损失：BCEWithLogitsLoss + per-graph `pos_weight = neg/pos`

### 5.3 训练配置

| 参数 | 值 |
|------|-----|
| latent_dim | 64 |
| num_mp_steps | 3 |
| head_hidden | 64 |
| lr | 1e-3 |
| optimizer | Adam |
| scheduler | ReduceLROnPlateau (patience=10) |

---

## 6. 评估体系

### 6.1 离线指标（test_correction.py）

- **top-k overlap**：GNN 概率最高的 k 个 cell 与 teacher_mark 的交集 / k
- **AUC**：ROC-AUC（GNN vs physics baseline）
- **gnn_beats_physics_ratio**：GNN 优于 physics 的样本比例

### 6.2 Rollout 评估（eval_correction.py）

从初始 mesh 出发，**无 teacher_eta，无残差标记，无中间 FEM solve**。每步只用分数决定加密单元，最终一次 FEM solve 计算误差。

三种方法：

| 方法 | 打分 | 说明 |
|------|------|------|
| `physics` | `physics_score` | 纯物理先验 |
| `neural` | `model(features + physics_score)` | 纯 GNN |
| `hybrid` | `α·zscore(physics) + β·zscore(neural)` | 混合（默认 α=β=0.5） |

### 6.3 误差指标

**aw_rel**（面积加权相对误差）：
```
aw_rel = √( Σ err²_tri · area / Σ ref²_tri · area )
```

**max_err**（最大逐点误差）：
```
max_err = max |Re(u_fem) − Re(u_mie)|
```

---

## 7. 可视化

### 7.1 amr 模式

GNN 驱动完整 AMR，每步 FEM solve，展示网格和场演变。

输出：
- `amr_overview.png` — 所有步骤总览
- `amr_steps/step{XX}.png` — 每步 3 面板（FEM 场 / Mie 参考 / 误差）
- `ground_truth.png` — 高保真参考解（`--compare` 时）
- `compare.png` — Physics vs GNN vs Eta 对比（`--compare` 时）

### 7.2 step 模式

重建指定 AMR 步的 mesh，对比 GNN / teacher / physics 标记。

输出：
- `marks_*.png` — 2×2 对比图（teacher / GNN / physics / TP/FP/FN/TN）
- `field_gnn_*.png` — GNN 加密后 3 面板图
- `field_eta_*.png` — 传统 η 加密后 3 面板图

---

## 8. 使用方法

### 8.1 数据生成

```bash
python outlook/src/gen.py \
    --n-samples 100 \
    --max-elements 12000 \
    --mark-fraction 0.03 \
    --output-dir outlook/data_correction
```

### 8.2 训练

```bash
python outlook/src/train_correction.py \
    --data-dir outlook/data_correction \
    --epochs 100 \
    --batch-size 32 \
    --lr 1e-3 \
    --device cuda \
    --checkpoint-out outlook/ckpt/correction.pt
```

输出：
- `correction.pt` — 最终模型
- `correction_best.pt` — val_loss 最低 checkpoint
- `correction_train_log.json` — 逐 epoch 日志

### 8.3 离线指标评估

```bash
python outlook/src/test_correction.py \
    --checkpoint outlook/ckpt/correction.pt \
    --data-dir outlook/data_correction \
    --device cuda \
    --visualize \
    --output-dir outlook/result/correction/test
```

### 8.4 Rollout 评估

```bash
python outlook/src/eval_correction.py \
    --checkpoint outlook/ckpt/correction.pt \
    --data-dir outlook/data_correction \
    --target-elements 2000,4000,8000,12000 \
    --mark-fraction 0.03 \
    --max-steps 40 \
    --methods physics,neural,hybrid \
    --alpha 0.5 --beta 0.5 \
    --device cuda \
    --output-dir outlook/result/correction/rollout
```

输出：
- `eval_results.json` — 详细结果 + 聚合统计 + 改善比例
- `summary.csv` — 每行一个 (method, target) 汇总
- `aw_rel_vs_elements.png` / `max_err_vs_elements.png`

### 8.5 可视化

```bash
# 端到端 AMR
python outlook/src/viz_correction.py \
    --checkpoint outlook/ckpt/correction.pt \
    --data-dir outlook/data_correction \
    --mode amr --sample-id 0 --max-elements 5000 \
    --device cuda \
    --output-dir outlook/result/correction/viz

# 单步对比
python outlook/src/viz_correction.py \
    --checkpoint outlook/ckpt/correction.pt \
    --data-dir outlook/data_correction \
    --mode step --sample-id 0 --step 0 \
    --device cuda \
    --output-dir outlook/result/correction/viz

# OOD 评估（自定义物理参数，需同时指定 k, eps-r, cx, cy, radius）
python outlook/src/viz_correction.py \
    --checkpoint outlook/ckpt/correction.pt \
    --data-dir outlook/data_correction \
    --mode amr \
    --k 30 --eps-r 4 --cx 0.5 --cy 0.5 --radius 0.15 \
    --max-elements 6000 --compare \
    --device cuda \
    --output-dir outlook/result/correction/viz
```

---

## 9. 目录结构

```
outlook/
├── README.md                          # 本文档
├── train.sh                           # SLURM 训练脚本（可配置 batch_size，默认 1）
├── analyze_budget_teacher.py          # Budget teacher 数据集分析
├── check_correction_data.py           # 数据质量校验
│
├── src/
│   ├── gen.py                        # 数据生成
│   ├── train_correction.py           # 训练
│   ├── test_correction.py            # 离线指标评估
│   ├── eval_correction.py            # Rollout 评估
│   ├── viz_correction.py             # 可视化
│   ├── rollout.py                    # 统一 AMR rollout（核心循环 + 共享 refinement 工具）
│   ├── gnn.py                        # DensityGNN 模型
│   ├── feat.py                       # 特征提取
│   ├── graph.py                      # mesh → PyG graph（含 build_edge_index_np）
│   ├── mesh.py                       # score → refined mesh
│   ├── metrics.py                    # aw_rel / max_err
│   ├── problem.py                    # PDE 参数 → HelmholtzProblem
│   └── amr.py                        # 残差 AMR teacher
│
├── ckpt/                             # checkpoint
├── data_correction/                  # 训练数据
└── result/                           # 评估结果
```

---

## 10. 辅助模块

| 模块 | 职责 |
|------|------|
| `rollout.py` | 统一 AMR rollout：`run_rollout_to_budget()` 驱动完整加密循环，支持 physics/neural/hybrid/eta 四种打分方法 |
| `feat.py` | 构建 15 维基础特征 + budget_code |
| `graph.py` | mesh → PyG Data（边特征 = phase_distance） |
| `mesh.py` | score → 迭代 top-k 加密（叶子继承初始单元 score） |
| `problem.py` | 参数字典 → HelmholtzProblem 实例 |
| `amr.py` | 纯残差驱动 AMR teacher（无网络） |
| `metrics.py` | `compute_mie_error`（aw_rel + max_err） |

---

## 11. 测试结果

### 11.1 训练配置

- 数据集：100 样本，~19 步/样本，1888 个 step 文件
- 训练/验证划分：80/20（按 sample_id，seed=42）
- 训练图：1513 个，验证图：375 个
- 正样本比例：3.0%（teacher_mark top-3%）
- 训练 100 epoch，最佳 epoch 57（val_loss=0.5873）

### 11.2 离线指标（test_correction.py）

在 20 个验证样本的 375 个 step 图上评估：

| 指标 | GNN | Physics Baseline |
|------|-----|-----------------|
| **AUC** | **0.9412** | 0.0000 |
| **top-k overlap 均值** | **0.440** | 0.131 |
| GNN beats physics | **372/375 (99.2%)** | — |

> Physics AUC=0.0 是因为 physics_score 在均匀网格上只有 2 个离散值，无法区分排序。GNN 通过学习 η 的空间分布模式，top-k overlap 提升 **3.4 倍**。

### 11.3 Rollout 评估（eval_correction.py）

20 个验证样本 × 3 种方法 × 3 个目标预算，每次 rollout 最终做 1 次 FEM solve：

**aw_rel（面积加权相对误差）**：

| 方法 | target=2000 | target=4000 | target=8000 |
|------|------------|------------|------------|
| physics | 19.19% | 14.04% | 13.25% |
| neural | 15.96% | 13.72% | 12.98% |
| **hybrid** | **15.78%** | **13.57%** | **12.94%** |

**max_err（最大逐点误差）**：

| 方法 | target=2000 | target=4000 | target=8000 |
|------|------------|------------|------------|
| physics | 0.3301 | 0.2514 | 0.2418 |
| neural | 0.2741 | 0.2431 | 0.2376 |
| **hybrid** | **0.2686** | **0.2430** | **0.2383** |

**相对 physics-only 的改善**：

| 方法 | target | aw_rel 改善 | max_err 改善 |
|------|--------|-----------|-------------|
| hybrid | 2000 | **+9.8%** | **+6.9%** |
| neural | 2000 | +9.9% | +6.2% |
| hybrid | 4000 | +1.9% | +1.5% |
| neural | 4000 | +1.1% | +1.4% |
| hybrid | 8000 | +1.9% | +1.3% |
| neural | 8000 | +1.8% | +1.5% |

### 11.4 关键发现

1. **GNN 显著优于 physics baseline**：离线 top-k overlap 从 0.131 提升到 0.440（3.4×），99.2% 的验证图上 GNN 胜出
2. **低预算改善最大**：target=2000 时 aw_rel 改善 ~10%，max_err 改善 ~7%；高预算时改善收窄到 ~2%（因为预算充足时 physics 也够用）
3. **hybrid 略优于 neural**：z-score 混合策略在多数场景下比纯 GNN 更稳定
4. **GNN 推理效率**：neural 方法比 physics 方法少用 ~12% 的 refinement 步数达到相同预算（8.9 vs 9.4 步 @2000），因为 GNN 的标记更精准

### 11.5 OOD 评估（k 超出训练范围 [3,15]）

固定 `eps_r=4, cx=cy=0.5, radius=0.15`，在 k=20/30/50/80 上做三方法对比（target=5000/10000/20000）：

**aw_rel (%)**：

| k | 方法 | target=5000 | target=10000 | target=20000 |
|---|------|------------|-------------|-------------|
| 20 | physics | 32.17 | 24.55 | 10.61 |
| 20 | neural | 32.56 | **17.85** | 10.68 |
| 20 | hybrid | **31.73** | 19.57 | 14.03 |
| 30 | physics | 33.94 | 15.15 | 10.14 |
| 30 | neural | **23.52** | 13.40 | 8.38 |
| 30 | hybrid | 27.94 | **13.34** | **7.96** |
| 50 | physics | 94.00 | 73.56 | 33.00 |
| 50 | neural | **88.16** | **52.92** | **27.99** |
| 50 | hybrid | 90.97 | 68.47 | 29.27 |
| 80 | physics | 122.06 | 135.78 | 101.64 |
| 80 | neural | 139.84 | **127.65** | **92.67** |
| 80 | hybrid | 139.84 | 127.68 | 105.94 |

**相对 physics-only 的改善 (%)**：

| k | 方法 | target=5000 | target=10000 | target=20000 |
|---|------|------------|-------------|-------------|
| 20 | neural | −1.2 | **+27.3** | −0.7 |
| 20 | hybrid | +1.4 | +20.3 | −32.3 |
| 30 | neural | **+30.7** | +11.6 | +17.4 |
| 30 | hybrid | +17.7 | +12.0 | +21.5 |
| 50 | neural | +6.2 | **+28.1** | **+15.2** |
| 50 | hybrid | +3.2 | +6.9 | +11.3 |
| 80 | neural | −14.6 | +6.0 | **+8.8** |
| 80 | hybrid | −14.6 | +6.0 | −4.2 |

**OOD 关键发现**：

1. **neural 在远 OOD（k=50）优势最大**：target=10000 时 aw_rel 从 73.56% 降至 52.92%（改善 +28.1%），hybrid 仅改善 +6.9%
2. **hybrid 的 physics 先验在 OOD 时成为拖累**：z-score 混合中 physics_score 的 `max_edge/h_target` 在高 k 下不再准确，导致 hybrid 的标记不如纯 neural
3. **neural 的泛化来自几何特征**：GNN 学到的是"界面附近 + 介质内 → 高密度"的空间模式，这一模式在 k 超出训练范围时仍然成立；而 physics_score 依赖 `lambda_eff = 2π/(k√ε_r)` 的具体数值
4. **极端 OOD（k=80）两者都差**：误差超过 100%，说明 20000 个单元完全不够分辨 k=80 的短波结构（λ_eff≈0.028，需要 ~70000+ 单元）
5. **k=20/30 时 hybrid 反而更好**：接近训练分布时 physics 先验有价值，hybrid 通过混合两种信号获得更稳定的标记