Atomizer/examples/test_journal_optimization.py

"""
Test: Complete Optimization with Journal-Based NX Solver

This tests the complete workflow:
1. Update model parameters in .prt
2. Solve via journal (using running NX GUI)
3. Extract results from OP2
4. Run optimization loop

REQUIREMENTS:
- Simcenter3D must be open (but no files need to be loaded)
- test_env conda environment activated
"""

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
)


# Global variable to store current design variables for the simulation runner
_current_design_vars = {}


def bracket_model_updater(design_vars: dict):
    """
    Store design variables for the simulation runner.

    Note: We no longer directly update the .prt file here.
    Instead, design variables are passed to the journal which applies them in NX.
    """
    global _current_design_vars
    _current_design_vars = design_vars.copy()

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


def bracket_simulation_runner() -> Path:
    """
    Run NX solver via journal on running NX GUI session.

    This connects to the running Simcenter3D GUI and:
    1. Opens the .sim file
    2. Applies expression updates in the journal
    3. Updates geometry and FEM
    4. Solves the simulation
    5. Returns path to .op2 file
    """
    global _current_design_vars
    sim_file = project_root / "examples/bracket/Bracket_sim1.sim"

    print("\n[SIMULATION] Running via journal on NX GUI...")
    print(f"  SIM file: {sim_file.name}")
    if _current_design_vars:
        print(f"  Expression updates: {_current_design_vars}")

    try:
        # Run solver via journal (connects to running NX GUI)
        # Pass expression updates directly to the journal
        op2_file = run_nx_simulation(
            sim_file=sim_file,
            nastran_version="2412",
            timeout=300,  # 5 minute timeout
            cleanup=True,  # Clean up temp files
            use_journal=True,  # Use journal mode (requires NX GUI open)
            expression_updates=_current_design_vars  # Pass design vars to journal
        )

        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("JOURNAL-BASED OPTIMIZATION TEST")
    print("="*60)
    print("\nREQUIREMENTS:")
    print("- Simcenter3D must be OPEN (no files need to be loaded)")
    print("- Will run 50 optimization trials (~3-4 minutes)")
    print("- Strategy: 20 random trials (exploration) + 30 TPE trials (exploitation)")
    print("- Each trial: update params -> solve via journal -> extract results")
    print("="*60)

    response = input("\nIs Simcenter3D open? (yes/no): ")
    if response.lower() not in ['yes', 'y']:
        print("Please open Simcenter3D and try again.")
        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,  # Journal-based solver!
        result_extractors={
            'stress_extractor': stress_extractor,
            'displacement_extractor': displacement_extractor
        }
    )

    # Use the configured number of trials (50 by default)
    n_trials = runner.config['optimization_settings']['n_trials']

    print("\n" + "="*60)
    print(f"Starting optimization with {n_trials} trials")
    print("Objective: Minimize max von Mises stress")
    print("Constraint: Max displacement <= 1.0 mm")
    print("Solver: Journal-based (using running NX GUI)")
    print(f"Sampler: TPE (20 random startup + {n_trials-20} TPE)")
    print("="*60)

    try:
        study = runner.run(study_name="journal_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 optimization progress!")

    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("  - Simcenter3D is open and running")
        print("  - .sim file is valid and solvable")
        print("  - No other processes are locking the files")