160 lines
5.1 KiB
Python
Executable File
160 lines
5.1 KiB
Python
Executable File
"""
|
|
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
|
|
),
|
|
}
|