From 4e463e731598a46e4e92a673ac0f55b733cccfc3 Mon Sep 17 00:00:00 2001 From: Ben Tudor Price Date: Wed, 29 Jul 2026 20:28:38 -0700 Subject: [PATCH] fix(ml): don't require CCD leaving groups when substituting residues In _map_to_closest_canonical_residue we use a template canonical residue to check if an unwanted residue might be substituted for a canonical. The template canonical residue includes leaving groups, causing it to always fail for residues in a polymer when checking if it subsets the atoms of the residue to be substituted. This fix removes leaving groups from the comparison logic when required and adds regression tests demonstrating both the in-chain and chain-terminus cases. --- src/atomworks/ml/transforms/filters.py | 24 +++++- .../test_undesired_res_token_leaving_atoms.py | 82 +++++++++++++++++++ 2 files changed, 104 insertions(+), 2 deletions(-) create mode 100644 tests/ml/transforms/test_undesired_res_token_leaving_atoms.py diff --git a/src/atomworks/ml/transforms/filters.py b/src/atomworks/ml/transforms/filters.py index b0451652..3e6d0749 100644 --- a/src/atomworks/ml/transforms/filters.py +++ b/src/atomworks/ml/transforms/filters.py @@ -13,6 +13,7 @@ from atomworks.common import exists, not_isin from atomworks.constants import HYDROGEN_LIKE_SYMBOLS from atomworks.enums import ChainType, ChainTypeInfo +from atomworks.io.utils.ccd import get_chem_comp_leaving_atom_names from atomworks.io.utils.query import QueryExpression from atomworks.io.utils.selection import get_annotation from atomworks.io.utils.sequence import get_1_from_3_letter_code, get_3_from_1_letter_code @@ -400,6 +401,19 @@ def __init__(self, min_occupancy: float = 0.5): super().__init__(func=lambda arr: remove_unresolved_atoms(arr, min_occupancy)) +def _leaving_atom_names(res_name: str) -> frozenset[str]: + """Atom names the CCD declares as displaced when `res_name` forms an inter-residue bond. + + `OXT`/`HXT` for amino acids, `OP3`/`HOP3` for nucleotides. A residue inside a polymer has + already lost these during parsing (see `atomworks.io.template`), so they must not be + required when matching an observed residue against a free-monomer CCD template. + """ + names: set[str] = set() + for displaced_atom_names in get_chem_comp_leaving_atom_names(res_name).values(): + names |= {str(name) for name in displaced_atom_names} + return frozenset(names) + + class HandleUndesiredResTokens(Transform): """ Remove, or otherwise handle, undesired residue tokens from the AtomArray. @@ -458,9 +472,15 @@ def _map_to_closest_canonical_residue( if not has_hydrogens: canonical_res = canonical_res[not_isin(canonical_res.element, HYDROGEN_LIKE_SYMBOLS)] - # If canonical residue is a strict subset of the original residue, + # Leaving atoms are absent from every residue inside a polymer, so they must not + # be required else no residue in a chain can ever match. + required_atom_names = canonical_res.atom_name[ + not_isin(canonical_res.atom_name, list(_leaving_atom_names(canonical_res_name))) + ] + + # If the canonical residue is a subset of the original residue, # keep all matching atom names and delete the rest - if np.all(np.isin(canonical_res.atom_name, atom_name)): + if np.all(np.isin(required_atom_names, atom_name)): to_keep = np.isin(atom_name, canonical_res.atom_name) # ... if we match without `force_unknown` break loop early return to_keep, canonical_res_name diff --git a/tests/ml/transforms/test_undesired_res_token_leaving_atoms.py b/tests/ml/transforms/test_undesired_res_token_leaving_atoms.py new file mode 100644 index 00000000..a2664fbf --- /dev/null +++ b/tests/ml/transforms/test_undesired_res_token_leaving_atoms.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import biotite.structure as struc +import numpy as np +import pytest +from biotite.structure import AtomArray + +from atomworks.enums import ChainType +from atomworks.io.utils.ccd import get_chem_comp_leaving_atom_names +from atomworks.ml.transforms.filters import HandleUndesiredResTokens + +# Non-canonical residues found inside protein chains whose canonical parent is unambiguous +# once the leaving group is disregarded. +MAPPABLE_RESIDUES = [ + pytest.param("ABA", "ALA", id="ABA-to-ALA"), + pytest.param("SAR", "GLY", id="SAR-to-GLY"), + pytest.param("PTR", "TYR", id="PTR-to-TYR"), + pytest.param("SEP", "SER", id="SEP-to-SER"), +] + + +def _make_polymer_residue(res_name: str, *, keep_leaving_group: bool) -> AtomArray: + """Build a single protein polymer residue with the annotations the transform requires. + + Args: + res_name: CCD component id, e.g. ``"PTR"``. + keep_leaving_group: If ``False``, drop the CCD-declared leaving atoms, i.e. the state + of a residue inside a chain. If ``True``, keep them, i.e. a chain terminus. + + Returns: + AtomArray of shape ``[n_atoms]``, annotated with `is_polymer`, `pn_unit_iid`, + `chain_type` and `atomize`. + """ + residue = struc.info.residue(res_name) + residue = residue[residue.element != "H"] + + if not keep_leaving_group: + leaving = {str(name) for names in get_chem_comp_leaving_atom_names(res_name).values() for name in names} + residue = residue[~np.isin(residue.atom_name, list(leaving))] + + n_atoms = residue.array_length() + residue.res_id = np.full(n_atoms, 5) + residue.set_annotation("is_polymer", np.ones(n_atoms, dtype=bool)) + residue.set_annotation("pn_unit_iid", np.full(n_atoms, -1, dtype=int)) + residue.set_annotation("chain_type", np.full(n_atoms, int(ChainType.POLYPEPTIDE_L))) + residue.set_annotation("atomize", np.zeros(n_atoms, dtype=bool)) + return residue + + +def _apply(residue: AtomArray, res_name: str) -> AtomArray: + transform = HandleUndesiredResTokens(undesired_res_tokens=[res_name]) + return transform.forward({"atom_array": residue})["atom_array"] + + +@pytest.mark.parametrize(("res_name", "expected_canonical"), MAPPABLE_RESIDUES) +def test_residue_inside_a_polymer_maps_to_closest_canonical(res_name: str, expected_canonical: str): + """A residue in a chain is substituted even though it has lost its leaving group. + + Before the fix these were atomized, because the canonical template demanded an `OXT` + that polymerisation had already removed. + """ + result = _apply(_make_polymer_residue(res_name, keep_leaving_group=False), res_name) + + assert set(map(str, result.res_name)) == { + expected_canonical + }, f"{res_name} should map to {expected_canonical}, got {sorted(set(map(str, result.res_name)))}" + assert not result.atomize.any(), f"{res_name} should not be atomized after a successful substitution" + + +@pytest.mark.parametrize(("res_name", "expected_canonical"), MAPPABLE_RESIDUES) +def test_chain_terminus_retains_its_leaving_group(res_name: str, expected_canonical: str): + """Leaving groups are not *required*, but they are still *kept* when present. + + A chain terminus has a real, observed `OXT`; dropping it would discard experimental + density. Only the required-atom set excludes leaving groups; the kept-atom set does not. + """ + result = _apply(_make_polymer_residue(res_name, keep_leaving_group=True), res_name) + + assert set(map(str, result.res_name)) == {expected_canonical} + assert "OXT" in set( + map(str, result.atom_name) + ), f"{res_name} at a chain terminus should keep its OXT, got {sorted(set(map(str, result.atom_name)))}"