Atomizer/docs/02_ARCHITECTURE.md

# Atomizer Architecture - Complete System Overview

## Overview

Atomizer consists of three major architectural components:

1. **Hook System** - Unified lifecycle hooks for FEA workflow automation
2. **Neural Acceleration** - Graph Neural Network surrogates for fast predictions
3. **Dashboard** - Real-time monitoring and visualization

This document covers the complete system architecture.

---

## Part 1: Neural Network Architecture (AtomizerField)

### System Overview

```
┌─────────────────────────────────────────────────────────┐
│                  AtomizerField System                   │
├─────────────────────────────────────────────────────────┤
│                                                          │
│  ┌─────────────┐   ┌─────────────┐   ┌─────────────┐   │
│  │ BDF/OP2     │   │    GNN      │   │  Inference  │   │
│  │ Parser      │──>│  Training   │──>│   Engine    │   │
│  │ (Phase 1)   │   │  (Phase 2)  │   │  (Phase 2)  │   │
│  └─────────────┘   └─────────────┘   └─────────────┘   │
│        │                 │                  │           │
│        ▼                 ▼                  ▼           │
│  ┌─────────────────────────────────────────────────┐   │
│  │              Neural Model Types                  │   │
│  ├─────────────────────────────────────────────────┤   │
│  │ • Field Predictor GNN (displacement + stress)   │   │
│  │ • Parametric GNN (all 4 objectives directly)    │   │
│  │ • Ensemble models for uncertainty               │   │
│  └─────────────────────────────────────────────────┘   │
│                                                          │
└─────────────────────────────────────────────────────────┘
```

### Integration with Optimization

```
┌───────────────────────────┬─────────────────────────────┐
│     Traditional Path      │      Neural Path            │
├───────────────────────────┼─────────────────────────────┤
│  NX Solver (via Journals) │   AtomizerField GNN         │
│  ~10-30 min per eval      │   ~4.5 ms per eval          │
│  Full physics fidelity    │   Physics-informed learning │
└───────────────────────────┴─────────────────────────────┘
                          ↕
┌─────────────────────────────────────────────────────────┐
│         Hybrid Decision Engine                          │
│  Confidence-based switching • Uncertainty quantification│
│  Automatic FEA validation • Online learning             │
└─────────────────────────────────────────────────────────┘
```

### Key Neural Components

| Component | File | Purpose |
|-----------|------|---------|
| **BDF/OP2 Parser** | `atomizer-field/neural_field_parser.py` | Convert NX Nastran → neural format |
| **Field Predictor** | `atomizer-field/neural_models/field_predictor.py` | GNN for displacement/stress fields |
| **Parametric GNN** | `atomizer-field/neural_models/parametric_predictor.py` | Direct objective prediction |
| **Physics Loss** | `atomizer-field/neural_models/physics_losses.py` | Physics-informed training |
| **Neural Surrogate** | `optimization_engine/neural_surrogate.py` | Integration layer |
| **Neural Runner** | `optimization_engine/runner_with_neural.py` | Optimization with NN |

### Neural Data Flow

```
Training Data Collection:
  FEA Run → BDF/OP2 Export → Parser → HDF5+JSON → Dataset

Model Training:
  Dataset → DataLoader → GNN → Physics Loss → Optimizer → Checkpoint

Inference (Production):
  Design Params → Normalize → GNN → Denormalize → Predictions (4.5ms)
```

### Performance Metrics

| Metric | FEA Only | Neural Only | Hybrid |
|--------|----------|-------------|--------|
| Time per trial | 10-30 min | 4.5 ms | 0.5s avg |
| Speedup | 1x | 2,200x | 20x |
| Accuracy | Baseline | <5% error | <3% error |

**See [GNN_ARCHITECTURE.md](GNN_ARCHITECTURE.md) for technical details.**

---

## Part 2: Hook Architecture - Unified Lifecycle System

Atomizer uses a **unified lifecycle hook system** where all hooks - whether system plugins or auto-generated post-processing scripts - integrate seamlessly through the `HookManager`.

## Hook Types

### 1. Lifecycle Hooks (Phase 1 - System Plugins)

Located in: `optimization_engine/plugins/<hook_point>/`

**Purpose**: Plugin system for FEA workflow automation

**Hook Points**:
```
pre_mesh          → Before meshing
post_mesh         → After meshing, before solve
pre_solve         → Before FEA solver execution
post_solve        → After solve, before extraction
post_extraction   → After result extraction
post_calculation  → After inline calculations (NEW in Phase 2.9)
custom_objective  → Custom objective functions
```

**Example**: System logging, state management, file operations

### 2. Generated Post-Processing Hooks (Phase 2.9)

Located in: `optimization_engine/plugins/post_calculation/` (by default)

**Purpose**: Auto-generated custom calculations on extracted data

**Can be placed at ANY hook point** for maximum flexibility!

**Types**:
- Weighted objectives
- Custom formulas
- Constraint checks
- Comparisons (ratios, differences, percentages)

## Complete Optimization Workflow

```
Optimization Trial N
    ↓
┌─────────────────────────────────────┐
│ PRE-SOLVE HOOKS                     │
│ - Log trial parameters              │
│ - Validate design variables         │
│ - Backup model files                │
└─────────────────────────────────────┘
    ↓
┌─────────────────────────────────────┐
│ RUN NX NASTRAN SOLVE                │
└─────────────────────────────────────┘
    ↓
┌─────────────────────────────────────┐
│ POST-SOLVE HOOKS                    │
│ - Check solution convergence        │
│ - Log solve completion              │
└─────────────────────────────────────┘
    ↓
┌─────────────────────────────────────┐
│ EXTRACT RESULTS (OP2/F06)           │
│ - Read stress, displacement, etc.   │
└─────────────────────────────────────┘
    ↓
┌─────────────────────────────────────┐
│ POST-EXTRACTION HOOKS               │
│ - Log extracted values              │
│ - Validate result ranges            │
└─────────────────────────────────────┘
    ↓
┌─────────────────────────────────────┐
│ INLINE CALCULATIONS (Phase 2.8)     │
│ - avg_stress = sum(stresses) / len │
│ - norm_stress = avg_stress / 200   │
│ - norm_disp = max_disp / 5          │
└─────────────────────────────────────┘
    ↓
┌─────────────────────────────────────┐
│ POST-CALCULATION HOOKS (Phase 2.9)  │
│ - weighted_objective()              │
│ - safety_factor()                   │
│ - constraint_check()                │
└─────────────────────────────────────┘
    ↓
┌─────────────────────────────────────┐
│ REPORT TO OPTUNA                    │
│ - Return objective value(s)         │
└─────────────────────────────────────┘
    ↓
Next Trial
```

## Directory Structure

```
optimization_engine/plugins/
├── hooks.py                          # HookPoint enum, Hook dataclass
├── hook_manager.py                   # HookManager class
├── pre_mesh/                         # Pre-meshing hooks
├── post_mesh/                        # Post-meshing hooks
├── pre_solve/                        # Pre-solve hooks
│   ├── detailed_logger.py
│   └── optimization_logger.py
├── post_solve/                       # Post-solve hooks
│   └── log_solve_complete.py
├── post_extraction/                  # Post-extraction hooks
│   ├── log_results.py
│   └── optimization_logger_results.py
└── post_calculation/                 # Post-calculation hooks (NEW!)
    ├── weighted_objective_test.py    # Generated by Phase 2.9
    ├── safety_factor_hook.py         # Generated by Phase 2.9
    └── min_to_avg_ratio_hook.py      # Generated by Phase 2.9
```

## Hook Format

All hooks follow the same interface:

```python
def my_hook(context: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """
    Hook function.

    Args:
        context: Dictionary containing relevant data:
            - trial_number: Current optimization trial
            - design_variables: Current design variable values
            - results: Extracted FEA results (post-extraction)
            - calculations: Inline calculation results (post-calculation)

    Returns:
        Optional dictionary with results to add to context
    """
    # Hook logic here
    return {'my_result': value}


def register_hooks(hook_manager):
    """Register this hook with the HookManager."""
    hook_manager.register_hook(
        hook_point='post_calculation',  # or any other HookPoint
        function=my_hook,
        description="My custom hook",
        name="my_hook",
        priority=100,
        enabled=True
    )
```

## Hook Generation (Phase 2.9)

### Standalone Scripts (Original)

Generated as independent Python scripts with JSON I/O:

```python
from optimization_engine.hook_generator import HookGenerator

generator = HookGenerator()

hook_spec = {
    "action": "weighted_objective",
    "description": "Combine stress and displacement",
    "params": {
        "inputs": ["norm_stress", "norm_disp"],
        "weights": [0.7, 0.3]
    }
}

# Generate standalone script
hook = generator.generate_from_llm_output(hook_spec)
generator.save_hook_to_file(hook, "generated_hooks/")
```

**Use case**: Independent execution, debugging, external tools

### Lifecycle Hooks (Integrated)

Generated as lifecycle-compatible plugins:

```python
from optimization_engine.hook_generator import HookGenerator

generator = HookGenerator()

hook_spec = {
    "action": "weighted_objective",
    "description": "Combine stress and displacement",
    "params": {
        "inputs": ["norm_stress", "norm_disp"],
        "weights": [0.7, 0.3]
    }
}

# Generate lifecycle hook
hook_content = generator.generate_lifecycle_hook(
    hook_spec,
    hook_point='post_calculation'  # or pre_solve, post_extraction, etc.
)

# Save to plugins directory
output_file = Path("optimization_engine/plugins/post_calculation/weighted_objective.py")
with open(output_file, 'w') as f:
    f.write(hook_content)

# HookManager automatically discovers and loads it!
```

**Use case**: Integration with optimization workflow, automatic execution

## Flexibility: Hooks Can Be Placed Anywhere!

The beauty of the lifecycle system is that **generated hooks can be placed at ANY hook point**:

### Example 1: Pre-Solve Validation

```python
# Generate a constraint check to run BEFORE solving
constraint_spec = {
    "action": "constraint_check",
    "description": "Ensure wall thickness is reasonable",
    "params": {
        "inputs": ["wall_thickness", "max_thickness"],
        "condition": "wall_thickness / max_thickness",
        "threshold": 1.0,
        "constraint_name": "thickness_check"
    }
}

hook_content = generator.generate_lifecycle_hook(
    constraint_spec,
    hook_point='pre_solve'  # Run BEFORE solve!
)
```

###Example 2: Post-Extraction Safety Factor

```python
# Generate safety factor calculation right after extraction
safety_spec = {
    "action": "custom_formula",
    "description": "Calculate safety factor from extracted stress",
    "params": {
        "inputs": ["max_stress", "yield_strength"],
        "formula": "yield_strength / max_stress",
        "output_name": "safety_factor"
    }
}

hook_content = generator.generate_lifecycle_hook(
    safety_spec,
    hook_point='post_extraction'  # Run right after extraction!
)
```

### Example 3: Pre-Mesh Parameter Validation

```python
# Generate parameter check before meshing
validation_spec = {
    "action": "comparison",
    "description": "Check if thickness exceeds maximum",
    "params": {
        "inputs": ["requested_thickness", "max_allowed"],
        "operation": "ratio",
        "output_name": "thickness_ratio"
    }
}

hook_content = generator.generate_lifecycle_hook(
    validation_spec,
    hook_point='pre_mesh'  # Run before meshing!
)
```

## Hook Manager Usage

```python
from optimization_engine.plugins.hook_manager import HookManager

# Create manager
hook_manager = HookManager()

# Auto-load all plugins from directory structure
hook_manager.load_plugins_from_directory(
    Path("optimization_engine/plugins")
)

# Execute hooks at specific point
context = {
    'trial_number': 42,
    'results': {'max_stress': 150.5},
    'calculations': {'norm_stress': 0.75, 'norm_disp': 0.64}
}

results = hook_manager.execute_hooks('post_calculation', context)

# Get summary
summary = hook_manager.get_summary()
print(f"Total hooks: {summary['total_hooks']}")
print(f"Hooks at post_calculation: {summary['by_hook_point']['post_calculation']}")
```

## Integration with Optimization Runner

The optimization runner will be updated to call hooks at appropriate lifecycle points:

```python
# In optimization_engine/runner.py

def run_trial(self, trial_number, design_variables):
    # Create context
    context = {
        'trial_number': trial_number,
        'design_variables': design_variables,
        'working_dir': self.working_dir
    }

    # Pre-solve hooks
    self.hook_manager.execute_hooks('pre_solve', context)

    # Run solve
    self.nx_solver.run(...)

    # Post-solve hooks
    self.hook_manager.execute_hooks('post_solve', context)

    # Extract results
    results = self.extractor.extract(...)
    context['results'] = results

    # Post-extraction hooks
    self.hook_manager.execute_hooks('post_extraction', context)

    # Inline calculations (Phase 2.8)
    calculations = self.inline_calculator.calculate(...)
    context['calculations'] = calculations

    # Post-calculation hooks (Phase 2.9)
    hook_results = self.hook_manager.execute_hooks('post_calculation', context)

    # Merge hook results into context
    for result in hook_results:
        if result:
            context.update(result)

    # Return final objective
    return context.get('weighted_objective') or results['stress']
```

## Benefits of Unified System

1. **Consistency**: All hooks use same interface, same registration, same execution
2. **Flexibility**: Generated hooks can be placed at any lifecycle point
3. **Discoverability**: HookManager auto-loads from directory structure
4. **Extensibility**: Easy to add new hook points or new hook types
5. **Debugging**: All hooks have logging, history tracking, enable/disable
6. **Priority Control**: Hooks execute in priority order
7. **Error Handling**: Configurable fail-fast or continue-on-error

## Example: Complete CBAR Optimization

**User Request:**
> "Extract CBAR element forces in Z direction, calculate average and minimum, create objective that minimizes min/avg ratio, optimize CBAR stiffness X with genetic algorithm"

**Phase 2.7 LLM Analysis:**
```json
{
  "engineering_features": [
    {"action": "extract_1d_element_forces", "domain": "result_extraction"},
    {"action": "update_cbar_stiffness", "domain": "fea_properties"}
  ],
  "inline_calculations": [
    {"action": "calculate_average", "params": {"input": "forces_z"}},
    {"action": "find_minimum", "params": {"input": "forces_z"}}
  ],
  "post_processing_hooks": [
    {
      "action": "comparison",
      "params": {
        "inputs": ["min_force", "avg_force"],
        "operation": "ratio",
        "output_name": "min_to_avg_ratio"
      }
    }
  ]
}
```

**Phase 2.8 Generated (Inline):**
```python
avg_forces_z = sum(forces_z) / len(forces_z)
min_forces_z = min(forces_z)
```

**Phase 2.9 Generated (Lifecycle Hook):**
```python
# optimization_engine/plugins/post_calculation/min_to_avg_ratio_hook.py

def min_to_avg_ratio_hook(context):
    calculations = context.get('calculations', {})

    min_force = calculations.get('min_forces_z')
    avg_force = calculations.get('avg_forces_z')

    result = min_force / avg_force

    return {'min_to_avg_ratio': result, 'objective': result}

def register_hooks(hook_manager):
    hook_manager.register_hook(
        hook_point='post_calculation',
        function=min_to_avg_ratio_hook,
        description="Compare min force to average",
        name="min_to_avg_ratio_hook"
    )
```

**Execution:**
```
Trial 1:
  pre_solve hooks → log trial
  solve → NX Nastran
  post_solve hooks → check convergence
  post_extraction hooks → validate results

  Extract: forces_z = [10.5, 12.3, 8.9, 11.2, 9.8]

  Inline calculations:
    avg_forces_z = 10.54
    min_forces_z = 8.9

  post_calculation hooks → min_to_avg_ratio_hook
    min_to_avg_ratio = 8.9 / 10.54 = 0.844

  Report to Optuna: objective = 0.844
```

**All code auto-generated! Zero manual scripting!** 🚀

## Future Enhancements

1. **Hook Dependencies**: Hooks can declare dependencies on other hooks
2. **Conditional Execution**: Hooks can have conditions (e.g., only run if stress > threshold)
3. **Hook Composition**: Combine multiple hooks into pipelines
4. **Study-Specific Hooks**: Hooks stored in `studies/<study_name>/plugins/`
5. **Hook Marketplace**: Share hooks between projects/users

## Summary

The unified lifecycle hook system provides:
- ✅ Single consistent interface for all hooks
- ✅ Generated hooks integrate seamlessly with system hooks
- ✅ Hooks can be placed at ANY lifecycle point
- ✅ Auto-discovery and loading
- ✅ Priority control and error handling
- ✅ Maximum flexibility for optimization workflows

**Phase 2.9 hooks are now true lifecycle hooks, usable anywhere in the FEA workflow!**