feat: Add Zernike GNN surrogate module and M1 mirror V12/V13 studies
This commit introduces the GNN-based surrogate for Zernike mirror optimization and the M1 mirror study progression from V12 (GNN validation) to V13 (pure NSGA-II). ## GNN Surrogate Module (optimization_engine/gnn/) New module for Graph Neural Network surrogate prediction of mirror deformations: - `polar_graph.py`: PolarMirrorGraph - fixed 3000-node polar grid structure - `zernike_gnn.py`: ZernikeGNN with design-conditioned message passing - `differentiable_zernike.py`: GPU-accelerated Zernike fitting and objectives - `train_zernike_gnn.py`: ZernikeGNNTrainer with multi-task loss - `gnn_optimizer.py`: ZernikeGNNOptimizer for turbo mode (~900k trials/hour) - `extract_displacement_field.py`: OP2 to HDF5 field extraction - `backfill_field_data.py`: Extract fields from existing FEA trials Key innovation: Design-conditioned convolutions that modulate message passing based on structural design parameters, enabling accurate field prediction. ## M1 Mirror Studies ### V12: GNN Field Prediction + FEA Validation - Zernike GNN trained on V10/V11 FEA data (238 samples) - Turbo mode: 5000 GNN predictions → top candidates → FEA validation - Calibration workflow for GNN-to-FEA error correction - Scripts: run_gnn_turbo.py, validate_gnn_best.py, compute_full_calibration.py ### V13: Pure NSGA-II FEA (Ground Truth) - Seeds 217 FEA trials from V11+V12 - Pure multi-objective NSGA-II without any surrogate - Establishes ground-truth Pareto front for GNN accuracy evaluation - Narrowed blank_backface_angle range to [4.0, 5.0] ## Documentation Updates - SYS_14: Added Zernike GNN section with architecture diagrams - CLAUDE.md: Added GNN module reference and quick start - V13 README: Study documentation with seeding strategy 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Antoine
@@ -167,7 +167,127 @@ Strategy:
|
||||
|
||||
---
|
||||
|
||||
## GNN Field Predictor (Advanced)
|
||||
## Zernike GNN (Mirror Optimization)
|
||||
|
||||
### Overview
|
||||
|
||||
The **Zernike GNN** is a specialized Graph Neural Network for mirror surface optimization. Unlike the MLP surrogate that predicts objectives directly, the Zernike GNN predicts the full displacement field, then computes Zernike coefficients and objectives via differentiable layers.
|
||||
|
||||
**Why GNN over MLP for Zernike?**
|
||||
1. **Spatial awareness**: GNN learns smooth deformation fields via message passing
|
||||
2. **Correct relative computation**: Predicts fields, then subtracts (like FEA)
|
||||
3. **Multi-task learning**: Field + objective supervision
|
||||
4. **Physics-informed**: Edge structure respects mirror geometry
|
||||
|
||||
### Architecture
|
||||
|
||||
```
|
||||
Design Variables [11]
|
||||
│
|
||||
▼
|
||||
Design Encoder [11 → 128]
|
||||
│
|
||||
└──────────────────┐
|
||||
│
|
||||
Node Features │
|
||||
[r, θ, x, y] │
|
||||
│ │
|
||||
▼ │
|
||||
Node Encoder │
|
||||
[4 → 128] │
|
||||
│ │
|
||||
└─────────┬────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────┐
|
||||
│ Design-Conditioned │
|
||||
│ Message Passing (× 6) │
|
||||
│ │
|
||||
│ • Polar-aware edges │
|
||||
│ • Design modulates messages │
|
||||
│ • Residual connections │
|
||||
└─────────────┬───────────────┘
|
||||
│
|
||||
▼
|
||||
Per-Node Decoder [128 → 4]
|
||||
│
|
||||
▼
|
||||
Z-Displacement Field [3000, 4]
|
||||
(one value per node per subcase)
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────┐
|
||||
│ DifferentiableZernikeFit │
|
||||
│ (GPU-accelerated) │
|
||||
└─────────────┬───────────────┘
|
||||
│
|
||||
▼
|
||||
Zernike Coefficients → Objectives
|
||||
```
|
||||
|
||||
### Module Structure
|
||||
|
||||
```
|
||||
optimization_engine/gnn/
|
||||
├── __init__.py # Public API
|
||||
├── polar_graph.py # PolarMirrorGraph - fixed polar grid
|
||||
├── zernike_gnn.py # ZernikeGNN model (design-conditioned conv)
|
||||
├── differentiable_zernike.py # GPU Zernike fitting & objective layers
|
||||
├── extract_displacement_field.py # OP2 → HDF5 field extraction
|
||||
├── train_zernike_gnn.py # ZernikeGNNTrainer pipeline
|
||||
├── gnn_optimizer.py # ZernikeGNNOptimizer for turbo mode
|
||||
└── backfill_field_data.py # Extract fields from existing trials
|
||||
```
|
||||
|
||||
### Training Workflow
|
||||
|
||||
```bash
|
||||
# Step 1: Extract displacement fields from FEA trials
|
||||
python -m optimization_engine.gnn.backfill_field_data V11
|
||||
|
||||
# Step 2: Train GNN on extracted data
|
||||
python -m optimization_engine.gnn.train_zernike_gnn V11 V12 --epochs 200
|
||||
|
||||
# Step 3: Run GNN-accelerated optimization
|
||||
python run_gnn_turbo.py --trials 5000
|
||||
```
|
||||
|
||||
### Key Classes
|
||||
|
||||
| Class | Purpose |
|
||||
|-------|---------|
|
||||
| `PolarMirrorGraph` | Fixed 3000-node polar grid for mirror surface |
|
||||
| `ZernikeGNN` | Main model with design-conditioned message passing |
|
||||
| `DifferentiableZernikeFit` | GPU-accelerated Zernike coefficient computation |
|
||||
| `ZernikeObjectiveLayer` | Compute rel_rms objectives from coefficients |
|
||||
| `ZernikeGNNTrainer` | Complete training pipeline with multi-task loss |
|
||||
| `ZernikeGNNOptimizer` | Turbo optimization with GNN predictions |
|
||||
|
||||
### Calibration
|
||||
|
||||
GNN predictions require calibration against FEA ground truth. Use the full FEA dataset (not just validation samples) for robust calibration:
|
||||
|
||||
```python
|
||||
# compute_full_calibration.py
|
||||
# Computes calibration factors: GNN_pred * factor ≈ FEA_truth
|
||||
calibration_factors = {
|
||||
'rel_filtered_rms_40_vs_20': 1.15, # GNN underpredicts by ~15%
|
||||
'rel_filtered_rms_60_vs_20': 1.08,
|
||||
'mfg_90_optician_workload': 0.95, # GNN overpredicts by ~5%
|
||||
}
|
||||
```
|
||||
|
||||
### Performance
|
||||
|
||||
| Metric | FEA | Zernike GNN |
|
||||
|--------|-----|-------------|
|
||||
| Time per eval | 8-10 min | 4 ms |
|
||||
| Trials per hour | 6-7 | 900,000 |
|
||||
| Typical accuracy | Ground truth | 5-15% error |
|
||||
|
||||
---
|
||||
|
||||
## GNN Field Predictor (Generic)
|
||||
|
||||
### Core Components
|
||||
|
||||
@@ -560,5 +680,6 @@ optimization_engine/
|
||||
|
||||
| Version | Date | Changes |
|
||||
|---------|------|---------|
|
||||
| 2.1 | 2025-12-10 | Added Zernike GNN section for mirror optimization |
|
||||
| 2.0 | 2025-12-06 | Added MLP Surrogate with Turbo Mode |
|
||||
| 1.0 | 2025-12-05 | Initial consolidation from neural docs |
|
||||
|
||||
Reference in New Issue
Block a user