fix(synthetic): drop setup_entities() from atomarray_to_gemmi so written cifs round-trip chain_info

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:

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.