""" optimization_interface.py Bridge between AtomizerField neural network and Atomizer optimization platform AtomizerField Optimization Interface v2.1 Enables gradient-based optimization with neural field predictions. Key Features: - Drop-in replacement for FEA evaluation (1000× faster) - Gradient computation for sensitivity analysis - Field-aware optimization (knows WHERE stress occurs) - Uncertainty quantification (knows when to trust predictions) - Automatic FEA fallback for high-uncertainty cases """ import torch import torch.nn.functional as F import numpy as np from pathlib import Path import json import time from neural_models.field_predictor import AtomizerFieldModel from neural_models.data_loader import FEAMeshDataset class NeuralFieldOptimizer: """ Optimization interface for AtomizerField This class provides a simple API for optimization: - evaluate(parameters) → objectives (max_stress, max_disp, etc.) - get_sensitivities(parameters) → gradients for optimization - get_fields(parameters) → complete stress/displacement fields Usage: optimizer = NeuralFieldOptimizer('checkpoint_best.pt') results = optimizer.evaluate(parameters) print(f"Max stress: {results['max_stress']:.2f} MPa") """ def __init__( self, model_path, uncertainty_threshold=0.1, enable_gradients=True, device=None ): """ Initialize optimizer Args: model_path (str): Path to trained model checkpoint uncertainty_threshold (float): Uncertainty above which to recommend FEA enable_gradients (bool): Enable gradient computation device (str): Device to run on ('cuda' or 'cpu') """ if device is None: self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') else: self.device = torch.device(device) print(f"\nAtomizerField Optimization Interface v2.1") print(f"Device: {self.device}") # Load model print(f"Loading model from {model_path}...") checkpoint = torch.load(model_path, map_location=self.device) # Create model 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.config = checkpoint['config'] self.uncertainty_threshold = uncertainty_threshold self.enable_gradients = enable_gradients # Model info self.model_info = { 'version': checkpoint.get('epoch', 'unknown'), 'best_val_loss': checkpoint.get('best_val_loss', 'unknown'), 'training_config': checkpoint['config'] } print(f"Model loaded successfully!") print(f" Epoch: {checkpoint.get('epoch', 'N/A')}") print(f" Validation loss: {checkpoint.get('best_val_loss', 'N/A')}") # Statistics for tracking self.eval_count = 0 self.total_time = 0.0 def evaluate(self, graph_data, return_fields=False): """ Evaluate design using neural network (drop-in FEA replacement) Args: graph_data: PyTorch Geometric Data object with mesh graph return_fields (bool): Return complete fields or just objectives Returns: dict: Optimization objectives and optionally complete fields - max_stress: Maximum von Mises stress (MPa) - max_displacement: Maximum displacement (mm) - mass: Total mass (kg) if available - fields: Complete stress/displacement fields (if return_fields=True) - inference_time_ms: Prediction time - uncertainty: Prediction uncertainty (if ensemble enabled) """ start_time = time.time() # Move to device graph_data = graph_data.to(self.device) # Predict with torch.set_grad_enabled(self.enable_gradients): predictions = self.model(graph_data, return_stress=True) inference_time = (time.time() - start_time) * 1000 # ms # Extract objectives max_displacement = torch.max( torch.norm(predictions['displacement'][:, :3], dim=1) ).item() max_stress = torch.max(predictions['von_mises']).item() results = { 'max_stress': max_stress, 'max_displacement': max_displacement, 'inference_time_ms': inference_time, 'evaluation_count': self.eval_count } # Add complete fields if requested if return_fields: results['fields'] = { 'displacement': predictions['displacement'].cpu().detach().numpy(), 'stress': predictions['stress'].cpu().detach().numpy(), 'von_mises': predictions['von_mises'].cpu().detach().numpy() } # Update statistics self.eval_count += 1 self.total_time += inference_time return results def get_sensitivities(self, graph_data, objective='max_stress'): """ Compute gradients for gradient-based optimization This enables MUCH faster optimization than finite differences! Args: graph_data: PyTorch Geometric Data with requires_grad=True objective (str): Which objective to differentiate ('max_stress' or 'max_displacement') Returns: dict: Gradients with respect to input features - node_gradients: ∂objective/∂node_features - edge_gradients: ∂objective/∂edge_features """ if not self.enable_gradients: raise RuntimeError("Gradients not enabled. Set enable_gradients=True") # Enable gradients graph_data = graph_data.to(self.device) graph_data.x.requires_grad_(True) if graph_data.edge_attr is not None: graph_data.edge_attr.requires_grad_(True) # Forward pass predictions = self.model(graph_data, return_stress=True) # Compute objective if objective == 'max_stress': obj = torch.max(predictions['von_mises']) elif objective == 'max_displacement': disp_mag = torch.norm(predictions['displacement'][:, :3], dim=1) obj = torch.max(disp_mag) else: raise ValueError(f"Unknown objective: {objective}") # Backward pass obj.backward() # Extract gradients gradients = { 'node_gradients': graph_data.x.grad.cpu().numpy(), 'objective_value': obj.item() } if graph_data.edge_attr is not None and graph_data.edge_attr.grad is not None: gradients['edge_gradients'] = graph_data.edge_attr.grad.cpu().numpy() return gradients def batch_evaluate(self, graph_data_list, return_fields=False): """ Evaluate multiple designs in batch (even faster!) Args: graph_data_list (list): List of graph data objects return_fields (bool): Return complete fields Returns: list: List of evaluation results """ results = [] for graph_data in graph_data_list: result = self.evaluate(graph_data, return_fields=return_fields) results.append(result) return results def needs_fea_validation(self, uncertainty): """ Determine if FEA validation is recommended Args: uncertainty (float): Prediction uncertainty Returns: bool: True if FEA is recommended """ return uncertainty > self.uncertainty_threshold def compare_with_fea(self, graph_data, fea_results): """ Compare neural predictions with FEA ground truth Args: graph_data: Mesh graph fea_results (dict): FEA results with 'max_stress', 'max_displacement' Returns: dict: Comparison metrics """ # Neural prediction pred = self.evaluate(graph_data) # Compute errors stress_error = abs(pred['max_stress'] - fea_results['max_stress']) stress_rel_error = stress_error / (fea_results['max_stress'] + 1e-8) disp_error = abs(pred['max_displacement'] - fea_results['max_displacement']) disp_rel_error = disp_error / (fea_results['max_displacement'] + 1e-8) comparison = { 'neural_prediction': pred, 'fea_results': fea_results, 'errors': { 'stress_error_abs': stress_error, 'stress_error_rel': stress_rel_error, 'displacement_error_abs': disp_error, 'displacement_error_rel': disp_rel_error }, 'within_tolerance': stress_rel_error < 0.1 and disp_rel_error < 0.1 } return comparison def get_statistics(self): """ Get optimizer usage statistics Returns: dict: Statistics about predictions """ avg_time = self.total_time / self.eval_count if self.eval_count > 0 else 0 return { 'total_evaluations': self.eval_count, 'total_time_ms': self.total_time, 'average_time_ms': avg_time, 'model_info': self.model_info } def reset_statistics(self): """Reset usage statistics""" self.eval_count = 0 self.total_time = 0.0 class ParametricOptimizer: """ Optimizer for parametric designs This wraps NeuralFieldOptimizer and adds parameter → mesh conversion. Enables direct optimization over design parameters (thickness, radius, etc.) """ def __init__( self, model_path, parameter_names, parameter_bounds, mesh_generator_fn ): """ Initialize parametric optimizer Args: model_path (str): Path to trained model parameter_names (list): Names of design parameters parameter_bounds (dict): Bounds for each parameter mesh_generator_fn: Function that converts parameters → graph_data """ self.neural_optimizer = NeuralFieldOptimizer(model_path) self.parameter_names = parameter_names self.parameter_bounds = parameter_bounds self.mesh_generator = mesh_generator_fn print(f"\nParametric Optimizer initialized") print(f"Design parameters: {parameter_names}") def evaluate_parameters(self, parameters): """ Evaluate design from parameters Args: parameters (dict): Design parameters Returns: dict: Objectives (max_stress, max_displacement, etc.) """ # Generate mesh from parameters graph_data = self.mesh_generator(parameters) # Evaluate results = self.neural_optimizer.evaluate(graph_data) # Add parameters to results results['parameters'] = parameters return results def optimize( self, initial_parameters, objectives, constraints, method='gradient', max_iterations=100 ): """ Run optimization Args: initial_parameters (dict): Starting point objectives (list): Objectives to minimize/maximize constraints (list): Constraint functions method (str): Optimization method ('gradient' or 'genetic') max_iterations (int): Maximum iterations Returns: dict: Optimal parameters and results """ # This would integrate with scipy.optimize or genetic algorithms # Placeholder for now print(f"\nStarting optimization with {method} method...") print(f"Initial parameters: {initial_parameters}") print(f"Objectives: {objectives}") print(f"Max iterations: {max_iterations}") # TODO: Implement optimization loop # For gradient-based: # 1. Evaluate at current parameters # 2. Compute sensitivities # 3. Update parameters using gradients # 4. Repeat until convergence raise NotImplementedError("Full optimization loop coming in next update!") def create_optimizer(model_path, config=None): """ Factory function to create optimizer Args: model_path (str): Path to trained model config (dict): Optimizer configuration Returns: NeuralFieldOptimizer instance """ if config is None: config = {} return NeuralFieldOptimizer(model_path, **config) if __name__ == "__main__": # Example usage print("AtomizerField Optimization Interface") print("=" * 60) print("\nThis module provides fast optimization with neural field predictions.") print("\nExample usage:") print(""" # Create optimizer optimizer = NeuralFieldOptimizer('checkpoint_best.pt') # Evaluate design results = optimizer.evaluate(graph_data) print(f"Max stress: {results['max_stress']:.2f} MPa") print(f"Inference time: {results['inference_time_ms']:.1f} ms") # Get sensitivities for gradient-based optimization gradients = optimizer.get_sensitivities(graph_data, objective='max_stress') # Batch evaluation (test 1000 designs in seconds!) all_results = optimizer.batch_evaluate(design_variants) """) print("\nOptimization interface ready!")