Atomizer/studies/Simple_Bracket/bracket_pareto_3obj/run_optimization.py

"""
bracket_pareto_3obj - Optimization Script
============================================================

Three-objective Pareto optimization: minimize mass, minimize stress, maximize stiffness

Protocol: Multi-Objective NSGA-II

Staged Workflow:
----------------
1. DISCOVER: python run_optimization.py --discover
2. VALIDATE: python run_optimization.py --validate
3. TEST: python run_optimization.py --test
4. RUN: python run_optimization.py --run --trials 100

Generated by StudyWizard on 2025-12-06 14:43
"""

from pathlib import Path
import sys
import json
import argparse
from datetime import datetime
from typing import Optional, Tuple, List

# Add parent directory to path
project_root = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(project_root))

import optuna
from optuna.samplers import NSGAIISampler

# Core imports
from optimization_engine.nx.solver import NXSolver
from optimization_engine.utils.logger import get_logger

# Extractor imports
from optimization_engine.extractors.bdf_mass_extractor import extract_mass_from_bdf
from optimization_engine.extractors.extract_displacement import extract_displacement
from optimization_engine.extractors.extract_von_mises_stress import extract_solid_stress


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


def clean_nastran_files(model_dir: Path, logger) -> List[Path]:
    """Remove old Nastran solver output files."""
    patterns = ['*.op2', '*.f06', '*.log', '*.f04', '*.pch', '*.DBALL', '*.MASTER', '_temp*.txt']
    deleted = []

    for pattern in patterns:
        for f in model_dir.glob(pattern):
            try:
                f.unlink()
                deleted.append(f)
                logger.info(f"  Deleted: {f.name}")
            except Exception as e:
                logger.warning(f"  Failed to delete {f.name}: {e}")

    return deleted


def objective(trial: optuna.Trial, config: dict, nx_solver: NXSolver,
              model_dir: Path, logger) -> Tuple[float, float, float]:
    """
    Objective function for optimization.

    Returns tuple of objectives for multi-objective optimization.
    """
    # Sample design variables
    design_vars = {}
    for var in config['design_variables']:
        param_name = var['parameter']
        bounds = var['bounds']
        design_vars[param_name] = trial.suggest_float(param_name, bounds[0], bounds[1])

    logger.trial_start(trial.number, design_vars)

    try:
        # Get file paths
        sim_file = model_dir / config['simulation']['sim_file']

        # Run FEA simulation
        result = nx_solver.run_simulation(
            sim_file=sim_file,
            working_dir=model_dir,
            expression_updates=design_vars,
            solution_name=config['simulation'].get('solution_name'),
            cleanup=True
        )

        if not result['success']:
            logger.trial_failed(trial.number, f"Simulation failed: {result.get('error', 'Unknown')}")
            return (float('inf'), float('inf'), float('inf'))

        op2_file = result['op2_file']
        dat_file = model_dir / config['simulation']['dat_file']

        # Extract results
        obj_mass = extract_mass_from_bdf(str(dat_file))
        logger.info(f'  mass: {obj_mass}')

        stress_result = extract_solid_stress(op2_file, subcase=1, element_type='chexa')
        obj_stress = stress_result.get('max_von_mises', float('inf')) / 1000.0  # kPa -> MPa
        logger.info(f'  stress: {obj_stress:.2f} MPa')

        disp_result = extract_displacement(op2_file, subcase=1)
        max_displacement = disp_result['max_displacement']
        # For stiffness maximization, use inverse of displacement
        applied_force = 1000.0  # N - adjust based on your model
        obj_stiffness = -applied_force / max(abs(max_displacement), 1e-6)
        logger.info(f'  stiffness: {obj_stiffness}')


        # Check constraints
        feasible = True
        constraint_results = {}
        # Check stress_limit (stress from OP2 is in kPa for mm/kg units, convert to MPa)
        const_stress_limit = extract_solid_stress(op2_file, element_type='chexa')
        stress_mpa = const_stress_limit.get('max_von_mises', float('inf')) / 1000.0  # kPa -> MPa
        constraint_results['stress_limit'] = stress_mpa
        if stress_mpa > 300:
            feasible = False
            logger.warning(f'  Constraint violation: stress_limit = {stress_mpa:.1f} MPa vs 300 MPa')


        # Set user attributes
        trial.set_user_attr('mass', obj_mass)
        trial.set_user_attr('stress', obj_stress)
        trial.set_user_attr('stiffness', obj_stiffness)
        trial.set_user_attr('feasible', feasible)

        objectives = {'mass': obj_mass, 'stress': obj_stress, 'stiffness': obj_stiffness}
        logger.trial_complete(trial.number, objectives, constraint_results, feasible)

        return (obj_mass, obj_stress, obj_stiffness)

    except Exception as e:
        logger.trial_failed(trial.number, str(e))
        return (float('inf'), float('inf'), float('inf'))


def main():
    """Main optimization workflow."""
    parser = argparse.ArgumentParser(description='bracket_pareto_3obj')

    stage_group = parser.add_mutually_exclusive_group()
    stage_group.add_argument('--discover', action='store_true', help='Discover model outputs')
    stage_group.add_argument('--validate', action='store_true', help='Run single validation trial')
    stage_group.add_argument('--test', action='store_true', help='Run 3-trial test')
    stage_group.add_argument('--run', action='store_true', help='Run optimization')

    parser.add_argument('--trials', type=int, default=100, help='Number of trials')
    parser.add_argument('--resume', action='store_true', help='Resume existing study')
    parser.add_argument('--clean', action='store_true', help='Clean old files first')

    args = parser.parse_args()

    if not any([args.discover, args.validate, args.test, args.run]):
        print("No stage specified. Use --discover, --validate, --test, or --run")
        return 1

    # Setup paths
    study_dir = Path(__file__).parent
    config_path = study_dir / "1_setup" / "optimization_config.json"
    model_dir = study_dir / "1_setup" / "model"
    results_dir = study_dir / "2_results"
    results_dir.mkdir(exist_ok=True)

    study_name = "bracket_pareto_3obj"

    # Initialize
    logger = get_logger(study_name, study_dir=results_dir)
    config = load_config(config_path)
    nx_solver = NXSolver(nastran_version="2506")

    if args.clean:
        clean_nastran_files(model_dir, logger)

    # Run appropriate stage
    if args.discover or args.validate or args.test:
        # Run limited trials for these stages
        n = 1 if args.discover or args.validate else 3
        storage = f"sqlite:///{results_dir / 'study_test.db'}"

        study = optuna.create_study(
            study_name=f"{study_name}_test",
            storage=storage,
            sampler=NSGAIISampler(population_size=5, seed=42),
            directions=['minimize'] * 3,
            load_if_exists=False
        )

        study.optimize(
            lambda trial: objective(trial, config, nx_solver, model_dir, logger),
            n_trials=n,
            show_progress_bar=True
        )

        logger.info(f"Completed {len(study.trials)} trial(s)")
        return 0

    # Full optimization run
    storage = f"sqlite:///{results_dir / 'study.db'}"

    if args.resume:
        study = optuna.load_study(
            study_name=study_name,
            storage=storage,
            sampler=NSGAIISampler(population_size=20, seed=42)
        )
    else:
        study = optuna.create_study(
            study_name=study_name,
            storage=storage,
            sampler=NSGAIISampler(population_size=20, seed=42),
            directions=['minimize'] * 3,
            load_if_exists=True
        )

    logger.study_start(study_name, args.trials, "NSGAIISampler")

    study.optimize(
        lambda trial: objective(trial, config, nx_solver, model_dir, logger),
        n_trials=args.trials,
        show_progress_bar=True
    )

    n_complete = len([t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE])
    logger.study_complete(study_name, len(study.trials), n_complete)

    # Report results
    pareto_trials = study.best_trials
    logger.info(f"\nOptimization Complete!")
    logger.info(f"Total trials: {len(study.trials)}")
    logger.info(f"Successful: {n_complete}")

    return 0


if __name__ == "__main__":
    exit(main())