Atomizer/examples/run_optimization.py

"""
Example: Running Complete Optimization

This example demonstrates the complete optimization workflow:
1. Load optimization configuration
2. Update NX model parameters
3. Run simulation (dummy for now - would call NX solver)
4. Extract results from OP2
5. Optimize with Optuna

For a real run, you would need:
- pyNastran installed for OP2 extraction
- NX solver accessible to run simulations
"""

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


# ==================================================
# 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.

    In a real implementation, this would:
    1. Open NX (or use batch mode)
    2. Update the .sim file
    3. Run the solver
    4. Wait for completion
    5. Return path to .op2 file

    For now, we return the path to existing results.
    """
    print("\n[SIMULATION] Running NX Nastran solver...")
    print("[SIMULATION] (Using existing results for demonstration)")

    # In real use, this would run the actual solver
    # For now, 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


# ==================================================
# STEP 3: Define result extractors (dummy versions)
# ==================================================
def dummy_mass_extractor(result_path: Path) -> dict:
    """
    Dummy mass extractor.
    In real use, would call: from optimization_engine.result_extractors.extractors import mass_extractor
    """
    import random
    # Simulate varying mass based on a simple model
    # In reality, this would extract from OP2
    base_mass = 0.45  # kg
    variation = random.uniform(-0.05, 0.05)

    return {
        'total_mass': base_mass + variation,
        'cg_x': 0.0,
        'cg_y': 0.0,
        'cg_z': 0.0,
        'units': 'kg'
    }


def dummy_stress_extractor(result_path: Path) -> dict:
    """
    Dummy stress extractor.
    In real use, would call: from optimization_engine.result_extractors.extractors import stress_extractor
    """
    import random
    # Simulate stress results
    base_stress = 180.0  # MPa
    variation = random.uniform(-30.0, 30.0)

    return {
        'max_von_mises': base_stress + variation,
        'stress_type': 'von_mises',
        'element_id': 1234,
        'units': 'MPa'
    }


def dummy_displacement_extractor(result_path: Path) -> dict:
    """
    Dummy displacement extractor.
    In real use, would call: from optimization_engine.result_extractors.extractors import displacement_extractor
    """
    import random
    # Simulate displacement results
    base_disp = 0.9  # mm
    variation = random.uniform(-0.2, 0.2)

    return {
        'max_displacement': base_disp + variation,
        'max_node_id': 5678,
        'dx': 0.0,
        'dy': 0.0,
        'dz': base_disp + variation,
        'units': 'mm'
    }


# ==================================================
# MAIN: Run optimization
# ==================================================
if __name__ == "__main__":
    print("="*60)
    print("ATOMIZER - OPTIMIZATION EXAMPLE")
    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}")

    # Create result extractors dict
    extractors = {
        'mass_extractor': dummy_mass_extractor,
        'stress_extractor': dummy_stress_extractor,
        'displacement_extractor': dummy_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 (use fewer trials for demo)
    print("\n" + "="*60)
    print("Starting optimization with 10 trials (demo)")
    print("For full optimization, modify n_trials in config")
    print("="*60)

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

    # Run!
    study = runner.run(study_name="bracket_optimization_demo")

    print("\n" + "="*60)
    print("OPTIMIZATION RESULTS")
    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. Install pyNastran: conda install -c conda-forge pynastran")
    print("2. Replace dummy extractors with real OP2 extractors")
    print("3. Integrate with NX solver (batch mode or NXOpen)")
    print("4. Run full optimization with n_trials=100+")
    print("="*60)