diff --git a/src/sampleworks/synthetic/synthetic_utils.py b/src/sampleworks/synthetic/synthetic_utils.py index 2b9a97db..8eba166e 100644 --- a/src/sampleworks/synthetic/synthetic_utils.py +++ b/src/sampleworks/synthetic/synthetic_utils.py @@ -2,7 +2,9 @@ import math import traceback -from collections.abc import Iterator +from collections import Counter +from collections.abc import Hashable, Iterable, Iterator +from itertools import pairwise from pathlib import Path import gemmi @@ -26,6 +28,10 @@ ) +# How many explicit duplicates to show in an error message when converting atomarray to gemmi. +MAX_REPORTED_DUPLICATES = 3 + + def resolve_parallel_jobs(device: torch.device | str, n_jobs: int) -> int: """Choose a safe job count for synthetic calculations on a device. @@ -310,28 +316,195 @@ def _resolve_altlocs_for_gemmi(atom_array: AtomArray) -> list[str]: return ["\x00" if a in BLANK_ALTLOC_IDS else a for a in atom_array.altloc_id] -def _residue_group_bounds(atom_array: AtomArray) -> Iterator[tuple[int, int]]: - """Yield the atom-index spans, one per residue. +def _check_keys_unique(keys: Iterable[Hashable], *, level: str, identity: str) -> None: + """Require the key gemmi uses to identify one hierarchy level to be unique. + + Parameters + ---------- + keys : Iterable of Hashable + Keys about the structure's hierarchy level to be checked for uniqueness. + level : str + Singular noun for what is identified (``"chain"``, ``"residue"``, ``"atom"``). Used + in the error message. + identity : str + The key's fields spelled out, so a reader can map a reported tuple positionally. Used + in the error message. + + Raises + ------ + ValueError + If any key appears more than once. + """ + # Counter keeps first-occurrence order for error message + repeats = [f"{key!r} x{count}" for key, count in Counter(keys).items() if count > 1] + if repeats: + shown = ", ".join(repeats[:MAX_REPORTED_DUPLICATES]) + if len(repeats) > MAX_REPORTED_DUPLICATES: + shown += f", ... and {len(repeats) - MAX_REPORTED_DUPLICATES} more" + raise ValueError( + f"gemmi identifies each {level} by {identity}, so duplicates would be " + f"indistinguishable: {shown}." + ) + + +def _check_residue_fields_homogeneous( + atom_array: AtomArray, residue_boundary_mask: np.ndarray +) -> None: + """Require the fields read from a residue's first atom to hold across all atoms in the + residue. + + ``res_id`` and ``chain_id`` need no check: they are the grouping key. Parameters ---------- atom_array : AtomArray - Structure whose atoms are grouped into residues. Atoms of a residue are - assumed contiguous (true for arrays loaded in file order). + Structure to check. + residue_boundary_mask : np.ndarray + ``(n_atoms - 1,)`` bool residue boundaries of ``atom_array``; element ``i`` is True + when atom ``i + 1`` starts a new residue. - Yields + Raises ------ - tuple of int + ValueError + If a span disagrees on ``res_name`` or ``hetero``. + """ + chain_id, res_id = atom_array.chain_id, atom_array.res_id + for field in ("res_name", "hetero"): + values = atom_array.get_annotation(field) + # a change with no boundary at that position is a change *inside* a span + changed = (values[1:] != values[:-1]) & ~residue_boundary_mask # (n_atoms - 1,) bool + if changed.any(): + idx = int(np.flatnonzero(changed)[0]) + 1 + raise ValueError( + f"Atoms of residue (chain {chain_id[idx]!r}, res_id {res_id[idx]}) disagree " + f"on {field}: atom {idx - 1} has {values[idx - 1]!r} but atom {idx} has " + f"{values[idx]!r}. atomarray_to_gemmi reads {field} from each residue's " + f"first atom, so the differing value would be silently dropped." + ) + + +def _check_no_repeated_atoms(atom_array: AtomArray, altlocs: list[str]) -> None: + """Require each ``(atom_name, altloc)`` pair to be unique within its residue. + + gemmi (0.6.7) identifies an atom within a residue by that pair (seqid.hpp:124-141), + so a repeat yields two indistinguishable atoms. + + Keyed on the full ``(chain_id, res_id, atom_name, altloc)`` for informative error + message. Assumes that each ``(chain_id, res_id)`` occupies exactly one span, which + should have been established by ``_check_no_repeated_residues``. + + Parameters + ---------- + atom_array : AtomArray + Structure to check. + altlocs : list of str + Per-atom altloc labels in gemmi convention. + + Raises + ------ + ValueError + If any ``(atom_name, altloc)`` pair repeats within a residue. + """ + _check_keys_unique( + zip( + atom_array.chain_id.tolist(), + atom_array.res_id.tolist(), + atom_array.atom_name.tolist(), + altlocs, + ), + level="atom", + identity="(chain_id, res_id, atom_name, altloc)", + ) + + +def _check_no_repeated_residues(atom_array: AtomArray, residue_span_start_idx: np.ndarray) -> None: + """Require each ``(chain_id, res_id)`` key to occupy exactly one contiguous block. + + Parameters + ---------- + atom_array : AtomArray + Structure to check. + residue_span_start_idx : np.ndarray + ``(n_residues,)`` int; first atom index of each residue of ``atom_array``. + + Raises + ------ + ValueError + If a ``(chain_id, res_id)`` key spans more than one residue block. + """ + residue_keys = zip( + atom_array.chain_id[residue_span_start_idx].tolist(), + atom_array.res_id[residue_span_start_idx].tolist(), + ) + _check_keys_unique(residue_keys, level="residue", identity="(chain_id, res_id)") + + +def _check_no_repeated_chains(atom_array: AtomArray) -> None: + """Require each ``chain_id`` to occupy exactly one contiguous block. + + Parameters + ---------- + atom_array : AtomArray + Structure to check. + + Raises + ------ + ValueError + If a ``chain_id`` spans more than one chain block. + """ + chain_id = atom_array.chain_id + chain_starts = np.flatnonzero(np.concatenate([[True], chain_id[1:] != chain_id[:-1]])) + _check_keys_unique(chain_id[chain_starts].tolist(), level="chain", identity="chain_id") + + +def _prepare_residue_spans( + atom_array: AtomArray, + altlocs: list[str], +) -> Iterator[tuple[int, int]]: + """Validate an atom array's residue spans and return the spans for building gemmi + Structure hierarchically. + + ``_build_gemmi_residue`` reads per-residue fields from each span's first atom, and + the chain loop in ``atomarray_to_gemmi`` assumes contiguous chains and residues. + This function checks both assumptions and raises an error if they are violated. + + Residues are keyed on ``(chain_id, res_id)``. ``ins_code`` is ignored until issue + #306 is resolved. + + Parameters + ---------- + atom_array : AtomArray + Structure to validate for conversion to gemmi. Must be non-empty. + altlocs : list of str + Per-atom altloc labels in gemmi convention. + + Returns + ------- + Iterator of tuple of int One ``(start_idx, stop_idx)`` per residue, covering the atoms ``atom_array[start_idx:stop_idx]`` that share the same ``(chain_id, res_id)``. + + Raises + ------ + ValueError + If ``atom_array`` is malformed by having duplicate atoms, residues, chains, or + atoms within a residue do not share the same per-residue fields. """ - if len(atom_array) == 0: - return chain_id, res_id = atom_array.chain_id, atom_array.res_id - # boundary shows where a new residue begins (i.e., chain or res_id changed). - boundary = (chain_id[1:] != chain_id[:-1]) | (res_id[1:] != res_id[:-1]) - start_indices = [0, *(np.flatnonzero(boundary) + 1).tolist(), len(atom_array)] - yield from zip(start_indices[:-1], start_indices[1:]) + # residue_boundary_mask[i] is True when atom i + 1 starts a new residue; atom 0 always + # marks a new start, which is why the flatnonzero below prepends True. + # (n_atoms - 1,) bool + residue_boundary_mask = (chain_id[1:] != chain_id[:-1]) | (res_id[1:] != res_id[:-1]) + # (n_residues,) int + residue_span_start_idx = np.flatnonzero(np.concatenate([[True], residue_boundary_mask])) + + _check_no_repeated_chains(atom_array) + _check_no_repeated_residues(atom_array, residue_span_start_idx) + _check_residue_fields_homogeneous(atom_array, residue_boundary_mask) + _check_no_repeated_atoms(atom_array, altlocs) + + residue_span_idx = [*residue_span_start_idx.tolist(), len(atom_array)] # append the end + return pairwise(residue_span_idx) def _build_gemmi_residue( @@ -345,7 +518,8 @@ def _build_gemmi_residue( Structure supplying per-atom annotations. start_idx : int Inclusive atom index of the residue's first atom; per-residue fields - (name, seqid, subchain) are read from this atom. + (name, seqid, subchain, het_flag) are read from this atom, which assumes the + span agrees on them -- ``_prepare_residue_spans`` enforces that upstream. stop_idx : int Exclusive atom index marking the end of the residue's atom span. altlocs : list of str @@ -412,12 +586,13 @@ def atomarray_to_gemmi( raise ValueError("Cannot convert an empty AtomArray to a gemmi.Structure.") altlocs = _resolve_altlocs_for_gemmi(atom_array) + residue_spans = _prepare_residue_spans(atom_array, altlocs) - # Group atoms into residues up front, then walk residues into chains. Contiguous - # grouping guarantees the hierarchy is well-formed by construction. + # Group atoms into residues up front, then walk residues into chains. Validated, + # contiguous grouping guarantees the hierarchy is well-formed by construction. model = gemmi.Model("1") # numeric name -> valid mmCIF pdbx_PDB_model_num current_chain: gemmi.Chain | None = None - for start_idx, stop_idx in _residue_group_bounds(atom_array): + for start_idx, stop_idx in residue_spans: chain_id = atom_array.chain_id[start_idx] if current_chain is None or chain_id != current_chain.name: if current_chain is not None: diff --git a/tests/synthetic/test_generate_synthetic_sf.py b/tests/synthetic/test_generate_synthetic_sf.py index cc73db9b..0b37627e 100644 --- a/tests/synthetic/test_generate_synthetic_sf.py +++ b/tests/synthetic/test_generate_synthetic_sf.py @@ -81,6 +81,40 @@ def _compute_fprotein(gemmi_structure: gemmi.Structure, device: torch.device) -> return assert_numpy(sfc.Fprotein_asu) +def _build_atom_array( + *, + chain_id: list[str], + res_id: list[int], + res_name: list[str], + atom_name: list[str], + hetero: list[bool] | None = None, + altloc_id: list[str] | None = None, +) -> AtomArray: + """Build a minimal AtomArray for residue-span validation tests. + + Only the fields a test varies need to be passed; coords are distinct per atom and + element/b_factor/occupancy are uniform, since none of them participate in residue + grouping or span validation. ``hetero`` defaults to biotite's all-False, and + ``altloc_id`` is left unset (not blank) when omitted, exercising the + missing-annotation path in ``_resolve_altlocs_for_gemmi``. + """ + n = len(atom_name) + arr = AtomArray(n) + arr.coord = np.arange(n * 3, dtype=np.float32).reshape(n, 3) + arr.chain_id = np.array(chain_id) + arr.res_id = np.array(res_id) + arr.res_name = np.array(res_name) + arr.atom_name = np.array(atom_name) + arr.element = np.array(["C"] * n) + if hetero is not None: + arr.hetero = np.array(hetero) + if altloc_id is not None: + arr.set_annotation("altloc_id", np.array(altloc_id)) + arr.set_annotation("b_factor", np.full(n, 20.0)) + arr.set_annotation("occupancy", np.ones(n)) + return arr + + def _dataset_with_columns(columns: dict[str, MTZDtype]) -> rs.DataSet: """Build a minimal rs.DataSet with each column cast to its MTZ dtype. @@ -94,30 +128,7 @@ def _dataset_with_columns(columns: dict[str, MTZDtype]) -> rs.DataSet: class TestAtomArrayToGemmi: - """Tests for atomarray_to_gemmi using the 6b8x structure.""" - - @pytest.fixture - def multichain_shared_resid_array(self) -> AtomArray: - """Two chains A (residues 1, 2) and B (residues 2, 3) that collide at the boundary. - - Chain A's last residue and chain B's first residue both have res_id 2 and are - adjacent in atom order, so residue grouping keyed on res_id alone would merge them - into a single residue across the chain boundary. This collision at the chain - boundary is exactly what exercises the chain_id term of the grouping predicate -- - a fixture whose res_id merely changes at the boundary (e.g. 2 -> 1) would split - correctly with or without that term and so would not guard it. - """ - n = 4 - arr = AtomArray(n) - arr.coord = np.arange(n * 3, dtype=np.float32).reshape(n, 3) - arr.chain_id = np.array(["A", "A", "B", "B"]) - arr.res_id = np.array([1, 2, 2, 3]) - arr.res_name = np.array(["ALA", "GLY", "GLY", "ALA"]) - arr.atom_name = np.array(["CA", "CA", "CA", "CA"]) - arr.element = np.array(["C", "C", "C", "C"]) - arr.set_annotation("b_factor", np.full(n, 20.0)) - arr.set_annotation("occupancy", np.ones(n)) - return arr + """Tests that the atomarray→gemmi conversion is faithful, using the 6b8x structure.""" def test_cell_matches_pdb(self, gemmi_structure_from_atomarray, stripped_gemmi): """Unit cell parameters are preserved through the biotite→gemmi conversion.""" @@ -223,16 +234,57 @@ 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_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. - """ - arr = multichain_shared_resid_array - model = atomarray_to_gemmi(arr)[0] + def test_fprotein_changes_with_occupancy(self, stripped_atom_array, stripped_gemmi, device): + """Fprotein amplitudes differ when occupancies changes from uniform to custom values.""" + altloc_info = detect_altlocs(stripped_atom_array) + + arr_uniform = assign_occupancies(stripped_atom_array, altloc_info, "uniform") + f_uniform = _compute_fprotein( + atomarray_to_gemmi(arr_uniform, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm), + device, + ) + + arr_custom = assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.2, 0.8, 0.0]) + f_custom = _compute_fprotein( + atomarray_to_gemmi(arr_custom, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm), + device, + ) + + assert not np.allclose(np.abs(f_uniform), np.abs(f_custom), atol=1e-3) + + +class TestGemmiHierarchyValidation: + """Tests that atomarray_to_gemmi rejects atom arrays whose hierarchy gemmi cannot + represent faithfully, and accepts the legitimate shapes that resemble them. + + A gemmi structure has the following hierarchy: first model is structure[0], first chain + is model[0], first residue is chain[0], first atom is residue[0]. + """ + + def test_empty_atom_array_raises(self): + """An empty AtomArray fails fast rather than yielding a chain-less structure.""" + with pytest.raises(ValueError, match="empty AtomArray"): + atomarray_to_gemmi(AtomArray(0)) + + def test_single_atom_array_converts(self): + """A one-atom array converts: the span arithmetic holds when the per-atom + difference arrays it is built from are empty.""" + arr = _build_atom_array(chain_id=["A"], res_id=[1], res_name=["ALA"], atom_name=["CA"]) + chains = list(atomarray_to_gemmi(arr)[0]) + assert [(res.seqid.num, atom.name) for res in chains[0] for atom in res] == [(1, "CA")] - chains = list(model) - expected_chain_ids = list(dict.fromkeys(arr.chain_id.tolist())) # unique, input order - assert [chain.name for chain in chains] == expected_chain_ids + def test_multichain_shared_res_ids_not_merged(self): + """Chains A (residues 1, 2) and B (residues 2, 3) share a res_id 2 but should not + be merged. + """ + arr = _build_atom_array( + chain_id=["A", "A", "B", "B"], + res_id=[1, 2, 2, 3], + res_name=["ALA", "GLY", "GLY", "ALA"], + atom_name=["CA", "CA", "CA", "CA"], + ) + chains = list(atomarray_to_gemmi(arr)[0]) + assert [chain.name for chain in chains] == ["A", "B"] for chain in chains: mask = arr.chain_id == chain.name residues = list(chain) @@ -244,10 +296,101 @@ def test_multichain_shared_res_ids_not_merged_in_gemmi(self, multichain_shared_r assert [res.name for res in residues] == arr.res_name[mask].tolist() assert [atom.name for res in residues for atom in res] == arr.atom_name[mask].tolist() - def test_empty_atom_array_raises(self): - """An empty AtomArray fails fast rather than yielding a chain-less structure.""" - with pytest.raises(ValueError, match="empty AtomArray"): - atomarray_to_gemmi(AtomArray(0)) + def test_same_atom_name_with_distinct_altlocs_is_valid(self): + """Atoms that differ only by altloc are treated as distinct.""" + arr = _build_atom_array( + chain_id=["A", "A"], + res_id=[1, 1], + res_name=["ALA", "ALA"], + atom_name=["CA", "CA"], + altloc_id=["A", "B"], + ) + chains = list(atomarray_to_gemmi(arr)[0]) + assert [chain.name for chain in chains] == ["A"] + residues = list(chains[0]) + assert len(residues) == 1 + assert [atom.altloc for atom in residues[0]] == ["A", "B"] + + def test_span_with_mixed_res_name_raises(self): + """Atoms in one residue disagreeing on res_name are rejected.""" + arr = _build_atom_array( + chain_id=["A", "A"], + res_id=[1, 1], + res_name=["ALA", "GLY"], + atom_name=["N", "CA"], + ) + with pytest.raises(ValueError, match="disagree on res_name"): + atomarray_to_gemmi(arr) + + def test_span_with_mixed_hetero_raises(self): + """Atoms in one residue disagreeing on hetero are rejected.""" + arr = _build_atom_array( + chain_id=["A", "A"], + res_id=[1, 1], + res_name=["MSE", "MSE"], + atom_name=["N", "SE"], + hetero=[False, True], + ) + with pytest.raises(ValueError, match="disagree on hetero"): + atomarray_to_gemmi(arr) + + def test_hetero_change_between_residues_is_valid(self): + """hetero may change at a residue boundary, and reaches gemmi as the het_flag that + marks the residue HETATM ('H') or ATOM ('A').""" + arr = _build_atom_array( + chain_id=["A", "A", "A"], + res_id=[1, 1, 2], + res_name=["ALA", "ALA", "MSE"], + atom_name=["N", "CA", "SE"], + hetero=[False, False, True], + ) + residues = list(atomarray_to_gemmi(arr)[0][0]) + assert [res.het_flag for res in residues] == ["A", "H"] + + @pytest.mark.parametrize( + "altloc_id", + [None, ["", "."]], + ids=["annotation_absent", "distinct_blank_altloc_ids"], + ) + def test_duplicate_atom_name_in_residue_raises(self, altloc_id): + """Atoms in a residue sharing the name and with no or blank altloc are rejected.""" + arr = _build_atom_array( + chain_id=["A", "A"], + res_id=[1, 1], + res_name=["ALA", "ALA"], + atom_name=["CA", "CA"], + altloc_id=altloc_id, + ) + with pytest.raises(ValueError, match="identifies each atom"): + atomarray_to_gemmi(arr) + + def test_noncontiguous_residue_raises(self): + """Atoms of a residue must be contiguous: a res_id that reappears after another + residue intervenes is rejected.""" + arr = _build_atom_array( + chain_id=["A", "A", "A"], + res_id=[1, 2, 1], + res_name=["ALA", "GLY", "ALA"], + atom_name=["CA", "CA", "CB"], + ) + with pytest.raises(ValueError, match="identifies each residue"): + atomarray_to_gemmi(arr) + + def test_noncontiguous_chain_raises(self): + """Atoms of a chain must be contiguous: a chain_id that reappears after another + chain intervenes is rejected.""" + arr = _build_atom_array( + chain_id=["A", "B", "A"], + res_id=[1, 2, 3], + res_name=["ALA", "GLY", "SER"], + atom_name=["CA", "CA", "CA"], + ) + with pytest.raises(ValueError, match="identifies each chain"): + atomarray_to_gemmi(arr) + + +class TestAssignOccupancies: + """Tests for assign_occupancies value handling, using the 6b8x altlocs.""" def test_occupancy_warns_on_extra_values(self, stripped_atom_array, caplog): """A warning is logged when more occupancy values are provided than there are altlocs.""" @@ -275,24 +418,6 @@ def test_occupancy_raises_on_bad_sum(self, stripped_atom_array): with pytest.raises(ValueError, match="sum to 1.0"): assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.3, 0.3, 0.3]) - def test_fprotein_changes_with_occupancy(self, stripped_atom_array, stripped_gemmi, device): - """Fprotein amplitudes differ when occupancies changes from uniform to custom values.""" - altloc_info = detect_altlocs(stripped_atom_array) - - arr_uniform = assign_occupancies(stripped_atom_array, altloc_info, "uniform") - f_uniform = _compute_fprotein( - atomarray_to_gemmi(arr_uniform, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm), - device, - ) - - arr_custom = assign_occupancies(stripped_atom_array, altloc_info, "custom", [0.2, 0.8, 0.0]) - f_custom = _compute_fprotein( - atomarray_to_gemmi(arr_custom, stripped_gemmi.cell, stripped_gemmi.spacegroup_hm), - device, - ) - - assert not np.allclose(np.abs(f_uniform), np.abs(f_custom), atol=1e-3) - class TestResolveMtzColumn: """Tests for resolve_mtz_column column-selection logic."""