diff --git a/src/sampleworks/synthetic/synthetic_utils.py b/src/sampleworks/synthetic/synthetic_utils.py index 2b9a97db..2c0c3dba 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" @@ -392,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 @@ -429,7 +436,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..5de9e3fe 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,27 @@ 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. + + 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) + @pytest.fixture(scope="module") def stripped_gemmi(resources_dir: Path) -> gemmi.Structure: @@ -223,6 +246,48 @@ 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() 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 the first model + + 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" + ) + + # 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) + 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 in the Gemmi Structure object.