""" Neural-Enhanced Optimization Runner Extends the base OptimizationRunner with neural network surrogate capabilities from AtomizerField for super-efficient optimization. Features: - Automatic neural surrogate integration when models are available - Hybrid optimization with smart FEA/NN switching - Confidence-based fallback to FEA - Speedup tracking and reporting """ from pathlib import Path from typing import Dict, Any, List, Optional, Callable, Tuple import json import time import logging import numpy as np from datetime import datetime import optuna from optimization_engine.runner import OptimizationRunner from optimization_engine.neural_surrogate import ( create_surrogate_from_config, create_hybrid_optimizer_from_config, NeuralSurrogate, HybridOptimizer ) logger = logging.getLogger(__name__) class NeuralOptimizationRunner(OptimizationRunner): """ Extended optimization runner with neural network surrogate support. Seamlessly integrates AtomizerField neural models to achieve 600x-500,000x speedup over traditional FEA-based optimization. """ def __init__( self, config_path: Path, model_updater: Callable, simulation_runner: Callable, result_extractors: Dict[str, Callable] ): """ Initialize neural-enhanced optimization runner. Args: config_path: Path to optimization_config.json model_updater: Function to update NX model parameters simulation_runner: Function to run FEA simulation result_extractors: Dictionary of result extraction functions """ # Initialize base class super().__init__(config_path, model_updater, simulation_runner, result_extractors) # Initialize neural surrogate components self.neural_surrogate: Optional[NeuralSurrogate] = None self.hybrid_optimizer: Optional[HybridOptimizer] = None self.neural_speedup_tracker = [] # Try to initialize neural components self._initialize_neural_components() def _initialize_neural_components(self): """Initialize neural surrogate and hybrid optimizer if configured.""" try: # Create neural surrogate from config self.neural_surrogate = create_surrogate_from_config(self.config) if self.neural_surrogate: logger.info("✓ Neural surrogate initialized successfully") logger.info(f" Model: {self.neural_surrogate.model_checkpoint}") logger.info(f" Confidence threshold: {self.neural_surrogate.confidence_threshold}") # Create hybrid optimizer for smart FEA/NN switching self.hybrid_optimizer = create_hybrid_optimizer_from_config(self.config) if self.hybrid_optimizer: logger.info("✓ Hybrid optimizer initialized") logger.info(f" Exploration trials: {self.hybrid_optimizer.exploration_trials}") logger.info(f" Validation frequency: {self.hybrid_optimizer.validation_frequency}") else: logger.info("Neural surrogate not configured - using standard FEA optimization") except Exception as e: logger.warning(f"Could not initialize neural components: {e}") logger.info("Falling back to standard FEA optimization") def _objective_function_with_neural(self, trial: optuna.Trial) -> float: """ Enhanced objective function with neural network surrogate support. Attempts to use neural network for fast prediction, falls back to FEA when confidence is low or validation is needed. Args: trial: Optuna trial object Returns: Objective value (float) """ # Sample design variables (same as base class) design_vars = self._sample_design_variables(trial) # Decide whether to use neural network or FEA use_neural = False nn_prediction = None nn_confidence = 0.0 if self.neural_surrogate and self.hybrid_optimizer: # Check if hybrid optimizer recommends using NN if self.hybrid_optimizer.should_use_nn(trial.number): # Try neural network prediction start_time = time.time() try: # Get case data for the current model case_data = self._prepare_case_data(design_vars) # Get neural network prediction predictions, confidence, used_nn = self.neural_surrogate.predict( design_vars, case_data=case_data ) if used_nn and predictions is not None: # Successfully used neural network nn_time = time.time() - start_time use_neural = True nn_prediction = predictions nn_confidence = confidence logger.info(f"Trial {trial.number}: Used neural network (confidence: {confidence:.2%}, time: {nn_time:.3f}s)") # Track speedup self.neural_speedup_tracker.append({ 'trial': trial.number, 'nn_time': nn_time, 'confidence': confidence }) else: logger.info(f"Trial {trial.number}: Neural confidence too low ({confidence:.2%}), using FEA") except Exception as e: logger.warning(f"Trial {trial.number}: Neural prediction failed: {e}, using FEA") # Execute hooks and get results if use_neural and nn_prediction is not None: # Use neural network results extracted_results = self._process_neural_results(nn_prediction, design_vars) # Skip model update and simulation since we used NN result_path = None else: # Fall back to standard FEA (using base class method) return super()._objective_function(trial) # Process constraints and objectives (same as base class) return self._evaluate_objectives_and_constraints( trial, design_vars, extracted_results, result_path ) def _sample_design_variables(self, trial: optuna.Trial) -> Dict[str, float]: """ Sample design variables from trial (extracted from base class). Args: trial: Optuna trial object Returns: Dictionary of design variable values """ design_vars = {} # Handle both dict and list formats for design_variables if isinstance(self.config['design_variables'], dict): for var_name, var_info in self.config['design_variables'].items(): if var_info['type'] == 'continuous': value = trial.suggest_float( var_name, var_info['min'], var_info['max'] ) precision = self._get_precision(var_name, var_info.get('units', '')) design_vars[var_name] = round(value, precision) elif var_info['type'] in ['discrete', 'integer']: design_vars[var_name] = trial.suggest_int( var_name, int(var_info['min']), int(var_info['max']) ) else: # Old format for dv in self.config['design_variables']: if dv['type'] == 'continuous': value = trial.suggest_float( dv['name'], dv['bounds'][0], dv['bounds'][1] ) precision = self._get_precision(dv['name'], dv.get('units', '')) design_vars[dv['name']] = round(value, precision) elif dv['type'] == 'discrete': design_vars[dv['name']] = trial.suggest_int( dv['name'], int(dv['bounds'][0]), int(dv['bounds'][1]) ) return design_vars def _prepare_case_data(self, design_vars: Dict[str, float]) -> Optional[Dict[str, Any]]: """ Prepare case-specific data for neural network prediction. This includes mesh file paths, boundary conditions, loads, etc. Args: design_vars: Current design variable values Returns: Case data dictionary or None """ try: case_data = { 'fem_file': self.config.get('fem_file', ''), 'sim_file': self.config.get('sim_file', ''), 'design_variables': design_vars, # Add any case-specific data needed by the neural network } # Add boundary conditions if specified if 'boundary_conditions' in self.config: case_data['boundary_conditions'] = self.config['boundary_conditions'] # Add loads if specified if 'loads' in self.config: case_data['loads'] = self.config['loads'] return case_data except Exception as e: logger.warning(f"Could not prepare case data: {e}") return None def _process_neural_results( self, nn_prediction: Dict[str, Any], design_vars: Dict[str, float] ) -> Dict[str, float]: """ Process neural network predictions into extracted results format. Args: nn_prediction: Raw neural network predictions design_vars: Current design variable values Returns: Dictionary of extracted results matching objective/constraint names """ extracted_results = {} # Map neural network outputs to objective/constraint names for obj in self.config['objectives']: obj_name = obj['name'] # Try to find matching prediction if obj_name in nn_prediction: value = nn_prediction[obj_name] elif 'metric' in obj and obj['metric'] in nn_prediction: value = nn_prediction[obj['metric']] else: # Try common mappings metric_mappings = { 'max_stress': ['max_von_mises_stress', 'stress', 'von_mises'], 'max_displacement': ['max_displacement', 'displacement', 'disp'], 'mass': ['mass', 'weight'], 'volume': ['volume'], 'compliance': ['compliance', 'strain_energy'], 'frequency': ['frequency', 'natural_frequency', 'freq'] } value = None for mapped_names in metric_mappings.get(obj_name, []): if mapped_names in nn_prediction: value = nn_prediction[mapped_names] break if value is None: raise ValueError(f"Could not find neural prediction for objective '{obj_name}'") # Apply appropriate precision precision = self._get_precision(obj_name, obj.get('units', '')) extracted_results[obj_name] = round(float(value), precision) # Process constraints similarly for const in self.config.get('constraints', []): const_name = const['name'] if const_name in nn_prediction: value = nn_prediction[const_name] elif 'metric' in const and const['metric'] in nn_prediction: value = nn_prediction[const['metric']] else: # Try to reuse objective values if constraint uses same metric if const_name in extracted_results: value = extracted_results[const_name] else: raise ValueError(f"Could not find neural prediction for constraint '{const_name}'") precision = self._get_precision(const_name, const.get('units', '')) extracted_results[const_name] = round(float(value), precision) return extracted_results def _evaluate_objectives_and_constraints( self, trial: optuna.Trial, design_vars: Dict[str, float], extracted_results: Dict[str, float], result_path: Optional[Path] ) -> float: """ Evaluate objectives and constraints (extracted from base class). Args: trial: Optuna trial object design_vars: Design variable values extracted_results: Extracted simulation/NN results result_path: Path to result files (None if using NN) Returns: Total objective value """ # Export training data if using FEA if self.training_data_exporter and result_path: self._export_training_data(trial.number, design_vars, extracted_results, result_path) # Evaluate constraints for const in self.config.get('constraints', []): value = extracted_results[const['name']] limit = const['limit'] if const['type'] == 'upper_bound' and value > limit: logger.info(f"Constraint violated: {const['name']} = {value:.4f} > {limit:.4f}") raise optuna.TrialPruned() elif const['type'] == 'lower_bound' and value < limit: logger.info(f"Constraint violated: {const['name']} = {value:.4f} < {limit:.4f}") raise optuna.TrialPruned() # Calculate weighted objective total_objective = 0.0 for obj in self.config['objectives']: value = extracted_results[obj['name']] weight = obj.get('weight', 1.0) direction = obj.get('direction', 'minimize') if direction == 'minimize': total_objective += weight * value else: # maximize total_objective -= weight * value # Store in history history_entry = { 'trial_number': trial.number, 'timestamp': datetime.now().isoformat(), 'design_variables': design_vars, 'objectives': {obj['name']: extracted_results[obj['name']] for obj in self.config['objectives']}, 'constraints': {const['name']: extracted_results[const['name']] for const in self.config.get('constraints', [])}, 'total_objective': total_objective, 'used_neural': result_path is None # Track if NN was used } self.history.append(history_entry) # Save history self._save_history() logger.info(f"Trial {trial.number} completed:") logger.info(f" Design vars: {design_vars}") logger.info(f" Objectives: {history_entry['objectives']}") logger.info(f" Total objective: {total_objective:.6f}") if history_entry.get('used_neural'): logger.info(f" Method: Neural Network") return total_objective def run( self, study_name: Optional[str] = None, n_trials: Optional[int] = None, resume: bool = False ) -> optuna.Study: """ Run neural-enhanced optimization. Args: study_name: Optional study name n_trials: Number of trials to run resume: Whether to resume existing study Returns: Completed Optuna study """ # Override objective function if neural surrogate is available if self.neural_surrogate: # Temporarily replace objective function original_objective = self._objective_function self._objective_function = self._objective_function_with_neural try: # Run optimization using base class study = super().run(study_name, n_trials, resume) # Print neural speedup summary if applicable if self.neural_speedup_tracker: self._print_speedup_summary() return study finally: # Restore original objective function if replaced if self.neural_surrogate: self._objective_function = original_objective def _print_speedup_summary(self): """Print summary of neural network speedup achieved.""" if not self.neural_speedup_tracker: return nn_trials = len(self.neural_speedup_tracker) total_trials = len(self.history) nn_percentage = (nn_trials / total_trials) * 100 avg_nn_time = np.mean([t['nn_time'] for t in self.neural_speedup_tracker]) avg_confidence = np.mean([t['confidence'] for t in self.neural_speedup_tracker]) # Estimate FEA time (rough estimate if not tracked) estimated_fea_time = 30 * 60 # 30 minutes in seconds estimated_speedup = estimated_fea_time / avg_nn_time print("\n" + "="*60) print("NEURAL NETWORK SPEEDUP SUMMARY") print("="*60) print(f"Trials using neural network: {nn_trials}/{total_trials} ({nn_percentage:.1f}%)") print(f"Average NN inference time: {avg_nn_time:.3f} seconds") print(f"Average NN confidence: {avg_confidence:.1%}") print(f"Estimated speedup: {estimated_speedup:.0f}x") print(f"Time saved: ~{(estimated_fea_time - avg_nn_time) * nn_trials / 3600:.1f} hours") print("="*60) def update_neural_model(self, new_checkpoint: Path): """ Update the neural network model checkpoint. Useful for updating to a newly trained model during optimization. Args: new_checkpoint: Path to new model checkpoint """ if self.neural_surrogate: try: self.neural_surrogate.load_model(new_checkpoint) logger.info(f"Updated neural model to: {new_checkpoint}") except Exception as e: logger.error(f"Failed to update neural model: {e}") def train_neural_model(self, training_data_dir: Path, epochs: int = 100): """ Train a new neural model on collected data. Args: training_data_dir: Directory containing training data epochs: Number of training epochs """ if self.hybrid_optimizer: try: model_path = self.hybrid_optimizer.train_surrogate_model(training_data_dir, epochs) # Update to use the newly trained model if model_path and self.neural_surrogate: self.update_neural_model(model_path) except Exception as e: logger.error(f"Failed to train neural model: {e}") def create_neural_runner( config_path: Path, model_updater: Callable, simulation_runner: Callable, result_extractors: Dict[str, Callable] ) -> NeuralOptimizationRunner: """ Factory function to create a neural-enhanced optimization runner. Args: config_path: Path to optimization configuration model_updater: Function to update model parameters simulation_runner: Function to run simulation result_extractors: Dictionary of result extraction functions Returns: NeuralOptimizationRunner instance """ return NeuralOptimizationRunner( config_path, model_updater, simulation_runner, result_extractors )