Atomizer/examples/test_optimization_with_solver.py

"""
Test: Complete Optimization with Real NX Solver

This runs the complete optimization loop:
1. Update model parameters
2. Run NX solver (REAL simulation)
3. Extract results from OP2
4. Optimize with Optuna

WARNING: This will run NX solver for each trial!
For 5 trials, expect ~5-10 minutes depending on solver speed.
"""

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.nx_solver import run_nx_simulation
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:
    """
    Run NX Nastran solver and return path to OP2 file.

    This is the key difference from the test version -
    it actually runs the solver for each trial!
    """
    sim_file = project_root / "examples/bracket/Bracket_sim1.sim"

    print("\n[SIMULATION] Running NX Nastran solver...")
    print(f"  SIM file: {sim_file.name}")

    try:
        # Run solver (this will take ~1-2 minutes per trial)
        op2_file = run_nx_simulation(
            sim_file=sim_file,
            nastran_version="2412",
            timeout=600,  # 10 minute timeout
            cleanup=True  # Clean up temp files
        )

        print(f"[SIMULATION] Complete! Results: {op2_file.name}")
        return op2_file

    except Exception as e:
        print(f"[SIMULATION] FAILED: {e}")
        raise


if __name__ == "__main__":
    print("="*60)
    print("REAL OPTIMIZATION WITH NX SOLVER")
    print("="*60)
    print("\n⚠️  WARNING ⚠️")
    print("This will run NX Nastran solver for each trial!")
    print("For 3 trials, expect ~5-10 minutes total.")
    print("\nMake sure:")
    print("  - NX 2412 is installed and licensed")
    print("  - No NX GUI sessions are open")
    print("  - Bracket.prt and Bracket_sim1.sim are accessible")
    print("="*60)

    response = input("\nContinue? (yes/no): ")
    if response.lower() not in ['yes', 'y']:
        print("Cancelled.")
        sys.exit(0)

    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,  # REAL SOLVER!
        result_extractors={
            'stress_extractor': stress_extractor,
            'displacement_extractor': displacement_extractor
        }
    )

    # Run just 3 trials for testing (change to 20-50 for real optimization)
    runner.config['optimization_settings']['n_trials'] = 3

    print("\n" + "="*60)
    print("Starting optimization with 3 trials")
    print("Objective: Minimize max von Mises stress")
    print("Constraint: Max displacement <= 1.0 mm")
    print("="*60)

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

        print("\n" + "="*60)
        print("OPTIMIZATION COMPLETE!")
        print("="*60)
        print(f"\nBest stress: {study.best_value:.2f} MPa")
        print(f"\nBest parameters:")
        for param, value in study.best_params.items():
            print(f"  {param}: {value:.4f}")

        print(f"\nResults saved to: {runner.output_dir}")
        print("\nCheck history.csv to see how stress changed with parameters!")

    except Exception as e:
        print(f"\n{'='*60}")
        print("ERROR DURING OPTIMIZATION")
        print("="*60)
        print(f"{e}")
        import traceback
        traceback.print_exc()
        print("\nMake sure:")
        print("  - NX Nastran is properly installed")
        print("  - License is available")
        print("  - .sim file is valid and solvable")