diff --git a/scripts/eval/rscc_grid_search_script.py b/scripts/eval/rscc_grid_search_script.py index 4b13179d..0fe75712 100644 --- a/scripts/eval/rscc_grid_search_script.py +++ b/scripts/eval/rscc_grid_search_script.py @@ -113,7 +113,10 @@ def process_group( try: # Load base map for canonical unit cell, # don't overwrite the base map with selection map--we'll use the full map later too. - base_xmap = protein_config.load_map(base_map_path) + base_xmap = protein_config.load_map( + base_map_path, + resolution=trials[0].resolution, + ) if base_xmap is None: raise ValueError(f"Failed to load base map from {base_map_path}") @@ -123,7 +126,9 @@ def process_group( # Load the reference structure (used to align refined structures so the calculated # maps line up with the base map, for a correct RSCC calculation). - ref_path = protein_config.get_reference_structure_path(trials[0].altloc_occupancies) + ref_path = trials[0].input_structure_path or protein_config.get_reference_structure_path( + trials[0].altloc_occupancies + ) if ref_path is None: raise ValueError( f"Could not find reference structure for occupancy {trials[0].altloc_occupancies}" @@ -294,7 +299,7 @@ def main(args: argparse.Namespace): # Sort so all trials sharing a (protein, occ_key) are contiguous, then build groups. # Resolve protein name once per group and slice ref_coords for each protein. groups: list[tuple[str, list[Trial], Path, dict[str, np.ndarray]]] = [] - group_index: dict[tuple[str, OccKey], int] = {} + group_index: dict[tuple[str, OccKey, Path | None], int] = {} for trial in sorted(all_trials, key=lambda t: (t.protein, t.occ_key)): if trial.protein in protein_configs: protein = trial.protein @@ -303,11 +308,13 @@ def main(args: argparse.Namespace): else: logger.warning(f"Skipping protein with no configuration: {trial.protein}") continue - key = (protein, trial.occ_key) + key = (protein, trial.occ_key, trial.density_path) idx = group_index.get(key) if idx is None: protein_config = protein_configs[protein] - base_map_path = protein_config.get_base_map_path_for_occupancy(trial.altloc_occupancies) + base_map_path = trial.density_path or protein_config.get_base_map_path_for_occupancy( + trial.altloc_occupancies + ) if base_map_path is None: logger.warning( f"Skipping group {protein}/{trial.altloc_occupancies}: base map not found" diff --git a/src/sampleworks/eval/eval_dataclasses.py b/src/sampleworks/eval/eval_dataclasses.py index 349d83ac..6353ecc9 100644 --- a/src/sampleworks/eval/eval_dataclasses.py +++ b/src/sampleworks/eval/eval_dataclasses.py @@ -21,6 +21,9 @@ class Trial: trial_dir: Path refined_cif_path: Path protein_dir_name: str + input_structure_path: Path | None = None + density_path: Path | None = None + resolution: float | None = None rscc: float = np.nan # these last three are placeholders for RSCC calculations. base_map_path: Path | None = None error: Exception | None = None @@ -88,9 +91,15 @@ def get_base_map_path_for_occupancy(self, altloc_occupancies: dict[str, float]) return None def load_map( - self, map_path: Path, canonical_unit_cell=True, selection_coords=None, padding=0.0 + self, + map_path: Path, + canonical_unit_cell=True, + selection_coords=None, + padding=0.0, + resolution: float | None = None, ) -> XMap | None: - xmap = XMap.fromfile(str(map_path), resolution=self.resolution) + """Load a map using recorded trial resolution when available.""" + xmap = XMap.fromfile(str(map_path), resolution=resolution or self.resolution) if canonical_unit_cell: xmap = xmap.canonical_unit_cell() if selection_coords is not None: diff --git a/src/sampleworks/eval/grid_search_eval_utils.py b/src/sampleworks/eval/grid_search_eval_utils.py index 9c917981..98c00832 100644 --- a/src/sampleworks/eval/grid_search_eval_utils.py +++ b/src/sampleworks/eval/grid_search_eval_utils.py @@ -4,6 +4,7 @@ """ import argparse +import json import re import sys from importlib.resources import files @@ -17,6 +18,117 @@ from sampleworks.utils.guidance_constants import StructurePredictor +def _metadata_value( + metadata: dict[str, object], + key: str, + fallback: object, + aliases: tuple[str, ...] = (), +) -> object: + """Return the first non-empty metadata value, otherwise a fallback. + + Parameters + ---------- + metadata + Parsed job metadata. + key + Preferred field name. + fallback + Value used when the preferred field and aliases are absent or empty. + aliases + Older or alternate field names checked after ``key``. + + Returns + ------- + object + Resolved metadata value or ``fallback``. + """ + for field_name in (key, *aliases): + value = metadata.get(field_name) + if value is not None and value != "": + return value + return fallback + + +def _metadata_float( + metadata: dict[str, object], + key: str, + fallback: int | float | None, + aliases: tuple[str, ...] = (), +) -> float | None: + """Resolve a metadata field as a float, falling back when invalid.""" + value = _metadata_value(metadata, key, fallback, aliases) + if value is None: + return None + try: + return float(str(value)) + except ValueError: + logger.warning(f"Invalid numeric metadata {key}={value!r}; using {fallback!r}") + return float(fallback) if fallback is not None else None + + +def _metadata_int( + metadata: dict[str, object], + key: str, + fallback: int | None, + aliases: tuple[str, ...] = (), +) -> int | None: + """Resolve a metadata field as an integer, falling back when invalid.""" + value = _metadata_value(metadata, key, fallback, aliases) + if value is None: + return None + try: + return int(float(str(value))) + except ValueError: + logger.warning(f"Invalid integer metadata {key}={value!r}; using {fallback!r}") + return fallback + + +def _metadata_occupancies( + metadata: dict[str, object], + fallback: dict[str, float], +) -> dict[str, float]: + """Resolve explicit altloc occupancies, falling back to path metadata.""" + value = _metadata_value(metadata, "altloc_occupancies", fallback) + if not isinstance(value, dict): + logger.warning("altloc_occupancies metadata must be a JSON object; using path fallback") + return fallback + try: + return {str(label).upper(): float(str(occupancy)) for label, occupancy in value.items()} + except ValueError: + logger.warning( + "altloc_occupancies metadata contains a non-numeric value; using path fallback" + ) + return fallback + + +def load_job_metadata(trial_dir: Path) -> dict[str, object] | None: + """Load a trial's ``job_metadata.json`` when it is present and valid. + + Parameters + ---------- + trial_dir + Directory containing one grid-search trial. + + Returns + ------- + dict[str, object] | None + Parsed metadata, or ``None`` when no usable metadata file exists. + """ + metadata_path = trial_dir / "job_metadata.json" + if not metadata_path.is_file(): + return None + try: + with open(metadata_path) as handle: + metadata = json.load(handle) + except (OSError, json.JSONDecodeError) as exc: + logger.warning(f"Could not load trial metadata from {metadata_path}: {exc}") + return None + if not isinstance(metadata, dict): + logger.warning(f"Trial metadata must be a JSON object: {metadata_path}") + return None + return metadata + + def resolve_cif_path(row: pd.Series, cif_root: Path | None) -> Path: """Resolve a CIF path from a row, preferring ``structure`` then ``structure_pattern``. @@ -149,28 +261,54 @@ def scan_grid_search_results( # Check if we found a refined.cif file in the current directory refined_cif = current_directory / target_filename if current_depth == target_depth and refined_cif.exists(): - # Reconstruct metadata from path structure + metadata = load_job_metadata(current_directory) or {} + + # Retain path parsing as a compatibility fallback for historical runs + # that predate job_metadata.json. # Expected structure: .../protein_dir/model_dir/scaler_dir/trial_dir/refined.cif trial_dir = current_directory scaler_dir = trial_dir.parent model_dir = scaler_dir.parent protein_dir = model_dir.parent - protein, altloc_occupancies = extract_protein_and_occupancy(protein_dir.name) - method, model = get_method_and_model_name(model_dir.name) + path_protein, path_altloc_occupancies = extract_protein_and_occupancy(protein_dir.name) + method, path_model = get_method_and_model_name(model_dir.name) + protein_value = _metadata_value(metadata, "protein", path_protein) + protein = str(protein_value) if protein_value is not None else None + altloc_occupancies = _metadata_occupancies(metadata, path_altloc_occupancies) + model_value = _metadata_value(metadata, "model_name", path_model, aliases=("model",)) + model = str(model_value) + method_value = _metadata_value(metadata, "method", method) + method = str(method_value) if method_value is not None else None + scaler = str(_metadata_value(metadata, "guidance_type", scaler_dir.name)) params = parse_trial_dir(trial_dir) - guidance_weight = None - if params["guidance_weight"] is not None: - guidance_weight = float(params["guidance_weight"]) - gd_steps = int(params["gd_steps"]) if params["gd_steps"] is not None else None + guidance_weight = _metadata_float( + metadata, + "guidance_weight", + params["guidance_weight"], + aliases=("step_size",), + ) + gd_steps = _metadata_int( + metadata, + "num_gd_steps", + int(params["gd_steps"]) if params["gd_steps"] is not None else None, + ) + ensemble_size = _metadata_int( + metadata, + "ensemble_size", + int(params["ensemble_size"]) if params["ensemble_size"] is not None else None, + ) + input_structure_value = _metadata_value(metadata, "structure", "") + density_value = _metadata_value(metadata, "density", "") + resolution = _metadata_float(metadata, "resolution", None) # Validate parameters to satisfy ty if ( protein is None or not altloc_occupancies or (model == StructurePredictor.BOLTZ_2 and method is None) - or params["ensemble_size"] is None + or ensemble_size is None or (guidance_weight is None and gd_steps is None) ): logger.warning(f"Skipping trial in {trial_dir} due to missing metadata") @@ -182,13 +320,18 @@ def scan_grid_search_results( altloc_occupancies=altloc_occupancies, model=model, method=method, - scaler=scaler_dir.name, - ensemble_size=int(params["ensemble_size"]), + scaler=scaler, + ensemble_size=ensemble_size, guidance_weight=guidance_weight, gd_steps=gd_steps, trial_dir=trial_dir, refined_cif_path=refined_cif, protein_dir_name=protein_dir.name, + input_structure_path=( + Path(str(input_structure_value)) if input_structure_value else None + ), + density_path=Path(str(density_value)) if density_value else None, + resolution=resolution, ) ) diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index f477ab73..a575951a 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -29,6 +29,7 @@ NoiseSpaceDPSScaler, NoScalingScaler, ) +from sampleworks.eval.occupancy_utils import extract_protein_and_occupancy from sampleworks.utils.cif_utils import add_category_to_cif, resolve_mixed_hetatm_atom_altlocs from sampleworks.utils.guidance_constants import ( GuidanceType, @@ -632,6 +633,9 @@ def _write_job_metadata( """ metadata = args.as_dict() metadata.update(job_result.as_dict()) + _, altloc_occupancies = extract_protein_and_occupancy(str(args.protein)) + if altloc_occupancies: + metadata["altloc_occupancies"] = altloc_occupancies output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) with open(output_dir / "job_metadata.json", "w") as fp: diff --git a/tests/eval/test_eval_dataclasses.py b/tests/eval/test_eval_dataclasses.py index bcb90127..cd039378 100644 --- a/tests/eval/test_eval_dataclasses.py +++ b/tests/eval/test_eval_dataclasses.py @@ -79,6 +79,9 @@ def test_dict_copy_contains_all_fields(self, trial: Trial): "trial_dir", "refined_cif_path", "protein_dir_name", + "input_structure_path", + "density_path", + "resolution", "rscc", "base_map_path", "error", diff --git a/tests/eval/test_grid_search_trial_metadata.py b/tests/eval/test_grid_search_trial_metadata.py new file mode 100644 index 00000000..cb92f301 --- /dev/null +++ b/tests/eval/test_grid_search_trial_metadata.py @@ -0,0 +1,66 @@ +"""Tests for metadata-backed grid-search trial discovery.""" + +import json +from pathlib import Path + +from sampleworks.eval.grid_search_eval_utils import load_job_metadata, scan_grid_search_results + + +def test_scan_grid_search_results_prefers_job_metadata(tmp_path) -> None: + """Trial identity and input paths come from recorded job metadata.""" + trial_dir = tmp_path / "1ABC_0.5occA_0.5occB" / "wrong_model" / "wrong_scaler" / "ens1_gw0.1" + trial_dir.mkdir(parents=True) + (trial_dir / "refined.cif").write_text("data_test") + metadata = { + "protein": "1ABC", + "model_name": "boltz2", + "method": "MD", + "guidance_type": "pure_guidance", + "ensemble_size": 8, + "step_size": 0.25, + "altloc_occupancies": {"A": 0.25, "B": 0.75}, + "structure": "/inputs/1abc.cif", + "density": "/inputs/1abc.ccp4", + "resolution": 1.8, + } + (trial_dir / "job_metadata.json").write_text(json.dumps(metadata)) + + trials = scan_grid_search_results(trial_dir, current_depth=4, target_depth=4) + + assert len(trials) == 1 + trial = trials[0] + assert trial.protein == "1ABC" + assert trial.model == "boltz2" + assert trial.method == "MD" + assert trial.scaler == "pure_guidance" + assert trial.ensemble_size == 8 + assert trial.guidance_weight == 0.25 + assert trial.altloc_occupancies == {"A": 0.25, "B": 0.75} + assert trial.input_structure_path == Path("/inputs/1abc.cif") + assert trial.density_path == Path("/inputs/1abc.ccp4") + assert trial.resolution == 1.8 + + +def test_load_job_metadata_rejects_non_object_json(tmp_path) -> None: + """Metadata arrays are ignored rather than breaking trial discovery.""" + (tmp_path / "job_metadata.json").write_text("[]") + + assert load_job_metadata(tmp_path) is None + + +def test_scan_uses_path_fallback_for_empty_metadata_values(tmp_path) -> None: + """Null or empty metadata fields do not replace usable path metadata.""" + trial_dir = tmp_path / "1ABC_1.0occA" / "boltz2_MD" / "pure_guidance" / "ens8_gw0.1" + trial_dir.mkdir(parents=True) + (trial_dir / "refined.cif").write_text("data_test") + (trial_dir / "job_metadata.json").write_text( + json.dumps({"protein": None, "model": None, "method": "", "ensemble_size": "8.0"}) + ) + + trials = scan_grid_search_results(trial_dir, current_depth=4, target_depth=4) + + assert len(trials) == 1 + assert trials[0].protein == "1abc" + assert trials[0].model == "boltz2" + assert trials[0].method == "MD" + assert trials[0].ensemble_size == 8 diff --git a/tests/utils/test_guidance_script_utils.py b/tests/utils/test_guidance_script_utils.py index a91a0fa5..c8a2ec25 100644 --- a/tests/utils/test_guidance_script_utils.py +++ b/tests/utils/test_guidance_script_utils.py @@ -162,6 +162,26 @@ def test_write_job_metadata_creates_missing_output_dir( assert (nested / "job_metadata.json").exists() +def test_write_job_metadata_records_altloc_occupancies( + tmp_path: Path, guidance_job_result: JobResult +): + """Metadata stores occupancies explicitly instead of relying on directory names.""" + args = GuidanceConfig( + protein="1l63_0.25occA_0.75occB", + structure=Path("dummy"), + density=Path("dummy"), + model="boltz2", + guidance_type="pure_guidance", + log_path="dummy", + output_dir=str(tmp_path), + ) + + _write_job_metadata(tmp_path, args, guidance_job_result) + + metadata = json.loads((tmp_path / "job_metadata.json").read_text()) + assert metadata["altloc_occupancies"] == {"A": 0.25, "B": 0.75} + + def test_write_job_metadata_remaps_job_result_paths_to_host( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ):