auto: daily sync

projects/isogrid-dev-plate/studies/01_v1_tpe/trial_retention.py

@@ -0,0 +1,159 @@
+"""
+Trial file retention policy for isogrid optimization.
+
+Rules:
+  - NEVER delete:  *.png  (density, mesh, rib figures — full history always kept)
+  - NEVER delete:  *.json (params, results, rib profiles)
+  - HEAVY files:   NX model copies (.prt, .fem, .sim, .afm, .afem)
+                   + Nastran outputs (.op2, .f06, .dat, .log)
+  - KEEP heavy:    last KEEP_RECENT trials  +  best KEEP_BEST trials (by objective)
+  - STRIP heavy:   all other trial folders
+
+Usage in run_optimization.py:
+    rm = TrialRetentionManager(ITER_DIR, keep_recent=10, keep_best=5)
+    # after each trial:
+    rm.register(trial_number, trial_dir, objective=obj, mass_kg=mass, feasible=ok)
+    rm.apply()
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+# Extensions considered "heavy" (copied once, stripped when not in keep set)
+# NX model copies + Nastran outputs — everything needed to reproduce / re-open a trial
+HEAVY_EXTENSIONS = {".prt", ".fem", ".sim", ".afm", ".afem", ".op2", ".f06", ".dat", ".log"}
+
+# Extensions that are NEVER deleted regardless of retention policy
+SAFE_EXTENSIONS = {".png", ".json"}
+
+
+@dataclass
+class _TrialRecord:
+    number: int
+    path: Path
+    objective: float = float("inf")
+    mass_kg: float = float("inf")
+    feasible: bool = False
+    has_heavy: bool = True
+
+
+class TrialRetentionManager:
+    """
+    Manages heavy-file retention across trial folders.
+
+    After each trial:
+      1. Call register() with the trial's outcome
+      2. Call apply() to enforce the keep-recent + keep-best policy
+    """
+
+    def __init__(
+        self,
+        iter_dir: Path,
+        keep_recent: int = 10,
+        keep_best: int = 5,
+    ):
+        self.iter_dir = iter_dir
+        self.keep_recent = keep_recent
+        self.keep_best = keep_best
+        self._records: dict[int, _TrialRecord] = {}
+
+    def register(
+        self,
+        trial_number: int,
+        trial_dir: Path,
+        objective: float,
+        mass_kg: float,
+        feasible: bool,
+    ) -> None:
+        """Register a completed trial so the retention policy can track it."""
+        # Detect whether any heavy files currently exist
+        has_heavy = False
+        if trial_dir.exists():
+            has_heavy = any(
+                f.is_file() and f.suffix in HEAVY_EXTENSIONS
+                for f in trial_dir.iterdir()
+            )
+        self._records[trial_number] = _TrialRecord(
+            number=trial_number,
+            path=trial_dir,
+            objective=objective,
+            mass_kg=mass_kg,
+            feasible=feasible,
+            has_heavy=has_heavy,
+        )
+
+    def apply(self) -> list[int]:
+        """
+        Enforce retention policy.
+
+        Returns list of trial numbers whose heavy files were stripped.
+        """
+        if not self._records:
+            return []
+
+        all_nums = sorted(self._records.keys())
+
+        # Set 1: most recent N trials
+        recent_set: set[int] = set(all_nums[-self.keep_recent :])
+
+        # Set 2: best K trials — feasible first, then lowest objective
+        sorted_by_quality = sorted(
+            self._records.values(),
+            key=lambda r: (0 if r.feasible else 1, r.objective),
+        )
+        best_set: set[int] = {r.number for r in sorted_by_quality[: self.keep_best]}
+
+        keep_set = recent_set | best_set
+
+        stripped: list[int] = []
+        for num, record in self._records.items():
+            if num not in keep_set and record.has_heavy:
+                n_removed = self._strip_heavy(record)
+                if n_removed > 0:
+                    stripped.append(num)
+
+        return stripped
+
+    def _strip_heavy(self, record: _TrialRecord) -> int:
+        """
+        Remove heavy files from a trial folder.
+
+        PNGs and JSONs are NEVER touched.
+        Returns the number of files removed.
+        """
+        if not record.path.exists():
+            record.has_heavy = False
+            return 0
+
+        removed = 0
+        for f in list(record.path.iterdir()):
+            if f.is_file() and f.suffix in HEAVY_EXTENSIONS:
+                f.unlink()
+                removed += 1
+
+        record.has_heavy = False
+        return removed
+
+    def summary(self) -> dict:
+        """Return a brief status summary."""
+        all_nums = sorted(self._records.keys())
+        recent_set = set(all_nums[-self.keep_recent :])
+        sorted_by_quality = sorted(
+            self._records.values(),
+            key=lambda r: (0 if r.feasible else 1, r.objective),
+        )
+        best_set = {r.number for r in sorted_by_quality[: self.keep_best]}
+        keep_set = recent_set | best_set
+
+        return {
+            "total_trials": len(self._records),
+            "keep_recent": self.keep_recent,
+            "keep_best": self.keep_best,
+            "currently_kept": sorted(keep_set),
+            "stripped": sorted(
+                n for n, r in self._records.items()
+                if not r.has_heavy and n not in keep_set
+            ),
+        }