96b196de58
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>
70 lines
2.0 KiB
Python
70 lines
2.0 KiB
Python
"""
|
|
GNN (Graph Neural Network) Surrogate Module for Atomizer
|
|
=========================================================
|
|
|
|
This module provides Graph Neural Network-based surrogates for FEA optimization,
|
|
particularly designed for Zernike-based mirror optimization where spatial structure
|
|
matters.
|
|
|
|
Key Components:
|
|
- PolarMirrorGraph: Fixed polar grid graph structure for mirror surface
|
|
- ZernikeGNN: GNN model for predicting displacement fields
|
|
- DifferentiableZernikeFit: GPU-accelerated Zernike fitting
|
|
- ZernikeObjectiveLayer: Compute objectives from displacement fields
|
|
- ZernikeGNNTrainer: Complete training pipeline
|
|
|
|
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
|
|
|
|
Usage:
|
|
# Training
|
|
python -m optimization_engine.gnn.train_zernike_gnn V11 V12 --epochs 200
|
|
|
|
# API
|
|
from optimization_engine.gnn import PolarMirrorGraph, ZernikeGNN, ZernikeGNNTrainer
|
|
"""
|
|
|
|
__version__ = "1.0.0"
|
|
|
|
# Core components
|
|
from .polar_graph import PolarMirrorGraph, create_mirror_dataset
|
|
from .zernike_gnn import ZernikeGNN, ZernikeGNNLite, create_model, load_model
|
|
from .differentiable_zernike import (
|
|
DifferentiableZernikeFit,
|
|
ZernikeObjectiveLayer,
|
|
ZernikeRMSLoss,
|
|
build_zernike_matrix,
|
|
)
|
|
from .extract_displacement_field import (
|
|
extract_displacement_field,
|
|
save_field,
|
|
load_field,
|
|
)
|
|
from .train_zernike_gnn import ZernikeGNNTrainer, MirrorDataset
|
|
|
|
__all__ = [
|
|
# Polar Graph
|
|
'PolarMirrorGraph',
|
|
'create_mirror_dataset',
|
|
# GNN Model
|
|
'ZernikeGNN',
|
|
'ZernikeGNNLite',
|
|
'create_model',
|
|
'load_model',
|
|
# Zernike Layers
|
|
'DifferentiableZernikeFit',
|
|
'ZernikeObjectiveLayer',
|
|
'ZernikeRMSLoss',
|
|
'build_zernike_matrix',
|
|
# Field Extraction
|
|
'extract_displacement_field',
|
|
'save_field',
|
|
'load_field',
|
|
# Training
|
|
'ZernikeGNNTrainer',
|
|
'MirrorDataset',
|
|
]
|