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

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:

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:

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

# 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

# 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

# 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

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:

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

{
  "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):

avg_forces_z = sum(forces_z) / len(forces_z)
min_forces_z = min(forces_z)

Phase 2.9 Generated (Lifecycle Hook):

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