Files
Atomizer/studies/m1_mirror_adaptive_V13/run_optimization.py
Antoine 96b196de58 feat: Add Zernike GNN surrogate module and M1 mirror V12/V13 studies
This commit introduces the GNN-based surrogate for Zernike mirror optimization
and the M1 mirror study progression from V12 (GNN validation) to V13 (pure NSGA-II).

## GNN Surrogate Module (optimization_engine/gnn/)

New module for Graph Neural Network surrogate prediction of mirror deformations:

- `polar_graph.py`: PolarMirrorGraph - fixed 3000-node polar grid structure
- `zernike_gnn.py`: ZernikeGNN with design-conditioned message passing
- `differentiable_zernike.py`: GPU-accelerated Zernike fitting and objectives
- `train_zernike_gnn.py`: ZernikeGNNTrainer with multi-task loss
- `gnn_optimizer.py`: ZernikeGNNOptimizer for turbo mode (~900k trials/hour)
- `extract_displacement_field.py`: OP2 to HDF5 field extraction
- `backfill_field_data.py`: Extract fields from existing FEA trials

Key innovation: Design-conditioned convolutions that modulate message passing
based on structural design parameters, enabling accurate field prediction.

## M1 Mirror Studies

### V12: GNN Field Prediction + FEA Validation
- Zernike GNN trained on V10/V11 FEA data (238 samples)
- Turbo mode: 5000 GNN predictions → top candidates → FEA validation
- Calibration workflow for GNN-to-FEA error correction
- Scripts: run_gnn_turbo.py, validate_gnn_best.py, compute_full_calibration.py

### V13: Pure NSGA-II FEA (Ground Truth)
- Seeds 217 FEA trials from V11+V12
- Pure multi-objective NSGA-II without any surrogate
- Establishes ground-truth Pareto front for GNN accuracy evaluation
- Narrowed blank_backface_angle range to [4.0, 5.0]

## Documentation Updates

- SYS_14: Added Zernike GNN section with architecture diagrams
- CLAUDE.md: Added GNN module reference and quick start
- V13 README: Study documentation with seeding strategy

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2025-12-10 08:44:04 -05:00

568 lines
20 KiB
Python

#!/usr/bin/env python3
"""
M1 Mirror Pure NSGA-II FEA Optimization V13
=============================================
Pure multi-objective optimization with NSGA-II sampler and FEA only.
No neural surrogate - every trial is a real FEA evaluation.
Key Features:
1. NSGA-II sampler for true multi-objective Pareto optimization
2. Seeds from V11 + V12 FEA trials (~110+ prior trials)
3. No surrogate bias - ground truth only
4. 3 objectives: rel_rms_40_vs_20, rel_rms_60_vs_20, mfg_90
Usage:
python run_optimization.py --start
python run_optimization.py --start --trials 50
python run_optimization.py --start --trials 50 --resume
For 8-hour overnight run (~55 trials at 8-9 min/trial):
python run_optimization.py --start --trials 55
"""
import sys
import os
import json
import time
import argparse
import logging
import sqlite3
import shutil
import re
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass, field
from datetime import datetime
import numpy as np
# Add parent directories to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
import optuna
from optuna.samplers import NSGAIISampler
# Atomizer imports
from optimization_engine.nx_solver import NXSolver
from optimization_engine.utils import ensure_nx_running
from optimization_engine.extractors import ZernikeExtractor
# ============================================================================
# Paths
# ============================================================================
STUDY_DIR = Path(__file__).parent
SETUP_DIR = STUDY_DIR / "1_setup"
ITERATIONS_DIR = STUDY_DIR / "2_iterations"
RESULTS_DIR = STUDY_DIR / "3_results"
CONFIG_PATH = SETUP_DIR / "optimization_config.json"
# Source studies for seeding
V11_DB = STUDY_DIR.parent / "m1_mirror_adaptive_V11" / "3_results" / "study.db"
V12_DB = STUDY_DIR.parent / "m1_mirror_adaptive_V12" / "3_results" / "study.db"
# Ensure directories exist
ITERATIONS_DIR.mkdir(exist_ok=True)
RESULTS_DIR.mkdir(exist_ok=True)
# Logging
LOG_FILE = RESULTS_DIR / "optimization.log"
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s | %(levelname)-8s | %(message)s',
handlers=[
logging.StreamHandler(sys.stdout),
logging.FileHandler(LOG_FILE, mode='a')
]
)
logger = logging.getLogger(__name__)
# ============================================================================
# Objective names
# ============================================================================
OBJ_NAMES = [
'rel_filtered_rms_40_vs_20',
'rel_filtered_rms_60_vs_20',
'mfg_90_optician_workload'
]
DESIGN_VAR_NAMES = [
'lateral_inner_angle', 'lateral_outer_angle', 'lateral_outer_pivot',
'lateral_inner_pivot', 'lateral_middle_pivot', 'lateral_closeness',
'whiffle_min', 'whiffle_outer_to_vertical', 'whiffle_triangle_closeness',
'blank_backface_angle', 'inner_circular_rib_dia'
]
# ============================================================================
# Prior Data Loader
# ============================================================================
def load_fea_trials_from_db(db_path: Path, label: str) -> List[Dict]:
"""Load FEA trials from an Optuna database."""
if not db_path.exists():
logger.warning(f"{label} database not found: {db_path}")
return []
fea_data = []
conn = sqlite3.connect(str(db_path))
try:
cursor = conn.cursor()
cursor.execute('''
SELECT trial_id, number FROM trials
WHERE state = 'COMPLETE'
''')
trials = cursor.fetchall()
for trial_id, trial_num in trials:
# Get user attributes
cursor.execute('''
SELECT key, value_json FROM trial_user_attributes
WHERE trial_id = ?
''', (trial_id,))
attrs = {row[0]: json.loads(row[1]) for row in cursor.fetchall()}
# Check if FEA trial (source contains 'FEA')
source = attrs.get('source', 'FEA')
if 'FEA' not in source:
continue # Skip NN trials
# Get params
cursor.execute('''
SELECT param_name, param_value FROM trial_params
WHERE trial_id = ?
''', (trial_id,))
params = {name: float(value) for name, value in cursor.fetchall()}
if not params:
continue
# Get objectives (stored as individual attributes or in 'objectives')
objectives = {}
if 'objectives' in attrs:
objectives = attrs['objectives']
else:
# Try individual attributes
for obj_name in OBJ_NAMES:
if obj_name in attrs:
objectives[obj_name] = attrs[obj_name]
if all(k in objectives for k in OBJ_NAMES):
fea_data.append({
'trial_num': trial_num,
'params': params,
'objectives': objectives,
'source': f'{label}_{source}'
})
except Exception as e:
logger.error(f"Error loading {label} data: {e}")
finally:
conn.close()
logger.info(f"Loaded {len(fea_data)} FEA trials from {label}")
return fea_data
def load_all_prior_fea_data() -> List[Dict]:
"""Load FEA trials from V11 and V12."""
all_data = []
# V11 data
v11_data = load_fea_trials_from_db(V11_DB, "V11")
all_data.extend(v11_data)
# V12 data
v12_data = load_fea_trials_from_db(V12_DB, "V12")
all_data.extend(v12_data)
logger.info(f"Total prior FEA trials: {len(all_data)}")
return all_data
# ============================================================================
# FEA Runner
# ============================================================================
class FEARunner:
"""Runs actual FEA simulations."""
def __init__(self, config: Dict[str, Any]):
self.config = config
self.nx_solver = None
self.nx_manager = None
self.master_model_dir = SETUP_DIR / "model"
def setup(self):
"""Setup NX and solver."""
logger.info("Setting up NX session...")
study_name = self.config.get('study_name', 'm1_mirror_adaptive_V13')
try:
self.nx_manager, nx_was_started = ensure_nx_running(
session_id=study_name,
auto_start=True,
start_timeout=120
)
logger.info("NX session ready" + (" (started)" if nx_was_started else " (existing)"))
except Exception as e:
logger.error(f"Failed to setup NX: {e}")
raise
# Initialize solver
nx_settings = self.config.get('nx_settings', {})
nx_install_dir = nx_settings.get('nx_install_path', 'C:\\Program Files\\Siemens\\NX2506')
version_match = re.search(r'NX(\d+)', nx_install_dir)
nastran_version = version_match.group(1) if version_match else "2506"
self.nx_solver = NXSolver(
master_model_dir=str(self.master_model_dir),
nx_install_dir=nx_install_dir,
nastran_version=nastran_version,
timeout=nx_settings.get('simulation_timeout_s', 600),
use_iteration_folders=True,
study_name="m1_mirror_adaptive_V13"
)
def run_fea(self, params: Dict[str, float], trial_num: int) -> Optional[Dict]:
"""Run FEA and extract objectives."""
if self.nx_solver is None:
self.setup()
logger.info(f" [FEA {trial_num}] Running simulation...")
expressions = {var['expression_name']: params[var['name']]
for var in self.config['design_variables']}
iter_folder = self.nx_solver.create_iteration_folder(
iterations_base_dir=ITERATIONS_DIR,
iteration_number=trial_num,
expression_updates=expressions
)
try:
nx_settings = self.config.get('nx_settings', {})
sim_file = iter_folder / nx_settings.get('sim_file', 'ASSY_M1_assyfem1_sim1.sim')
t_start = time.time()
result = self.nx_solver.run_simulation(
sim_file=sim_file,
working_dir=iter_folder,
expression_updates=expressions,
solution_name=nx_settings.get('solution_name', 'Solution 1'),
cleanup=False
)
solve_time = time.time() - t_start
if not result['success']:
logger.error(f" [FEA {trial_num}] Solve failed: {result.get('error')}")
return None
logger.info(f" [FEA {trial_num}] Solved in {solve_time:.1f}s")
# Extract objectives
op2_path = Path(result['op2_file'])
objectives = self._extract_objectives(op2_path)
if objectives is None:
return None
logger.info(f" [FEA {trial_num}] 40-20: {objectives['rel_filtered_rms_40_vs_20']:.2f} nm")
logger.info(f" [FEA {trial_num}] 60-20: {objectives['rel_filtered_rms_60_vs_20']:.2f} nm")
logger.info(f" [FEA {trial_num}] Mfg: {objectives['mfg_90_optician_workload']:.2f} nm")
return {
'trial_num': trial_num,
'params': params,
'objectives': objectives,
'source': 'FEA',
'solve_time': solve_time
}
except Exception as e:
logger.error(f" [FEA {trial_num}] Error: {e}")
import traceback
traceback.print_exc()
return None
def _extract_objectives(self, op2_path: Path) -> Optional[Dict[str, float]]:
"""Extract objectives using ZernikeExtractor."""
try:
zernike_settings = self.config.get('zernike_settings', {})
extractor = ZernikeExtractor(
op2_path,
bdf_path=None,
displacement_unit=zernike_settings.get('displacement_unit', 'mm'),
n_modes=zernike_settings.get('n_modes', 50),
filter_orders=zernike_settings.get('filter_low_orders', 4)
)
ref = zernike_settings.get('reference_subcase', '2')
rel_40 = extractor.extract_relative("3", ref)
rel_60 = extractor.extract_relative("4", ref)
rel_90 = extractor.extract_relative("1", ref)
return {
'rel_filtered_rms_40_vs_20': rel_40['relative_filtered_rms_nm'],
'rel_filtered_rms_60_vs_20': rel_60['relative_filtered_rms_nm'],
'mfg_90_optician_workload': rel_90['relative_rms_filter_j1to3']
}
except Exception as e:
logger.error(f"Zernike extraction failed: {e}")
return None
def cleanup(self):
"""Cleanup NX session."""
if self.nx_manager:
if self.nx_manager.can_close_nx():
self.nx_manager.close_nx_if_allowed()
self.nx_manager.cleanup()
# ============================================================================
# NSGA-II Optimizer
# ============================================================================
class NSGA2Optimizer:
"""Pure FEA multi-objective optimizer with NSGA-II."""
def __init__(self, config: Dict[str, Any]):
self.config = config
self.fea_runner = FEARunner(config)
# Load prior data for seeding
self.prior_data = load_all_prior_fea_data()
# Database
self.db_path = RESULTS_DIR / "study.db"
self.storage = optuna.storages.RDBStorage(f'sqlite:///{self.db_path}')
# State
self.trial_count = 0
self.best_pareto = []
def _get_next_trial_number(self) -> int:
"""Get the next trial number based on existing iterations."""
existing = list(ITERATIONS_DIR.glob("iter*"))
if not existing:
return 1
max_num = max(int(p.name.replace("iter", "")) for p in existing)
return max_num + 1
def seed_from_prior(self, study: optuna.Study):
"""Seed the study with prior FEA trials."""
if not self.prior_data:
logger.warning("No prior data to seed from")
return
logger.info(f"Seeding study with {len(self.prior_data)} prior FEA trials...")
for i, d in enumerate(self.prior_data):
try:
# Create a trial with the prior data
distributions = {}
for var in self.config['design_variables']:
if var.get('enabled', False):
distributions[var['name']] = optuna.distributions.FloatDistribution(
var['min'], var['max']
)
# Create frozen trial
frozen_trial = optuna.trial.create_trial(
params=d['params'],
distributions=distributions,
values=[
d['objectives']['rel_filtered_rms_40_vs_20'],
d['objectives']['rel_filtered_rms_60_vs_20'],
d['objectives']['mfg_90_optician_workload']
],
user_attrs={
'source': d.get('source', 'prior_FEA'),
'rel_filtered_rms_40_vs_20': d['objectives']['rel_filtered_rms_40_vs_20'],
'rel_filtered_rms_60_vs_20': d['objectives']['rel_filtered_rms_60_vs_20'],
'mfg_90_optician_workload': d['objectives']['mfg_90_optician_workload'],
}
)
study.add_trial(frozen_trial)
except Exception as e:
logger.warning(f"Failed to seed trial {i}: {e}")
logger.info(f"Seeded {len(study.trials)} trials")
def run(self, n_trials: int = 50, resume: bool = False):
"""Run NSGA-II optimization."""
logger.info("\n" + "=" * 70)
logger.info("M1 MIRROR NSGA-II PURE FEA OPTIMIZATION V13")
logger.info("=" * 70)
logger.info(f"Prior FEA trials: {len(self.prior_data)}")
logger.info(f"New trials to run: {n_trials}")
logger.info(f"Objectives: {OBJ_NAMES}")
start_time = time.time()
# Create or load study
sampler = NSGAIISampler(
population_size=self.config.get('nsga2_settings', {}).get('population_size', 20),
crossover_prob=self.config.get('nsga2_settings', {}).get('crossover_prob', 0.9),
mutation_prob=self.config.get('nsga2_settings', {}).get('mutation_prob', 0.1),
seed=42
)
study = optuna.create_study(
study_name="v13_nsga2",
storage=self.storage,
directions=['minimize', 'minimize', 'minimize'], # 3 objectives
sampler=sampler,
load_if_exists=resume
)
# Seed with prior data if starting fresh
if not resume or len(study.trials) == 0:
self.seed_from_prior(study)
self.trial_count = self._get_next_trial_number()
logger.info(f"Starting from trial {self.trial_count}")
# Run optimization
def objective(trial: optuna.Trial) -> Tuple[float, float, float]:
# Sample parameters
params = {}
for var in self.config['design_variables']:
if var.get('enabled', False):
params[var['name']] = trial.suggest_float(var['name'], var['min'], var['max'])
# Run FEA
result = self.fea_runner.run_fea(params, self.trial_count)
self.trial_count += 1
if result is None:
# Return worst-case values for failed trials
return (1000.0, 1000.0, 1000.0)
# Store objectives as user attributes
trial.set_user_attr('source', 'FEA')
trial.set_user_attr('rel_filtered_rms_40_vs_20', result['objectives']['rel_filtered_rms_40_vs_20'])
trial.set_user_attr('rel_filtered_rms_60_vs_20', result['objectives']['rel_filtered_rms_60_vs_20'])
trial.set_user_attr('mfg_90_optician_workload', result['objectives']['mfg_90_optician_workload'])
trial.set_user_attr('solve_time', result.get('solve_time', 0))
return (
result['objectives']['rel_filtered_rms_40_vs_20'],
result['objectives']['rel_filtered_rms_60_vs_20'],
result['objectives']['mfg_90_optician_workload']
)
# Run
try:
study.optimize(
objective,
n_trials=n_trials,
show_progress_bar=True,
gc_after_trial=True
)
except KeyboardInterrupt:
logger.info("\nOptimization interrupted by user")
finally:
self.fea_runner.cleanup()
# Print results
elapsed = time.time() - start_time
self._print_results(study, elapsed)
def _print_results(self, study: optuna.Study, elapsed: float):
"""Print optimization results."""
logger.info("\n" + "=" * 70)
logger.info("OPTIMIZATION COMPLETE")
logger.info("=" * 70)
logger.info(f"Time: {elapsed/60:.1f} min ({elapsed/3600:.2f} hours)")
logger.info(f"Total trials: {len(study.trials)}")
# Get Pareto front
pareto_trials = study.best_trials
logger.info(f"Pareto-optimal trials: {len(pareto_trials)}")
# Print Pareto front
logger.info("\nPareto Front:")
logger.info("-" * 70)
logger.info(f"{'Trial':>6} {'40-20 (nm)':>12} {'60-20 (nm)':>12} {'Mfg (nm)':>12}")
logger.info("-" * 70)
pareto_data = []
for trial in sorted(pareto_trials, key=lambda t: t.values[0]):
logger.info(f"{trial.number:>6} {trial.values[0]:>12.2f} {trial.values[1]:>12.2f} {trial.values[2]:>12.2f}")
pareto_data.append({
'trial': trial.number,
'params': trial.params,
'objectives': {
'rel_filtered_rms_40_vs_20': trial.values[0],
'rel_filtered_rms_60_vs_20': trial.values[1],
'mfg_90_optician_workload': trial.values[2]
}
})
# Save results
results = {
'summary': {
'total_trials': len(study.trials),
'pareto_size': len(pareto_trials),
'elapsed_hours': elapsed / 3600
},
'pareto_front': pareto_data
}
with open(RESULTS_DIR / 'final_results.json', 'w') as f:
json.dump(results, f, indent=2)
logger.info(f"\nResults saved to {RESULTS_DIR / 'final_results.json'}")
# ============================================================================
# Main
# ============================================================================
def main():
parser = argparse.ArgumentParser(description='M1 Mirror NSGA-II V13')
parser.add_argument('--start', action='store_true', help='Start optimization')
parser.add_argument('--trials', type=int, default=50, help='Number of new FEA trials')
parser.add_argument('--resume', action='store_true', help='Resume from existing study')
args = parser.parse_args()
if not args.start:
print("M1 Mirror NSGA-II Pure FEA Optimization V13")
print("=" * 50)
print("\nUsage:")
print(" python run_optimization.py --start")
print(" python run_optimization.py --start --trials 55")
print(" python run_optimization.py --start --trials 55 --resume")
print("\nFor 8-hour overnight run (~55 trials at 8-9 min/trial):")
print(" python run_optimization.py --start --trials 55")
print("\nThis will:")
print(f" 1. Load ~{107} FEA trials from V11 database")
print(f" 2. Load additional FEA trials from V12 database")
print(" 3. Seed NSGA-II with all prior FEA data")
print(" 4. Run pure FEA multi-objective optimization")
print(" 5. No surrogate - every trial is real FEA")
return
with open(CONFIG_PATH, 'r') as f:
config = json.load(f)
optimizer = NSGA2Optimizer(config)
optimizer.run(n_trials=args.trials, resume=args.resume)
if __name__ == "__main__":
main()