Atomizer/docs/04_USER_GUIDES/NEURAL_FEATURES_COMPLETE.md

# Atomizer Neural Features - Complete Guide

**Version**: 1.0.0
**Last Updated**: 2025-11-25
**Status**: Production Ready

---

## Executive Summary

AtomizerField brings **Graph Neural Network (GNN) acceleration** to Atomizer, enabling:

| 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** |

This guide covers all neural network features, architectures, and integration points.

---

## Table of Contents

1. [Overview](#overview)
2. [Neural Model Types](#neural-model-types)
3. [Architecture Deep Dive](#architecture-deep-dive)
4. [Training Pipeline](#training-pipeline)
5. [Integration Layer](#integration-layer)
6. [Loss Functions](#loss-functions)
7. [Uncertainty Quantification](#uncertainty-quantification)
8. [Pre-trained Models](#pre-trained-models)
9. [Configuration Reference](#configuration-reference)
10. [Performance Benchmarks](#performance-benchmarks)

---

## Overview

### What is AtomizerField?

AtomizerField is a neural network system that learns to predict FEA simulation results. Instead of solving physics equations numerically (expensive), it uses trained neural networks to predict results instantly.

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

Neural Workflow:
Design → Neural Network (4.5 ms) → Results → Objective
```

### Core Components

| Component | File | Purpose |
|-----------|------|---------|
| **BDF/OP2 Parser** | `neural_field_parser.py` | Converts NX Nastran files to neural format |
| **Data Validator** | `validate_parsed_data.py` | Physics and quality checks |
| **Field Predictor** | `field_predictor.py` | GNN for displacement/stress fields |
| **Parametric Predictor** | `parametric_predictor.py` | GNN for direct objective prediction |
| **Physics Loss** | `physics_losses.py` | Physics-informed training |
| **Neural Surrogate** | `neural_surrogate.py` | Integration with Atomizer |
| **Neural Runner** | `runner_with_neural.py` | Optimization with neural acceleration |

---

## Neural Model Types

### 1. Field Predictor GNN

**Purpose**: Predicts complete displacement and stress fields across the entire mesh.

**Architecture**:
```
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)

Edge Features (5D per edge):
├── Edge length
├── Direction vector (3D)
└── Element type indicator

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 parameters

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

### 2. Parametric Predictor GNN (Recommended)

**Purpose**: Predicts all 4 optimization objectives directly from design parameters.

**Architecture**:
```
Design Parameters (4D):
├── beam_half_core_thickness
├── beam_face_thickness
├── holes_diameter
└── hole_count

Design Encoder (MLP):
├── Linear(4 → 64)
├── ReLU
├── Linear(64 → 128)
└── ReLU

GNN Backbone (4 layers):
├── Design-conditioned message passing
├── Hidden channels: 128
└── Global pooling: Mean + Max

Scalar Heads (MLP):
├── Linear(384 → 128)
├── ReLU
├── Linear(128 → 64)
├── ReLU
└── Linear(64 → 4)

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

**Parameters**: ~500,000 trainable parameters

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

### 3. Ensemble Models

**Purpose**: Uncertainty quantification through multiple model predictions.

**How it works**:
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

---

## Architecture Deep Dive

### Graph Construction from FEA Mesh

The neural network treats the FEA mesh as a graph:

```
FEA Mesh                    →    Neural Graph
─────────────────────────────────────────────
Nodes (grid points)         →    Graph nodes
Elements (CTETRA, CQUAD)    →    Graph edges
Node properties             →    Node features
Element connectivity        →    Edge connections
```

### Message Passing

The GNN learns physics through message passing:

```python
# Simplified message passing
for layer in gnn_layers:
    # Aggregate neighbor information
    messages = aggregate(
        node_features,
        edge_features,
        adjacency
    )

    # Update node features
    node_features = update(
        node_features,
        messages
    )
```

This is analogous to how FEA distributes forces through elements.

### Design Conditioning (Parametric GNN)

The parametric model conditions the GNN on design parameters:

```python
# Design parameters are encoded
design_encoding = design_encoder(design_params)  # [batch, 128]

# Broadcast to all nodes
node_features = node_features + design_encoding.unsqueeze(1)

# GNN processes with design context
for layer in gnn_layers:
    node_features = layer(node_features, edge_index, edge_attr)
```

---

## Training Pipeline

### Step 1: Collect Training Data

Run optimization with training data export:

```python
# In workflow_config.json
{
  "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
```

### Step 2: Parse to Neural Format

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

Creates HDF5 + JSON files:
```
trial_0001/
├── neural_field_data.json    # Metadata, structure
└── neural_field_data.h5      # Mesh coordinates, field results
```

### Step 3: Train Model

**Field Predictor**:
```bash
python train.py \
  --train_dir ../training_data/parsed \
  --epochs 200 \
  --model FieldPredictorGNN \
  --hidden_channels 128 \
  --num_layers 6 \
  --physics_loss_weight 0.3
```

**Parametric Predictor** (recommended):
```bash
python train_parametric.py \
  --train_dir ../training_data/parsed \
  --val_dir ../validation_data/parsed \
  --epochs 200 \
  --hidden_channels 128 \
  --num_layers 4
```

### Step 4: Validate

```bash
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

```python
# In workflow_config.json
{
  "neural_surrogate": {
    "enabled": true,
    "model_checkpoint": "atomizer-field/runs/my_model/checkpoint_best.pt",
    "confidence_threshold": 0.85
  }
}
```

---

## Integration Layer

### NeuralSurrogate Class

```python
from optimization_engine.neural_surrogate import NeuralSurrogate

# Load trained model
surrogate = NeuralSurrogate(
    model_path="atomizer-field/runs/model/checkpoint_best.pt",
    device="cuda",
    confidence_threshold=0.85
)

# Predict
results, confidence, used_nn = surrogate.predict(
    design_variables={"thickness": 5.0, "width": 10.0},
    bdf_template="model.bdf"
)

if used_nn:
    print(f"Predicted with {confidence:.1%} confidence")
else:
    print("Fell back to FEA (low confidence)")
```

### ParametricSurrogate Class (Recommended)

```python
from optimization_engine.neural_surrogate import create_parametric_surrogate_for_study

# Auto-detect and load
surrogate = create_parametric_surrogate_for_study()

# Predict all 4 objectives
results = surrogate.predict({
    "beam_half_core_thickness": 7.0,
    "beam_face_thickness": 2.5,
    "holes_diameter": 35.0,
    "hole_count": 10.0
})

print(f"Mass: {results['mass']:.2f} g")
print(f"Frequency: {results['frequency']:.2f} Hz")
print(f"Max displacement: {results['max_displacement']:.6f} mm")
print(f"Max stress: {results['max_stress']:.2f} MPa")
print(f"Inference time: {results['inference_time_ms']:.2f} ms")
```

### HybridOptimizer Class

```python
from optimization_engine.neural_surrogate import create_hybrid_optimizer_from_config

# Smart FEA/NN switching
hybrid = create_hybrid_optimizer_from_config(config_path)

for trial in range(1000):
    if hybrid.should_use_nn(trial):
        result = hybrid.predict_with_nn(design_vars)
    else:
        result = hybrid.run_fea(design_vars)
        hybrid.add_training_sample(design_vars, result)
```

---

## Loss Functions

### 1. MSE Loss (Standard)

```python
loss = mean_squared_error(predicted, target)
```

Equal weighting of all outputs. Simple but effective.

### 2. Relative Loss

```python
loss = mean(|predicted - target| / |target|)
```

Better for multi-scale outputs (stress in MPa, displacement in mm).

### 3. Physics-Informed Loss

```python
loss = mse_loss + lambda_physics * physics_loss

physics_loss = (
    equilibrium_violation +      # F = ma
    constitutive_violation +     # σ = Eε
    boundary_condition_violation # u = 0 at supports
)
```

Enforces physical laws during training. Improves generalization.

### 4. Max Error Loss

```python
loss = max(|predicted - target|)
```

Penalizes worst predictions. Critical for safety-critical applications.

### When to Use Each

| Loss Function | Use Case |
|--------------|----------|
| MSE | General training, balanced errors |
| Relative | Multi-scale outputs |
| Physics | Better generalization, extrapolation |
| Max Error | Safety-critical, avoid outliers |

---

## Uncertainty Quantification

### Ensemble Method

```python
from atomizer_field.neural_models.uncertainty import EnsemblePredictor

# Load 5 models
ensemble = EnsemblePredictor([
    "model_fold_1.pt",
    "model_fold_2.pt",
    "model_fold_3.pt",
    "model_fold_4.pt",
    "model_fold_5.pt"
])

# Predict with uncertainty
mean_pred, std_pred = ensemble.predict_with_uncertainty(design_vars)
confidence = 1.0 / (1.0 + std_pred / mean_pred)

if confidence < 0.85:
    # Trigger FEA validation
    fea_result = run_fea(design_vars)
```

### Monte Carlo Dropout

```python
# Enable dropout at inference
model.train()  # Keeps dropout active

predictions = []
for _ in range(10):
    pred = model(input_data)
    predictions.append(pred)

mean_pred = np.mean(predictions)
std_pred = np.std(predictions)
```

---

## Pre-trained Models

### Available Models

| Model | Location | Design Variables | Objectives |
|-------|----------|------------------|------------|
| UAV Arm (Parametric) | `runs/parametric_uav_arm_v2/` | 4 | 4 |
| UAV Arm (Field) | `runs/uav_arm_model/` | 4 | 2 fields |

### Using Pre-trained Models

```python
from optimization_engine.neural_surrogate import create_parametric_surrogate_for_study

# Auto-detects model in atomizer-field/runs/
surrogate = create_parametric_surrogate_for_study()

# Immediate predictions - no training needed!
result = surrogate.predict({
    "beam_half_core_thickness": 7.0,
    "beam_face_thickness": 2.5,
    "holes_diameter": 35.0,
    "hole_count": 10.0
})
```

---

## Configuration Reference

### Complete workflow_config.json

```json
{
  "study_name": "neural_optimization_study",

  "neural_surrogate": {
    "enabled": true,
    "model_checkpoint": "atomizer-field/runs/parametric_uav_arm_v2/checkpoint_best.pt",
    "confidence_threshold": 0.85,
    "device": "cuda",
    "cache_predictions": true,
    "cache_size": 10000
  },

  "hybrid_optimization": {
    "enabled": true,
    "exploration_trials": 30,
    "validation_frequency": 20,
    "retrain_frequency": 100,
    "drift_threshold": 0.15,
    "retrain_on_drift": true
  },

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

### Parameter Reference

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `enabled` | bool | false | Enable neural surrogate |
| `model_checkpoint` | str | - | Path to trained model |
| `confidence_threshold` | float | 0.85 | Min confidence for NN |
| `device` | str | "cuda" | "cuda" or "cpu" |
| `cache_predictions` | bool | true | Cache repeated designs |
| `exploration_trials` | int | 30 | Initial FEA trials |
| `validation_frequency` | int | 20 | FEA validation interval |
| `retrain_frequency` | int | 100 | Retrain interval |
| `drift_threshold` | float | 0.15 | Max error before retrain |

---

## Performance Benchmarks

### UAV Arm Study (4 design variables, 4 objectives)

| Metric | FEA Only | Neural Only | Hybrid |
|--------|----------|-------------|--------|
| Time per trial | 10.2s | 4.5ms | 0.5s avg |
| Total time (1000 trials) | 2.8 hours | 4.5 seconds | 8 minutes |
| Prediction error | - | 2.3% | 1.8% |
| Speedup | 1x | 2,267x | 21x |

### Accuracy by Objective

| Objective | MAE | MAPE | R² |
|-----------|-----|------|-----|
| Mass | 0.5g | 0.8% | 0.998 |
| Frequency | 2.1 Hz | 1.2% | 0.995 |
| Max Displacement | 0.001mm | 2.8% | 0.987 |
| Max Stress | 3.2 MPa | 3.5% | 0.981 |

### GPU vs CPU

| Device | Inference Time | Throughput |
|--------|---------------|------------|
| CPU (i7-12700) | 45ms | 22/sec |
| GPU (RTX 3080) | 4.5ms | 220/sec |
| Speedup | 10x | 10x |

---

## Quick Reference

### Files and Locations

```
atomizer-field/
├── neural_field_parser.py      # Parse BDF/OP2
├── batch_parser.py             # Batch processing
├── validate_parsed_data.py     # Data validation
├── train.py                    # Train field predictor
├── train_parametric.py         # Train parametric model
├── predict.py                  # Inference engine
├── neural_models/
│   ├── field_predictor.py      # GNN architecture
│   ├── parametric_predictor.py # Parametric GNN
│   ├── physics_losses.py       # Loss functions
│   ├── uncertainty.py          # Uncertainty quantification
│   └── data_loader.py          # PyTorch dataset
├── runs/                       # Trained models
│   └── parametric_uav_arm_v2/
│       └── checkpoint_best.pt
└── tests/                      # 18 comprehensive tests
```

### Common Commands

```bash
# Parse training data
python batch_parser.py ../training_data

# Train parametric model
python train_parametric.py --train_dir ../data --epochs 200

# Validate model
python validate.py --checkpoint runs/model/checkpoint_best.pt

# Run tests
python -m pytest tests/ -v
```

### Python API

```python
# Quick start
from optimization_engine.neural_surrogate import create_parametric_surrogate_for_study

surrogate = create_parametric_surrogate_for_study()
result = surrogate.predict({"param1": 1.0, "param2": 2.0})
```

---

## See Also

- [Neural Workflow Tutorial](NEURAL_WORKFLOW_TUTORIAL.md) - Step-by-step guide
- [GNN Architecture](GNN_ARCHITECTURE.md) - Technical deep dive
- [Physics Loss Guide](PHYSICS_LOSS_GUIDE.md) - Loss function selection
- [Atomizer-Field Integration Plan](ATOMIZER_FIELD_INTEGRATION_PLAN.md) - Implementation details

---

**AtomizerField**: Revolutionizing structural optimization through neural field learning.

*Built with PyTorch Geometric, designed for the future of engineering.*