Atomizer/examples/test_stress_displacement_optimization.py

"""
Test: Stress + Displacement Optimization

Tests the complete pipeline with:
- Objective: Minimize max von Mises stress
- Constraint: Max displacement <= 1.0 mm
"""

from pathlib import Path
import sys

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

from optimization_engine.runner import OptimizationRunner
from optimization_engine.nx_updater import update_nx_model
from optimization_engine.result_extractors.extractors import (
    stress_extractor,
    displacement_extractor
)


def bracket_model_updater(design_vars: dict):
    """Update bracket model parameters."""
    prt_file = project_root / "examples/bracket/Bracket.prt"
    print(f"\n[MODEL UPDATE] {prt_file.name}")
    for name, value in design_vars.items():
        print(f"  {name} = {value:.4f}")
    update_nx_model(prt_file, design_vars, backup=False)


def bracket_simulation_runner() -> Path:
    """Return existing OP2 (no re-solve for now)."""
    print("\n[SIMULATION] Using existing OP2")
    return project_root / "examples/bracket/bracket_sim1-solution_1.op2"


if __name__ == "__main__":
    print("="*60)
    print("STRESS + DISPLACEMENT OPTIMIZATION TEST")
    print("="*60)

    config_path = project_root / "examples/bracket/optimization_config_stress_displacement.json"

    runner = OptimizationRunner(
        config_path=config_path,
        model_updater=bracket_model_updater,
        simulation_runner=bracket_simulation_runner,
        result_extractors={
            'stress_extractor': stress_extractor,
            'displacement_extractor': displacement_extractor
        }
    )

    # Run 5 trials to test
    runner.config['optimization_settings']['n_trials'] = 5

    print("\nRunning 5 test trials...")
    print("Objective: Minimize max von Mises stress")
    print("Constraint: Max displacement <= 1.0 mm")
    print("="*60)

    try:
        study = runner.run(study_name="stress_displacement_test")

        print("\n" + "="*60)
        print("SUCCESS! Complete pipeline works!")
        print("="*60)
        print(f"Best stress: {study.best_value:.2f} MPa")
        print(f"Best parameters: {study.best_params}")
        print(f"\nResults in: {runner.output_dir}")

        # Show summary
        print("\n" + "="*60)
        print("EXTRACTED VALUES (from OP2):")
        print("="*60)

        # Read the last trial results
        import json
        history_file = runner.output_dir / "history.json"
        if history_file.exists():
            with open(history_file, 'r') as f:
                history = json.load(f)
                if history:
                    last_trial = history[-1]
                    print(f"Max stress: {last_trial['results'].get('max_von_mises', 'N/A')} MPa")
                    print(f"Max displacement: {last_trial['results'].get('max_displacement', 'N/A')} mm")
                    print(f"Stress element: {last_trial['results'].get('element_id', 'N/A')}")
                    print(f"Displacement node: {last_trial['results'].get('max_node_id', 'N/A')}")

    except Exception as e:
        print(f"\nERROR: {e}")
        import traceback
        traceback.print_exc()