Atomizer/examples/run_optimization_real.py

"""
Example: Running Complete Optimization WITH REAL OP2 EXTRACTION

This version uses real pyNastran extractors instead of dummy data.

Requirements:
- conda activate test_env (with pyNastran and optuna installed)

What this does:
1. Updates NX model parameters in the .prt file
2. Uses existing OP2 results (simulation step skipped for now)
3. Extracts REAL mass, stress, displacement from OP2
4. Runs Optuna optimization

Note: Since we're using the same OP2 file for all trials (no re-solving),
      the results will be constant. This is just to test the pipeline.
      For real optimization, you'd need to run NX solver for each trial.
"""

from pathlib import Path
import sys

# Add project root to path
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 (
    mass_extractor,
    stress_extractor,
    displacement_extractor
)


# ==================================================
# STEP 1: Define model updater function
# ==================================================
def bracket_model_updater(design_vars: dict):
    """
    Update the bracket model with new design variable values.

    Args:
        design_vars: Dict like {'tip_thickness': 22.5, 'support_angle': 35.0}
    """
    prt_file = project_root / "examples/bracket/Bracket.prt"

    print(f"\n[MODEL UPDATE] Updating {prt_file.name} with:")
    for name, value in design_vars.items():
        print(f"  {name} = {value:.4f}")

    # Update the .prt file with new parameter values
    update_nx_model(prt_file, design_vars, backup=False)

    print("[MODEL UPDATE] Complete")


# ==================================================
# STEP 2: Define simulation runner function
# ==================================================
def bracket_simulation_runner() -> Path:
    """
    Run NX simulation and return path to result files.

    For this demo, we just return the existing OP2 file.
    In production, this would:
    1. Run NX solver with updated model
    2. Wait for completion
    3. Return path to new OP2 file
    """
    print("\n[SIMULATION] Running NX Nastran solver...")
    print("[SIMULATION] (Using existing OP2 for demo - no actual solve)")

    # Return path to existing OP2 file
    result_file = project_root / "examples/bracket/bracket_sim1-solution_1.op2"

    if not result_file.exists():
        raise FileNotFoundError(f"Result file not found: {result_file}")

    print(f"[SIMULATION] Results: {result_file.name}")
    return result_file


# ==================================================
# MAIN: Run optimization
# ==================================================
if __name__ == "__main__":
    print("="*60)
    print("ATOMIZER - REAL OPTIMIZATION TEST")
    print("="*60)

    # Path to optimization configuration
    config_path = project_root / "examples/bracket/optimization_config.json"

    if not config_path.exists():
        print(f"Error: Configuration file not found: {config_path}")
        print("Please run the MCP build_optimization_config tool first.")
        sys.exit(1)

    print(f"\nConfiguration: {config_path}")

    # Use REAL extractors
    print("\nUsing REAL OP2 extractors (pyNastran)")
    extractors = {
        'mass_extractor': mass_extractor,
        'stress_extractor': stress_extractor,
        'displacement_extractor': displacement_extractor
    }

    # Create optimization runner
    runner = OptimizationRunner(
        config_path=config_path,
        model_updater=bracket_model_updater,
        simulation_runner=bracket_simulation_runner,
        result_extractors=extractors
    )

    # Run optimization with just 5 trials for testing
    print("\n" + "="*60)
    print("Starting optimization with 5 trials (test mode)")
    print("="*60)
    print("\nNOTE: Since we're using the same OP2 file for all trials")
    print("(not re-running solver), results will be constant.")
    print("This is just to test the pipeline integration.")
    print("="*60)

    # Override n_trials for demo
    runner.config['optimization_settings']['n_trials'] = 5

    try:
        # Run!
        study = runner.run(study_name="bracket_real_extraction_test")

        print("\n" + "="*60)
        print("TEST COMPLETE - PIPELINE WORKS!")
        print("="*60)
        print(f"\nBest parameters found:")
        for param, value in study.best_params.items():
            print(f"  {param}: {value:.4f}")

        print(f"\nBest objective value: {study.best_value:.6f}")

        print(f"\nResults saved to: {runner.output_dir}")
        print("  - history.csv (all trials)")
        print("  - history.json (detailed results)")
        print("  - optimization_summary.json (best results)")

        print("\n" + "="*60)
        print("NEXT STEPS:")
        print("="*60)
        print("1. Check the history.csv to see extracted values")
        print("2. Integrate NX solver execution (batch mode)")
        print("3. Run real optimization with solver re-runs")
        print("="*60)

    except Exception as e:
        print(f"\n{'='*60}")
        print("ERROR DURING OPTIMIZATION")
        print("="*60)
        print(f"Error: {e}")
        print("\nMake sure you're running in test_env with:")
        print("  - pyNastran installed")
        print("  - optuna installed")
        print("  - pandas installed")
        import traceback
        traceback.print_exc()