Atomizer/docs/protocols/extensions/templates/hook_template.py

"""
{Hook Name} - Lifecycle Hook Plugin

This is a template for creating new lifecycle hooks.
Copy this file to optimization_engine/plugins/{hook_point}/{hook_name}.py

Available hook points:
- pre_mesh: Before meshing
- post_mesh: After meshing
- pre_solve: Before solver execution
- post_solve: After solver completion
- post_extraction: After result extraction
- post_calculation: After objective calculation
- custom_objective: Custom objective functions

Author: {Your Name}
Created: {Date}
Version: 1.0
Hook Point: {hook_point}
"""

from typing import Dict, Any, Optional
from pathlib import Path
import json
from datetime import datetime


def {hook_name}_hook(context: Dict[str, Any]) -> Dict[str, Any]:
    """
    {Description of what this hook does}.

    This hook runs at the {hook_point} stage of the optimization trial.

    Args:
        context: Dictionary containing trial context:
            - trial_number (int): Current trial number
            - design_params (dict): Current design parameter values
            - config (dict): Optimization configuration
            - working_dir (Path): Study working directory

            For post_solve and later:
            - op2_file (Path): Path to OP2 results file
            - solve_success (bool): Whether solve succeeded
            - solve_time (float): Solve duration in seconds

            For post_extraction and later:
            - results (dict): Extracted results so far

            For post_calculation:
            - objectives (dict): Computed objective values
            - constraints (dict): Constraint values

    Returns:
        Dictionary with computed values or modifications.
        These values are added to the trial context.
        Return empty dict {} if no modifications needed.

    Raises:
        Exception: Any exception will be logged but won't stop the trial
                   unless you want it to (raise optuna.TrialPruned instead)

    Example:
        >>> context = {'trial_number': 1, 'design_params': {'x': 5.0}}
        >>> result = {hook_name}_hook(context)
        >>> print(result)
        {{'{computed_key}': 123.45}}
    """
    # =========================================
    # Access context values
    # =========================================
    trial_num = context.get('trial_number', 0)
    design_params = context.get('design_params', {})
    config = context.get('config', {})
    working_dir = context.get('working_dir', Path('.'))

    # For post_solve hooks and later:
    # op2_file = context.get('op2_file')
    # solve_success = context.get('solve_success', False)

    # For post_extraction hooks and later:
    # results = context.get('results', {})

    # For post_calculation hooks:
    # objectives = context.get('objectives', {})
    # constraints = context.get('constraints', {})

    # =========================================
    # Your hook logic here
    # =========================================

    # Example: Log trial start (pre_solve hook)
    # print(f"[Hook] Trial {trial_num} starting with params: {design_params}")

    # Example: Compute derived quantity (post_extraction hook)
    # max_stress = results.get('max_von_mises', 0)
    # yield_strength = config.get('material', {}).get('yield_strength', 250)
    # safety_factor = yield_strength / max(max_stress, 1e-6)

    # Example: Write log file (post_calculation hook)
    # log_entry = {
    #     'trial': trial_num,
    #     'timestamp': datetime.now().isoformat(),
    #     'objectives': context.get('objectives', {}),
    # }
    # with open(working_dir / 'trial_log.jsonl', 'a') as f:
    #     f.write(json.dumps(log_entry) + '\n')

    # =========================================
    # Return computed values
    # =========================================

    # Values returned here are added to the context
    # and can be accessed by later hooks or the optimizer

    return {
        # '{computed_key}': computed_value,
    }


def register_hooks(hook_manager) -> None:
    """
    Register this hook with the hook manager.

    This function is called automatically when plugins are discovered.
    It must be named exactly 'register_hooks' and take one argument.

    Args:
        hook_manager: The HookManager instance from optimization_engine
    """
    hook_manager.register_hook(
        hook_point='{hook_point}',  # pre_mesh, post_mesh, pre_solve, etc.
        function={hook_name}_hook,
        name='{hook_name}_hook',
        description='{Brief description of what this hook does}',
        priority=100,  # Lower number = runs earlier (1-200 typical range)
        enabled=True   # Set to False to disable by default
    )


# =========================================
# Optional: Helper functions
# =========================================

def _helper_function(data: Any) -> Any:
    """
    Private helper function for the hook.

    Keep hook logic clean by extracting complex operations
    into helper functions.
    """
    pass


# =========================================
# After creating your hook:
#
# 1. Place in correct directory:
#    optimization_engine/plugins/{hook_point}/{hook_name}.py
#
# 2. Hook is auto-discovered - no __init__.py changes needed
#
# 3. Test the hook:
#    python -c "
#    from optimization_engine.plugins.hook_manager import HookManager
#    hm = HookManager()
#    hm.discover_plugins()
#    print(hm.list_hooks())
#    "
#
# 4. Update documentation if significant:
#    - Add to EXT_02_CREATE_HOOK.md examples section
# =========================================


# =========================================
# Example hooks for reference
# =========================================

def example_logger_hook(context: Dict[str, Any]) -> Dict[str, Any]:
    """Example: Simple trial logger for pre_solve."""
    trial = context.get('trial_number', 0)
    params = context.get('design_params', {})
    print(f"[LOG] Trial {trial} starting: {params}")
    return {}


def example_safety_factor_hook(context: Dict[str, Any]) -> Dict[str, Any]:
    """Example: Safety factor calculator for post_extraction."""
    results = context.get('results', {})
    max_stress = results.get('max_von_mises', 0)

    if max_stress > 0:
        safety_factor = 250.0 / max_stress  # Assuming 250 MPa yield
    else:
        safety_factor = float('inf')

    return {'safety_factor': safety_factor}


def example_validator_hook(context: Dict[str, Any]) -> Dict[str, Any]:
    """Example: Result validator for post_solve."""
    import optuna

    solve_success = context.get('solve_success', False)
    op2_file = context.get('op2_file')

    if not solve_success:
        raise optuna.TrialPruned("Solve failed")

    if op2_file and not Path(op2_file).exists():
        raise optuna.TrialPruned("OP2 file not generated")

    return {'validation_passed': True}