""" LLM-Enhanced Optimization Runner - Phase 3.2 Flexible LLM-enhanced optimization runner that integrates: - Phase 2.7: LLM workflow analysis - Phase 2.8: Inline code generation (optional) - Phase 2.9: Post-processing hook generation (optional) - Phase 3.0: pyNastran research agent (optional) - Phase 3.1: Extractor orchestration (optional) This runner enables users to describe optimization goals in natural language and choose to leverage automated code generation, manual coding, or a hybrid approach. Author: Atomizer Development Team Version: 0.1.0 (Phase 3.2) Last Updated: 2025-01-16 """ from pathlib import Path from typing import Dict, Any, List, Optional import json import logging import optuna from datetime import datetime from optimization_engine.extractor_orchestrator import ExtractorOrchestrator from optimization_engine.inline_code_generator import InlineCodeGenerator from optimization_engine.hook_generator import HookGenerator from optimization_engine.plugins.hook_manager import HookManager logger = logging.getLogger(__name__) class LLMOptimizationRunner: """ LLM-enhanced optimization runner with flexible automation options. This runner empowers users to leverage LLM-assisted code generation for: - OP2 result extractors (Phase 3.1) - optional - Inline calculations (Phase 2.8) - optional - Post-processing hooks (Phase 2.9) - optional Users can describe goals in natural language and choose automated generation, manual coding, or a hybrid approach based on their needs. """ def __init__(self, llm_workflow: Dict[str, Any], model_updater: callable, simulation_runner: callable, study_name: str = "llm_optimization", output_dir: Optional[Path] = None): """ Initialize LLM-driven optimization runner. Args: llm_workflow: Output from Phase 2.7 LLM analysis with: - engineering_features: List of FEA operations - inline_calculations: List of simple math operations - post_processing_hooks: List of custom calculations - optimization: Dict with algorithm, design_variables, etc. model_updater: Function(design_vars: Dict) -> None Updates NX expressions in the CAD model and saves changes. simulation_runner: Function(design_vars: Dict) -> Path Runs FEM simulation with updated design variables. Returns path to OP2 results file. study_name: Name for Optuna study output_dir: Directory for results """ self.llm_workflow = llm_workflow self.model_updater = model_updater self.simulation_runner = simulation_runner self.study_name = study_name if output_dir is None: output_dir = Path.cwd() / "optimization_results" / study_name self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) # Initialize automation components self._initialize_automation() # Optuna study self.study = None self.history = [] logger.info(f"LLMOptimizationRunner initialized for study: {study_name}") def _initialize_automation(self): """Initialize all automation components from LLM workflow.""" logger.info("Initializing automation components...") # Phase 3.1: Extractor Orchestrator logger.info(" - Phase 3.1: Extractor Orchestrator") self.orchestrator = ExtractorOrchestrator( extractors_dir=self.output_dir / "generated_extractors" ) # Generate extractors from LLM workflow self.extractors = self.orchestrator.process_llm_workflow(self.llm_workflow) logger.info(f" Generated {len(self.extractors)} extractor(s)") # Phase 2.8: Inline Code Generator logger.info(" - Phase 2.8: Inline Code Generator") self.inline_generator = InlineCodeGenerator() self.inline_code = [] for calc in self.llm_workflow.get('inline_calculations', []): generated = self.inline_generator.generate_from_llm_output(calc) self.inline_code.append(generated.code) logger.info(f" Generated {len(self.inline_code)} inline calculation(s)") # Phase 2.9: Hook Generator logger.info(" - Phase 2.9: Hook Generator") self.hook_generator = HookGenerator() # Generate lifecycle hooks from post_processing_hooks hook_dir = self.output_dir / "generated_hooks" hook_dir.mkdir(exist_ok=True) for hook_spec in self.llm_workflow.get('post_processing_hooks', []): hook_content = self.hook_generator.generate_lifecycle_hook( hook_spec, hook_point='post_calculation' ) # Save hook hook_name = hook_spec.get('action', 'custom_hook') hook_file = hook_dir / f"{hook_name}.py" with open(hook_file, 'w') as f: f.write(hook_content) logger.info(f" Generated hook: {hook_name}") # Phase 1: Hook Manager logger.info(" - Phase 1: Hook Manager") self.hook_manager = HookManager() # Load generated hooks if hook_dir.exists(): self.hook_manager.load_plugins_from_directory(hook_dir) # Load system hooks system_hooks_dir = Path(__file__).parent / 'plugins' if system_hooks_dir.exists(): self.hook_manager.load_plugins_from_directory(system_hooks_dir) summary = self.hook_manager.get_summary() logger.info(f" Loaded {summary['enabled_hooks']} hook(s)") logger.info("Automation components initialized successfully!") def _create_optuna_study(self) -> optuna.Study: """Create Optuna study from LLM workflow optimization settings.""" opt_config = self.llm_workflow.get('optimization', {}) # Determine direction (minimize or maximize) direction = opt_config.get('direction', 'minimize') # Create study study = optuna.create_study( study_name=self.study_name, direction=direction, storage=f"sqlite:///{self.output_dir / f'{self.study_name}.db'}", load_if_exists=True ) logger.info(f"Created Optuna study: {self.study_name} (direction: {direction})") return study def _objective(self, trial: optuna.Trial) -> float: """ Optuna objective function - LLM-enhanced with flexible automation! This function leverages LLM workflow analysis with user-configurable automation: 1. Suggests design variables from LLM analysis 2. Updates model 3. Runs simulation 4. Extracts results (using generated or manual extractors) 5. Executes inline calculations (generated or manual) 6. Executes post-calculation hooks (generated or manual) 7. Returns objective value Args: trial: Optuna trial Returns: Objective value """ trial_number = trial.number logger.info(f"\n{'='*80}") logger.info(f"Trial {trial_number} starting...") logger.info(f"{'='*80}") # ==================================================================== # STEP 1: Suggest Design Variables # ==================================================================== design_vars_config = self.llm_workflow.get('optimization', {}).get('design_variables', []) design_vars = {} for var_config in design_vars_config: var_name = var_config['parameter'] var_min = var_config.get('min', 0.0) var_max = var_config.get('max', 1.0) # Suggest value using Optuna design_vars[var_name] = trial.suggest_float(var_name, var_min, var_max) logger.info(f"Design variables: {design_vars}") # Execute pre-solve hooks self.hook_manager.execute_hooks('pre_solve', { 'trial_number': trial_number, 'design_variables': design_vars }) # ==================================================================== # STEP 2: Update Model # ==================================================================== logger.info("Updating model...") self.model_updater(design_vars) # ==================================================================== # STEP 3: Run Simulation # ==================================================================== logger.info("Running simulation...") # Pass design_vars to simulation_runner so NX journal can update expressions op2_file = self.simulation_runner(design_vars) logger.info(f"Simulation complete: {op2_file}") # Execute post-solve hooks self.hook_manager.execute_hooks('post_solve', { 'trial_number': trial_number, 'op2_file': op2_file }) # ==================================================================== # STEP 4: Extract Results (Phase 3.1 - Auto-Generated Extractors) # ==================================================================== logger.info("Extracting results...") results = {} for extractor in self.extractors: try: extraction_result = self.orchestrator.execute_extractor( extractor.name, Path(op2_file), subcase=1 ) results.update(extraction_result) logger.info(f" {extractor.name}: {list(extraction_result.keys())}") except Exception as e: logger.error(f"Extraction failed for {extractor.name}: {e}") # Continue with other extractors # Execute post-extraction hooks self.hook_manager.execute_hooks('post_extraction', { 'trial_number': trial_number, 'results': results }) # ==================================================================== # STEP 5: Inline Calculations (Phase 2.8 - Auto-Generated Code) # ==================================================================== logger.info("Executing inline calculations...") calculations = {} calc_namespace = {**results, **calculations} # Make results available for calc_code in self.inline_code: try: exec(calc_code, calc_namespace) # Extract newly created variables for key, value in calc_namespace.items(): if key not in results and not key.startswith('_'): calculations[key] = value logger.info(f" Executed: {calc_code[:50]}...") except Exception as e: logger.error(f"Inline calculation failed: {e}") logger.info(f"Calculations: {calculations}") # ==================================================================== # STEP 6: Post-Calculation Hooks (Phase 2.9 - Auto-Generated Hooks) # ==================================================================== logger.info("Executing post-calculation hooks...") hook_results = self.hook_manager.execute_hooks('post_calculation', { 'trial_number': trial_number, 'design_variables': design_vars, 'results': results, 'calculations': calculations }) # Merge hook results final_context = {**results, **calculations} for hook_result in hook_results: if hook_result: final_context.update(hook_result) logger.info(f"Hook results: {hook_results}") # ==================================================================== # STEP 7: Extract Objective Value # ==================================================================== # Try to get objective from hooks first objective = None # Check hook results for 'objective' or 'weighted_objective' for hook_result in hook_results: if hook_result: if 'objective' in hook_result: objective = hook_result['objective'] break elif 'weighted_objective' in hook_result: objective = hook_result['weighted_objective'] break # Fallback: use first extracted result if objective is None: # Try common objective names for key in ['max_displacement', 'max_stress', 'max_von_mises']: if key in final_context: objective = final_context[key] logger.warning(f"No explicit objective found, using: {key}") break if objective is None: raise ValueError("Could not determine objective value from results/calculations/hooks") logger.info(f"Objective value: {objective}") # Save trial history trial_data = { 'trial_number': trial_number, 'design_variables': design_vars, 'results': results, 'calculations': calculations, 'objective': objective } self.history.append(trial_data) # Incremental save - write history after each trial # This allows monitoring progress in real-time self._save_incremental_history() return float(objective) def run_optimization(self, n_trials: int = 50) -> Dict[str, Any]: """ Run LLM-enhanced optimization with flexible automation. Args: n_trials: Number of optimization trials Returns: Dict with: - best_params: Best design variable values - best_value: Best objective value - history: Complete trial history """ logger.info(f"\n{'='*80}") logger.info(f"Starting LLM-Driven Optimization") logger.info(f"{'='*80}") logger.info(f"Study: {self.study_name}") logger.info(f"Trials: {n_trials}") logger.info(f"Output: {self.output_dir}") logger.info(f"{'='*80}\n") # Create study self.study = self._create_optuna_study() # Run optimization self.study.optimize(self._objective, n_trials=n_trials) # Get results best_trial = self.study.best_trial results = { 'best_params': best_trial.params, 'best_value': best_trial.value, 'best_trial_number': best_trial.number, 'history': self.history } # Save results self._save_results(results) logger.info(f"\n{'='*80}") logger.info("Optimization Complete!") logger.info(f"{'='*80}") logger.info(f"Best value: {results['best_value']}") logger.info(f"Best params: {results['best_params']}") logger.info(f"Results saved to: {self.output_dir}") logger.info(f"{'='*80}\n") return results def _save_incremental_history(self): """ Save trial history incrementally after each trial. This allows real-time monitoring of optimization progress. """ history_file = self.output_dir / "optimization_history_incremental.json" # Convert history to JSON-serializable format serializable_history = [] for trial in self.history: trial_copy = trial.copy() # Convert any numpy types to native Python types for key in ['results', 'calculations', 'design_variables']: if key in trial_copy: trial_copy[key] = {k: float(v) if isinstance(v, (int, float)) else v for k, v in trial_copy[key].items()} if 'objective' in trial_copy: trial_copy['objective'] = float(trial_copy['objective']) serializable_history.append(trial_copy) # Write to file with open(history_file, 'w') as f: json.dump(serializable_history, f, indent=2, default=str) def _save_results(self, results: Dict[str, Any]): """Save optimization results to file.""" results_file = self.output_dir / "optimization_results.json" # Make history JSON serializable serializable_results = { 'best_params': results['best_params'], 'best_value': results['best_value'], 'best_trial_number': results['best_trial_number'], 'timestamp': datetime.now().isoformat(), 'study_name': self.study_name, 'n_trials': len(results['history']) } with open(results_file, 'w') as f: json.dump(serializable_results, f, indent=2) logger.info(f"Results saved to: {results_file}") def main(): """Test LLM-driven optimization runner.""" print("=" * 80) print("Phase 3.2: LLM-Driven Optimization Runner Test") print("=" * 80) print() # Example LLM workflow (from Phase 2.7) llm_workflow = { "engineering_features": [ { "action": "extract_displacement", "domain": "result_extraction", "description": "Extract displacement from OP2", "params": {"result_type": "displacement"} } ], "inline_calculations": [ { "action": "normalize", "params": { "input": "max_displacement", "reference": "max_allowed_disp", "value": 5.0 }, "code_hint": "norm_disp = max_displacement / 5.0" } ], "post_processing_hooks": [ { "action": "weighted_objective", "params": { "inputs": ["norm_disp"], "weights": [1.0], "objective": "minimize" } } ], "optimization": { "algorithm": "TPE", "direction": "minimize", "design_variables": [ { "parameter": "wall_thickness", "min": 3.0, "max": 8.0, "type": "continuous" } ] } } print("LLM Workflow Configuration:") print(f" Engineering features: {len(llm_workflow['engineering_features'])}") print(f" Inline calculations: {len(llm_workflow['inline_calculations'])}") print(f" Post-processing hooks: {len(llm_workflow['post_processing_hooks'])}") print(f" Design variables: {len(llm_workflow['optimization']['design_variables'])}") print() # Dummy functions for testing def dummy_model_updater(design_vars): print(f" [Dummy] Updating model with: {design_vars}") def dummy_simulation_runner(): print(" [Dummy] Running simulation...") # Return path to test OP2 return Path("tests/bracket_sim1-solution_1.op2") # Initialize runner print("Initializing LLM-driven optimization runner...") runner = LLMOptimizationRunner( llm_workflow=llm_workflow, model_updater=dummy_model_updater, simulation_runner=dummy_simulation_runner, study_name="test_llm_optimization" ) print() print("=" * 80) print("Runner initialized successfully!") print("Ready to run optimization with auto-generated code!") print("=" * 80) if __name__ == '__main__': main()