Atomizer/examples/study_management_example.py

"""
Study Management Example

This script demonstrates how to use the study management features:
1. Create a new study
2. Resume an existing study to add more trials
3. List all available studies
4. Create a new study after topology/configuration changes
"""

import sys
from pathlib import Path

# 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_solver import run_nx_simulation
from optimization_engine.result_extractors import (
    extract_stress_from_op2,
    extract_displacement_from_op2
)


def bracket_model_updater(design_vars: dict):
    """Update bracket model with new design variable values."""
    from integration.nx_expression_updater import update_expressions_from_file

    sim_file = Path('examples/bracket/Bracket_sim1.sim')

    # Map design variables to NX expressions
    expressions = {
        'tip_thickness': design_vars['tip_thickness'],
        'support_angle': design_vars['support_angle']
    }

    update_expressions_from_file(
        sim_file=sim_file,
        expressions=expressions
    )


def bracket_simulation_runner() -> Path:
    """Run bracket simulation using journal-based NX solver."""
    sim_file = Path('examples/bracket/Bracket_sim1.sim')

    op2_file = run_nx_simulation(
        sim_file=sim_file,
        nastran_version='2412',
        timeout=300,
        cleanup=False,
        use_journal=True
    )

    return op2_file


def stress_extractor(result_path: Path) -> dict:
    """Extract stress results from OP2."""
    results = extract_stress_from_op2(result_path)
    return results


def displacement_extractor(result_path: Path) -> dict:
    """Extract displacement results from OP2."""
    results = extract_displacement_from_op2(result_path)
    return results


def example_1_new_study():
    """
    Example 1: Create a new optimization study with 20 trials
    """
    print("\n" + "="*70)
    print("EXAMPLE 1: Creating a New Study")
    print("="*70)

    config_path = Path('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
        }
    )

    # Create a new study with a specific name
    # This uses the config's n_trials (50) unless overridden
    study = runner.run(
        study_name="bracket_optimization_v1",
        n_trials=20,  # Override to 20 trials for this example
        resume=False  # Create new study
    )

    print("\nStudy completed successfully!")
    print(f"Database saved to: {runner._get_study_db_path('bracket_optimization_v1')}")


def example_2_resume_study():
    """
    Example 2: Resume an existing study to add more trials
    """
    print("\n" + "="*70)
    print("EXAMPLE 2: Resuming an Existing Study")
    print("="*70)

    config_path = Path('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
        }
    )

    # Resume the study created in example 1
    # Add 30 more trials (bringing total to 50)
    study = runner.run(
        study_name="bracket_optimization_v1",
        n_trials=30,  # Additional trials to run
        resume=True   # Resume existing study
    )

    print("\nStudy resumed and expanded successfully!")
    print(f"Total trials: {len(study.trials)}")


def example_3_list_studies():
    """
    Example 3: List all available studies
    """
    print("\n" + "="*70)
    print("EXAMPLE 3: Listing All Studies")
    print("="*70)

    config_path = Path('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
        }
    )

    studies = runner.list_studies()

    if not studies:
        print("No studies found.")
    else:
        print(f"\nFound {len(studies)} studies:\n")
        for study in studies:
            print(f"Study: {study['study_name']}")
            print(f"  Created: {study['created_at']}")
            print(f"  Total trials: {study.get('total_trials', 0)}")
            print(f"  Resume count: {study.get('resume_count', 0)}")
            print(f"  Config hash: {study.get('config_hash', 'N/A')[:8]}...")
            print()


def example_4_new_study_after_change():
    """
    Example 4: Create a new study after topology/configuration changes

    This demonstrates what to do when:
    - Geometry topology has changed significantly
    - Design variables have been added/removed
    - Objectives have changed

    In these cases, the surrogate model from the previous study is no longer valid,
    so you should create a NEW study rather than resume.
    """
    print("\n" + "="*70)
    print("EXAMPLE 4: New Study After Configuration Change")
    print("="*70)
    print("\nScenario: Bracket topology was modified, added new design variable")
    print("Old surrogate is invalid -> Create NEW study with different name\n")

    config_path = Path('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
        }
    )

    # Create a NEW study with a different name
    # Version number (v2) indicates this is a different geometry/configuration
    study = runner.run(
        study_name="bracket_optimization_v2",  # Different name!
        n_trials=50,
        resume=False  # New study, not resuming
    )

    print("\nNew study created for modified configuration!")
    print("Old study (v1) remains unchanged in database.")


if __name__ == "__main__":
    print("="*70)
    print("STUDY MANAGEMENT DEMONSTRATION")
    print("="*70)
    print("\nThis script demonstrates study management features:")
    print("1. Create new study")
    print("2. Resume existing study (add more trials)")
    print("3. List all studies")
    print("4. Create new study after topology change")
    print("\nREQUIREMENT: Simcenter3D must be OPEN")
    print("="*70)

    response = input("\nIs Simcenter3D open? (yes/no): ")
    if response.lower() not in ['yes', 'y']:
        print("Please open Simcenter3D and try again.")
        sys.exit(0)

    print("\n" + "="*70)
    print("Which example would you like to run?")
    print("="*70)
    print("1. Create a new study (20 trials)")
    print("2. Resume existing study 'bracket_optimization_v1' (+30 trials)")
    print("3. List all available studies")
    print("4. Create new study after topology change (50 trials)")
    print("0. Exit")
    print("="*70)

    choice = input("\nEnter choice (0-4): ").strip()

    try:
        if choice == '1':
            example_1_new_study()
        elif choice == '2':
            example_2_resume_study()
        elif choice == '3':
            example_3_list_studies()
        elif choice == '4':
            example_4_new_study_after_change()
        elif choice == '0':
            print("Exiting.")
        else:
            print("Invalid choice.")

    except Exception as e:
        print("\n" + "="*70)
        print("ERROR")
        print("="*70)
        print(f"{e}")
        import traceback
        traceback.print_exc()