516 lines
19 KiB
Python
516 lines
19 KiB
Python
|
"""
|
||
|
|
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
|
||
|
|
)
|