From 92fdf386edf8ba196eaa9daeed7434e83d75787f Mon Sep 17 00:00:00 2001 From: Doris Mai Date: Mon, 3 Aug 2026 17:15:23 -0700 Subject: [PATCH 1/2] fix(synthetic): drop setup_entities() from atomarray_to_gemmi so written cifs round-trip chain_info --- src/sampleworks/synthetic/synthetic_utils.py | 19 +++++-- tests/synthetic/test_generate_synthetic_sf.py | 49 +++++++++++++++++++ 2 files changed, 65 insertions(+), 3 deletions(-) diff --git a/src/sampleworks/synthetic/synthetic_utils.py b/src/sampleworks/synthetic/synthetic_utils.py index 2b9a97db..738274a7 100644 --- a/src/sampleworks/synthetic/synthetic_utils.py +++ b/src/sampleworks/synthetic/synthetic_utils.py @@ -362,8 +362,9 @@ def _build_gemmi_residue( residue.name = atom_array.res_name[start_idx] residue.seqid = gemmi.SeqId(str(res_id)) # writes auth_seq_id residue.label_seq = res_id # writes label_seq_id, important for saving mmCIF - # if the subchain id is not set, gemmi's setup_entities() will set it to multi-char, - # which is rejected by SFcalculator's PDB-header step. + # writes label_asym_id; nothing else assigns it, since atomarray_to_gemmi + # deliberately skips setup_entities(). Must stay single-char -- SFcalculator's + # PDB-header step rejects the multi-char subchain ids setup_entities() invents. residue.subchain = atom_array.chain_id[start_idx] # biotite's bool `hetero` -> gemmi's single-char het_flag ('H' HETATM / 'A' ATOM) residue.het_flag = "H" if bool(atom_array.hetero[start_idx]) else "A" @@ -407,6 +408,14 @@ def atomarray_to_gemmi( ------- gemmi.Structure Structure ready to be wrapped by SFC_Torch.io.PDBParser + + Notes + ----- + No entities are assigned, so a cif written from the result has no entity block + and ``_atom_site.label_entity_id`` is ``.``. Sequences are then inferred from + ``_atom_site`` on reload, which is what model wrappers need; the cost is that + ``chain_info`` loses ``rcsb_entity`` (only Protenix reads it, and it falls back + to the chain id). """ if len(atom_array) == 0: raise ValueError("Cannot convert an empty AtomArray to a gemmi.Structure.") @@ -429,7 +438,11 @@ def atomarray_to_gemmi( structure = gemmi.Structure() structure.add_model(model) - structure.setup_entities() # SFcalculator/PDBParser expects entities assigned + # No setup_entities(): it fabricates entities with an empty full_sequence, so the + # written cif carries _entity/_entity_poly but no _entity_poly_seq. Atomworks reads + # this partial block which leads to KeyError when model wrapper accessing fields like + # `processed_entity_canonical_sequence`. When there are no entities, gemmi writes no + # entity block and atomworks infers the sequence from _atom_site instead. if unit_cell is not None: structure.cell = unit_cell if space_group is not None: diff --git a/tests/synthetic/test_generate_synthetic_sf.py b/tests/synthetic/test_generate_synthetic_sf.py index cc73db9b..609db7ce 100644 --- a/tests/synthetic/test_generate_synthetic_sf.py +++ b/tests/synthetic/test_generate_synthetic_sf.py @@ -9,9 +9,11 @@ import pytest import reciprocalspaceship as rs import torch +from atomworks import parse from atomworks.io.transforms.atom_array import remove_waters from biotite.structure import AtomArray from reciprocalspaceship.dtypes.base import MTZDtype +from sampleworks.eval.structure_utils import get_asym_unit_from_structure from sampleworks.synthetic.synthetic_utils import ( assign_occupancies, atomarray_to_gemmi, @@ -33,6 +35,15 @@ DMIN = 2.0 +# chain_info fields the model wrappers read: chain_type (all), the canonical sequence +# (Boltz/Protpardelle polymer YAML), and res_name (Boltz ligand CCD code). +CROSS_MODEL_CHAIN_INFO_KEYS = ("chain_type", "processed_entity_canonical_sequence", "res_name") + + +def _parse_at_production_kwargs(path: Path) -> dict: + """Parse a structure with the kwargs ``run_guidance`` uses.""" + return parse(path, hydrogen_policy="remove", add_missing_atoms=False, ccd_mirror_path=None) + @pytest.fixture(scope="module") def stripped_gemmi(resources_dir: Path) -> gemmi.Structure: @@ -223,6 +234,44 @@ def test_saved_structure_round_trips_annotations( np.testing.assert_allclose(loaded.b_factor, ref.b_factor, atol=1e-2) np.testing.assert_allclose(loaded.occupancy, ref.occupancy, atol=1e-2) + def test_saved_structure_round_trips_chain_info(self, resources_dir, stripped_gemmi, tmp_path): + """Test that a cif written by atomarray_to_gemmi parses back to the same chain_info + the source file does, so generated cifs can feed model featurization. + + The load_any test above covers _atom_site fidelity. This one covers the layer above + it: atomworks derives chain_info from the entity block when one is present and from + _atom_site when it is not. Emitting a partial entity block (an _entity/_entity_poly + with no _entity_poly_seq, which gemmi's setup_entities() produces) puts atomworks on + the first path with nothing to read, so chain_info comes back carrying + unprocessed_entity_canonical_sequence and every wrapper that reads + processed_entity_canonical_sequence raises KeyError. + """ + source = _parse_at_production_kwargs(resources_dir / "6b8x" / "6b8x_final.pdb") + ref = get_asym_unit_from_structure(source, 0) # parse returns a stack; take model 0 + + save_cif_path = tmp_path / "saved.cif" + gemmi_structure = atomarray_to_gemmi(ref, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm) + gemmi_structure.make_mmcif_document().write_file(str(save_cif_path)) + written = _parse_at_production_kwargs(save_cif_path) + + source_info, written_info = source["chain_info"], written["chain_info"] + assert set(written_info) == set(source_info) + for chain_id, expected in source_info.items(): + actual = written_info[chain_id] + for key in CROSS_MODEL_CHAIN_INFO_KEYS: + assert key in expected, f"source chain {chain_id} has no {key!r}" + assert key in actual, f"round-tripped chain {chain_id} lost {key!r}" + assert np.array_equal(np.asarray(actual[key]), np.asarray(expected[key])), ( + f"chain {chain_id} field {key!r} did not round-trip" + ) + + # chain_info only means something if the same atoms came back: parse's processing + # (altloc resolution, chain splitting, hydrogen policy) can drop or renumber atoms + # without disturbing the sequence. + loaded = get_asym_unit_from_structure(written, 0) + assert len(loaded) == len(ref) + assert np.array_equal(loaded.res_id, ref.res_id) + def test_multichain_shared_res_ids_not_merged_in_gemmi(self, multichain_shared_resid_array): """Test that atomarray_to_gemmi splits shared res_ids into separate residues per chain in the Gemmi Structure object. From b9278b3e9523a97d1a77a74191e902d3fe3a53cd Mon Sep 17 00:00:00 2001 From: Doris Mai Date: Sat, 15 Aug 2026 10:25:53 -0700 Subject: [PATCH 2/2] test(synthetic): assert atom identity in the chain_info round-trip --- src/sampleworks/synthetic/synthetic_utils.py | 14 ++++---- tests/synthetic/test_generate_synthetic_sf.py | 32 ++++++++++++++----- 2 files changed, 30 insertions(+), 16 deletions(-) diff --git a/src/sampleworks/synthetic/synthetic_utils.py b/src/sampleworks/synthetic/synthetic_utils.py index 738274a7..2c0c3dba 100644 --- a/src/sampleworks/synthetic/synthetic_utils.py +++ b/src/sampleworks/synthetic/synthetic_utils.py @@ -393,6 +393,12 @@ def atomarray_to_gemmi( the atom array has no ``altloc_id`` annotation (e.g. arrays reconstructed by a model wrapper), all altlocs default to blank. + No entities are assigned, so a cif written from the result has no entity block + and ``_atom_site.label_entity_id`` is ``.``. Sequences are then inferred from + ``_atom_site`` on reload, which is what model wrappers need; the cost is that + ``chain_info`` loses ``rcsb_entity`` (only Protenix reads it, and it falls back + to the chain id). + Parameters ---------- atom_array @@ -408,14 +414,6 @@ def atomarray_to_gemmi( ------- gemmi.Structure Structure ready to be wrapped by SFC_Torch.io.PDBParser - - Notes - ----- - No entities are assigned, so a cif written from the result has no entity block - and ``_atom_site.label_entity_id`` is ``.``. Sequences are then inferred from - ``_atom_site`` on reload, which is what model wrappers need; the cost is that - ``chain_info`` loses ``rcsb_entity`` (only Protenix reads it, and it falls back - to the chain id). """ if len(atom_array) == 0: raise ValueError("Cannot convert an empty AtomArray to a gemmi.Structure.") diff --git a/tests/synthetic/test_generate_synthetic_sf.py b/tests/synthetic/test_generate_synthetic_sf.py index 609db7ce..5de9e3fe 100644 --- a/tests/synthetic/test_generate_synthetic_sf.py +++ b/tests/synthetic/test_generate_synthetic_sf.py @@ -41,7 +41,19 @@ def _parse_at_production_kwargs(path: Path) -> dict: - """Parse a structure with the kwargs ``run_guidance`` uses.""" + """Parse a structure with the kwargs ``run_guidance`` uses. + + Parameters + ---------- + path : Path + Path to the structure file to parse. + + Returns + ------- + dict + The parsed Atomworks structure, keyed by ``"asym_unit"``, ``"chain_info"``, and the + rest of the parse metadata. + """ return parse(path, hydrogen_policy="remove", add_missing_atoms=False, ccd_mirror_path=None) @@ -241,13 +253,13 @@ def test_saved_structure_round_trips_chain_info(self, resources_dir, stripped_ge The load_any test above covers _atom_site fidelity. This one covers the layer above it: atomworks derives chain_info from the entity block when one is present and from _atom_site when it is not. Emitting a partial entity block (an _entity/_entity_poly - with no _entity_poly_seq, which gemmi's setup_entities() produces) puts atomworks on - the first path with nothing to read, so chain_info comes back carrying + with no _entity_poly_seq, which gemmi's setup_entities() would produce) leads to the + first path with nothing to read, so chain_info comes back carrying unprocessed_entity_canonical_sequence and every wrapper that reads processed_entity_canonical_sequence raises KeyError. """ source = _parse_at_production_kwargs(resources_dir / "6b8x" / "6b8x_final.pdb") - ref = get_asym_unit_from_structure(source, 0) # parse returns a stack; take model 0 + ref = get_asym_unit_from_structure(source, 0) # parse returns a stack; take the first model save_cif_path = tmp_path / "saved.cif" gemmi_structure = atomarray_to_gemmi(ref, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm) @@ -265,12 +277,16 @@ def test_saved_structure_round_trips_chain_info(self, resources_dir, stripped_ge f"chain {chain_id} field {key!r} did not round-trip" ) - # chain_info only means something if the same atoms came back: parse's processing - # (altloc resolution, chain splitting, hydrogen policy) can drop or renumber atoms - # without disturbing the sequence. + # The canonical sequence is per-residue, but parse's processing could drop or renumber + # atoms without disturbing the sequence. Here we check that on a per-atom level the + # identity is preserved. Other annotation fields (b_factor/occupancy/element) are checked + # in the annotations round-trip test above. loaded = get_asym_unit_from_structure(written, 0) assert len(loaded) == len(ref) - assert np.array_equal(loaded.res_id, ref.res_id) + for category in ("chain_id", "res_id", "atom_name"): + assert np.array_equal(loaded.get_annotation(category), ref.get_annotation(category)), ( + f"annotation {category!r} did not survive the parse round-trip" + ) def test_multichain_shared_res_ids_not_merged_in_gemmi(self, multichain_shared_resid_array): """Test that atomarray_to_gemmi splits shared res_ids into separate residues per chain