feat(ciffiles): Write Sampleworks metadata into CIF files produced by grid search and patching script; resolves #208

Author: marcuscollins

scripts/patch_output_cif_files.py

@@ -1,5 +1,6 @@
 # utility script to put all header information from original PDB entry into our CIF files
 import fnmatch
+import json
 import re
 from argparse import ArgumentParser
 from pathlib import Path
@@ -12,6 +13,7 @@
 from biotite.structure.io.pdbx import CIFColumn, CIFFile, set_structure
 from loguru import logger
 from sampleworks.utils.atom_array_utils import remove_atoms_with_any_nan_coords
+from sampleworks.utils.cif_utils import add_category_to_cif
 
 
 SAMPLEWORKS_CACHE = Path("~/.sampleworks/rcsb").expanduser()
@@ -73,7 +75,11 @@ def parse_args():
     )
     parser.add_argument("--grid-search-input-dir", required=True)
     parser.add_argument(
-        "--input-pdb-pattern", default="{pdb_id}/{pdb_id}_single_001_density_input.cif"
+        "--input-pdb-pattern",
+        default="{pdb_id}/{pdb_id}_single_001_density_input.cif",
+        help="Pattern used by fnmatch/glob for input pdb files. The complete path of the input "
+        "pdb must match f'{grid-search-input-dir}/{input-pdb-pattern}'. Defaults to "
+        "'{pdb_id}/{pdb_id}_single_001_density_input.cif'",
     )
     args = parser.parse_args()
     return args
@@ -159,9 +165,23 @@ def patch_individual_cif_file(
     atom_keys = list(zip(asym_unit.chain_id.tolist(), asym_unit.res_id.tolist()))
     asym_unit.res_id = np.array([mapping[k] for k in atom_keys], dtype=asym_unit.res_id.dtype)
 
-    # load the actual PDB, we'll copy the new coordinates to it.
+    # load the actual PDB, we'll copy the new coordinates and metadata into it.
     template = CIFFile.read(rcsb_path)
 
+    # Write sampleworks trial metadata to the CIF file, if we can find it
+    cif_data = CIFFile.read(cif_path)
+    if "sampleworks" in cif_data.block:
+        template.block["sampleworks"] = cif_data.block["sampleworks"]
+    elif (metadata_path := cif_path.parent / "job_metadata.json").exists():
+        with open(metadata_path, "r") as fp:
+            metadata = json.load(fp)
+        if metadata is not None:
+            add_category_to_cif(template, metadata, "sampleworks")
+        else:
+            logger.warning(f"Sampleworks metadata file at {metadata_path} is empty")
+    else:
+        logger.warning(f"No sampleworks metadata found for {cif_path}")
+
     # remove any atoms with nan coordinates--these seem to come in because we sometimes use parse
     # (from AtomWorks) which creates them. Still, we'll do this here just in case.
     asym_unit = remove_atoms_with_any_nan_coords(asym_unit)

src/sampleworks/utils/cif_utils.py

@@ -297,8 +297,7 @@ def add_category_to_cif(
             raise ValueError("CIFFile has no blocks. Cannot add category.")
         elif len(blocks) > 1:
             raise ValueError(
-                f"CIFFile has multiple blocks: {blocks}. "
-                "Please specify block_name parameter."
+                f"CIFFile has multiple blocks: {blocks}. Please specify block_name parameter."
             )
         block = ciffile[blocks[0]]
     else:

src/sampleworks/utils/guidance_script_utils.py

@@ -30,7 +30,7 @@
     NoiseSpaceDPSScaler,
     NoScalingScaler,
 )
-from sampleworks.utils.cif_utils import resolve_mixed_hetatm_atom_altlocs
+from sampleworks.utils.cif_utils import add_category_to_cif, resolve_mixed_hetatm_atom_altlocs
 from sampleworks.utils.guidance_constants import (
     GuidanceType,
     StructurePredictor,
@@ -265,7 +265,7 @@ def get_reward_function_and_structure(
 
 
 def save_everything(
-    output_dir: str | Path,
+    args: GuidanceConfig | argparse.Namespace,
     losses: list[Any],
     refined_structure: dict,
     traj_denoised: list[Any],
@@ -283,8 +283,10 @@ def save_everything(
 
     Parameters
     ----------
-    output_dir : str | Path
-        Directory to write all output files into. Created if it doesn't exist.
+    args : GuidanceConfig | argparse.Namespace
+        The arguments for the guidance run. This method directly uses args.output_dir,
+        and creates that directory if it does not exist. The result of args.as_dict() is
+        written to a JSON file in the same directory, and inserted into the output CIF file.
     losses : list[Any]
         Per-step loss values (may contain ``None`` entries for unguided steps).
     refined_structure : dict
@@ -304,7 +306,7 @@ def save_everything(
         Optional model-space atom template. When provided (mismatch runs),
         this template is used for final structure and trajectory saving.
     """
-    output_dir = Path(output_dir)
+    output_dir = Path(args.output_dir)
     output_dir.mkdir(parents=True, exist_ok=True)
 
     logger.info("Saving results")
@@ -327,10 +329,20 @@ def save_everything(
     else:
         atom_array = base_atom_array
 
+    metadata = args.as_dict()
+
     final_structure = CIFFile()
     set_structure(final_structure, atom_array)
+    add_category_to_cif(final_structure, metadata, category_name="sampleworks")
     final_structure.write(str(output_dir / "refined.cif"))
 
+    # write out the job parameters to a JSON file in the same directory as the refined.cif file
+    # Even though this is technically duplicated, keep it around as a backup in case metadata
+    # is lost in some CIF transform.
+    with open(Path(output_dir) / "job_metadata.json", "w") as fp:
+        # use the GuidanceConfig's as_dict() method to avoid serializing PosixPath objects
+        json.dump(metadata, fp)
+
     # Two calls to save_trajectory, very similar, but saving different trajectories!
     save_trajectory(
         scaler_type,
@@ -565,7 +577,7 @@ def _run_guidance(
     model_atom_array = result.metadata.get("model_atom_array") if result.metadata else None
 
     save_everything(
-        args.output_dir,
+        args,
         losses,
         refined_structure,
         traj_denoised,
@@ -641,10 +653,6 @@ def run_guidance_job_queue(job_queue_path: str) -> list[JobResult]:
         logger.info(f"Running job {i + 1}/{len(job_queue)}: {job}")
 
         job_result = run_guidance(job, job.guidance_type, model_wrapper, device)
-        # write out the job parameters to a JSON file in the same directory as the refined.cif file
-        with open(Path(job_result.output_dir) / "job_metadata.json", "w") as fp:
-            # use the GuidanceConfig's as_dict() method to avoid serializing PosixPath objects
-            json.dump(job.as_dict(), fp)
 
         job_results.append(job_result)
         torch.cuda.empty_cache()  # just in case

tests/integration/test_mismatch_integration.py

@@ -21,6 +21,7 @@
 from sampleworks.utils.atom_array_utils import make_normalized_atom_id
 from sampleworks.utils.atom_reconciler import AtomReconciler
 from sampleworks.utils.frame_transforms import apply_forward_transform
+from sampleworks.utils.guidance_script_arguments import GuidanceConfig
 from sampleworks.utils.guidance_script_utils import save_everything
 
 from tests.mocks import MismatchCase, MismatchCaseWrapper
@@ -1011,8 +1012,18 @@ def test_save_with_model_template(self, tmp_path: Path):
         refined = {"asym_unit": build_test_atom_array(n_atoms=n_struct)}
         model_atom_array = build_test_atom_array(n_atoms=n_model, with_occupancy=False)
 
+        args = GuidanceConfig(
+            protein="1l63",
+            structure=Path("dummy"),
+            density=Path("dummy"),
+            model="boltz2",
+            guidance_type="pure_guidance",
+            log_path="dummy",
+            output_dir=str(tmp_path),
+        )
+
         save_everything(
-            output_dir=tmp_path,
+            args,
             losses=[0.5, 0.3],
             refined_structure=refined,
             traj_denoised=[],

tests/utils/test_cif_utils.py

@@ -1,14 +1,13 @@
 """Tests for cif_utils module."""
 
 import logging
-import tempfile
 from pathlib import Path
 
 import numpy as np
 import pytest
 from atomworks.io.utils.io_utils import load_any
 from biotite.structure import array, Atom, AtomArray, AtomArrayStack
-from biotite.structure.io.pdbx.cif import CIFFile, CIFColumn
+from biotite.structure.io.pdbx.cif import CIFColumn, CIFFile
 from sampleworks.utils.atom_array_utils import save_structure_to_cif
 from sampleworks.utils.cif_utils import add_category_to_cif, resolve_mixed_hetatm_atom_altlocs
 

tests/utils/test_guidance_script_utils.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 
 import torch
+from sampleworks.utils.guidance_script_arguments import GuidanceConfig
 from sampleworks.utils.guidance_script_utils import save_everything
 
 from tests.utils.atom_array_builders import build_test_atom_array
@@ -15,8 +16,18 @@ def test_save_everything_uses_model_atom_array_for_mismatch(tmp_path: Path):
 
     final_state = torch.zeros((1, 5, 3), dtype=torch.float32)
 
+    args = GuidanceConfig(
+        protein="1l63",
+        structure=Path("dummy"),
+        density=Path("dummy"),
+        model="boltz2",
+        guidance_type="pure_guidance",
+        log_path="dummy",
+        output_dir=str(tmp_path),
+    )
+
     save_everything(
-        output_dir=tmp_path,
+        args,
         losses=[],
         refined_structure=refined_structure,
         traj_denoised=[],