Atomizer/studies/simple_beam_optimization/run_optimization.py

"""
Simple Beam Optimization Study
===============================

Multi-objective optimization:
- Minimize displacement (constraint: < 10mm)
- Minimize stress
- Minimize mass

Design Variables:
- beam_half_core_thickness: 10-40 mm
- beam_face_thickness: 10-40 mm
- holes_diameter: 150-450 mm
- hole_count: 5-15
"""

import sys
import json
import optuna
from pathlib import Path
from datetime import datetime
from typing import Dict

# Add parent directories to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from optimization_engine.nx_updater import NXParameterUpdater
from optimization_engine.nx_solver import NXSolver
from optimization_engine.result_extractors.generated.extract_displacement import extract_displacement
from optimization_engine.result_extractors.generated.extract_solid_stress import extract_solid_stress
from optimization_engine.result_extractors.generated.extract_expression import extract_expression


def print_section(title: str):
    """Print a section header."""
    print()
    print("=" * 80)
    print(f"  {title}")
    print("=" * 80)
    print()


def load_config(config_file: Path) -> dict:
    """Load JSON configuration."""
    with open(config_file, 'r') as f:
        return json.load(f)


def main():
    print_section("SIMPLE BEAM OPTIMIZATION STUDY")

    # File paths
    study_dir = Path(__file__).parent
    config_file = study_dir / "beam_optimization_config.json"
    prt_file = study_dir / "model" / "Beam.prt"
    sim_file = study_dir / "model" / "Beam_sim1.sim"

    if not config_file.exists():
        print(f"ERROR: Config file not found: {config_file}")
        sys.exit(1)

    if not prt_file.exists():
        print(f"ERROR: Part file not found: {prt_file}")
        sys.exit(1)

    if not sim_file.exists():
        print(f"ERROR: Simulation file not found: {sim_file}")
        sys.exit(1)

    # Load configuration
    config = load_config(config_file)

    print("Study Configuration:")
    print(f"  - Study: {config['study_name']}")
    print(f"  - Substudy: {config['substudy_name']}")
    print(f"  - Description: {config['description']}")
    print()
    print("Objectives:")
    for obj in config['objectives']:
        print(f"  - {obj['name']}: weight={obj['weight']}")
    print()
    print("Constraints:")
    for con in config['constraints']:
        print(f"  - {con['name']}: {con['type']} {con['value']} {con['units']}")
    print()
    print("Design Variables:")
    for var_name, var_info in config['design_variables'].items():
        print(f"  - {var_name}: {var_info['min']}-{var_info['max']} {var_info['units']}")
    print()
    print(f"Optimization Settings:")
    print(f"  - Algorithm: {config['optimization_settings']['algorithm']}")
    print(f"  - Trials: {config['optimization_settings']['n_trials']}")
    print(f"  - Sampler: {config['optimization_settings']['sampler']}")
    print()

    # Setup output directory
    output_dir = study_dir / "substudies" / config['substudy_name']
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Part file: {prt_file}")
    print(f"Simulation file: {sim_file}")
    print(f"Output directory: {output_dir}")
    print()

    # =========================================================================
    # DEFINE OBJECTIVE FUNCTION
    # =========================================================================

    def objective(trial: optuna.Trial) -> float:
        """
        Optuna objective function.

        Evaluates one design point:
        1. Updates geometry parameters
        2. Runs FEM simulation
        3. Extracts results
        4. Computes weighted multi-objective with penalties
        """
        trial_num = trial.number

        print(f"\n[Trial {trial_num}] Starting...")

        # Sample design variables
        design_vars = {}
        for var_name, var_info in config['design_variables'].items():
            if var_info['type'] == 'continuous':
                design_vars[var_name] = trial.suggest_float(
                    var_name,
                    var_info['min'],
                    var_info['max']
                )
            elif var_info['type'] == 'integer':
                design_vars[var_name] = trial.suggest_int(
                    var_name,
                    int(var_info['min']),
                    int(var_info['max'])
                )

        print(f"[Trial {trial_num}] Design variables:")
        for var_name, var_value in design_vars.items():
            print(f"  - {var_name}: {var_value:.3f}")

        # Create trial directory
        trial_dir = output_dir / f"trial_{trial_num:03d}"
        trial_dir.mkdir(exist_ok=True)

        # Copy all 4 files to trial directory (.prt, _i.prt, .fem, .sim)
        import shutil
        trial_prt = trial_dir / prt_file.name
        trial_sim = trial_dir / sim_file.name

        shutil.copy2(prt_file, trial_prt)
        shutil.copy2(sim_file, trial_sim)

        # Copy FEM file
        fem_file = prt_file.parent / f"{prt_file.stem}_fem1.fem"
        if fem_file.exists():
            trial_fem = trial_dir / fem_file.name
            shutil.copy2(fem_file, trial_fem)

        # Copy idealized geometry (_i.prt) - contains midsurface thickness data
        # Pattern: Beam_fem1_i.prt (derived from FEM file name)
        if fem_file.exists():
            prt_i_file = prt_file.parent / f"{fem_file.stem}_i.prt"
            if prt_i_file.exists():
                trial_prt_i = trial_dir / prt_i_file.name
                shutil.copy2(prt_i_file, trial_prt_i)

        try:
            # Update geometry
            print(f"[Trial {trial_num}] Updating geometry...")
            updater = NXParameterUpdater(trial_prt)
            updater.update_expressions(design_vars)

            # Run simulation
            print(f"[Trial {trial_num}] Running FEM simulation...")
            solver = NXSolver()
            result = solver.run_simulation(trial_sim)

            if not result['success']:
                raise RuntimeError(f"Simulation failed: {result}")

            op2_file = result['op2_file']

            print(f"[Trial {trial_num}] Extracting results...")

            # Extract displacement
            disp_result = extract_displacement(op2_file)
            max_disp = disp_result['max_displacement']

            # Extract stress
            stress_result = extract_solid_stress(op2_file)
            max_stress = stress_result['max_von_mises']

            # Extract mass
            mass_result = extract_expression(trial_prt, 'p173')
            mass = mass_result['p173']

            print(f"[Trial {trial_num}] Results:")
            print(f"  - Displacement: {max_disp:.3f} mm")
            print(f"  - Stress: {max_stress:.3f} MPa")
            print(f"  - Mass: {mass:.3f} kg")

            # Compute weighted multi-objective
            objective_value = 0.0

            for obj in config['objectives']:
                if obj['extractor'] == 'max_displacement':
                    value = max_disp
                elif obj['extractor'] == 'max_stress':
                    value = max_stress
                elif obj['extractor'] == 'mass':
                    value = mass
                else:
                    continue

                weight = obj['weight']
                objective_value += weight * value

            # Apply constraint penalties
            penalty = 0.0
            for constraint in config['constraints']:
                if constraint['extractor'] == 'max_displacement':
                    current_value = max_disp
                elif constraint['extractor'] == 'max_stress':
                    current_value = max_stress
                else:
                    continue

                if constraint['type'] == 'less_than':
                    if current_value > constraint['value']:
                        violation = (current_value - constraint['value']) / constraint['value']
                        penalty += 1000.0 * violation
                        print(f"[Trial {trial_num}] CONSTRAINT VIOLATED: {constraint['name']}")
                        print(f"  Current: {current_value:.3f}, Limit: {constraint['value']}")

            total_objective = objective_value + penalty

            print(f"[Trial {trial_num}] Objective: {objective_value:.3f}, Penalty: {penalty:.3f}, Total: {total_objective:.3f}")

            # Save trial results
            trial_results = {
                'trial_number': trial_num,
                'design_variables': design_vars,
                'results': {
                    'max_displacement': max_disp,
                    'max_stress': max_stress,
                    'mass': mass
                },
                'objective': objective_value,
                'penalty': penalty,
                'total_objective': total_objective,
                'timestamp': datetime.now().isoformat()
            }

            with open(trial_dir / "results.json", 'w') as f:
                json.dump(trial_results, f, indent=2)

            return total_objective

        except Exception as e:
            print(f"[Trial {trial_num}] FAILED: {e}")
            import traceback
            traceback.print_exc()
            return 1e10  # Return large penalty for failed trials

    # =========================================================================
    # RUN OPTIMIZATION
    # =========================================================================

    print_section("RUNNING OPTIMIZATION")

    # Create Optuna study
    study = optuna.create_study(
        direction='minimize',
        sampler=optuna.samplers.TPESampler() if config['optimization_settings']['sampler'] == 'TPE' else None
    )

    # Run optimization
    print(f"Starting {config['optimization_settings']['n_trials']} optimization trials...")
    print()

    study.optimize(
        objective,
        n_trials=config['optimization_settings']['n_trials'],
        show_progress_bar=True
    )

    # =========================================================================
    # SAVE RESULTS
    # =========================================================================

    print_section("SAVING RESULTS")

    # Save full study
    study_file = output_dir / "optuna_study.pkl"
    import pickle
    with open(study_file, 'wb') as f:
        pickle.dump(study, f)

    print(f"Study saved to: {study_file}")

    # Save best trial
    best_trial = study.best_trial
    best_results = {
        'best_trial_number': best_trial.number,
        'best_params': best_trial.params,
        'best_value': best_trial.value,
        'timestamp': datetime.now().isoformat()
    }

    best_file = output_dir / "best_trial.json"
    with open(best_file, 'w') as f:
        json.dump(best_results, f, indent=2)

    print(f"Best trial saved to: {best_file}")
    print()

    # =========================================================================
    # PRINT SUMMARY
    # =========================================================================

    print_section("OPTIMIZATION COMPLETE")

    print(f"Total trials: {len(study.trials)}")
    print(f"Best trial: {best_trial.number}")
    print(f"Best objective value: {best_trial.value:.6f}")
    print()
    print("Best design variables:")
    for var_name, var_value in best_trial.params.items():
        print(f"  - {var_name}: {var_value:.3f}")
    print()

    # Load best trial results to show performance
    best_trial_dir = output_dir / f"trial_{best_trial.number:03d}"
    best_results_file = best_trial_dir / "results.json"

    if best_results_file.exists():
        with open(best_results_file, 'r') as f:
            best_results = json.load(f)

        print("Best performance:")
        print(f"  - Displacement: {best_results['results']['max_displacement']:.3f} mm")
        print(f"  - Stress: {best_results['results']['max_stress']:.3f} MPa")
        print(f"  - Mass: {best_results['results']['mass']:.3f} kg")
        print()


if __name__ == "__main__":
    main()