"""Smart iteration folder management for Hydrotech Beam optimization. Manages iteration folders with intelligent retention: - Each iteration gets a full copy of model files (openable in NX for debug) - Last N iterations: keep full model files (rolling window) - Best K iterations: keep full model files (by objective value) - All others: strip model files, keep only solver outputs + params This gives debuggability (open any recent/best iteration in NX) while keeping disk usage bounded. References: CEO design brief (2026-02-11): "all models properly saved in their iteration folder, keep last 10, keep best 3, delete stacking models" """ from __future__ import annotations import json import logging import shutil from dataclasses import dataclass, field from pathlib import Path from typing import Optional logger = logging.getLogger(__name__) # NX model file extensions (copied to each iteration) MODEL_EXTENSIONS = {".prt", ".fem", ".sim"} # Solver output extensions (always kept, even after stripping) KEEP_EXTENSIONS = {".op2", ".f06", ".dat", ".log", ".json", ".txt", ".csv"} # Default retention policy DEFAULT_KEEP_RECENT = 10 # keep last N iterations with full models DEFAULT_KEEP_BEST = 3 # keep best K iterations with full models @dataclass class IterationInfo: """Metadata for a single iteration.""" number: int path: Path mass: float = float("inf") displacement: float = float("inf") stress: float = float("inf") feasible: bool = False has_models: bool = True # False after stripping @dataclass class IterationManager: """Manages iteration folders with smart retention. Usage: mgr = IterationManager(study_dir, master_model_dir) # Before each trial: iter_dir = mgr.prepare_iteration(iteration_number) # After trial completes: mgr.record_result(iteration_number, mass=..., displacement=..., stress=...) # Periodically or at study end: mgr.apply_retention() """ study_dir: Path master_model_dir: Path keep_recent: int = DEFAULT_KEEP_RECENT keep_best: int = DEFAULT_KEEP_BEST _iterations: dict[int, IterationInfo] = field(default_factory=dict, repr=False) def __post_init__(self) -> None: self.iterations_dir = self.study_dir / "iterations" self.iterations_dir.mkdir(parents=True, exist_ok=True) # Scan existing iterations (for resume support) for d in sorted(self.iterations_dir.iterdir()): if d.is_dir() and d.name.startswith("iter"): try: num = int(d.name.replace("iter", "")) info = IterationInfo(number=num, path=d) # Load results if available results_file = d / "results.json" if results_file.exists(): data = json.loads(results_file.read_text()) info.mass = data.get("mass_kg", float("inf")) info.displacement = data.get("tip_displacement_mm", float("inf")) info.stress = data.get("max_von_mises_mpa", float("inf")) info.feasible = ( info.displacement <= 10.0 and info.stress <= 130.0 ) # Check if model files are present info.has_models = any( f.suffix in MODEL_EXTENSIONS for f in d.iterdir() ) self._iterations[num] = info except (ValueError, json.JSONDecodeError): continue if self._iterations: logger.info( "Loaded %d existing iterations (resume support)", len(self._iterations), ) def prepare_iteration(self, iteration_number: int) -> Path: """Set up an iteration folder with fresh model copies. Copies all model files from master_model_dir to the iteration folder. All paths are resolved to absolute to avoid NX reference issues. Args: iteration_number: Trial number (1-indexed). Returns: Absolute path to the iteration folder. """ iter_dir = (self.iterations_dir / f"iter{iteration_number:03d}").resolve() # Clean up if exists (failed previous run) if iter_dir.exists(): shutil.rmtree(iter_dir) iter_dir.mkdir(parents=True) # Copy ALL model files (so NX can resolve references within the folder) master = self.master_model_dir.resolve() copied = 0 for ext in MODEL_EXTENSIONS: for src in master.glob(f"*{ext}"): shutil.copy2(src, iter_dir / src.name) copied += 1 logger.info( "Prepared iter%03d: copied %d model files to %s", iteration_number, copied, iter_dir, ) # Track iteration self._iterations[iteration_number] = IterationInfo( number=iteration_number, path=iter_dir, has_models=True, ) return iter_dir def record_result( self, iteration_number: int, mass: float, displacement: float, stress: float, ) -> None: """Record results for an iteration and run retention check. Args: iteration_number: Trial number. mass: Extracted mass in kg. displacement: Tip displacement in mm. stress: Max von Mises stress in MPa. """ if iteration_number in self._iterations: info = self._iterations[iteration_number] info.mass = mass info.displacement = displacement info.stress = stress info.feasible = displacement <= 10.0 and stress <= 130.0 # Apply retention every 5 iterations to keep disk in check if iteration_number % 5 == 0: self.apply_retention() def apply_retention(self) -> None: """Apply the smart retention policy. Keep full model files for: 1. Last `keep_recent` iterations (rolling window) 2. Best `keep_best` iterations (by mass, feasible first) Strip model files from everything else (keep solver outputs only). """ if not self._iterations: return all_nums = sorted(self._iterations.keys()) # Set 1: Last N iterations recent_set = set(all_nums[-self.keep_recent:]) # Set 2: Best K by objective (feasible first, then lowest mass) sorted_by_quality = sorted( self._iterations.values(), key=lambda info: ( 0 if info.feasible else 1, # feasible first info.mass, # then lowest mass ), ) best_set = {info.number for info in sorted_by_quality[:self.keep_best]} # Keep set = recent ∪ best keep_set = recent_set | best_set # Strip model files from everything NOT in keep set stripped = 0 for num, info in self._iterations.items(): if num not in keep_set and info.has_models: self._strip_models(info) stripped += 1 if stripped > 0: logger.info( "Retention: kept %d recent + %d best, stripped %d iterations", len(recent_set), len(best_set), stripped, ) def _strip_models(self, info: IterationInfo) -> None: """Remove model files from an iteration folder, keep solver outputs.""" if not info.path.exists(): return removed = 0 for f in info.path.iterdir(): if f.is_file() and f.suffix in MODEL_EXTENSIONS: f.unlink() removed += 1 info.has_models = False if removed > 0: logger.debug( "Stripped %d model files from iter%03d", removed, info.number, ) def get_best_iterations(self, n: int = 3) -> list[IterationInfo]: """Return the N best iterations (feasible first, then lowest mass).""" return sorted( self._iterations.values(), key=lambda info: ( 0 if info.feasible else 1, info.mass, ), )[:n]