From c1a8a23de6d54461f88c611621f32a7b907bdd25 Mon Sep 17 00:00:00 2001 From: Karson Chrispens <33336327+k-chrispens@users.noreply.github.com> Date: Tue, 21 Jul 2026 22:10:16 -0400 Subject: [PATCH] fix(eval): harden CIF output patching and RSCC parsing Strip mixed ATOM/HETATM altlocs and remap chain+residue per atom in patch_output_cif_files.py; pass add_missing_atoms=False in the RSCC grid search parse; extend .actlignore. --- .actlignore | 3 ++ scripts/eval/rscc_grid_search_script.py | 4 +-- scripts/patch_output_cif_files.py | 47 ++++++++++++++++++------- 3 files changed, 39 insertions(+), 15 deletions(-) diff --git a/.actlignore b/.actlignore index 167a601d..27dd16e2 100644 --- a/.actlignore +++ b/.actlignore @@ -13,3 +13,6 @@ release_data/ *.pt *.tar.gz *.tgz +*.tar +output/ +.jj/ diff --git a/scripts/eval/rscc_grid_search_script.py b/scripts/eval/rscc_grid_search_script.py index 0fe75712..17bb96cf 100644 --- a/scripts/eval/rscc_grid_search_script.py +++ b/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") diff --git a/scripts/patch_output_cif_files.py b/scripts/patch_output_cif_files.py index 0094df4f..d4d2850f 100644 --- a/scripts/patch_output_cif_files.py +++ b/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)