Atomizer/docs/protocols/system/SYS_14_NEURAL_ACCELERATION.md

SYS_14: Neural Network Acceleration

Overview

Atomizer provides neural network surrogate acceleration enabling 100-1000x faster optimization by replacing expensive FEA evaluations with instant neural predictions.

Two approaches available:

1. MLP Surrogate (Simple, integrated) - 4-layer MLP trained on FEA data, runs within study
2. GNN Field Predictor (Advanced) - Graph neural network for full field predictions

Key Innovation: Train once on FEA data, then explore 5,000-50,000+ designs in the time it takes to run 50 FEA trials.

---

When to Use

Trigger	Action
>50 trials needed	Consider neural acceleration
"neural", "surrogate", "NN" mentioned	Load this protocol
"fast", "acceleration", "speed" needed	Suggest neural acceleration
Training data available	Enable surrogate

---

Quick Reference

Performance Comparison:

Metric	Traditional FEA	Neural Network	Improvement
Time per evaluation	10-30 minutes	4.5 milliseconds	2,000-500,000x
Trials per hour	2-6	800,000+	1000x
Design exploration	~50 designs	~50,000 designs	1000x

Model Types:

Model	Purpose	Use When
MLP Surrogate	Direct objective prediction	Simple studies, quick setup
Field Predictor GNN	Full displacement/stress fields	Need field visualization
Parametric Predictor GNN	Direct objective prediction	Complex geometry, need accuracy
Ensemble	Uncertainty quantification	Need confidence bounds

---

MLP Surrogate (Recommended for Quick Start)

Overview

The MLP (Multi-Layer Perceptron) surrogate is a simple but effective neural network that predicts objectives directly from design parameters. It's integrated into the study workflow via run_nn_optimization.py.

Architecture

Input Layer (N design variables)
    ↓
Linear(N, 64) + ReLU + BatchNorm + Dropout(0.1)
    ↓
Linear(64, 128) + ReLU + BatchNorm + Dropout(0.1)
    ↓
Linear(128, 128) + ReLU + BatchNorm + Dropout(0.1)
    ↓
Linear(128, 64) + ReLU + BatchNorm + Dropout(0.1)
    ↓
Linear(64, M objectives)

Parameters: ~34,000 trainable

Workflow Modes

1. Standard Hybrid Mode (--all)

Run all phases sequentially:

python run_nn_optimization.py --all

Phases:

1. Export: Extract training data from existing FEA trials
2. Train: Train MLP surrogate (300 epochs default)
3. NN-Optimize: Run 1000 NN trials with NSGA-II
4. Validate: Validate top 10 candidates with FEA

2. Hybrid Loop Mode (--hybrid-loop)

Iterative refinement:

python run_nn_optimization.py --hybrid-loop --iterations 5 --nn-trials 500

Each iteration:

1. Train/retrain surrogate from current FEA data
2. Run NN optimization
3. Validate top candidates with FEA
4. Add validated results to training set
5. Repeat until convergence (max error < 5%)

3. Turbo Mode (--turbo) ⚡ RECOMMENDED

Aggressive single-best validation:

python run_nn_optimization.py --turbo --nn-trials 5000 --batch-size 100 --retrain-every 10

Strategy:

• Run NN in small batches (100 trials)
• Validate ONLY the single best candidate with FEA
• Add to training data immediately
• Retrain surrogate every N FEA validations
• Repeat until total NN budget exhausted

Example: 5,000 NN trials with batch=100 → 50 FEA validations in ~12 minutes

Configuration

{
  "neural_acceleration": {
    "enabled": true,
    "min_training_points": 50,
    "auto_train": true,
    "epochs": 300,
    "validation_split": 0.2,
    "nn_trials": 1000,
    "validate_top_n": 10,
    "model_file": "surrogate_best.pt",
    "separate_nn_database": true
  }
}

Important: separate_nn_database: true stores NN trials in nn_study.db instead of study.db to avoid overloading the dashboard with thousands of NN-only results.

Typical Accuracy

Objective	Expected Error
Mass	1-5%
Stress	1-4%
Stiffness	5-15%

Output Files

2_results/
├── study.db                    # Main FEA + validated results (dashboard)
├── nn_study.db                 # NN-only results (not in dashboard)
├── surrogate_best.pt           # Trained model weights
├── training_data.json          # Normalized training data
├── nn_optimization_state.json  # NN optimization state
├── nn_pareto_front.json        # NN-predicted Pareto front
├── validation_report.json      # FEA validation results
└── turbo_report.json           # Turbo mode results (if used)

---

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

# 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:

# 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

Component	File	Purpose
BDF/OP2 Parser	neural_field_parser.py	Convert NX files to neural format
Data Validator	validate_parsed_data.py	Physics and quality checks
Field Predictor	field_predictor.py	GNN for full field prediction
Parametric Predictor	parametric_predictor.py	GNN for direct objectives
Physics Loss	physics_losses.py	Physics-informed training
Neural Surrogate	neural_surrogate.py	Integration with Atomizer
Neural Runner	runner_with_neural.py	Optimization with NN acceleration

Workflow Diagram

Traditional:
Design → NX Model → Mesh → Solve (30 min) → Results → Objective

Neural (after training):
Design → Neural Network (4.5 ms) → Results → Objective

---

Neural Model Types

1. Field Predictor GNN

Use Case: When you need full field predictions (stress distribution, deformation shape).

Input Features (12D per node):
├── Node coordinates (x, y, z)
├── Material properties (E, nu, rho)
├── Boundary conditions (fixed/free per DOF)
└── Load information (force magnitude, direction)

GNN Layers (6 message passing):
├── MeshGraphConv (custom for FEA topology)
├── Layer normalization
├── ReLU activation
└── Dropout (0.1)

Output (per node):
├── Displacement (6 DOF: Tx, Ty, Tz, Rx, Ry, Rz)
└── Von Mises stress (1 value)

Parameters: ~718,221 trainable

2. Parametric Predictor GNN (Recommended)

Use Case: Direct optimization objective prediction (fastest option).

Design Parameters (ND) → Design Encoder (MLP) → GNN Backbone → Scalar Heads

Output (objectives):
├── mass (grams)
├── frequency (Hz)
├── max_displacement (mm)
└── max_stress (MPa)

Parameters: ~500,000 trainable

3. Ensemble Models

Use Case: Uncertainty quantification.

1. Train 3-5 models with different random seeds
2. At inference, run all models
3. Use mean for prediction, std for uncertainty
4. High uncertainty → trigger FEA validation

---

Training Pipeline

Step 1: Collect Training Data

Enable export in workflow config:

{
  "training_data_export": {
    "enabled": true,
    "export_dir": "atomizer_field_training_data/my_study"
  }
}

Output structure:

atomizer_field_training_data/my_study/
├── trial_0001/
│   ├── input/model.bdf       # Nastran input
│   ├── output/model.op2      # Binary results
│   └── metadata.json         # Design params + objectives
├── trial_0002/
│   └── ...
└── study_summary.json

Recommended: 100-500 FEA samples for good generalization.

Step 2: Parse to Neural Format

cd atomizer-field
python batch_parser.py ../atomizer_field_training_data/my_study

Creates HDF5 + JSON files per trial.

Step 3: Train Model

Parametric Predictor (recommended):

python train_parametric.py \
  --train_dir ../training_data/parsed \
  --val_dir ../validation_data/parsed \
  --epochs 200 \
  --hidden_channels 128 \
  --num_layers 4

Field Predictor:

python train.py \
  --train_dir ../training_data/parsed \
  --epochs 200 \
  --model FieldPredictorGNN \
  --hidden_channels 128 \
  --num_layers 6 \
  --physics_loss_weight 0.3

Step 4: Validate

python validate.py --checkpoint runs/my_model/checkpoint_best.pt

Expected output:

Validation Results:
├── Mean Absolute Error: 2.3% (mass), 1.8% (frequency)
├── R² Score: 0.987
├── Inference Time: 4.5ms ± 0.8ms
└── Physics Violations: 0.2%

Step 5: Deploy

{
  "neural_surrogate": {
    "enabled": true,
    "model_checkpoint": "atomizer-field/runs/my_model/checkpoint_best.pt",
    "confidence_threshold": 0.85
  }
}

---

Configuration

Full Neural Configuration Example

{
  "study_name": "bracket_neural_optimization",

  "surrogate_settings": {
    "enabled": true,
    "model_type": "parametric_gnn",
    "model_path": "models/bracket_surrogate.pt",
    "confidence_threshold": 0.85,
    "validation_frequency": 10,
    "fallback_to_fea": true
  },

  "training_data_export": {
    "enabled": true,
    "export_dir": "atomizer_field_training_data/bracket_study",
    "export_bdf": true,
    "export_op2": true,
    "export_fields": ["displacement", "stress"]
  },

  "neural_optimization": {
    "initial_fea_trials": 50,
    "neural_trials": 5000,
    "retraining_interval": 500,
    "uncertainty_threshold": 0.15
  }
}

Configuration Parameters

Parameter	Type	Default	Description
enabled	bool	false	Enable neural surrogate
model_type	string	"parametric_gnn"	Model architecture
model_path	string	-	Path to trained model
confidence_threshold	float	0.85	Min confidence for predictions
validation_frequency	int	10	FEA validation every N trials
fallback_to_fea	bool	true	Use FEA when uncertain

---

Hybrid FEA/Neural Workflow

Phase 1: FEA Exploration (50-100 trials)

• Run standard FEA optimization
• Export training data automatically
• Build landscape understanding

Phase 2: Neural Training

• Parse collected data
• Train parametric predictor
• Validate accuracy

Phase 3: Neural Acceleration (1000s of trials)

• Use neural network for rapid exploration
• Periodic FEA validation
• Retrain if distribution shifts

Phase 4: FEA Refinement (10-20 trials)

• Validate top candidates with FEA
• Ensure results are physically accurate
• Generate final Pareto front

---

Adaptive Iteration Loop

For complex optimizations, use iterative refinement:

┌─────────────────────────────────────────────────────────────────┐
│  Iteration 1:                                                    │
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────┐       │
│  │ Initial FEA  │ -> │ Train NN     │ -> │ NN Search    │       │
│  │ (50-100)     │    │ Surrogate    │    │ (1000 trials)│       │
│  └──────────────┘    └──────────────┘    └──────────────┘       │
│                                                 │                │
│  Iteration 2+:                                  ▼                │
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────┐       │
│  │ Validate Top │ -> │ Retrain NN   │ -> │ NN Search    │       │
│  │ NN with FEA  │    │ with new data│    │ (1000 trials)│       │
│  └──────────────┘    └──────────────┘    └──────────────┘       │
└─────────────────────────────────────────────────────────────────┘

Adaptive Configuration

{
  "adaptive_settings": {
    "enabled": true,
    "initial_fea_trials": 50,
    "nn_trials_per_iteration": 1000,
    "fea_validation_per_iteration": 5,
    "max_iterations": 10,
    "convergence_threshold": 0.01,
    "retrain_epochs": 100
  }
}

Convergence Criteria

Stop when:

• No improvement for 2-3 consecutive iterations
• Reached FEA budget limit
• Objective improvement < 1% threshold

Output Files

studies/my_study/3_results/
├── adaptive_state.json      # Current iteration state
├── surrogate_model.pt       # Trained neural network
└── training_history.json    # NN training metrics

---

Loss Functions

Data Loss (MSE)

Standard prediction error:

data_loss = MSE(predicted, target)

Physics Loss

Enforce physical constraints:

physics_loss = (
    equilibrium_loss +      # Force balance
    boundary_loss +         # BC satisfaction
    compatibility_loss      # Strain compatibility
)

Combined Training

total_loss = data_loss + 0.3 * physics_loss

Physics loss weight typically 0.1-0.5.

---

Uncertainty Quantification

Ensemble Method

# Run N models
predictions = [model_i(x) for model_i in ensemble]

# Statistics
mean_prediction = np.mean(predictions)
uncertainty = np.std(predictions)

# Decision
if uncertainty > threshold:
    # Use FEA instead
    result = run_fea(x)
else:
    result = mean_prediction

Confidence Thresholds

Uncertainty	Action
< 5%	Use neural prediction
5-15%	Use neural, flag for validation
> 15%	Fall back to FEA

---

Troubleshooting

Symptom	Cause	Solution
High prediction error	Insufficient training data	Collect more FEA samples
Out-of-distribution warnings	Design outside training range	Retrain with expanded range
Slow inference	Large mesh	Use parametric predictor instead
Physics violations	Low physics loss weight	Increase physics_loss_weight

---

Cross-References

• Depends On: SYS_10_IMSO for optimization framework
• Used By: OP_02_RUN_OPTIMIZATION, OP_05_EXPORT_TRAINING_DATA
• See Also: modules/neural-acceleration.md

---

Implementation Files

atomizer-field/
├── neural_field_parser.py       # BDF/OP2 parsing
├── field_predictor.py           # Field GNN
├── parametric_predictor.py      # Parametric GNN
├── train.py                     # Field training
├── train_parametric.py          # Parametric training
├── validate.py                  # Model validation
├── physics_losses.py            # Physics-informed loss
└── batch_parser.py              # Batch data conversion

optimization_engine/
├── neural_surrogate.py          # Atomizer integration
└── runner_with_neural.py        # Neural runner

---

Version History

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