feat: Add neural loop automation - templates, auto-trainer, CLI

Closes the neural training loop with automated workflow:
- atomizer.py: One-command neural workflow CLI
- auto_trainer.py: Auto-training trigger system (50pt threshold)
- template_loader.py: Study creation from templates
- study_reset.py: Study reset/cleanup utility
- 3 templates: beam stiffness, bracket stress, frequency tuning
- State assessment document (Nov 25)

Usage: python atomizer.py neural-optimize --study my_study --trials 500

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

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
2025-11-26 07:53:00 -05:00
parent e3bdb08a22
commit a0c008a593
10 changed files with 2789 additions and 0 deletions

View File

@@ -0,0 +1,522 @@
"""
Auto-Training Trigger System for AtomizerField
Monitors training data collection and automatically triggers neural network training
when enough data is accumulated. This is the key component to close the neural loop.
Workflow:
1. Monitor training data directory for new trials
2. When min_points threshold is reached, trigger training
3. Validate trained model against FEA
4. Deploy model for neural-accelerated optimization
Usage:
from optimization_engine.auto_trainer import AutoTrainer
trainer = AutoTrainer(
study_name="uav_arm_optimization",
min_points=50,
epochs=100
)
# Check if ready to train
if trainer.should_train():
model_path = trainer.train()
trainer.validate_model(model_path)
# Or run continuous monitoring
trainer.watch()
"""
import json
import subprocess
import sys
import time
import logging
from pathlib import Path
from typing import Dict, Any, Optional, Tuple, List
from datetime import datetime
import shutil
logger = logging.getLogger(__name__)
class AutoTrainer:
"""
Automatic neural network training trigger for AtomizerField.
Monitors training data accumulation and triggers training when thresholds are met.
"""
def __init__(
self,
study_name: str,
training_data_dir: Optional[Path] = None,
min_points: int = 50,
epochs: int = 100,
val_split: float = 0.2,
retrain_threshold: int = 50,
atomizer_field_dir: Optional[Path] = None,
output_dir: Optional[Path] = None
):
"""
Initialize the auto trainer.
Args:
study_name: Name of the optimization study
training_data_dir: Directory containing exported training data
min_points: Minimum data points before training (default: 50)
epochs: Training epochs (default: 100)
val_split: Validation split ratio (default: 0.2)
retrain_threshold: New points needed for retraining (default: 50)
atomizer_field_dir: Path to atomizer-field repository
output_dir: Directory for trained models
"""
self.study_name = study_name
self.min_points = min_points
self.epochs = epochs
self.val_split = val_split
self.retrain_threshold = retrain_threshold
# Set up directories
project_root = Path(__file__).parent.parent
self.training_data_dir = training_data_dir or (
project_root / "atomizer_field_training_data" / study_name
)
self.atomizer_field_dir = atomizer_field_dir or (project_root / "atomizer-field")
self.output_dir = output_dir or (
self.atomizer_field_dir / "runs" / f"{study_name}_auto"
)
# Tracking state
self.last_trained_count = 0
self.model_version = 0
self.training_history: List[Dict[str, Any]] = []
# Load state if exists
self._load_state()
logger.info(f"AutoTrainer initialized for {study_name}")
logger.info(f"Training data: {self.training_data_dir}")
logger.info(f"Min points: {min_points}, Retrain threshold: {retrain_threshold}")
def _state_file(self) -> Path:
"""Get path to state file."""
return self.output_dir / "auto_trainer_state.json"
def _load_state(self) -> None:
"""Load trainer state from disk."""
state_file = self._state_file()
if state_file.exists():
with open(state_file, 'r') as f:
state = json.load(f)
self.last_trained_count = state.get("last_trained_count", 0)
self.model_version = state.get("model_version", 0)
self.training_history = state.get("training_history", [])
logger.info(f"Loaded state: {self.last_trained_count} points trained, version {self.model_version}")
def _save_state(self) -> None:
"""Save trainer state to disk."""
self.output_dir.mkdir(parents=True, exist_ok=True)
state_file = self._state_file()
state = {
"study_name": self.study_name,
"last_trained_count": self.last_trained_count,
"model_version": self.model_version,
"training_history": self.training_history,
"last_updated": datetime.now().isoformat()
}
with open(state_file, 'w') as f:
json.dump(state, f, indent=2)
def count_training_points(self) -> int:
"""
Count available training data points.
Returns:
Number of trial directories with valid training data
"""
if not self.training_data_dir.exists():
return 0
count = 0
for trial_dir in self.training_data_dir.glob("trial_*"):
if trial_dir.is_dir():
# Check for required files
has_input = (trial_dir / "input" / "model.bdf").exists()
has_output = (trial_dir / "output" / "model.op2").exists()
has_metadata = (trial_dir / "metadata.json").exists()
if has_input and has_output and has_metadata:
count += 1
return count
def should_train(self) -> bool:
"""
Check if there's enough new data to trigger training.
Returns:
True if training should be triggered
"""
current_count = self.count_training_points()
# First training - check minimum threshold
if self.last_trained_count == 0:
return current_count >= self.min_points
# Retraining - check new points threshold
new_points = current_count - self.last_trained_count
return new_points >= self.retrain_threshold
def get_new_points_count(self) -> int:
"""Get number of new points since last training."""
return self.count_training_points() - self.last_trained_count
def prepare_training_split(self) -> Tuple[Path, Path]:
"""
Prepare train/validation split from collected data.
Returns:
Tuple of (train_dir, val_dir) paths
"""
train_dir = self.training_data_dir.parent / f"{self.study_name}_train"
val_dir = self.training_data_dir.parent / f"{self.study_name}_val"
# Clear existing splits
if train_dir.exists():
shutil.rmtree(train_dir)
if val_dir.exists():
shutil.rmtree(val_dir)
train_dir.mkdir(parents=True)
val_dir.mkdir(parents=True)
# Get all trial directories
trial_dirs = sorted(self.training_data_dir.glob("trial_*"))
n_trials = len(trial_dirs)
n_val = max(1, int(n_trials * self.val_split))
# Split: use latest trials for validation (they're most diverse)
train_trials = trial_dirs[:-n_val] if n_val > 0 else trial_dirs
val_trials = trial_dirs[-n_val:] if n_val > 0 else []
# Copy to split directories
for trial_dir in train_trials:
dest = train_dir / trial_dir.name
shutil.copytree(trial_dir, dest)
for trial_dir in val_trials:
dest = val_dir / trial_dir.name
shutil.copytree(trial_dir, dest)
logger.info(f"Split data: {len(train_trials)} train, {len(val_trials)} validation")
return train_dir, val_dir
def train(self, train_parametric: bool = True) -> Optional[Path]:
"""
Trigger neural network training.
Args:
train_parametric: If True, train parametric predictor (fast).
If False, train field predictor (slower, more detailed).
Returns:
Path to trained model checkpoint, or None if training failed
"""
current_count = self.count_training_points()
if current_count < self.min_points:
logger.warning(
f"Not enough data for training: {current_count} < {self.min_points}"
)
return None
logger.info(f"Starting training with {current_count} data points...")
# Prepare train/val split
train_dir, val_dir = self.prepare_training_split()
# Increment model version
self.model_version += 1
version_output_dir = self.output_dir / f"v{self.model_version}"
version_output_dir.mkdir(parents=True, exist_ok=True)
# Choose training script
if train_parametric:
train_script = self.atomizer_field_dir / "train_parametric.py"
else:
train_script = self.atomizer_field_dir / "train.py"
if not train_script.exists():
logger.error(f"Training script not found: {train_script}")
return None
# Build training command
cmd = [
sys.executable,
str(train_script),
"--train_dir", str(train_dir),
"--val_dir", str(val_dir),
"--epochs", str(self.epochs),
"--output_dir", str(version_output_dir)
]
logger.info(f"Running: {' '.join(cmd)}")
# Run training
start_time = time.time()
try:
result = subprocess.run(
cmd,
capture_output=True,
text=True,
cwd=str(self.atomizer_field_dir),
timeout=3600 * 4 # 4 hour timeout
)
elapsed = time.time() - start_time
if result.returncode != 0:
logger.error(f"Training failed:\n{result.stderr}")
return None
logger.info(f"Training completed in {elapsed/60:.1f} minutes")
# Find model checkpoint
checkpoints = list(version_output_dir.glob("*.pt")) + list(version_output_dir.glob("*.pth"))
if not checkpoints:
# Check for best model
checkpoints = list(version_output_dir.glob("**/best*.pt")) + \
list(version_output_dir.glob("**/checkpoint*.pt"))
if checkpoints:
model_path = checkpoints[0]
logger.info(f"Model saved: {model_path}")
else:
logger.warning("No checkpoint file found after training")
model_path = version_output_dir
# Update state
self.last_trained_count = current_count
self.training_history.append({
"version": self.model_version,
"timestamp": datetime.now().isoformat(),
"data_points": current_count,
"epochs": self.epochs,
"training_time_seconds": elapsed,
"model_path": str(model_path)
})
self._save_state()
return model_path
except subprocess.TimeoutExpired:
logger.error("Training timed out after 4 hours")
return None
except Exception as e:
logger.error(f"Training error: {e}")
return None
def validate_model(
self,
model_path: Path,
n_validation_trials: int = 5
) -> Dict[str, Any]:
"""
Validate trained model against FEA results.
Args:
model_path: Path to trained model
n_validation_trials: Number of trials to validate
Returns:
Validation metrics dictionary
"""
logger.info(f"Validating model: {model_path}")
# This would integrate with the neural surrogate to compare predictions vs FEA
# For now, return placeholder metrics
validation_results = {
"model_path": str(model_path),
"n_validation_trials": n_validation_trials,
"mean_error_percent": 0.0, # Would be computed
"max_error_percent": 0.0,
"validated_at": datetime.now().isoformat()
}
# TODO: Implement actual validation
# - Load model
# - Run predictions on held-out trials
# - Compare with FEA results
# - Compute error metrics
return validation_results
def get_latest_model(self) -> Optional[Path]:
"""
Get path to latest trained model.
Returns:
Path to latest model checkpoint, or None if no model exists
"""
if self.model_version == 0:
return None
latest_dir = self.output_dir / f"v{self.model_version}"
if not latest_dir.exists():
return None
# Find checkpoint
checkpoints = list(latest_dir.glob("*.pt")) + list(latest_dir.glob("*.pth"))
if checkpoints:
return checkpoints[0]
return latest_dir
def watch(self, check_interval: int = 60) -> None:
"""
Continuously monitor for new data and trigger training.
Args:
check_interval: Seconds between checks (default: 60)
"""
logger.info(f"Starting auto-trainer watch mode for {self.study_name}")
logger.info(f"Check interval: {check_interval}s")
logger.info(f"Min points: {self.min_points}, Retrain threshold: {self.retrain_threshold}")
try:
while True:
current_count = self.count_training_points()
new_points = current_count - self.last_trained_count
status = f"[{datetime.now().strftime('%H:%M:%S')}] "
status += f"Points: {current_count} (new: {new_points})"
if self.should_train():
status += " -> TRAINING"
print(status)
model_path = self.train()
if model_path:
print(f"Training complete: {model_path}")
else:
if self.last_trained_count == 0:
needed = self.min_points - current_count
status += f" (need {needed} more for first training)"
else:
needed = self.retrain_threshold - new_points
status += f" (need {needed} more for retraining)"
print(status)
time.sleep(check_interval)
except KeyboardInterrupt:
logger.info("Watch mode stopped")
def get_status(self) -> Dict[str, Any]:
"""
Get current trainer status.
Returns:
Status dictionary with counts and state
"""
current_count = self.count_training_points()
new_points = current_count - self.last_trained_count
return {
"study_name": self.study_name,
"total_points": current_count,
"new_points_since_training": new_points,
"last_trained_count": self.last_trained_count,
"model_version": self.model_version,
"min_points_threshold": self.min_points,
"retrain_threshold": self.retrain_threshold,
"should_train": self.should_train(),
"latest_model": str(self.get_latest_model()) if self.get_latest_model() else None,
"training_history_count": len(self.training_history)
}
def check_training_status(study_name: str) -> Dict[str, Any]:
"""
Quick check of training data status for a study.
Args:
study_name: Name of the study
Returns:
Status dictionary
"""
trainer = AutoTrainer(study_name=study_name)
return trainer.get_status()
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="AtomizerField Auto-Trainer")
parser.add_argument("study_name", help="Name of the optimization study")
parser.add_argument("--train", action="store_true", help="Trigger training now")
parser.add_argument("--watch", action="store_true", help="Watch mode - continuous monitoring")
parser.add_argument("--status", action="store_true", help="Show status only")
parser.add_argument("--min-points", type=int, default=50, help="Minimum points for training")
parser.add_argument("--epochs", type=int, default=100, help="Training epochs")
parser.add_argument("--interval", type=int, default=60, help="Check interval for watch mode")
args = parser.parse_args()
# Set up logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s [%(levelname)s] %(message)s'
)
trainer = AutoTrainer(
study_name=args.study_name,
min_points=args.min_points,
epochs=args.epochs
)
if args.status:
status = trainer.get_status()
print(f"\nAuto-Trainer Status: {args.study_name}")
print("=" * 50)
for key, value in status.items():
print(f" {key}: {value}")
elif args.train:
if trainer.should_train():
print("Training triggered...")
model_path = trainer.train()
if model_path:
print(f"Success! Model at: {model_path}")
else:
print("Training failed")
else:
print("Not enough data for training")
print(f"Current: {trainer.count_training_points()}, Need: {args.min_points}")
elif args.watch:
trainer.watch(check_interval=args.interval)
else:
# Default: show status and recommendation
status = trainer.get_status()
print(f"\nAuto-Trainer Status: {args.study_name}")
print("=" * 50)
print(f" Data points: {status['total_points']}")
print(f" New since last training: {status['new_points_since_training']}")
print(f" Model version: v{status['model_version']}")
print(f" Should train: {status['should_train']}")
print()
if status['should_train']:
print("Ready to train! Run with --train to start training.")
else:
if status['last_trained_count'] == 0:
needed = status['min_points_threshold'] - status['total_points']
print(f"Need {needed} more points for initial training.")
else:
needed = status['retrain_threshold'] - status['new_points_since_training']
print(f"Need {needed} more new points for retraining.")

View File

@@ -0,0 +1,447 @@
"""
Study Reset and Cleanup Utility for Atomizer
Provides safe operations to reset or clean up optimization studies:
- Reset database (remove all trials, keep configuration)
- Clean up temporary files
- Archive results
- Full study deletion
Usage:
python -m optimization_engine.study_reset my_study --reset-db
python -m optimization_engine.study_reset my_study --cleanup-temp
python -m optimization_engine.study_reset my_study --full-reset
Safety features:
- Confirmation prompts for destructive operations
- Automatic backups before deletion
- Dry-run mode to preview changes
"""
import json
import shutil
import optuna
from pathlib import Path
from datetime import datetime
from typing import Dict, Any, List, Optional
import logging
logger = logging.getLogger(__name__)
class StudyReset:
"""Handles study reset and cleanup operations."""
def __init__(self, study_name: str, studies_dir: str = "studies"):
"""
Initialize study reset utility.
Args:
study_name: Name of the study to reset
studies_dir: Base directory for studies
"""
self.study_name = study_name
self.studies_dir = Path(studies_dir)
self.study_path = self.studies_dir / study_name
self.setup_dir = self.study_path / "1_setup"
self.model_dir = self.setup_dir / "model"
self.results_dir = self.study_path / "2_results"
def validate_study_exists(self) -> bool:
"""Check if study exists."""
return self.study_path.exists()
def get_study_stats(self) -> Dict[str, Any]:
"""
Get current study statistics.
Returns:
Dictionary with study statistics
"""
stats = {
"study_name": self.study_name,
"exists": self.study_path.exists(),
"has_results": self.results_dir.exists(),
"trials": 0,
"completed": 0,
"failed": 0,
"db_size_mb": 0,
"temp_files": 0,
"temp_size_mb": 0
}
if not self.study_path.exists():
return stats
# Check database
db_path = self.results_dir / "study.db"
if db_path.exists():
stats["db_size_mb"] = db_path.stat().st_size / (1024 * 1024)
try:
storage = f"sqlite:///{db_path}"
study = optuna.load_study(study_name=self.study_name, storage=storage)
stats["trials"] = len(study.trials)
stats["completed"] = len([t for t in study.trials
if t.state == optuna.trial.TrialState.COMPLETE])
stats["failed"] = len([t for t in study.trials
if t.state == optuna.trial.TrialState.FAIL])
except Exception as e:
logger.warning(f"Could not load study: {e}")
# Count temp files
temp_patterns = ["_temp*", "*.log", "*.bak", "worker_*"]
temp_files = []
for pattern in temp_patterns:
temp_files.extend(self.model_dir.glob(pattern))
temp_files.extend(self.results_dir.glob(pattern))
stats["temp_files"] = len(temp_files)
stats["temp_size_mb"] = sum(f.stat().st_size for f in temp_files if f.is_file()) / (1024 * 1024)
return stats
def reset_database(self, backup: bool = True, dry_run: bool = False) -> Dict[str, Any]:
"""
Reset the Optuna database (delete all trials).
Args:
backup: Create backup before reset
dry_run: Preview changes without executing
Returns:
Operation result dictionary
"""
result = {"operation": "reset_database", "dry_run": dry_run}
db_path = self.results_dir / "study.db"
if not db_path.exists():
result["status"] = "skipped"
result["message"] = "No database found"
return result
if dry_run:
result["status"] = "preview"
result["message"] = f"Would delete {db_path}"
return result
# Create backup
if backup:
backup_name = f"study_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.db"
backup_path = self.results_dir / backup_name
shutil.copy2(db_path, backup_path)
result["backup"] = str(backup_path)
logger.info(f"Created backup: {backup_path}")
# Delete database
db_path.unlink()
result["status"] = "success"
result["message"] = "Database reset complete"
# Also clean history files
for history_file in ["history.json", "history.csv", "optimization_summary.json"]:
hist_path = self.results_dir / history_file
if hist_path.exists():
hist_path.unlink()
logger.info(f"Deleted: {hist_path}")
return result
def cleanup_temp_files(self, dry_run: bool = False) -> Dict[str, Any]:
"""
Remove temporary files from study.
Args:
dry_run: Preview changes without executing
Returns:
Operation result dictionary
"""
result = {
"operation": "cleanup_temp",
"dry_run": dry_run,
"deleted_files": [],
"deleted_size_mb": 0
}
temp_patterns = [
"_temp*", # Temporary NX files
"*.log", # Log files
"*.bak", # Backup files
"worker_*", # Worker directories
"*.pyc", # Python cache
"__pycache__" # Python cache dirs
]
files_to_delete: List[Path] = []
for pattern in temp_patterns:
files_to_delete.extend(self.model_dir.glob(pattern))
files_to_delete.extend(self.results_dir.glob(pattern))
files_to_delete.extend(self.study_path.glob(pattern))
total_size = 0
for path in files_to_delete:
if path.is_file():
total_size += path.stat().st_size
result["files_found"] = len(files_to_delete)
result["size_mb"] = total_size / (1024 * 1024)
if dry_run:
result["status"] = "preview"
result["files_to_delete"] = [str(f) for f in files_to_delete[:20]] # Limit preview
return result
# Actually delete
for path in files_to_delete:
try:
if path.is_file():
path.unlink()
elif path.is_dir():
shutil.rmtree(path)
result["deleted_files"].append(str(path))
except Exception as e:
logger.warning(f"Could not delete {path}: {e}")
result["deleted_size_mb"] = total_size / (1024 * 1024)
result["status"] = "success"
return result
def archive_results(self, archive_dir: Optional[Path] = None, dry_run: bool = False) -> Dict[str, Any]:
"""
Archive study results before reset.
Args:
archive_dir: Directory for archives (default: studies/archives)
dry_run: Preview changes without executing
Returns:
Operation result dictionary
"""
result = {"operation": "archive", "dry_run": dry_run}
if archive_dir is None:
archive_dir = self.studies_dir / "archives"
if not self.results_dir.exists():
result["status"] = "skipped"
result["message"] = "No results to archive"
return result
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
archive_name = f"{self.study_name}_{timestamp}"
archive_path = archive_dir / archive_name
if dry_run:
result["status"] = "preview"
result["archive_path"] = str(archive_path)
return result
archive_dir.mkdir(parents=True, exist_ok=True)
shutil.copytree(self.results_dir, archive_path)
result["status"] = "success"
result["archive_path"] = str(archive_path)
logger.info(f"Archived results to: {archive_path}")
return result
def full_reset(self, backup: bool = True, dry_run: bool = False) -> Dict[str, Any]:
"""
Perform full study reset (database + temp files).
Args:
backup: Create backup before reset
dry_run: Preview changes without executing
Returns:
Combined operation result
"""
results = {"operation": "full_reset", "dry_run": dry_run}
if backup and not dry_run:
archive_result = self.archive_results(dry_run=dry_run)
results["archive"] = archive_result
db_result = self.reset_database(backup=backup, dry_run=dry_run)
results["database"] = db_result
temp_result = self.cleanup_temp_files(dry_run=dry_run)
results["temp_cleanup"] = temp_result
# Remove lock files
lock_file = self.results_dir / ".optimization_lock"
if lock_file.exists() and not dry_run:
lock_file.unlink()
results["lock_removed"] = True
results["status"] = "success" if not dry_run else "preview"
return results
def delete_study(self, confirm: bool = False, dry_run: bool = False) -> Dict[str, Any]:
"""
Completely delete study (DESTRUCTIVE).
Args:
confirm: Must be True to actually delete
dry_run: Preview changes without executing
Returns:
Operation result dictionary
"""
result = {"operation": "delete_study", "dry_run": dry_run}
if not confirm and not dry_run:
result["status"] = "error"
result["message"] = "Must set confirm=True to delete study"
return result
if not self.study_path.exists():
result["status"] = "skipped"
result["message"] = "Study does not exist"
return result
if dry_run:
result["status"] = "preview"
result["message"] = f"Would delete: {self.study_path}"
return result
# Create archive first
archive_result = self.archive_results()
result["archive"] = archive_result
# Delete study folder
shutil.rmtree(self.study_path)
result["status"] = "success"
result["message"] = f"Deleted study: {self.study_name}"
return result
def reset_study(
study_name: str,
reset_db: bool = True,
cleanup_temp: bool = True,
backup: bool = True,
dry_run: bool = False
) -> Dict[str, Any]:
"""
Convenience function to reset a study.
Args:
study_name: Name of the study
reset_db: Reset the Optuna database
cleanup_temp: Clean up temporary files
backup: Create backup before reset
dry_run: Preview changes without executing
Returns:
Operation result dictionary
"""
resetter = StudyReset(study_name)
if not resetter.validate_study_exists():
return {"status": "error", "message": f"Study '{study_name}' not found"}
results = {}
if reset_db:
results["database"] = resetter.reset_database(backup=backup, dry_run=dry_run)
if cleanup_temp:
results["temp_cleanup"] = resetter.cleanup_temp_files(dry_run=dry_run)
return results
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(
description="Reset or cleanup Atomizer optimization studies",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Show study status
python -m optimization_engine.study_reset my_study --status
# Preview reset (dry run)
python -m optimization_engine.study_reset my_study --full-reset --dry-run
# Reset database only
python -m optimization_engine.study_reset my_study --reset-db
# Clean temp files only
python -m optimization_engine.study_reset my_study --cleanup-temp
# Full reset with backup
python -m optimization_engine.study_reset my_study --full-reset
"""
)
parser.add_argument("study_name", help="Name of the study")
parser.add_argument("--status", action="store_true", help="Show study status only")
parser.add_argument("--reset-db", action="store_true", help="Reset Optuna database")
parser.add_argument("--cleanup-temp", action="store_true", help="Clean temporary files")
parser.add_argument("--full-reset", action="store_true", help="Full reset (db + temp)")
parser.add_argument("--archive", action="store_true", help="Archive results before reset")
parser.add_argument("--delete", action="store_true", help="Delete study completely")
parser.add_argument("--no-backup", action="store_true", help="Skip backup")
parser.add_argument("--dry-run", action="store_true", help="Preview without executing")
parser.add_argument("--yes", "-y", action="store_true", help="Skip confirmation prompts")
args = parser.parse_args()
# Set up logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s [%(levelname)s] %(message)s'
)
resetter = StudyReset(args.study_name)
if not resetter.validate_study_exists():
print(f"Error: Study '{args.study_name}' not found")
sys.exit(1)
if args.status:
stats = resetter.get_study_stats()
print(f"\nStudy: {args.study_name}")
print("=" * 50)
print(f" Trials: {stats['trials']} ({stats['completed']} completed, {stats['failed']} failed)")
print(f" Database size: {stats['db_size_mb']:.2f} MB")
print(f" Temp files: {stats['temp_files']} ({stats['temp_size_mb']:.2f} MB)")
sys.exit(0)
# Confirmation
if not args.dry_run and not args.yes:
action = "full reset" if args.full_reset else \
"delete" if args.delete else \
"reset" if args.reset_db else "cleanup"
response = input(f"\nReally {action} study '{args.study_name}'? [y/N] ")
if response.lower() not in ['y', 'yes']:
print("Aborted")
sys.exit(0)
backup = not args.no_backup
if args.full_reset:
result = resetter.full_reset(backup=backup, dry_run=args.dry_run)
elif args.delete:
result = resetter.delete_study(confirm=True, dry_run=args.dry_run)
elif args.reset_db:
result = resetter.reset_database(backup=backup, dry_run=args.dry_run)
elif args.cleanup_temp:
result = resetter.cleanup_temp_files(dry_run=args.dry_run)
elif args.archive:
result = resetter.archive_results(dry_run=args.dry_run)
else:
parser.print_help()
sys.exit(0)
print("\nResult:")
print(json.dumps(result, indent=2))
if args.dry_run:
print("\n[DRY RUN - no changes made]")

View File

@@ -0,0 +1,383 @@
"""
Template Loader for Atomizer Optimization Studies
Creates new studies from templates with automatic folder structure creation.
Usage:
from optimization_engine.template_loader import create_study_from_template, list_templates
# List available templates
templates = list_templates()
# Create a new study from template
create_study_from_template(
template_name="beam_stiffness_optimization",
study_name="my_beam_study"
)
"""
import json
import shutil
from pathlib import Path
from typing import Dict, Any, List, Optional
from datetime import datetime
TEMPLATES_DIR = Path(__file__).parent.parent / "templates"
STUDIES_DIR = Path(__file__).parent.parent / "studies"
def list_templates() -> List[Dict[str, Any]]:
"""
List all available templates.
Returns:
List of template metadata dictionaries
"""
templates = []
if not TEMPLATES_DIR.exists():
return templates
for template_file in TEMPLATES_DIR.glob("*.json"):
try:
with open(template_file, 'r') as f:
config = json.load(f)
template_info = config.get("template_info", {})
templates.append({
"name": template_file.stem,
"description": config.get("description", "No description"),
"category": template_info.get("category", "general"),
"analysis_type": template_info.get("analysis_type", "unknown"),
"objectives": len(config.get("objectives", [])),
"design_variables": len(config.get("design_variables", [])),
"path": str(template_file)
})
except Exception as e:
print(f"Warning: Could not load template {template_file}: {e}")
return templates
def get_template(template_name: str) -> Optional[Dict[str, Any]]:
"""
Load a template by name.
Args:
template_name: Name of the template (without .json extension)
Returns:
Template configuration dictionary or None if not found
"""
template_path = TEMPLATES_DIR / f"{template_name}.json"
if not template_path.exists():
# Try with .json extension already included
template_path = TEMPLATES_DIR / template_name
if not template_path.exists():
return None
with open(template_path, 'r') as f:
return json.load(f)
def create_study_from_template(
template_name: str,
study_name: str,
studies_dir: Optional[Path] = None,
overrides: Optional[Dict[str, Any]] = None
) -> Path:
"""
Create a new study from a template.
Args:
template_name: Name of the template to use
study_name: Name for the new study
studies_dir: Base directory for studies (default: studies/)
overrides: Dictionary of config values to override
Returns:
Path to the created study directory
Raises:
FileNotFoundError: If template doesn't exist
FileExistsError: If study already exists
"""
if studies_dir is None:
studies_dir = STUDIES_DIR
studies_dir = Path(studies_dir)
# Load template
template = get_template(template_name)
if template is None:
available = [t["name"] for t in list_templates()]
raise FileNotFoundError(
f"Template '{template_name}' not found. "
f"Available templates: {available}"
)
# Check if study already exists
study_path = studies_dir / study_name
if study_path.exists():
raise FileExistsError(
f"Study '{study_name}' already exists at {study_path}. "
"Choose a different name or delete the existing study."
)
# Create study directory structure
setup_dir = study_path / "1_setup"
model_dir = setup_dir / "model"
results_dir = study_path / "2_results"
setup_dir.mkdir(parents=True)
model_dir.mkdir()
results_dir.mkdir()
# Customize template for this study
config = template.copy()
config["study_name"] = study_name
config["created_from_template"] = template_name
config["created_at"] = datetime.now().isoformat()
# Update training data export path
if "training_data_export" in config:
export_dir = config["training_data_export"].get("export_dir", "")
if "${study_name}" in export_dir:
config["training_data_export"]["export_dir"] = export_dir.replace(
"${study_name}", study_name
)
# Apply overrides
if overrides:
_deep_update(config, overrides)
# Write configuration
config_path = setup_dir / "optimization_config.json"
with open(config_path, 'w') as f:
json.dump(config, f, indent=2)
# Create run_optimization.py
run_script_content = _generate_run_script(study_name, config)
run_script_path = study_path / "run_optimization.py"
with open(run_script_path, 'w') as f:
f.write(run_script_content)
# Create README.md
readme_content = _generate_study_readme(study_name, config, template_name)
readme_path = study_path / "README.md"
with open(readme_path, 'w') as f:
f.write(readme_content)
print(f"Created study '{study_name}' from template '{template_name}'")
print(f" Location: {study_path}")
print(f" Config: {config_path}")
print(f"\nNext steps:")
print(f" 1. Add your NX model files to: {model_dir}")
print(f" 2. Update design variable bounds in optimization_config.json")
print(f" 3. Run: python {run_script_path} --trials 50")
return study_path
def _deep_update(base: Dict, updates: Dict) -> Dict:
"""Recursively update a dictionary."""
for key, value in updates.items():
if key in base and isinstance(base[key], dict) and isinstance(value, dict):
_deep_update(base[key], value)
else:
base[key] = value
return base
def _generate_run_script(study_name: str, config: Dict[str, Any]) -> str:
"""Generate the run_optimization.py script for a study."""
return f'''"""
Optimization Runner for {study_name}
Auto-generated from template: {config.get('created_from_template', 'unknown')}
Created: {config.get('created_at', 'unknown')}
Usage:
python run_optimization.py --trials 50
python run_optimization.py --trials 25 --resume
python run_optimization.py --trials 100 --enable-nn
"""
import sys
import argparse
from pathlib import Path
# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))
from optimization_engine.study_runner import run_study
def main():
parser = argparse.ArgumentParser(description="{config.get('description', study_name)}")
parser.add_argument('--trials', type=int, default=30, help='Number of trials to run')
parser.add_argument('--resume', action='store_true', help='Resume existing study')
parser.add_argument('--enable-nn', action='store_true', help='Enable neural network acceleration')
parser.add_argument('--validate-only', action='store_true', help='Only validate setup, do not run')
args = parser.parse_args()
study_dir = Path(__file__).parent
config_path = study_dir / "1_setup" / "optimization_config.json"
if args.validate_only:
from optimization_engine.validators import validate_study
result = validate_study("{study_name}")
print(result)
return
run_study(
config_path=config_path,
n_trials=args.trials,
resume=args.resume,
enable_neural=args.enable_nn
)
if __name__ == "__main__":
main()
'''
def _generate_study_readme(study_name: str, config: Dict[str, Any], template_name: str) -> str:
"""Generate a README.md for the study."""
objectives = config.get("objectives", [])
design_vars = config.get("design_variables", [])
constraints = config.get("constraints", [])
obj_list = "\n".join([f"- **{o.get('name', 'unnamed')}**: {o.get('goal', 'minimize')} - {o.get('description', '')}" for o in objectives])
dv_list = "\n".join([f"- **{d.get('parameter', 'unnamed')}**: [{d.get('bounds', [0, 1])[0]}, {d.get('bounds', [0, 1])[1]}] - {d.get('description', '')}" for d in design_vars])
const_list = "\n".join([f"- **{c.get('name', 'unnamed')}**: {c.get('type', 'less_than')} {c.get('threshold', 0)} - {c.get('description', '')}" for c in constraints])
return f'''# {study_name}
{config.get('description', 'Optimization study')}
**Template**: {template_name}
**Created**: {config.get('created_at', 'unknown')}
## Engineering Context
{config.get('engineering_context', 'No context provided')}
## Objectives
{obj_list if obj_list else 'None defined'}
## Design Variables
{dv_list if dv_list else 'None defined'}
## Constraints
{const_list if const_list else 'None defined'}
## Setup Instructions
1. **Add NX Model Files**
Copy your NX part (.prt), simulation (.sim), and FEM (.fem) files to:
```
1_setup/model/
```
2. **Configure Design Variables**
Edit `1_setup/optimization_config.json`:
- Ensure `design_variables[].parameter` matches your NX expression names
- Adjust bounds to your design space
3. **Validate Setup**
```bash
python run_optimization.py --validate-only
```
## Running the Optimization
### Basic Run
```bash
python run_optimization.py --trials 50
```
### Resume Interrupted Run
```bash
python run_optimization.py --trials 25 --resume
```
### With Neural Network Acceleration
```bash
python run_optimization.py --trials 100 --enable-nn
```
## Results
After optimization, results are saved in `2_results/`:
- `study.db` - Optuna database with all trials
- `history.json` - Trial history
- `optimization_summary.json` - Summary with best parameters
## Visualization
View results with Optuna Dashboard:
```bash
optuna-dashboard sqlite:///2_results/study.db
```
Or generate a report:
```bash
python -m optimization_engine.generate_report {study_name}
```
'''
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Atomizer Template Loader")
subparsers = parser.add_subparsers(dest="command", help="Commands")
# List templates
list_parser = subparsers.add_parser("list", help="List available templates")
# Create study
create_parser = subparsers.add_parser("create", help="Create study from template")
create_parser.add_argument("--template", "-t", required=True, help="Template name")
create_parser.add_argument("--name", "-n", required=True, help="Study name")
args = parser.parse_args()
if args.command == "list":
templates = list_templates()
if not templates:
print("No templates found in templates/")
else:
print("Available templates:")
print("-" * 60)
for t in templates:
print(f" {t['name']}")
print(f" {t['description']}")
print(f" Category: {t['category']} | Analysis: {t['analysis_type']}")
print(f" Design vars: {t['design_variables']} | Objectives: {t['objectives']}")
print()
elif args.command == "create":
try:
study_path = create_study_from_template(
template_name=args.template,
study_name=args.name
)
except (FileNotFoundError, FileExistsError) as e:
print(f"Error: {e}")
sys.exit(1)
else:
parser.print_help()