Atomizer/tests/test_bracket_llm_runner.py

"""
Test LLM Optimization Runner with Bracket Problem

Objective: Maximize displacement
Constraint: Stress must remain under safety factor of 4

Assuming Aluminum 6061-T6:
- Yield Strength: 276 MPa
- Safety Factor: 4
- Allowable Stress: 276 / 4 = 69 MPa

This test validates Phase 3.2 LLM-enhanced optimization with stress constraints.
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from optimization_engine.llm_optimization_runner import LLMOptimizationRunner

# LLM workflow for bracket optimization
# Goal: Maximize displacement while keeping stress below safety factor
llm_workflow = {
    'engineering_features': [
        {
            'action': 'extract_displacement',
            'domain': 'result_extraction',
            'description': 'Extract displacement results from OP2 file',
            'params': {'result_type': 'displacement'}
        },
        {
            'action': 'extract_solid_stress',
            'domain': 'result_extraction',
            'description': 'Extract von Mises stress from solid elements',
            'params': {
                'result_type': 'stress',
                'element_type': 'ctetra'
            }
        }
    ],
    'inline_calculations': [
        {
            'action': 'calculate_safety_factor',
            'params': {
                'input': 'max_von_mises',
                'yield_strength': 276.0,  # MPa for Aluminum 6061-T6
                'operation': 'divide'
            },
            'code_hint': 'safety_factor = 276.0 / max_von_mises'
        },
        {
            'action': 'negate_displacement',
            'params': {
                'input': 'max_displacement',
                'operation': 'negate'
            },
            'code_hint': 'neg_displacement = -max_displacement'
        }
    ],
    'post_processing_hooks': [
        {
            'action': 'weighted_objective',
            'params': {
                'inputs': ['neg_displacement'],
                'weights': [1.0],
                'objective': 'minimize'
            },
            'description': 'Minimize negative displacement (maximize actual displacement)'
        }
    ],
    'optimization': {
        'algorithm': 'TPE',
        'direction': 'minimize',  # Minimize negative displacement = maximize displacement
        'design_variables': [
            {
                'parameter': 'wall_thickness',
                'min': 3.0,
                'max': 8.0,
                'units': 'mm'
            }
        ]
    }
}


def dummy_updater(design_vars):
    """Dummy model updater for testing."""
    print(f"  [Test] Would update model with: {design_vars}")


def dummy_runner():
    """Dummy simulation runner for testing."""
    print("  [Test] Would run NX simulation...")
    # Return path to existing test OP2
    return Path('tests/bracket_sim1-solution_1.op2')


if __name__ == '__main__':
    print("=" * 80)
    print("Test: LLM Runner with Bracket Problem")
    print("=" * 80)
    print()
    print("Problem Definition:")
    print("  Objective: MAXIMIZE displacement")
    print("  Constraint: Stress must remain below yield/SF")
    print("  Material: Aluminum 6061-T6 (Yield = 276 MPa)")
    print("  Safety Factor: 4.0")
    print("  Allowable Stress: 69 MPa")
    print()
    print("LLM Workflow:")
    print(f"  Engineering features: {len(llm_workflow['engineering_features'])}")
    print(f"  Inline calculations: {len(llm_workflow['inline_calculations'])}")
    print(f"  Post-processing hooks: {len(llm_workflow['post_processing_hooks'])}")
    print(f"  Design variables: {len(llm_workflow['optimization']['design_variables'])}")
    print()

    # Initialize LLM runner
    print("Initializing LLM optimization runner...")
    try:
        runner = LLMOptimizationRunner(
            llm_workflow=llm_workflow,
            model_updater=dummy_updater,
            simulation_runner=dummy_runner,
            study_name='bracket_maximize_disp_constrained'
        )

        print()
        print("=" * 80)
        print("Runner Initialized Successfully!")
        print("=" * 80)
        print(f"  Extractors generated: {len(runner.extractors)}")
        print(f"  Inline calculations: {len(runner.inline_code)}")
        hook_summary = runner.hook_manager.get_summary()
        print(f"  Hooks loaded: {hook_summary['enabled_hooks']}")
        print()

        print("To run full optimization, call runner.run_optimization(n_trials=50)")

    except Exception as e:
        print(f"\nERROR during initialization: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)