""" Study Creator - Atomizer Optimization Study Management Creates and manages optimization studies with mandatory benchmarking workflow. Workflow: 1. Create study structure 2. User provides NX models 3. Run benchmarking (mandatory) 4. Create substudies (substudy_1, substudy_2, etc.) 5. Each substudy validates against benchmarking before running Author: Antoine Letarte Date: 2025-11-17 Version: 1.0.0 """ import json import shutil from pathlib import Path from typing import Dict, Any, Optional, List from datetime import datetime import logging from optimization_engine.benchmarking_substudy import BenchmarkingSubstudy, BenchmarkResults logger = logging.getLogger(__name__) class StudyCreator: """ Creates and manages Atomizer optimization studies. Enforces mandatory benchmarking workflow and provides study structure management. """ def __init__(self, studies_root: Path = None): """ Initialize study creator. Args: studies_root: Root directory for all studies (default: ./studies) """ if studies_root is None: studies_root = Path.cwd() / "studies" self.studies_root = Path(studies_root) self.studies_root.mkdir(parents=True, exist_ok=True) logger.info(f"StudyCreator initialized: {self.studies_root}") def create_study(self, study_name: str, description: str = "") -> Path: """ Create a new optimization study with standard structure. Args: study_name: Name of the study (will be folder name) description: Brief description of the study Returns: Path to created study directory """ study_dir = self.studies_root / study_name if study_dir.exists(): logger.warning(f"Study already exists: {study_name}") return study_dir logger.info(f"Creating new study: {study_name}") # Create directory structure (study_dir / "model").mkdir(parents=True) (study_dir / "substudies" / "benchmarking").mkdir(parents=True) (study_dir / "config").mkdir(parents=True) (study_dir / "plugins" / "post_calculation").mkdir(parents=True) (study_dir / "results").mkdir(parents=True) # Create study metadata metadata = { "study_name": study_name, "description": description, "created": datetime.now().isoformat(), "status": "created", "benchmarking_completed": False, "substudies": [] } metadata_file = study_dir / "study_metadata.json" with open(metadata_file, 'w') as f: json.dump(metadata, f, indent=2) # Create README readme_content = self._generate_study_readme(study_name, description) readme_file = study_dir / "README.md" with open(readme_file, 'w', encoding='utf-8') as f: f.write(readme_content) logger.info(f"Study created: {study_dir}") logger.info("") logger.info("Next steps:") logger.info(f" 1. Add NX model files to: {study_dir / 'model'}/") logger.info(f" 2. Run benchmarking: study.run_benchmarking()") logger.info("") return study_dir def run_benchmarking(self, study_dir: Path, prt_file: Path, sim_file: Path) -> BenchmarkResults: """ Run mandatory benchmarking for a study. This MUST be run before any optimization substudies. Args: study_dir: Study directory prt_file: Path to NX part file sim_file: Path to NX simulation file Returns: BenchmarkResults """ logger.info("=" * 80) logger.info(f"RUNNING BENCHMARKING FOR STUDY: {study_dir.name}") logger.info("=" * 80) logger.info("") # Create benchmarking substudy benchmark = BenchmarkingSubstudy(study_dir, prt_file, sim_file) # Run discovery results = benchmark.run_discovery() # Generate report report_content = benchmark.generate_report(results) report_file = study_dir / "substudies" / "benchmarking" / "BENCHMARK_REPORT.md" with open(report_file, 'w', encoding='utf-8') as f: f.write(report_content) logger.info(f"Benchmark report saved to: {report_file}") logger.info("") # Update metadata self._update_metadata(study_dir, { "benchmarking_completed": results.validation_passed, "last_benchmarking": datetime.now().isoformat(), "status": "benchmarked" if results.validation_passed else "benchmark_failed" }) if not results.validation_passed: logger.error("Benchmarking validation FAILED!") logger.error("Fix issues before creating substudies") else: logger.info("Benchmarking validation PASSED!") logger.info("Ready to create substudies") logger.info("") return results def create_substudy(self, study_dir: Path, substudy_name: Optional[str] = None, config: Optional[Dict[str, Any]] = None) -> Path: """ Create a new substudy. Automatically validates against benchmarking before proceeding. Args: study_dir: Study directory substudy_name: Name of substudy (if None, auto-generates substudy_N) config: Optional configuration dict Returns: Path to substudy directory """ # Check benchmarking completed metadata = self._load_metadata(study_dir) if not metadata.get('benchmarking_completed', False): raise ValueError( "Benchmarking must be completed before creating substudies!\n" f"Run: study.run_benchmarking(prt_file, sim_file)" ) # Auto-generate substudy name if not provided if substudy_name is None: existing_substudies = metadata.get('substudies', []) # Filter out benchmarking non_benchmark = [s for s in existing_substudies if s != 'benchmarking'] substudy_number = len(non_benchmark) + 1 substudy_name = f"substudy_{substudy_number}" substudy_dir = study_dir / "substudies" / substudy_name if substudy_dir.exists(): logger.warning(f"Substudy already exists: {substudy_name}") return substudy_dir logger.info(f"Creating substudy: {substudy_name}") # Create substudy directory substudy_dir.mkdir(parents=True, exist_ok=True) # Create substudy config if config is None: # Use template config = self._create_default_substudy_config(study_dir, substudy_name) config_file = substudy_dir / "config.json" with open(config_file, 'w') as f: json.dump(config, f, indent=2) # Update metadata substudies = metadata.get('substudies', []) if substudy_name not in substudies: substudies.append(substudy_name) self._update_metadata(study_dir, {'substudies': substudies}) logger.info(f"Substudy created: {substudy_dir}") logger.info(f"Config: {config_file}") logger.info("") return substudy_dir def _create_default_substudy_config(self, study_dir: Path, substudy_name: str) -> Dict[str, Any]: """Create default substudy configuration based on benchmarking.""" # Load benchmark results benchmark_file = study_dir / "substudies" / "benchmarking" / "benchmark_results.json" if not benchmark_file.exists(): raise FileNotFoundError(f"Benchmark results not found: {benchmark_file}") with open(benchmark_file, 'r') as f: benchmark_data = json.load(f) # Create config from benchmark proposals config = { "substudy_name": substudy_name, "description": f"Substudy {substudy_name}", "created": datetime.now().isoformat(), "optimization": { "algorithm": "TPE", "direction": "minimize", "n_trials": 20, "n_startup_trials": 10, "design_variables": [] }, "continuation": { "enabled": False }, "solver": { "nastran_version": "2412", "use_journal": True, "timeout": 300 } } # Add proposed design variables for var in benchmark_data.get('proposed_design_variables', []): config["optimization"]["design_variables"].append({ "parameter": var['parameter'], "min": 0.0, # User must fill "max": 0.0, # User must fill "units": var.get('units', ''), "comment": f"From benchmarking: {var.get('suggested_range', 'define range')}" }) return config def _load_metadata(self, study_dir: Path) -> Dict[str, Any]: """Load study metadata.""" metadata_file = study_dir / "study_metadata.json" if not metadata_file.exists(): return {} with open(metadata_file, 'r') as f: return json.load(f) def _update_metadata(self, study_dir: Path, updates: Dict[str, Any]): """Update study metadata.""" metadata = self._load_metadata(study_dir) metadata.update(updates) metadata_file = study_dir / "study_metadata.json" with open(metadata_file, 'w') as f: json.dump(metadata, f, indent=2) def _generate_study_readme(self, study_name: str, description: str) -> str: """Generate README for new study.""" readme = [] readme.append(f"# {study_name}") readme.append("") readme.append(f"**Description**: {description}") readme.append(f"**Created**: {datetime.now().strftime('%Y-%m-%d')}") readme.append("") readme.append("## Study Structure") readme.append("") readme.append("```") readme.append(f"{study_name}/") readme.append("├── model/ # NX model files (.prt, .sim)") readme.append("├── substudies/") readme.append("│ ├── benchmarking/ # Mandatory discovery & validation") readme.append("│ ├── substudy_1/ # First optimization campaign") readme.append("│ └── substudy_2/ # Additional campaigns") readme.append("├── config/ # Configuration templates") readme.append("├── plugins/ # Study-specific hooks") readme.append("├── results/ # Optimization results") readme.append("└── README.md # This file") readme.append("```") readme.append("") readme.append("## Workflow") readme.append("") readme.append("### 1. Add NX Models") readme.append("Place your `.prt` and `.sim` files in the `model/` directory.") readme.append("") readme.append("### 2. Run Benchmarking (Mandatory)") readme.append("```python") readme.append("from optimization_engine.study_creator import StudyCreator") readme.append("") readme.append("creator = StudyCreator()") readme.append(f"results = creator.run_benchmarking(") readme.append(f" study_dir=Path('studies/{study_name}'),") readme.append(" prt_file=Path('studies/{study_name}/model/YourPart.prt'),") readme.append(" sim_file=Path('studies/{study_name}/model/YourSim.sim')") readme.append(")") readme.append("```") readme.append("") readme.append("### 3. Review Benchmark Report") readme.append("Check `substudies/benchmarking/BENCHMARK_REPORT.md` for:") readme.append("- Discovered expressions") readme.append("- OP2 contents") readme.append("- Baseline performance") readme.append("- Configuration proposals") readme.append("") readme.append("### 4. Create Substudies") readme.append("```python") readme.append("# Auto-numbered: substudy_1, substudy_2, etc.") readme.append(f"substudy_dir = creator.create_substudy(Path('studies/{study_name}'))") readme.append("") readme.append("# Or custom name:") readme.append(f"substudy_dir = creator.create_substudy(") readme.append(f" Path('studies/{study_name}'), ") readme.append(" substudy_name='coarse_exploration'") readme.append(")") readme.append("```") readme.append("") readme.append("### 5. Configure & Run Optimization") readme.append("Edit `substudies/substudy_N/config.json` with:") readme.append("- Design variable ranges") readme.append("- Objectives and constraints") readme.append("- Number of trials") readme.append("") readme.append("Then run the optimization!") readme.append("") readme.append("## Status") readme.append("") readme.append("See `study_metadata.json` for current study status.") readme.append("") return "\n".join(readme) def list_studies(self) -> List[Dict[str, Any]]: """List all studies in the studies root.""" studies = [] for study_dir in self.studies_root.iterdir(): if not study_dir.is_dir(): continue metadata_file = study_dir / "study_metadata.json" if metadata_file.exists(): with open(metadata_file, 'r') as f: metadata = json.load(f) studies.append({ 'name': study_dir.name, 'path': study_dir, 'status': metadata.get('status', 'unknown'), 'created': metadata.get('created', 'unknown'), 'benchmarking_completed': metadata.get('benchmarking_completed', False), 'substudies_count': len(metadata.get('substudies', [])) - 1 # Exclude benchmarking }) return studies def main(): """Example usage of StudyCreator.""" print("=" * 80) print("Atomizer Study Creator") print("=" * 80) print() creator = StudyCreator() # List existing studies studies = creator.list_studies() print(f"Existing studies: {len(studies)}") for study in studies: status_icon = "✅" if study['benchmarking_completed'] else "⚠️" print(f" {status_icon} {study['name']} ({study['status']}) - {study['substudies_count']} substudies") print() print("To create a new study:") print(" creator.create_study('my_study_name', 'Brief description')") print() if __name__ == '__main__': main()