""" Neural network surrogate integration for Atomizer. This module provides the integration layer between Atomizer optimization framework and AtomizerField neural network models for fast FEA predictions. Key Features: - Load and manage AtomizerField trained models - Convert design variables to neural field format - Provide millisecond FEA predictions - Automatic fallback to FEA when confidence is low - Performance tracking and statistics Usage: from optimization_engine.neural_surrogate import NeuralSurrogate, create_surrogate_for_study # Create surrogate for UAV arm study surrogate = create_surrogate_for_study( model_path="atomizer-field/runs/uav_arm_model/checkpoint_best.pt", training_data_dir="atomizer_field_training_data/uav_arm_train" ) # Predict for new design results = surrogate.predict(design_params) print(f"Max displacement: {results['max_displacement']:.6f} mm") """ import sys import time import json import logging import h5py from pathlib import Path from typing import Dict, Any, Optional, Tuple, List import numpy as np logger = logging.getLogger(__name__) # Add atomizer-field to path for imports _atomizer_field_path = Path(__file__).parent.parent / 'atomizer-field' if str(_atomizer_field_path) not in sys.path: sys.path.insert(0, str(_atomizer_field_path)) try: import torch from torch_geometric.data import Data TORCH_AVAILABLE = True except ImportError: TORCH_AVAILABLE = False logger.warning("PyTorch not installed. Neural surrogate features will be limited.") # Import AtomizerField model ATOMIZER_FIELD_AVAILABLE = False PARAMETRIC_MODEL_AVAILABLE = False if TORCH_AVAILABLE: try: from neural_models.field_predictor import AtomizerFieldModel, create_model ATOMIZER_FIELD_AVAILABLE = True except ImportError as e: logger.warning(f"AtomizerField modules not found: {e}") try: from neural_models.parametric_predictor import ParametricFieldPredictor, create_parametric_model PARAMETRIC_MODEL_AVAILABLE = True except ImportError as e: logger.warning(f"Parametric predictor modules not found: {e}") class NeuralSurrogate: """ Neural surrogate for fast FEA predictions using trained AtomizerField model. This class loads a trained AtomizerField model and provides fast predictions of displacement fields, which can then be used to compute derived quantities like max displacement, estimated stress, etc. """ def __init__( self, model_path: Path, training_data_dir: Path, device: str = 'auto' ): """ Initialize neural surrogate. Args: model_path: Path to trained model checkpoint (.pt file) training_data_dir: Path to training data (for normalization stats and mesh) device: Computing device ('cuda', 'cpu', or 'auto') """ if not TORCH_AVAILABLE: raise ImportError("PyTorch required. Install: pip install torch torch-geometric") if not ATOMIZER_FIELD_AVAILABLE: raise ImportError("AtomizerField modules not found") self.model_path = Path(model_path) self.training_data_dir = Path(training_data_dir) # Set device if device == 'auto': self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') else: self.device = torch.device(device) logger.info(f"Neural Surrogate initializing on {self.device}") # Load model self._load_model() # Load normalization statistics self._load_normalization_stats() # Load reference mesh structure self._load_reference_mesh() # Performance tracking self.stats = { 'predictions': 0, 'total_time_ms': 0.0, 'fea_validations': 0 } logger.info(f"Neural Surrogate ready: {self.num_nodes} nodes, model loaded") def _load_model(self): """Load trained AtomizerField model.""" logger.info(f"Loading model from {self.model_path}") checkpoint = torch.load(self.model_path, map_location=self.device) # Create model with saved config model_config = checkpoint['config']['model'] self.model = AtomizerFieldModel(**model_config) self.model.load_state_dict(checkpoint['model_state_dict']) self.model = self.model.to(self.device) self.model.eval() self.model_config = checkpoint['config'] self.best_val_loss = checkpoint.get('best_val_loss', None) n_params = sum(p.numel() for p in self.model.parameters()) logger.info(f"Model loaded: {n_params:,} parameters, val_loss={self.best_val_loss:.4f}") def _load_normalization_stats(self): """Load normalization statistics from training data.""" case_dirs = sorted(self.training_data_dir.glob("trial_*")) if not case_dirs: logger.warning("No training cases found - using identity normalization") self.coord_mean = np.zeros(3) self.coord_std = np.ones(3) self.disp_mean = np.zeros(6) self.disp_std = np.ones(6) return # Compute stats from all training data all_coords = [] all_disp = [] for case_dir in case_dirs: h5_file = case_dir / "neural_field_data.h5" if h5_file.exists(): with h5py.File(h5_file, 'r') as f: all_coords.append(f['mesh/node_coordinates'][:]) all_disp.append(f['results/displacement'][:]) if all_coords: all_coords = np.concatenate(all_coords, axis=0) all_disp = np.concatenate(all_disp, axis=0) self.coord_mean = all_coords.mean(axis=0) self.coord_std = all_coords.std(axis=0) + 1e-8 self.disp_mean = all_disp.mean(axis=0) self.disp_std = all_disp.std(axis=0) + 1e-8 logger.info(f"Normalization stats from {len(case_dirs)} cases") def _load_reference_mesh(self): """Load reference mesh structure for building graphs.""" case_dirs = sorted(self.training_data_dir.glob("trial_*")) if not case_dirs: raise ValueError(f"No training cases in {self.training_data_dir}") first_case = case_dirs[0] json_file = first_case / "neural_field_data.json" h5_file = first_case / "neural_field_data.h5" # Load metadata with open(json_file, 'r') as f: self.reference_metadata = json.load(f) # Load mesh with h5py.File(h5_file, 'r') as f: self.reference_coords = f['mesh/node_coordinates'][:] self.num_nodes = self.reference_coords.shape[0] # Build edge index (constant for parametric optimization) self._build_graph_structure() def _build_graph_structure(self): """Build graph edge index and attributes from mesh.""" metadata = self.reference_metadata num_nodes = self.num_nodes edge_list = [] # Get material properties mat_props = [0.0] * 5 if 'materials' in metadata: for mat in metadata['materials']: if mat['type'] == 'MAT1': mat_props = [ mat.get('E', 0.0) / 1e6, mat.get('nu', 0.0), mat.get('rho', 0.0) * 1e6, mat.get('G', 0.0) / 1e6 if mat.get('G') else 0.0, mat.get('alpha', 0.0) * 1e6 if mat.get('alpha') else 0.0 ] break # Process elements to create edges if 'mesh' in metadata and 'elements' in metadata['mesh']: for elem_type in ['solid', 'shell', 'beam']: if elem_type in metadata['mesh']['elements']: for elem in metadata['mesh']['elements'][elem_type]: elem_nodes = elem['nodes'] for i in range(len(elem_nodes)): for j in range(i + 1, len(elem_nodes)): node_i = elem_nodes[i] - 1 node_j = elem_nodes[j] - 1 if node_i < num_nodes and node_j < num_nodes: edge_list.append([node_i, node_j]) edge_list.append([node_j, node_i]) if edge_list: self.edge_index = torch.tensor(edge_list, dtype=torch.long).t().to(self.device) num_edges = self.edge_index.shape[1] self.edge_attr = torch.tensor([mat_props] * num_edges, dtype=torch.float).to(self.device) else: self.edge_index = torch.zeros((2, 0), dtype=torch.long).to(self.device) self.edge_attr = torch.zeros((0, 5), dtype=torch.float).to(self.device) # Build BC mask and load features (constant for this study) self._build_bc_and_loads() def _build_bc_and_loads(self): """Build boundary condition mask and load features.""" metadata = self.reference_metadata num_nodes = self.num_nodes # BC mask self.bc_mask = torch.zeros(num_nodes, 6) if 'boundary_conditions' in metadata and 'spc' in metadata['boundary_conditions']: for spc in metadata['boundary_conditions']['spc']: node_id = spc['node'] if node_id <= num_nodes: dofs = spc['dofs'] for dof_char in str(dofs): if dof_char.isdigit(): dof_idx = int(dof_char) - 1 if 0 <= dof_idx < 6: self.bc_mask[node_id - 1, dof_idx] = 1.0 # Load features self.load_features = torch.zeros(num_nodes, 3) if 'loads' in metadata and 'point_forces' in metadata['loads']: for force in metadata['loads']['point_forces']: node_id = force['node'] if node_id <= num_nodes: magnitude = force['magnitude'] direction = force['direction'] force_vector = [magnitude * d for d in direction] self.load_features[node_id - 1] = torch.tensor(force_vector) self.bc_mask = self.bc_mask.to(self.device) self.load_features = self.load_features.to(self.device) def _build_node_features(self) -> torch.Tensor: """Build node features tensor for model input.""" # Normalized coordinates coords = torch.from_numpy(self.reference_coords).float() coords_norm = (coords - torch.from_numpy(self.coord_mean).float()) / \ torch.from_numpy(self.coord_std).float() coords_norm = coords_norm.to(self.device) # Concatenate: [coords(3) + bc_mask(6) + loads(3)] = 12 features node_features = torch.cat([coords_norm, self.bc_mask, self.load_features], dim=-1) return node_features def predict( self, design_params: Dict[str, float], return_fields: bool = False ) -> Dict[str, Any]: """ Predict FEA results using neural network. Args: design_params: Design parameter values (not used for prediction, but kept for API compatibility - mesh is constant) return_fields: If True, return complete displacement field Returns: dict with: - max_displacement: Maximum displacement magnitude (mm) - max_stress: Estimated maximum stress (approximate) - inference_time_ms: Prediction time - fields: Complete displacement field (if return_fields=True) """ start_time = time.time() # Build graph data node_features = self._build_node_features() graph_data = Data( x=node_features, edge_index=self.edge_index, edge_attr=self.edge_attr ) # Predict with torch.no_grad(): predictions = self.model(graph_data, return_stress=True) # Denormalize displacement displacement = predictions['displacement'].cpu().numpy() displacement = displacement * self.disp_std + self.disp_mean # Compute max values disp_magnitude = np.linalg.norm(displacement[:, :3], axis=1) max_displacement = float(np.max(disp_magnitude)) # Stress (approximate - model trained on displacement only) max_stress = float(torch.max(predictions['von_mises']).item()) inference_time = (time.time() - start_time) * 1000 results = { 'max_displacement': max_displacement, 'max_stress': max_stress, 'inference_time_ms': inference_time } if return_fields: results['displacement_field'] = displacement results['von_mises_field'] = predictions['von_mises'].cpu().numpy() # Update stats self.stats['predictions'] += 1 self.stats['total_time_ms'] += inference_time return results def get_statistics(self) -> Dict[str, Any]: """Get prediction statistics.""" avg_time = self.stats['total_time_ms'] / self.stats['predictions'] \ if self.stats['predictions'] > 0 else 0 return { 'total_predictions': self.stats['predictions'], 'total_time_ms': self.stats['total_time_ms'], 'average_time_ms': avg_time, 'model_path': str(self.model_path), 'best_val_loss': self.best_val_loss, 'device': str(self.device) } def needs_fea_validation(self, trial_number: int) -> bool: """ Determine if FEA validation is recommended. Args: trial_number: Current trial number Returns: True if FEA validation is recommended """ # Validate periodically if trial_number < 5: return True # First few always validate if trial_number % 20 == 0: return True # Periodic validation return False class ParametricSurrogate: """ Parametric neural surrogate that predicts ALL objectives from design parameters. Unlike NeuralSurrogate which only predicts displacement fields, ParametricSurrogate directly predicts: - mass - frequency - max_displacement - max_stress This is the "future-proof" solution using design-conditioned GNN. """ def __init__( self, model_path: Path, training_data_dir: Path = None, device: str = 'auto', num_nodes: int = 500 ): """ Initialize parametric surrogate. Args: model_path: Path to trained parametric model checkpoint (.pt file) training_data_dir: Path to training data (optional - for mesh loading) device: Computing device ('cuda', 'cpu', or 'auto') num_nodes: Number of nodes for synthetic reference graph (default: 500) """ if not TORCH_AVAILABLE: raise ImportError("PyTorch required. Install: pip install torch torch-geometric") if not PARAMETRIC_MODEL_AVAILABLE: raise ImportError("Parametric predictor modules not found") self.model_path = Path(model_path) self.training_data_dir = Path(training_data_dir) if training_data_dir else None self.num_nodes = num_nodes # Set device if device == 'auto': self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') else: self.device = torch.device(device) logger.info(f"Parametric Surrogate initializing on {self.device}") # Load model and normalization self._load_model() # Create reference graph structure (synthetic - matching training) self._create_reference_graph() # Performance tracking self.stats = { 'predictions': 0, 'total_time_ms': 0.0 } logger.info(f"Parametric Surrogate ready: {self.num_nodes} nodes, " f"predicts mass/freq/disp/stress") def _load_model(self): """Load trained parametric model and normalization stats.""" logger.info(f"Loading parametric model from {self.model_path}") checkpoint = torch.load(self.model_path, map_location=self.device) # Create model with saved config model_config = checkpoint['config'] self.model = create_parametric_model(model_config) self.model.load_state_dict(checkpoint['model_state_dict']) self.model = self.model.to(self.device) self.model.eval() self.model_config = model_config self.best_val_loss = checkpoint.get('best_val_loss', None) # Load normalization stats norm = checkpoint.get('normalization', {}) self.design_var_names = checkpoint.get('design_var_names', []) self.n_design_vars = len(self.design_var_names) self.design_mean = torch.tensor(norm.get('design_mean', [0.0] * self.n_design_vars), dtype=torch.float32, device=self.device) self.design_std = torch.tensor(norm.get('design_std', [1.0] * self.n_design_vars), dtype=torch.float32, device=self.device) self.coord_mean = np.array(norm.get('coord_mean', [0.0, 0.0, 0.0])) self.coord_std = np.array(norm.get('coord_std', [1.0, 1.0, 1.0])) self.disp_mean = np.array(norm.get('disp_mean', [0.0] * 6)) self.disp_std = np.array(norm.get('disp_std', [1.0] * 6)) # Scalar normalization stats (for denormalization) self.mass_mean = norm.get('mass_mean', 3500.0) self.mass_std = norm.get('mass_std', 700.0) self.freq_mean = norm.get('freq_mean', 18.0) self.freq_std = norm.get('freq_std', 2.0) self.max_disp_mean = norm.get('max_disp_mean', 0.025) self.max_disp_std = norm.get('max_disp_std', 0.005) self.max_stress_mean = norm.get('max_stress_mean', 200.0) self.max_stress_std = norm.get('max_stress_std', 50.0) n_params = sum(p.numel() for p in self.model.parameters()) logger.info(f"Parametric model loaded: {n_params:,} params, " f"val_loss={self.best_val_loss:.4f}") logger.info(f"Design vars: {self.design_var_names}") def _create_reference_graph(self): """ Create a synthetic reference graph structure for the GNN. This matches the create_reference_graph() function used during training in train_parametric.py. The model was trained on synthetic graphs, so we need to use the same structure for inference. """ num_nodes = self.num_nodes # Create simple node features (random, like training) # [coords(3) + bc_mask(6) + loads(3)] = 12 features x = torch.randn(num_nodes, 12, device=self.device) # Create grid-like connectivity (same as training) edges = [] grid_size = int(np.ceil(np.sqrt(num_nodes))) for i in range(num_nodes): row = i // grid_size col = i % grid_size # Right neighbor (same row) right = i + 1 if col < grid_size - 1 and right < num_nodes: edges.append([i, right]) edges.append([right, i]) # Bottom neighbor (next row) bottom = i + grid_size if bottom < num_nodes: edges.append([i, bottom]) edges.append([bottom, i]) # Ensure we have at least some edges if len(edges) == 0: # Fallback: partially connected for very small graphs for i in range(num_nodes): for j in range(i + 1, min(i + 5, num_nodes)): edges.append([i, j]) edges.append([j, i]) edge_index = torch.tensor(edges, dtype=torch.long, device=self.device).t().contiguous() edge_attr = torch.randn(edge_index.shape[1], 5, device=self.device) # Create reference graph data object self.reference_graph = Data(x=x, edge_index=edge_index, edge_attr=edge_attr) logger.info(f"Created reference graph: {num_nodes} nodes, {edge_index.shape[1]} edges") def predict( self, design_params: Dict[str, float], return_fields: bool = False ) -> Dict[str, Any]: """ Predict all FEA objectives using parametric neural network. Args: design_params: Design parameter values (e.g. support_angle, tip_thickness, etc.) return_fields: If True, return complete displacement field (not supported for synthetic graphs) Returns: dict with: - mass: Predicted mass (g) - frequency: Predicted fundamental frequency (Hz) - max_displacement: Maximum displacement magnitude (mm) - max_stress: Maximum von Mises stress (MPa) - inference_time_ms: Prediction time """ start_time = time.time() # Build design parameter tensor param_values = [design_params.get(name, 0.0) for name in self.design_var_names] design_tensor = torch.tensor(param_values, dtype=torch.float32, device=self.device) # Normalize design params design_tensor_norm = (design_tensor - self.design_mean) / self.design_std # Add batch dimension for design params design_batch = design_tensor_norm.unsqueeze(0) # [1, n_design_vars] # Predict using reference graph with torch.no_grad(): predictions = self.model(self.reference_graph, design_batch) # Extract scalar predictions and denormalize # Model outputs normalized values, so we need to convert back to original scale mass_norm = predictions['mass'].item() freq_norm = predictions['frequency'].item() disp_norm = predictions['max_displacement'].item() stress_norm = predictions['max_stress'].item() # Denormalize to original scale mass = mass_norm * self.mass_std + self.mass_mean frequency = freq_norm * self.freq_std + self.freq_mean max_displacement = disp_norm * self.max_disp_std + self.max_disp_mean max_stress = stress_norm * self.max_stress_std + self.max_stress_mean inference_time = (time.time() - start_time) * 1000 results = { 'mass': mass, 'frequency': frequency, 'max_displacement': max_displacement, 'max_stress': max_stress, 'inference_time_ms': inference_time } # Update stats self.stats['predictions'] += 1 self.stats['total_time_ms'] += inference_time return results def get_statistics(self) -> Dict[str, Any]: """Get prediction statistics.""" avg_time = self.stats['total_time_ms'] / self.stats['predictions'] \ if self.stats['predictions'] > 0 else 0 return { 'total_predictions': self.stats['predictions'], 'total_time_ms': self.stats['total_time_ms'], 'average_time_ms': avg_time, 'model_path': str(self.model_path), 'best_val_loss': self.best_val_loss, 'device': str(self.device), 'design_var_names': self.design_var_names, 'n_design_vars': self.n_design_vars } class HybridOptimizer: """ Intelligent optimizer that combines FEA and neural surrogates. Phases: 1. Exploration: Use FEA to explore design space 2. Training: Train neural network on FEA data 3. Exploitation: Use NN for fast optimization 4. Validation: Periodically validate with FEA """ def __init__(self, config: Dict[str, Any]): """ Initialize hybrid optimizer. Args: config: Configuration dictionary """ self.config = config self.phase = 'exploration' self.fea_samples = [] self.nn_surrogate = None self.trial_count = 0 # Phase transition parameters self.min_fea_samples = config.get('min_fea_samples', 20) self.validation_frequency = config.get('validation_frequency', 10) self.retrain_frequency = config.get('retrain_frequency', 50) self.confidence_threshold = config.get('confidence_threshold', 0.95) # Training data export directory self.training_data_dir = Path(config.get('training_data_dir', 'hybrid_training_data')) self.training_data_dir.mkdir(parents=True, exist_ok=True) logger.info("Hybrid optimizer initialized") def should_use_nn(self, trial_number: int) -> Tuple[bool, str]: """ Decide whether to use NN for this trial. Args: trial_number: Current trial number Returns: Tuple of (use_nn, reason) """ self.trial_count = trial_number if self.phase == 'exploration': # Initial FEA exploration if trial_number < self.min_fea_samples: return False, f"Exploration phase ({trial_number}/{self.min_fea_samples})" else: # Transition to training self.phase = 'training' self._train_surrogate() self.phase = 'exploitation' return True, "Switched to neural surrogate" elif self.phase == 'exploitation': # Check if validation needed if trial_number % self.validation_frequency == 0: return False, f"Periodic FEA validation (every {self.validation_frequency} trials)" # Check if retraining needed if trial_number % self.retrain_frequency == 0: self._retrain_surrogate() return True, "Using neural surrogate" return False, f"Unknown phase: {self.phase}" def _train_surrogate(self): """Train surrogate model on accumulated FEA data.""" logger.info(f"Training surrogate on {len(self.fea_samples)} FEA samples") # In practice, this would: # 1. Parse all FEA data using neural_field_parser # 2. Train AtomizerField model # 3. Load trained model # For now, try to load pre-trained model if available model_path = self.config.get('pretrained_model_path') if model_path and Path(model_path).exists(): self.nn_surrogate = NeuralSurrogate( model_path=Path(model_path), confidence_threshold=self.confidence_threshold ) logger.info(f"Loaded pre-trained model from {model_path}") else: logger.warning("No pre-trained model available, continuing with FEA") self.phase = 'exploration' def _retrain_surrogate(self): """Retrain surrogate with additional data.""" logger.info(f"Retraining surrogate with {len(self.fea_samples)} total samples") # Trigger retraining pipeline # This would integrate with AtomizerField training def add_fea_sample(self, design: Dict[str, float], results: Dict[str, float]): """ Add FEA result to training data. Args: design: Design variables results: FEA results """ self.fea_samples.append({ 'trial': self.trial_count, 'design': design, 'results': results, 'timestamp': time.time() }) def get_phase_info(self) -> Dict[str, Any]: """Get current phase information.""" return { 'phase': self.phase, 'trial_count': self.trial_count, 'fea_samples': len(self.fea_samples), 'has_surrogate': self.nn_surrogate is not None, 'min_fea_samples': self.min_fea_samples, 'validation_frequency': self.validation_frequency } def create_parametric_surrogate_for_study( model_path: str = None, training_data_dir: str = None, project_root: Path = None ) -> Optional[ParametricSurrogate]: """ Factory function to create parametric neural surrogate for UAV arm study. This is the recommended surrogate type - predicts all objectives (mass, freq, etc.) Args: model_path: Path to parametric model checkpoint (auto-detect if None) training_data_dir: Path to training data (auto-detect if None) project_root: Project root directory for auto-detection Returns: ParametricSurrogate instance or None if not available """ if not TORCH_AVAILABLE or not PARAMETRIC_MODEL_AVAILABLE: logger.warning("Parametric surrogate not available: PyTorch or ParametricPredictor missing") return None # Auto-detect project root if project_root is None: project_root = Path(__file__).parent.parent # Auto-detect parametric model path if model_path is None: default_model = project_root / "atomizer-field" / "runs" / "parametric_uav_arm_v2" / "checkpoint_best.pt" if not default_model.exists(): # Try older path default_model = project_root / "atomizer-field" / "runs" / "parametric_uav_arm" / "checkpoint_best.pt" if default_model.exists(): model_path = str(default_model) else: logger.warning(f"No trained parametric model found") return None else: model_path = str(model_path) # Auto-detect training data if training_data_dir is None: default_data = project_root / "atomizer_field_training_data" / "uav_arm_train" if default_data.exists(): training_data_dir = str(default_data) else: logger.warning(f"No training data found at {default_data}") return None else: training_data_dir = str(training_data_dir) try: return ParametricSurrogate( model_path=Path(model_path), training_data_dir=Path(training_data_dir) ) except Exception as e: logger.error(f"Failed to create parametric surrogate: {e}") import traceback traceback.print_exc() return None def create_surrogate_for_study( model_path: str = None, training_data_dir: str = None, project_root: Path = None, study_name: str = None ) -> Optional[ParametricSurrogate]: """ Factory function to create neural surrogate for any study. Automatically detects whether to use ParametricSurrogate or NeuralSurrogate based on available models. Args: model_path: Path to model checkpoint (auto-detect if None) training_data_dir: Path to training data (optional, no longer required) project_root: Project root directory for auto-detection study_name: Name of the study (for auto-detection) Returns: ParametricSurrogate or NeuralSurrogate instance, or None if not available """ if not TORCH_AVAILABLE: logger.warning("Neural surrogate not available: PyTorch missing") return None # Auto-detect project root if project_root is None: project_root = Path(__file__).parent.parent # Try ParametricSurrogate first (more capable) if PARAMETRIC_MODEL_AVAILABLE: # Search order for parametric models model_search_paths = [] if study_name: # Study-specific paths model_search_paths.append(project_root / "atomizer-field" / "runs" / study_name / "checkpoint_best.pt") # Common model names to try model_search_paths.extend([ project_root / "atomizer-field" / "runs" / "bracket_model" / "checkpoint_best.pt", project_root / "atomizer-field" / "runs" / "bracket_stiffness_optimization_atomizerfield" / "checkpoint_best.pt", project_root / "atomizer-field" / "runs" / "parametric_uav_arm_v2" / "checkpoint_best.pt", project_root / "atomizer-field" / "runs" / "parametric_uav_arm" / "checkpoint_best.pt", project_root / "atomizer-field" / "runs" / "uav_arm_model" / "checkpoint_best.pt", ]) # Use explicit path if provided if model_path is not None: model_search_paths = [Path(model_path)] # Find first existing model found_model = None for mp in model_search_paths: if mp.exists(): found_model = mp logger.info(f"Found model at: {found_model}") break if found_model: try: # ParametricSurrogate no longer requires training_data_dir # It creates a synthetic reference graph like during training return ParametricSurrogate( model_path=found_model, training_data_dir=None # Not required anymore ) except Exception as e: logger.warning(f"Failed to create ParametricSurrogate: {e}") import traceback traceback.print_exc() # Fall through to try NeuralSurrogate # Fall back to NeuralSurrogate if ParametricSurrogate not available if ATOMIZER_FIELD_AVAILABLE: if model_path is None: default_model = project_root / "atomizer-field" / "runs" / "uav_arm_model" / "checkpoint_best.pt" if default_model.exists(): model_path = str(default_model) else: logger.warning(f"No trained model found") return None else: model_path = str(model_path) if training_data_dir is None: default_data = project_root / "atomizer_field_training_data" / "uav_arm_train" if default_data.exists(): training_data_dir = str(default_data) else: logger.warning(f"No training data found (required for NeuralSurrogate)") return None else: training_data_dir = str(training_data_dir) try: return NeuralSurrogate( model_path=Path(model_path), training_data_dir=Path(training_data_dir) ) except Exception as e: logger.error(f"Failed to create neural surrogate: {e}") return None logger.warning("No suitable neural model modules available") return None def create_surrogate_from_config(config: Dict[str, Any]) -> Optional[NeuralSurrogate]: """ Factory function to create neural surrogate from workflow configuration. Args: config: Workflow configuration dictionary Returns: NeuralSurrogate instance if enabled, None otherwise """ if not config.get('neural_surrogate', {}).get('enabled', False): logger.info("Neural surrogate is disabled") return None surrogate_config = config['neural_surrogate'] model_path = surrogate_config.get('model_path') training_data_dir = surrogate_config.get('training_data_dir') if not model_path: logger.error("Neural surrogate enabled but model_path not specified") return None if not training_data_dir: logger.error("Neural surrogate enabled but training_data_dir not specified") return None try: surrogate = NeuralSurrogate( model_path=Path(model_path), training_data_dir=Path(training_data_dir), device=surrogate_config.get('device', 'auto') ) logger.info("Neural surrogate created successfully") return surrogate except Exception as e: logger.error(f"Failed to create neural surrogate: {e}") return None def create_hybrid_optimizer_from_config(config: Dict[str, Any]) -> Optional[HybridOptimizer]: """ Factory function to create hybrid optimizer from configuration. Args: config: Workflow configuration dictionary Returns: HybridOptimizer instance if enabled, None otherwise """ if not config.get('hybrid_optimization', {}).get('enabled', False): logger.info("Hybrid optimization is disabled") return None hybrid_config = config.get('hybrid_optimization', {}) try: optimizer = HybridOptimizer(hybrid_config) logger.info("Hybrid optimizer created successfully") return optimizer except Exception as e: logger.error(f"Failed to create hybrid optimizer: {e}") return None