refactor(config): distinguish model names from objects

GRID_SEARCH.md

@@ -159,7 +159,7 @@ see what trial is actually being run by looking for lines like this one:
 ```log
 2026-03-10 02:03:23.664 | INFO     | sampleworks.utils.guidance_script_utils:run_guidance_job_queue:589 - Running job 1/15: GuidanceConfig(protein='5I
 MV_1.0occB', structure='/data/inputs/processed/5IMV/5IMV_single_001_density_input.cif', density='/data/inputs/occ_sweeps/1.0occB/density_maps/5IMV_1.0
-occB_1.00A.ccp4', model='boltz2', guidance_type='pure_guidance', log_path='/data/results/5IMV_1.0occB/boltz2_X-RAY_DIFFRACTION/pure_guidance/ens8_gw0.
+occB_1.00A.ccp4', model_name='boltz2', guidance_type='pure_guidance', log_path='/data/results/5IMV_1.0occB/boltz2_X-RAY_DIFFRACTION/pure_guidance/ens8_gw0.
 1/run.log', output_dir='/data/results/5IMV_1.0occB/boltz2_X-RAY_DIFFRACTION/pure_guidance/ens8_gw0.1', partial_diffusion_step=120, loss_order=2, resol
 ution=1.0, device='cuda:0', gradient_normalization=True, em=False, guidance_start=-1, augmentation=True, align_to_input=True)
 ```

run_grid_search.py

@@ -27,7 +27,7 @@
 class GridSearchConfig:
     """Serializable summary of the grid-search dimensions and output location."""
 
-    model: str
+    model_name: str
     scalers: list[str]
     ensemble_sizes: list[int]
     gradient_weights: list[float]
@@ -180,7 +180,7 @@ def build_args_for_process_pool(
         protein=job.protein,
         structure=job.structure_path,
         density=job.density_path,
-        model=job.model,
+        model_name=job.model_name,
         guidance_type=job.scaler,
         log_path=job.log_path,
         output_dir=job.output_dir,
@@ -266,7 +266,7 @@ def run_grid_search(
         futures = {}
 
         for worker_num, job_queue_path in enumerate(job_queue_paths):
-            model = worker_job_queues[worker_num][0].model
+            model = worker_job_queues[worker_num][0].model_name
             future = executor.submit(
                 run_guidance_queue_script,
                 job_queue_path,
@@ -286,13 +286,13 @@ def run_grid_search(
                     if r.status == "success":
                         successful += 1
                         log.info(
-                            f"SUCCESS ({r.protein}, {r.model}, {r.method}, {r.scaler} "
+                            f"SUCCESS ({r.protein}, {r.model_name}, {r.method}, {r.scaler} "
                             f"{r.runtime_seconds:.1f}s): {r.log_path}"
                         )
                     else:
                         failed += 1
                         log.error(
-                            f"FAILED ({r.protein}, {r.model}, {r.method}, {r.scaler} "
+                            f"FAILED ({r.protein}, {r.model_name}, {r.method}, {r.scaler} "
                             f"exit={r.exit_code}): {r.log_path}"
                         )
             except Exception as e:
@@ -447,7 +447,7 @@ def main(args: argparse.Namespace):
         return
 
     config = GridSearchConfig(
-        model=args.model,
+        model_name=args.model,
         scalers=args.scalers.split(),
         ensemble_sizes=[int(x) for x in args.ensemble_sizes.split()],
         gradient_weights=[float(x) for x in args.gradient_weights.split()],
@@ -506,7 +506,7 @@ def generate_jobs(args: argparse.Namespace) -> list[JobConfig]:
                                     structure_path=structure,
                                     density_path=density,
                                     resolution=resolution,
-                                    model=model,
+                                    model_name=model,
                                     scaler=scaler,
                                     ensemble_size=ens,
                                     gradient_weight=gw,
@@ -533,7 +533,7 @@ def generate_jobs(args: argparse.Namespace) -> list[JobConfig]:
                                 structure_path=structure,
                                 density_path=density,
                                 resolution=resolution,
-                                model=model,
+                                model_name=model,
                                 scaler=scaler,
                                 ensemble_size=ens,
                                 gradient_weight=gw,
@@ -573,7 +573,7 @@ def save_results(
     new_run_keys = {
         (
             r.protein,
-            r.model,
+            r.model_name,
             r.method,
             r.scaler,
             r.ensemble_size,
@@ -585,17 +585,20 @@ def save_results(
 
     merged_runs = [asdict(r) for r in results]
     for existing_run in existing_runs:
+        normalized_run = existing_run.copy()
+        if "model" in normalized_run:
+            normalized_run.setdefault("model_name", normalized_run.pop("model"))
         key = (
-            existing_run.get("protein"),
-            existing_run.get("model"),
-            existing_run.get("method"),
-            existing_run.get("scaler"),
-            existing_run.get("ensemble_size"),
-            existing_run.get("gradient_weight"),
-            existing_run.get("gd_steps"),
+            normalized_run.get("protein"),
+            normalized_run.get("model_name"),
+            normalized_run.get("method"),
+            normalized_run.get("scaler"),
+            normalized_run.get("ensemble_size"),
+            normalized_run.get("gradient_weight"),
+            normalized_run.get("gd_steps"),
         )
         if key not in new_run_keys:
-            merged_runs.append(existing_run)
+            merged_runs.append(normalized_run)
 
     output = {
         "config": asdict(config),

src/sampleworks/cli/guidance.py

@@ -13,7 +13,7 @@ def main(argv: list[str] | None = None) -> int:
     device, model_wrapper = get_model_and_device(
         config.device,
         getattr(config, "model_checkpoint", None),
-        config.model,
+        config.model_name,
         method=getattr(config, "method", None),
     )
     result = run_guidance(config, config.guidance_type, model_wrapper, device)

src/sampleworks/utils/guidance_script_arguments.py

@@ -215,7 +215,7 @@ class GuidanceConfig:
     protein: str
     structure: Path | str  # actually a path to a structure file
     density: Path | str
-    model: str | StructurePredictor
+    model_name: str | StructurePredictor
     guidance_type: str | GuidanceType
     log_path: str
     output_dir: str = "output"
@@ -241,23 +241,23 @@ def add_argument(self, name: str, default: Any = None, **kwargs):
     def from_cli(
         cls,
         argv: list[str] | None = None,
-        model: str | None = None,
+        model_name: str | None = None,
         guidance_type: str | None = None,
     ) -> GuidanceConfig:
         """Parse CLI arguments and return a fully populated GuidanceConfig.
 
-        When *model* and *guidance_type* are provided (e.g. from legacy
+        When *model_name* and *guidance_type* are provided (e.g. from legacy
         scripts), they are used directly and ``--model`` / ``--guidance-type``
         are not required on the command line.  Otherwise they are parsed as
         required CLI arguments.
         """
         model_choices = [m.value for m in StructurePredictor]
         guidance_choices = [g.value for g in GuidanceType]
-        model_preset = model is not None
+        model_preset = model_name is not None
         guidance_preset = guidance_type is not None
 
-        if model_preset and model not in model_choices:
-            raise ValueError(f"Unknown model type: {model}")
+        if model_preset and model_name not in model_choices:
+            raise ValueError(f"Unknown model type: {model_name}")
         if guidance_preset and guidance_type not in guidance_choices:
             raise ValueError(f"Unknown guidance type: {guidance_type}")
 
@@ -267,6 +267,7 @@ def from_cli(
             if not model_preset:
                 pre.add_argument(
                     "--model",
+                    dest="model_name",
                     type=str,
                     required=True,
                     choices=model_choices,
@@ -281,17 +282,21 @@ def from_cli(
                     help="Guidance method",
                 )
             pre_args, _ = pre.parse_known_args(argv)
-            model = model or pre_args.model
+            model_name = model_name or pre_args.model_name
             guidance_type = guidance_type or pre_args.guidance_type
 
+        if model_name is None or guidance_type is None:
+            raise RuntimeError("CLI parsing did not resolve a model name and guidance type")
+
         # -- full parser -----------------------------------------------------
         parser = argparse.ArgumentParser(
-            description=f"Run {guidance_type} guidance with {model}",
+            description=f"Run {guidance_type} guidance with {model_name}",
         )
         parser.add_argument(
             "--model",
+            dest="model_name",
             type=str,
-            default=model,
+            default=model_name,
             choices=model_choices,
             help=argparse.SUPPRESS if model_preset else "Structure prediction model",
         )
@@ -309,14 +314,14 @@ def from_cli(
             help="Protein identifier (must match naming used in grid search / evaluation)",
         )
         add_generic_args(parser)
-        _MODEL_ARG_ADDERS[model](parser)
+        _MODEL_ARG_ADDERS[model_name](parser)
         _GUIDANCE_ARG_ADDERS[guidance_type](parser)
 
         args = parser.parse_args(argv)
 
-        if model_preset and args.model != model:
+        if model_preset and args.model_name != model_name:
             parser.error(
-                f"This script is fixed to --model {model}."
+                f"This script is fixed to --model {model_name}."
                 f" Use sampleworks-guidance for other models."
             )
         if guidance_preset and args.guidance_type != guidance_type:
@@ -329,7 +334,7 @@ def from_cli(
             protein=args.protein,
             structure=args.structure,
             density=args.density,
-            model=model,
+            model_name=model_name,
             guidance_type=guidance_type,
             log_path=getattr(args, "log_path", None) or "",
             output_dir=args.output_dir,
@@ -362,9 +367,9 @@ def __post_init__(self):
             raise ValueError(f"Unknown guidance type: {self.guidance_type}")
 
         try:
-            _MODEL_ARG_ADDERS[self.model](self)
+            _MODEL_ARG_ADDERS[self.model_name](self)
         except KeyError:
-            raise ValueError(f"Unknown model type: {self.model}")
+            raise ValueError(f"Unknown model type: {self.model_name}")
 
     def populate_config_for_guidance_type(self, job: JobConfig, args: argparse.Namespace):
         """Apply per-job grid-search values onto this guidance configuration."""
@@ -373,13 +378,13 @@ def populate_config_for_guidance_type(self, job: JobConfig, args: argparse.Names
             self.model_checkpoint = checkpoint
         elif not getattr(self, "model_checkpoint", None):
             # Auto-resolve from baked-in /checkpoints/ or legacy fallback paths
-            model_key = str(self.model).lower().replace("structurepredictor.", "")
+            model_key = str(self.model_name).lower().replace("structurepredictor.", "")
             self.model_checkpoint = _resolve_checkpoint(model_key)
 
-        if job.model == StructurePredictor.BOLTZ_2 and job.method:
+        if job.model_name == StructurePredictor.BOLTZ_2 and job.method:
             self.method = job.method
 
-        if job.model == StructurePredictor.RF3:
+        if job.model_name == StructurePredictor.RF3:
             self.disable_chiral_features = getattr(args, "disable_chiral_features", False)
             self.track_chiral_features = getattr(args, "track_chiral_features", False)
 
@@ -409,6 +414,13 @@ def as_dict(self) -> dict[str, Any]:
         output["log_path"] = _remap_container_path(str(self.log_path))
         return output
 
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Restore state while migrating legacy pickles from ``model``."""
+        migrated = state.copy()
+        if "model" in migrated:
+            migrated.setdefault("model_name", migrated.pop("model"))
+        self.__dict__.update(migrated)
+
 
 def add_generic_args(parser: argparse.ArgumentParser | GuidanceConfig):
     """Add CLI arguments shared by all models and guidance methods."""
@@ -627,7 +639,7 @@ class JobConfig:
     structure_path: Path | str
     density_path: Path | str
     resolution: float
-    model: str
+    model_name: str
     scaler: str
     ensemble_size: int
     gradient_weight: float
@@ -642,7 +654,7 @@ class JobResult:
     """Serializable status record produced after a guidance job finishes."""
 
     protein: str
-    model: str
+    model_name: str
     method: str | None
     scaler: str
     ensemble_size: int
@@ -673,3 +685,10 @@ def as_dict(self) -> dict[str, Any]:
         output["output_dir"] = _remap_container_path(str(self.output_dir))
         output["log_path"] = _remap_container_path(str(self.log_path))
         return output
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Restore state while migrating legacy pickles from ``model``."""
+        migrated = state.copy()
+        if "model" in migrated:
+            migrated.setdefault("model_name", migrated.pop("model"))
+        self.__dict__.update(migrated)

src/sampleworks/utils/guidance_script_utils.py

@@ -661,7 +661,7 @@ def get_job_result(
     end_time = epoch_seconds(ended_at)
     result = JobResult(
         protein=args.protein,
-        model=args.model,
+        model_name=args.model_name,
         method=getattr(args, "method", None),
         scaler=args.guidance_type,
         ensemble_size=args.ensemble_size,
@@ -686,7 +686,7 @@ def run_guidance_job_queue(job_queue_path: str) -> list[JobResult]:
     template_job = job_queue[0]
     if template_job.model_checkpoint is None or template_job.model_checkpoint == "":
         # Auto-resolve from baked-in /checkpoints/ or legacy fallback paths
-        model_key = str(template_job.model).lower().replace("structurepredictor.", "")
+        model_key = str(template_job.model_name).lower().replace("structurepredictor.", "")
         resolved = _resolve_checkpoint(model_key)  # will raise if not found
         template_job.model_checkpoint = resolved
         # Propagate to all jobs in the queue
@@ -697,7 +697,7 @@ def run_guidance_job_queue(job_queue_path: str) -> list[JobResult]:
     device, model_wrapper = get_model_and_device(
         str(template_job.device),
         template_job.model_checkpoint,
-        template_job.model,
+        template_job.model_name,
         method=template_job.method if hasattr(template_job, "method") else None,
     )
     job_results = []