feat(eval): discover trials from job metadata

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"

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:

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,
             )
         )
 

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:

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",

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