diff --git a/GRID_SEARCH.md b/GRID_SEARCH.md index 8b98400d..f258ebf9 100644 --- a/GRID_SEARCH.md +++ b/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) ``` diff --git a/run_grid_search.py b/run_grid_search.py index d3f77df7..4987239c 100755 --- a/run_grid_search.py +++ b/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, @@ -587,7 +587,7 @@ def save_results( for existing_run in existing_runs: key = ( existing_run.get("protein"), - existing_run.get("model"), + existing_run.get("model_name"), existing_run.get("method"), existing_run.get("scaler"), existing_run.get("ensemble_size"), diff --git a/src/sampleworks/cli/guidance.py b/src/sampleworks/cli/guidance.py index 737757b7..ef945a77 100644 --- a/src/sampleworks/cli/guidance.py +++ b/src/sampleworks/cli/guidance.py @@ -18,7 +18,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, config=config, ) result = run_guidance(config, config.guidance_type, model_wrapper, device) diff --git a/src/sampleworks/utils/guidance_script_arguments.py b/src/sampleworks/utils/guidance_script_arguments.py index c1f318cd..27f68457 100644 --- a/src/sampleworks/utils/guidance_script_arguments.py +++ b/src/sampleworks/utils/guidance_script_arguments.py @@ -216,7 +216,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" @@ -242,23 +242,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}") @@ -268,6 +268,7 @@ def from_cli( if not model_preset: pre.add_argument( "--model", + dest="model_name", type=str, required=True, choices=model_choices, @@ -282,17 +283,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", ) @@ -310,14 +315,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: @@ -330,7 +335,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, @@ -363,9 +368,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.""" @@ -374,13 +379,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) @@ -648,7 +653,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 @@ -663,7 +668,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 diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index 0f511a61..8b51f4c0 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -715,7 +715,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, @@ -740,7 +740,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 @@ -751,7 +751,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, # this is not actually the model, it's the model name, e.g. boltz2 + template_job.model_name, config=template_job, ) job_results = [] diff --git a/tests/cli/test_guidance_cli.py b/tests/cli/test_guidance_cli.py index ea415edb..a916a4da 100644 --- a/tests/cli/test_guidance_cli.py +++ b/tests/cli/test_guidance_cli.py @@ -29,7 +29,7 @@ class TestFromCliUnified: def test_boltz2_pure_guidance(self): argv = ["--model", "boltz2", "--guidance-type", "pure_guidance"] + COMMON_ARGS config = GuidanceConfig.from_cli(argv) - assert config.model == "boltz2" + assert config.model_name == "boltz2" assert config.guidance_type == "pure_guidance" assert config.protein == "1VME" assert config.structure == "test.cif" @@ -39,7 +39,7 @@ def test_boltz2_pure_guidance(self): def test_rf3_fk_steering(self): argv = ["--model", "rf3", "--guidance-type", "fk_steering"] + COMMON_ARGS config = GuidanceConfig.from_cli(argv) - assert config.model == "rf3" + assert config.model_name == "rf3" assert config.guidance_type == "fk_steering" assert hasattr(config, "num_particles") assert hasattr(config, "fk_lambda") @@ -123,10 +123,10 @@ class TestFromCliLegacyScripts: def test_preset_model_and_guidance_type(self): config = GuidanceConfig.from_cli( COMMON_ARGS, - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", ) - assert config.model == "boltz2" + assert config.model_name == "boltz2" assert config.guidance_type == "pure_guidance" assert config.protein == "1VME" @@ -134,10 +134,10 @@ def test_no_model_guidance_type_on_cli_required(self): """Legacy scripts should not require --model/--guidance-type on CLI.""" config = GuidanceConfig.from_cli( COMMON_ARGS, - model="protenix", + model_name="protenix", guidance_type="fk_steering", ) - assert config.model == "protenix" + assert config.model_name == "protenix" assert config.guidance_type == "fk_steering" @pytest.mark.parametrize("model", ["boltz1", "boltz2", "protenix", "rf3"]) @@ -145,10 +145,10 @@ def test_no_model_guidance_type_on_cli_required(self): def test_all_model_guidance_combos(self, model, guidance_type): config = GuidanceConfig.from_cli( COMMON_ARGS, - model=model, + model_name=model, guidance_type=guidance_type, ) - assert config.model == model + assert config.model_name == model assert config.guidance_type == guidance_type @@ -272,12 +272,12 @@ def test_msa_path_rejected_for_boltz1(self): def test_conflicting_model_in_preset_mode(self): argv = ["--model", "boltz2"] + COMMON_ARGS with pytest.raises(SystemExit): - GuidanceConfig.from_cli(argv, model="rf3", guidance_type="fk_steering") + GuidanceConfig.from_cli(argv, model_name="rf3", guidance_type="fk_steering") def test_conflicting_guidance_type_in_preset_mode(self): argv = ["--guidance-type", "fk_steering"] + COMMON_ARGS with pytest.raises(SystemExit): - GuidanceConfig.from_cli(argv, model="boltz1", guidance_type="pure_guidance") + GuidanceConfig.from_cli(argv, model_name="boltz1", guidance_type="pure_guidance") def test_fk_args_rejected_for_pure_guidance(self): argv = [ @@ -332,12 +332,12 @@ class TestPresetMatchingAccepted: def test_matching_model_accepted(self): argv = ["--model", "boltz1"] + COMMON_ARGS - config = GuidanceConfig.from_cli(argv, model="boltz1", guidance_type="fk_steering") - assert config.model == "boltz1" + config = GuidanceConfig.from_cli(argv, model_name="boltz1", guidance_type="fk_steering") + assert config.model_name == "boltz1" def test_matching_guidance_type_accepted(self): argv = ["--guidance-type", "pure_guidance"] + COMMON_ARGS - config = GuidanceConfig.from_cli(argv, model="boltz2", guidance_type="pure_guidance") + config = GuidanceConfig.from_cli(argv, model_name="boltz2", guidance_type="pure_guidance") assert config.guidance_type == "pure_guidance" @@ -583,11 +583,11 @@ def test_help_exits_zero(self): def test_invalid_preset_model_raises(self): with pytest.raises(ValueError, match="Unknown model type"): - GuidanceConfig.from_cli(COMMON_ARGS, model="typo", guidance_type="fk_steering") + GuidanceConfig.from_cli(COMMON_ARGS, model_name="typo", guidance_type="fk_steering") def test_invalid_preset_guidance_type_raises(self): with pytest.raises(ValueError, match="Unknown guidance type"): - GuidanceConfig.from_cli(COMMON_ARGS, model="boltz1", guidance_type="typo") + GuidanceConfig.from_cli(COMMON_ARGS, model_name="boltz1", guidance_type="typo") class TestAlignmentReverseDiffusion: diff --git a/tests/integration/test_mismatch_integration.py b/tests/integration/test_mismatch_integration.py index 01393d4b..9daf8f7b 100644 --- a/tests/integration/test_mismatch_integration.py +++ b/tests/integration/test_mismatch_integration.py @@ -1016,7 +1016,7 @@ def test_save_with_model_template(self, tmp_path: Path): protein="1l63", structure=Path("dummy"), density=Path("dummy"), - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", log_path="dummy", output_dir=str(tmp_path), diff --git a/tests/utils/conftest.py b/tests/utils/conftest.py index d626d2fb..eb86734c 100644 --- a/tests/utils/conftest.py +++ b/tests/utils/conftest.py @@ -15,7 +15,7 @@ def guidance_job_result(tmp_path: Path) -> JobResult: """Successful JobResult populated with paths under ``tmp_path``.""" return JobResult( protein="1l63", - model="boltz2", + model_name="boltz2", method=None, scaler="pure_guidance", ensemble_size=8, diff --git a/tests/utils/test_guidance_script_arguments.py b/tests/utils/test_guidance_script_arguments.py index f18d5aa1..7b3c9258 100644 --- a/tests/utils/test_guidance_script_arguments.py +++ b/tests/utils/test_guidance_script_arguments.py @@ -54,7 +54,7 @@ def _build_job(model: StructurePredictor) -> JobConfig: structure_path="/tmp/structure.cif", density_path="/tmp/density.mrc", resolution=2.0, - model=model, + model_name=model, scaler=GuidanceType.PURE_GUIDANCE, ensemble_size=1, gradient_weight=0.1, @@ -75,7 +75,7 @@ def test_populate_config_resolves_checkpoint_when_none_provided(_mock_resolve, m protein="protein", structure="/tmp/structure.cif", density="/tmp/density.mrc", - model=model_wrapper_type, + model_name=model_wrapper_type, guidance_type=GuidanceType.PURE_GUIDANCE, log_path="/tmp/output/run.log", ) @@ -98,7 +98,7 @@ def test_populate_config_uses_model_checkpoint_argument(model_wrapper_type): protein="protein", structure="/tmp/structure.cif", density="/tmp/density.mrc", - model=model_wrapper_type, + model_name=model_wrapper_type, guidance_type=GuidanceType.PURE_GUIDANCE, log_path="/tmp/output/run.log", ) @@ -272,7 +272,7 @@ def test_as_dict_remaps_all_four_path_fields(monkeypatch): protein="1abc", structure="/data/inputs/structures/1abc.cif", density="/data/inputs/maps/1abc.ccp4", - model=StructurePredictor.BOLTZ_2, + model_name=StructurePredictor.BOLTZ_2, guidance_type=GuidanceType.PURE_GUIDANCE, log_path="/data/results/1abc/boltz2/run.log", output_dir="/data/results/1abc/boltz2", @@ -297,7 +297,7 @@ def test_as_dict_unchanged_when_no_env_vars(monkeypatch): protein="1abc", structure="/data/inputs/structures/1abc.cif", density="/data/inputs/maps/1abc.ccp4", - model=StructurePredictor.BOLTZ_2, + model_name=StructurePredictor.BOLTZ_2, guidance_type=GuidanceType.PURE_GUIDANCE, log_path="/data/results/1abc/run.log", output_dir="/data/results/1abc", @@ -317,7 +317,7 @@ def test_as_dict_remaps_with_catchall_host_dir(monkeypatch): protein="1abc", structure="/data/inputs/structures/1abc.cif", density="/data/inputs/maps/1abc.ccp4", - model=StructurePredictor.BOLTZ_2, + model_name=StructurePredictor.BOLTZ_2, guidance_type=GuidanceType.PURE_GUIDANCE, log_path="/data/results/1abc/run.log", output_dir="/data/results/1abc", diff --git a/tests/utils/test_guidance_script_utils.py b/tests/utils/test_guidance_script_utils.py index a91a0fa5..f09176c0 100644 --- a/tests/utils/test_guidance_script_utils.py +++ b/tests/utils/test_guidance_script_utils.py @@ -44,7 +44,7 @@ def test_save_everything_uses_model_atom_array_for_mismatch(tmp_path: Path): protein="1l63", structure=Path("dummy"), density=Path("dummy"), - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", log_path="dummy", output_dir=str(tmp_path), @@ -122,7 +122,7 @@ def test_write_job_metadata_with_job_result_appends_timing_and_status( protein="1l63", structure=Path("dummy"), density=Path("dummy"), - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", log_path="dummy", output_dir=str(tmp_path), @@ -134,6 +134,8 @@ def test_write_job_metadata_with_job_result_appends_timing_and_status( # GuidanceConfig keys are preserved assert metadata["protein"] == "1l63" assert metadata["guidance_type"] == "pure_guidance" + assert metadata["model_name"] == "boltz2" + assert "model" not in metadata # JobResult-only keys are appended assert metadata["started_at"] == "2026-05-05T10:00:00" assert metadata["finished_at"] == "2026-05-05T10:00:12.340000" @@ -151,7 +153,7 @@ def test_write_job_metadata_creates_missing_output_dir( protein="1l63", structure=Path("dummy"), density=Path("dummy"), - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", log_path="dummy", output_dir=str(nested), @@ -182,14 +184,14 @@ def test_write_job_metadata_remaps_job_result_paths_to_host( protein="1l63", structure=Path("dummy"), density=Path("dummy"), - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", log_path=container_log, output_dir=container_output, ) job_result = JobResult( protein="1l63", - model="boltz2", + model_name="boltz2", method=None, scaler="pure_guidance", ensemble_size=8,