feat: Add AtomizerField training data export and intelligent model discovery
Major additions: - Training data export system for AtomizerField neural network training - Bracket stiffness optimization study with 50+ training samples - Intelligent NX model discovery (auto-detect solutions, expressions, mesh) - Result extractors module for displacement, stress, frequency, mass - User-generated NX journals for advanced workflows - Archive structure for legacy scripts and test outputs - Protocol documentation and dashboard launcher 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
163
archive/scripts/run_nn_optimization.py
Normal file
163
archive/scripts/run_nn_optimization.py
Normal file
@@ -0,0 +1,163 @@
|
||||
"""
|
||||
Neural Network Only Optimization
|
||||
|
||||
This script runs multi-objective optimization using ONLY the neural network
|
||||
surrogate (no FEA). This demonstrates the speed improvement from NN predictions.
|
||||
|
||||
Objectives:
|
||||
- Minimize mass
|
||||
- Maximize frequency (minimize -frequency)
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import time
|
||||
import json
|
||||
import optuna
|
||||
from optuna.samplers import NSGAIISampler
|
||||
import numpy as np
|
||||
|
||||
# Add project paths
|
||||
project_root = Path(__file__).parent
|
||||
sys.path.insert(0, str(project_root))
|
||||
sys.path.insert(0, str(project_root / 'atomizer-field'))
|
||||
|
||||
from optimization_engine.simple_mlp_surrogate import SimpleSurrogate
|
||||
|
||||
|
||||
def main():
|
||||
print("="*60)
|
||||
print("Neural Network Only Optimization (Simple MLP)")
|
||||
print("="*60)
|
||||
|
||||
# Load surrogate
|
||||
print("\n[1] Loading neural surrogate...")
|
||||
model_path = project_root / "simple_mlp_surrogate.pt"
|
||||
|
||||
if not model_path.exists():
|
||||
print(f"ERROR: Model not found at {model_path}")
|
||||
print("Run 'python optimization_engine/simple_mlp_surrogate.py' first to train")
|
||||
return
|
||||
|
||||
surrogate = SimpleSurrogate.load(model_path)
|
||||
|
||||
if not surrogate:
|
||||
print("ERROR: Could not load neural surrogate")
|
||||
return
|
||||
|
||||
print(f" Design variables: {surrogate.design_var_names}")
|
||||
|
||||
# Define bounds (from UAV arm study)
|
||||
bounds = {
|
||||
'beam_half_core_thickness': (1.0, 5.0),
|
||||
'beam_face_thickness': (0.5, 3.0),
|
||||
'holes_diameter': (0.5, 5.0),
|
||||
'hole_count': (0.0, 6.0)
|
||||
}
|
||||
|
||||
print(f" Bounds: {bounds}")
|
||||
|
||||
# Create Optuna study
|
||||
print("\n[2] Creating Optuna study...")
|
||||
storage_path = project_root / "nn_only_optimization_study.db"
|
||||
|
||||
# Remove old study if exists
|
||||
if storage_path.exists():
|
||||
storage_path.unlink()
|
||||
|
||||
storage = optuna.storages.RDBStorage(f"sqlite:///{storage_path}")
|
||||
|
||||
study = optuna.create_study(
|
||||
study_name="nn_only_optimization",
|
||||
storage=storage,
|
||||
directions=["minimize", "minimize"], # mass, -frequency (minimize both)
|
||||
sampler=NSGAIISampler()
|
||||
)
|
||||
|
||||
# Track stats
|
||||
start_time = time.time()
|
||||
trial_times = []
|
||||
|
||||
def objective(trial: optuna.Trial):
|
||||
trial_start = time.time()
|
||||
|
||||
# Suggest parameters
|
||||
params = {}
|
||||
for name, (low, high) in bounds.items():
|
||||
if name == 'hole_count':
|
||||
params[name] = trial.suggest_int(name, int(low), int(high))
|
||||
else:
|
||||
params[name] = trial.suggest_float(name, low, high)
|
||||
|
||||
# Predict with NN
|
||||
results = surrogate.predict(params)
|
||||
|
||||
mass = results['mass']
|
||||
frequency = results['frequency']
|
||||
|
||||
trial_time = (time.time() - trial_start) * 1000
|
||||
trial_times.append(trial_time)
|
||||
|
||||
# Log progress every 100 trials
|
||||
if trial.number % 100 == 0:
|
||||
print(f" Trial {trial.number}: mass={mass:.1f}g, freq={frequency:.2f}Hz, time={trial_time:.1f}ms")
|
||||
|
||||
# Return objectives: minimize mass, minimize -frequency (= maximize frequency)
|
||||
return mass, -frequency
|
||||
|
||||
# Run optimization
|
||||
n_trials = 1000 # Much faster with NN!
|
||||
print(f"\n[3] Running {n_trials} trials...")
|
||||
|
||||
study.optimize(objective, n_trials=n_trials, show_progress_bar=True)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
|
||||
# Results
|
||||
print("\n" + "="*60)
|
||||
print("RESULTS")
|
||||
print("="*60)
|
||||
|
||||
print(f"\nTotal time: {total_time:.1f}s for {n_trials} trials")
|
||||
print(f"Average time per trial: {np.mean(trial_times):.1f}ms")
|
||||
print(f"Trials per second: {n_trials/total_time:.1f}")
|
||||
|
||||
# Get Pareto front
|
||||
pareto_front = study.best_trials
|
||||
print(f"\nPareto front size: {len(pareto_front)} designs")
|
||||
|
||||
print("\nTop 5 Pareto-optimal designs:")
|
||||
for i, trial in enumerate(pareto_front[:5]):
|
||||
mass = trial.values[0]
|
||||
freq = -trial.values[1] # Convert back to positive
|
||||
print(f" {i+1}. Mass={mass:.1f}g, Freq={freq:.2f}Hz")
|
||||
print(f" Params: {trial.params}")
|
||||
|
||||
# Save results
|
||||
results_file = project_root / "nn_optimization_results.json"
|
||||
results = {
|
||||
'n_trials': n_trials,
|
||||
'total_time_s': total_time,
|
||||
'avg_trial_time_ms': np.mean(trial_times),
|
||||
'trials_per_second': n_trials/total_time,
|
||||
'pareto_front_size': len(pareto_front),
|
||||
'pareto_designs': [
|
||||
{
|
||||
'mass': t.values[0],
|
||||
'frequency': -t.values[1],
|
||||
'params': t.params
|
||||
}
|
||||
for t in pareto_front
|
||||
]
|
||||
}
|
||||
|
||||
with open(results_file, 'w') as f:
|
||||
json.dump(results, f, indent=2)
|
||||
|
||||
print(f"\nResults saved to: {results_file}")
|
||||
print(f"Study database: {storage_path}")
|
||||
print("\nView in Optuna dashboard:")
|
||||
print(f" optuna-dashboard sqlite:///{storage_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user