fix(evals): cif patching regression to allow protenix outputs to run on all evals, and remove hydrogens in RSCC parsing

.actlignore

@@ -13,3 +13,6 @@ release_data/
 *.pt
 *.tar.gz
 *.tgz
+*.tar
+output/
+.jj/

scripts/eval/rscc_grid_search_script.py

@@ -134,7 +134,7 @@ def process_group(
                 f"Could not find reference structure for occupancy {trials[0].altloc_occupancies}"
             )
         # parse() returns only the first altloc.
-        ref_structure = parse(ref_path, ccd_mirror_path=None)
+        ref_structure = parse(ref_path, ccd_mirror_path=None, add_missing_atoms=False)
         ref_atom_array = get_asym_unit_from_structure(ref_structure)
         ref_atom_array = remove_atoms_with_any_nan_coords(ref_atom_array)
     except (FileNotFoundError, OSError, ValueError, RuntimeError, AttributeError, TypeError) as e:
@@ -157,7 +157,7 @@ def process_group(
     # parse refined, align, and compute density once per trial.
     for trial in trials:
         try:
-            structure = parse(trial.refined_cif_path, ccd_mirror_path=None)
+            structure = parse(trial.refined_cif_path, ccd_mirror_path=None, add_missing_atoms=False)
             atom_array = get_asym_unit_from_structure(structure)
             if not hasattr(atom_array, "coord") or atom_array.coord is None:
                 raise AttributeError("AtomArray | AtomArrayStack is missing coordinates")

scripts/patch_output_cif_files.py

@@ -13,7 +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
+from sampleworks.utils.cif_utils import add_category_to_cif, resolve_mixed_hetatm_atom_altlocs
 
 
 SAMPLEWORKS_CACHE = Path("~/.sampleworks/rcsb").expanduser()
@@ -245,7 +245,14 @@ def patch_individual_cif_file(
         # fetch only downloads the file if it isn't already present.
         rcsb_path = fetch(rcsb_id, format="cif", target_path=str(SAMPLEWORKS_CACHE))
 
-        reference = load_any(reference_path)
+        # Mirror the fix from guidance_script_utils - strip mixed ATOM/HETATM altlocs
+        # at the same residue position (e.g. 6NI5/6 CYS/CSO) so the reference matches
+        # what is generated by guidance. Returns the original path if nothing to fix.
+        safe_reference_path = resolve_mixed_hetatm_atom_altlocs(reference_path)
+        reference = load_any(safe_reference_path)
+        if safe_reference_path != reference_path:
+            safe_reference_path.unlink()
+
         asym_unit = load_any(cif_file)
         asym_unit = ensure_atom_array_stack(asym_unit)
     except Exception:
@@ -269,20 +276,34 @@ def patch_individual_cif_file(
         logger.error(msg)
         return msg
 
-    # patch the residue numbers to match the original pdb
-    mapping = {}
-    for cif_key, ref_key in zip(cif_keys, ref_keys, strict=True):
-        if cif_key[0] != ref_key[0]:
-            msg = f"Chain mismatch while remapping residues for {cif_path} vs {reference_path}"
-            logger.error(msg)
-            # return msg
-            # TODO: fix chain mismatches upstream (protenix json creation needs update)
-            # this breaks multi-chain stuff for now
+    # TODO: uncomment below region after fixing chain mismatches upstream
+    # (protenix json creation needs update)
+    # which breaks the commented multi-chain handling below
 
-        mapping[cif_key] = ref_key[1]
+    # patch the residue numbers to match the original pdb
+    # mapping = {}
+    # for cif_key, ref_key in zip(cif_keys, ref_keys, strict=True):
+    #     if cif_key[0] != ref_key[0]:
+    #         msg = f"Chain mismatch while remapping residues for {cif_path} vs {reference_path}"
+    #         logger.error(msg)
+    #         # return msg
+    #         # TODO: fix chain mismatches upstream (protenix json creation needs update)
+    #         # this breaks multi-chain stuff for now
+
+    #     mapping[cif_key] = ref_key[1]
+
+    # TODO: delete the line below after uncommenting the mapping above
+    # Some models currently relabel chains and/or renumber residues, so we
+    # remap both per atom. ``mapping`` is keyed by the predicted (chain, res) and
+    # returns the reference (chain, res) at the same position
+    mapping = dict(zip(cif_keys, ref_keys, strict=True))
 
     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)
+    # TODO: uncomment after fixing chain mismatches upstream
+    # asym_unit.res_id = np.array([mapping[k] for k in atom_keys], dtype=asym_unit.res_id.dtype)
+    # TODO: delete these two lines after fixing chain mismatches upstream
+    asym_unit.chain_id = np.array([mapping[k][0] for k in atom_keys])
+    asym_unit.res_id = np.array([mapping[k][1] for k in atom_keys], dtype=asym_unit.res_id.dtype)
 
     # load the actual PDB, we'll copy the new coordinates and metadata into it.
     template = CIFFile.read(rcsb_path)