Add C18/C19 steroid labels, tests, and a small demo

Author: kzubko

chebi_utils/extract_properties.py

@@ -238,12 +238,41 @@ def get_rings(mol: Chem.Mol) -> dict[str, list]:
     return atom_extensions
 
 
+def _add_angular_methyls(
+    mol: Chem.Mol,
+    atom_extensions: dict[str, list],
+    iupac_to_atom: dict[int, int],
+) -> None:
+    """Label angular methyls C18 (on C13) and C19 (on C10) when present.
+
+    Looks for the unique carbon neighbour of the attachment atom that is not
+    part of the gonane core (C1–C17). Estrogens without C19 simply yield no
+    ``steroid_19`` predicate.
+    """
+    core_atom_indices = set(iupac_to_atom.values())
+    for attachment_position, methyl_position in ((13, 18), (10, 19)):
+        attachment_atom_idx = iupac_to_atom.get(attachment_position)
+        if attachment_atom_idx is None:
+            continue
+        methyl_candidates = [
+            neighbor.GetIdx()
+            for neighbor in mol.GetAtomWithIdx(attachment_atom_idx).GetNeighbors()
+            if neighbor.GetAtomicNum() == 6 and neighbor.GetIdx() not in core_atom_indices
+        ]
+        if len(methyl_candidates) == 1:
+            atom_extensions.setdefault(f"steroid_{methyl_position}", []).append(
+                methyl_candidates[0]
+            )
+
+
 def get_steroid_positions(mol: Chem.Mol) -> dict[str, list]:
     """Extract steroid-nucleus position predicates.
 
     Matches the molecule against the gonane core and, on a match, labels the
     ring atoms with their IUPAC steroid position as predicates ``steroid_1`` …
-    ``steroid_17``. Molecules without a gonane core yield no predicates.
+    ``steroid_17``. When present, angular methyls are added as ``steroid_18``
+    (on C13) and ``steroid_19`` (on C10). Molecules without a gonane core yield
+    no predicates.
 
     Parameters
     ----------
@@ -258,11 +287,17 @@ def get_steroid_positions(mol: Chem.Mol) -> dict[str, list]:
     """
     atom_extensions: dict[str, list] = {}
     steroid_match = mol.GetSubstructMatch(_GONANE_PATTERN, useChirality=False)
-    if steroid_match:
-        for pat_idx, atom_idx in enumerate(steroid_match):
-            iupac = _GONANE_IDX_TO_IUPAC.get(pat_idx)
-            if iupac is not None:
-                atom_extensions.setdefault(f"steroid_{iupac}", []).append(atom_idx)
+    if not steroid_match:
+        return atom_extensions
+
+    iupac_to_atom: dict[int, int] = {}
+    for pat_idx, atom_idx in enumerate(steroid_match):
+        iupac = _GONANE_IDX_TO_IUPAC.get(pat_idx)
+        if iupac is not None:
+            atom_extensions.setdefault(f"steroid_{iupac}", []).append(atom_idx)
+            iupac_to_atom[iupac] = atom_idx
+
+    _add_angular_methyls(mol, atom_extensions, iupac_to_atom)
     return atom_extensions
 
 
@@ -289,3 +324,25 @@ def get_numerical_facts(mol: Chem.Mol) -> dict[str, list]:
     for ring in mol.GetRingInfo().AtomRings():
         atom_extensions.setdefault("ring_size", []).append(len(ring))
     return atom_extensions
+
+
+"""Manual check for steroid numbering (not part of the library API).
+
+Run: python -m chebi_utils.extract_properties
+Expect: cholesterol -> steroid_1..19, estrone -> steroid_1..18, benzene -> []
+"""
+if __name__ == "__main__":
+    from chebi_utils.read_molecule import smiles_or_inchi_to_mol
+
+    for name, smiles in {
+        "cholesterol": (
+            "C[C@H](CCCC(C)C)[C@H]1CC[C@@H]2[C@@]1(CC[C@H]3[C@H]2CC=C4[C@@]3(CC[C@@H](C4)O)C)C"
+        ),
+        "estrone": "C[C@]12CC[C@H]3[C@H]([C@@H]1CCC2=O)CCc4c3ccc(O)c4",
+        "benzene": "c1ccccc1",
+    }.items():
+        keys = sorted(
+            get_steroid_positions(smiles_or_inchi_to_mol(smiles)),
+            key=lambda k: int(k.split("_")[1]),
+        )
+        print(name, keys)

tests/test_extract_properties.py

@@ -0,0 +1,70 @@
+"""Tests for chebi_utils.extract_properties steroid numbering."""
+
+from __future__ import annotations
+
+import pytest
+
+from chebi_utils.extract_properties import get_steroid_positions, mol_to_fol_atoms
+from chebi_utils.read_molecule import smiles_or_inchi_to_mol
+
+# Cholesterol: gonane core + C18 and C19 angular methyls
+CHOLESTEROL_SMILES = (
+    "C[C@H](CCCC(C)C)[C@H]1CC[C@@H]2[C@@]1(CC[C@H]3[C@H]2CC=C4[C@@]3(CC[C@@H](C4)O)C)C"
+)
+# Estrone: aromatic ring A, C18 present, C19 absent
+ESTRONE_SMILES = "C[C@]12CC[C@H]3[C@H]([C@@H]1CCC2=O)CCc4c3ccc(O)c4"
+
+CORE_POSITIONS = {f"steroid_{n}" for n in range(1, 18)}
+
+
+def _steroid_keys(smiles: str) -> set[str]:
+    mol = smiles_or_inchi_to_mol(smiles)
+    assert mol is not None
+    return set(get_steroid_positions(mol))
+
+
+class TestSteroidPositions:
+    def test_non_steroid_has_no_positions(self):
+        assert get_steroid_positions(smiles_or_inchi_to_mol("c1ccccc1")) == {}
+
+    def test_cholesterol_has_core_and_angular_methyls(self):
+        keys = _steroid_keys(CHOLESTEROL_SMILES)
+        assert CORE_POSITIONS <= keys
+        assert "steroid_18" in keys
+        assert "steroid_19" in keys
+
+    def test_estrone_has_c18_but_not_c19(self):
+        keys = _steroid_keys(ESTRONE_SMILES)
+        assert CORE_POSITIONS <= keys
+        assert "steroid_18" in keys
+        assert "steroid_19" not in keys
+
+    @pytest.mark.parametrize(
+        "smiles,attachment,methyl",
+        [
+            (CHOLESTEROL_SMILES, "steroid_13", "steroid_18"),
+            (CHOLESTEROL_SMILES, "steroid_10", "steroid_19"),
+            (ESTRONE_SMILES, "steroid_13", "steroid_18"),
+        ],
+    )
+    def test_angular_methyl_is_bonded_to_attachment_atom(self, smiles, attachment, methyl):
+        mol = smiles_or_inchi_to_mol(smiles)
+        positions = get_steroid_positions(mol)
+        attachment_idx = positions[attachment][0]
+        methyl_idx = positions[methyl][0]
+        attachment_atom = mol.GetAtomWithIdx(attachment_idx)
+        neighbor_indices = {neighbor.GetIdx() for neighbor in attachment_atom.GetNeighbors()}
+        assert methyl_idx in neighbor_indices
+        assert mol.GetAtomWithIdx(methyl_idx).GetAtomicNum() == 6
+
+    def test_mol_to_fol_atoms_includes_steroid_predicates(self):
+        mol = smiles_or_inchi_to_mol(CHOLESTEROL_SMILES)
+        atom_facts, _mol_facts = mol_to_fol_atoms(mol, with_steroids=True)
+        assert "steroid_1" in atom_facts
+        assert "steroid_18" in atom_facts
+        assert "steroid_19" in atom_facts
+
+    def test_mol_to_fol_atoms_can_skip_steroids(self):
+        mol = smiles_or_inchi_to_mol(CHOLESTEROL_SMILES)
+        atom_facts, _mol_facts = mol_to_fol_atoms(mol, with_steroids=False)
+        assert not any(key.startswith("steroid_") for key in atom_facts)