""" uncertainty.py Uncertainty quantification for neural field predictions AtomizerField Uncertainty Quantification v2.1 Know when to trust predictions and when to run FEA! Key Features: - Ensemble-based uncertainty estimation - Confidence intervals for predictions - Automatic FEA recommendation - Online calibration """ import torch import torch.nn as nn import numpy as np from copy import deepcopy from .field_predictor import AtomizerFieldModel class UncertainFieldPredictor(nn.Module): """ Ensemble of models for uncertainty quantification Uses multiple models trained with different initializations to estimate prediction uncertainty. When uncertainty is high → Recommend FEA validation When uncertainty is low → Trust neural prediction """ def __init__(self, base_model_config, n_ensemble=5): """ Initialize ensemble Args: base_model_config (dict): Configuration for base model n_ensemble (int): Number of models in ensemble """ super().__init__() print(f"\nCreating ensemble with {n_ensemble} models...") # Create ensemble of models self.models = nn.ModuleList([ AtomizerFieldModel(**base_model_config) for _ in range(n_ensemble) ]) self.n_ensemble = n_ensemble # Initialize each model differently for i, model in enumerate(self.models): self._init_weights(model, seed=i) print(f"Ensemble created with {n_ensemble} models") def _init_weights(self, model, seed): """Initialize model weights with different seed""" torch.manual_seed(seed) def init_fn(m): if isinstance(m, nn.Linear): nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu') if m.bias is not None: nn.init.constant_(m.bias, 0) model.apply(init_fn) def forward(self, data, return_uncertainty=True, return_all_predictions=False): """ Forward pass through ensemble Args: data: Input graph data return_uncertainty (bool): Return uncertainty estimates return_all_predictions (bool): Return all individual predictions Returns: dict: Predictions with uncertainty - displacement: Mean prediction - stress: Mean prediction - von_mises: Mean prediction - displacement_std: Standard deviation (if return_uncertainty) - stress_std: Standard deviation (if return_uncertainty) - von_mises_std: Standard deviation (if return_uncertainty) - all_predictions: List of all predictions (if return_all_predictions) """ # Get predictions from all models all_predictions = [] for model in self.models: with torch.no_grad(): pred = model(data, return_stress=True) all_predictions.append(pred) # Stack predictions displacement_stack = torch.stack([p['displacement'] for p in all_predictions]) stress_stack = torch.stack([p['stress'] for p in all_predictions]) von_mises_stack = torch.stack([p['von_mises'] for p in all_predictions]) # Compute mean predictions results = { 'displacement': displacement_stack.mean(dim=0), 'stress': stress_stack.mean(dim=0), 'von_mises': von_mises_stack.mean(dim=0) } # Compute uncertainty (standard deviation across ensemble) if return_uncertainty: results['displacement_std'] = displacement_stack.std(dim=0) results['stress_std'] = stress_stack.std(dim=0) results['von_mises_std'] = von_mises_stack.std(dim=0) # Overall uncertainty metrics results['max_displacement_uncertainty'] = results['displacement_std'].max().item() results['max_stress_uncertainty'] = results['von_mises_std'].max().item() # Uncertainty as percentage of prediction results['displacement_rel_uncertainty'] = ( results['displacement_std'] / (torch.abs(results['displacement']) + 1e-8) ).mean().item() results['stress_rel_uncertainty'] = ( results['von_mises_std'] / (results['von_mises'] + 1e-8) ).mean().item() # Return all predictions if requested if return_all_predictions: results['all_predictions'] = all_predictions return results def needs_fea_validation(self, predictions, threshold=0.1): """ Determine if FEA validation is recommended Args: predictions (dict): Output from forward() with uncertainty threshold (float): Relative uncertainty threshold Returns: dict: Recommendation and reasons """ reasons = [] # Check displacement uncertainty if predictions['displacement_rel_uncertainty'] > threshold: reasons.append( f"High displacement uncertainty: " f"{predictions['displacement_rel_uncertainty']*100:.1f}% > {threshold*100:.1f}%" ) # Check stress uncertainty if predictions['stress_rel_uncertainty'] > threshold: reasons.append( f"High stress uncertainty: " f"{predictions['stress_rel_uncertainty']*100:.1f}% > {threshold*100:.1f}%" ) recommend_fea = len(reasons) > 0 return { 'recommend_fea': recommend_fea, 'reasons': reasons, 'displacement_uncertainty': predictions['displacement_rel_uncertainty'], 'stress_uncertainty': predictions['stress_rel_uncertainty'] } def get_confidence_intervals(self, predictions, confidence=0.95): """ Compute confidence intervals for predictions Args: predictions (dict): Output from forward() with uncertainty confidence (float): Confidence level (0.95 = 95% confidence) Returns: dict: Confidence intervals """ # For normal distribution, 95% CI is ±1.96 std # For 90% CI is ±1.645 std z_score = {0.90: 1.645, 0.95: 1.96, 0.99: 2.576}.get(confidence, 1.96) intervals = {} # Displacement intervals intervals['displacement_lower'] = predictions['displacement'] - z_score * predictions['displacement_std'] intervals['displacement_upper'] = predictions['displacement'] + z_score * predictions['displacement_std'] # Stress intervals intervals['von_mises_lower'] = predictions['von_mises'] - z_score * predictions['von_mises_std'] intervals['von_mises_upper'] = predictions['von_mises'] + z_score * predictions['von_mises_std'] # Max values with confidence intervals max_vm = predictions['von_mises'].max() max_vm_std = predictions['von_mises_std'].max() intervals['max_stress_estimate'] = max_vm.item() intervals['max_stress_lower'] = (max_vm - z_score * max_vm_std).item() intervals['max_stress_upper'] = (max_vm + z_score * max_vm_std).item() return intervals class OnlineLearner: """ Online learning from FEA runs during optimization As optimization progresses and you run FEA for validation, this module can quickly update the model to improve predictions. This creates a virtuous cycle: 1. Use neural network for fast exploration 2. Run FEA on promising designs 3. Update neural network with new data 4. Neural network gets better → need less FEA """ def __init__(self, model, learning_rate=0.0001): """ Initialize online learner Args: model: Neural network model learning_rate (float): Learning rate for updates """ self.model = model self.optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) self.replay_buffer = [] self.update_count = 0 print(f"\nOnline learner initialized") print(f"Learning rate: {learning_rate}") def add_fea_result(self, graph_data, fea_results): """ Add new FEA result to replay buffer Args: graph_data: Mesh graph fea_results (dict): FEA results (displacement, stress) """ self.replay_buffer.append({ 'graph_data': graph_data, 'fea_results': fea_results }) print(f"Added FEA result to buffer (total: {len(self.replay_buffer)})") def quick_update(self, steps=10): """ Quick fine-tuning on recent FEA results Args: steps (int): Number of gradient steps """ if len(self.replay_buffer) == 0: print("No data in replay buffer") return print(f"\nQuick update: {steps} steps on {len(self.replay_buffer)} samples") self.model.train() for step in range(steps): total_loss = 0.0 # Train on all samples in buffer for sample in self.replay_buffer: graph_data = sample['graph_data'] fea_results = sample['fea_results'] # Forward pass predictions = self.model(graph_data, return_stress=True) # Compute loss disp_loss = nn.functional.mse_loss( predictions['displacement'], fea_results['displacement'] ) if 'stress' in fea_results: stress_loss = nn.functional.mse_loss( predictions['stress'], fea_results['stress'] ) loss = disp_loss + stress_loss else: loss = disp_loss # Backward pass self.optimizer.zero_grad() loss.backward() self.optimizer.step() total_loss += loss.item() if step % 5 == 0: avg_loss = total_loss / len(self.replay_buffer) print(f" Step {step}/{steps}: Loss = {avg_loss:.6f}") self.model.eval() self.update_count += 1 print(f"Update complete (total updates: {self.update_count})") def clear_buffer(self): """Clear replay buffer""" self.replay_buffer = [] print("Replay buffer cleared") def create_uncertain_predictor(model_config, n_ensemble=5): """ Factory function to create uncertain predictor Args: model_config (dict): Model configuration n_ensemble (int): Ensemble size Returns: UncertainFieldPredictor instance """ return UncertainFieldPredictor(model_config, n_ensemble) if __name__ == "__main__": # Test uncertainty quantification print("Testing Uncertainty Quantification...\n") # Create ensemble model_config = { 'node_feature_dim': 12, 'edge_feature_dim': 5, 'hidden_dim': 64, 'num_layers': 4, 'dropout': 0.1 } ensemble = UncertainFieldPredictor(model_config, n_ensemble=3) print(f"\nEnsemble created with {ensemble.n_ensemble} models") print("Uncertainty quantification ready!") print("\nUsage:") print(""" # Get predictions with uncertainty predictions = ensemble(graph_data, return_uncertainty=True) # Check if FEA validation needed recommendation = ensemble.needs_fea_validation(predictions, threshold=0.1) if recommendation['recommend_fea']: print("Recommendation: Run FEA for validation") for reason in recommendation['reasons']: print(f" - {reason}") else: print("Prediction confident - no FEA needed!") """)