New shared functionality for sample filtering & property extraction

chebi_utils/extract_properties.py

@@ -0,0 +1,177 @@
+# extract basic (and not so basic) properties from molecules. This is used to construct
+# FOL structures for reasoning tasks on molecules.
+
+import logging
+
+from rdkit import Chem
+from rdkit.Chem import Descriptors
+
+MAX_RING_SIZE = 8
+
+# Gonane core SMARTS (C1–C17, IUPAC steroid numbering).
+# [#6] matches any carbon; ~ matches any bond — handles unsaturated steroids
+# (Δ4, Δ5), ketones, and estrogens (aromatic ring A) without modification.
+# C18/C19 methyls and side chains are intentionally excluded so the pattern
+# matches all steroid sub-classes, not just fully-saturated ones.
+_GONANE_PATTERN = Chem.MolFromSmarts(
+    "[#6:13]12~[#6:12]~[#6:11]~[#6:9]3~[#6:10]4~"
+    "[#6:1]~[#6:2]~[#6:3]~[#6:4]~[#6:5]4~"
+    "[#6:6]~[#6:7]~[#6:8]3~[#6:14]2~"
+    "[#6:15]~[#6:16]~[#6:17]1"
+)
+_GONANE_IDX_TO_IUPAC: dict[int, int] = {
+    atom.GetIdx(): atom.GetAtomMapNum()
+    for atom in _GONANE_PATTERN.GetAtoms()
+    if atom.GetAtomMapNum() > 0
+}
+
+
+def mol_to_fol_atoms(
+    mol: Chem.Mol, with_rings=True, with_steroids=True
+) -> tuple[dict[str, list], set[str]]:
+    """Convert an RDKit ``Mol`` into a first-order logic model at the atom level.
+
+    Returns ``(atom_extensions, mol_extensions)`` where:
+    - ``atom_extensions`` is a ``dict[str, list]``: unary predicates map to
+      ``list[int]`` of atom indices; binary predicates map to
+      ``list[tuple[int, int]]`` of (left, right) index pairs.
+    - ``mol_extensions`` is a ``set[str]`` of molecule-level predicate names
+      that hold for this molecule (e.g. ``net_charge_positive``).
+    """
+    atom_extensions: dict[str, list] = {}
+
+    atom_extensions.update(get_atom_properties(mol))
+    atom_extensions.update(get_bond_properties(mol))
+
+    if with_rings:
+        atom_extensions.update(get_rings(mol))
+
+    if with_steroids:
+        atom_extensions.update(get_steroid_positions(mol))
+
+    mol_extensions = get_molecule_level_properties(mol)
+
+    return atom_extensions, mol_extensions
+
+
+def get_atom_properties(mol: Chem.Mol) -> dict[str, list]:
+    try:
+        Chem.rdCIPLabeler.AssignCIPLabels(mol)
+    except Exception as e:
+        logging.error(
+            "Failed to assign CIP labels to molecule, skipping chirality-related extensions: %s",
+            e,
+        )
+
+    atom_extensions: dict[str, list] = {}
+
+    # For each atom: element symbol, charge, hydrogen counts, chirality
+    for atom in mol.GetAtoms():
+        atom_idx = atom.GetIdx()
+        atom_symbol = atom.GetSymbol().lower()
+        atom_extensions.setdefault(atom_symbol, []).append(atom_idx)
+
+        charge = atom.GetFormalCharge()
+        if charge != 0:
+            for pred in [
+                f"charge_{'n' if charge < 0 else 'p'}",
+                f"charge{'_m' + str(-charge) if charge < 0 else str(charge)}",
+            ]:
+                atom_extensions.setdefault(pred, []).append(atom_idx)
+        else:
+            atom_extensions.setdefault("charge0", []).append(atom_idx)
+
+        num_hs = atom.GetTotalNumHs(includeNeighbors=True)
+        for pred in [f"has_{num_hs}_hs"] + [f"has_at_least_{n}_hs" for n in range(1, num_hs + 1)]:
+            atom_extensions.setdefault(pred, []).append(atom_idx)
+
+        if atom.HasProp("_CIPCode"):
+            chiral_code = f"cip_code_{atom.GetProp('_CIPCode')}"
+            atom_extensions.setdefault(chiral_code, []).append(atom_idx)
+
+    return atom_extensions
+
+
+def get_bond_properties(mol: Chem.Mol) -> dict[str, list]:
+    # Bond predicates (symmetric)
+    atom_extensions: dict[str, list] = {}
+    for bond in mol.GetBonds():
+        left = bond.GetBeginAtomIdx()
+        right = bond.GetEndAtomIdx()
+
+        bond_pred = f"b{bond.GetBondType()}"
+        atom_extensions.setdefault(bond_pred, []).extend([(left, right), (right, left)])
+        atom_extensions.setdefault("has_bond_to", []).extend([(left, right), (right, left)])
+
+        if bond.GetStereo() != Chem.BondStereo.STEREONONE:
+            stereo_pred = f"b{bond.GetStereo().name}"
+            atom_extensions.setdefault(stereo_pred, []).extend([(left, right), (right, left)])
+    return atom_extensions
+
+
+def get_molecule_level_properties(mol: Chem.Mol) -> set[str]:
+    # Molecule-level (global) properties (either true or false for the whole molecule)
+    mol_extensions: set[str] = set()
+    net_charge = Chem.GetFormalCharge(mol)
+    if net_charge > 0:
+        mol_extensions.add("net_charge_positive")
+    elif net_charge < 0:
+        mol_extensions.add("net_charge_negative")
+    else:
+        mol_extensions.add("net_charge_neutral")
+    # aliphatic vs aromatic (defined as having at least one aromatic atom)
+    if len(list(mol.GetAromaticAtoms())) > 0:
+        mol_extensions.add("aromatic")
+    else:
+        mol_extensions.add("aliphatic")
+    return mol_extensions
+
+
+def get_rings(mol: Chem.Mol) -> dict[str, list]:
+    # Rings have two predicates. One for the atom-ring relation and one for the ring itself
+    #   ring{N}(A1, …, AN) – A1…AN form an N-membered ring (all permutations)
+    #                        Only for N <= MAX_RING_SIZE.
+    #   in_ring{N}(A)      – A belongs to some N-membered ring (N <= MAX_RING_SIZE).
+    #   in_ring(A)         – A belongs to some ring of any size.
+    atom_extensions: dict[str, list] = {}
+    in_ring_atoms: set[int] = set()
+    in_ringN_atoms: dict[int, set[int]] = {}
+    for ring in mol.GetRingInfo().AtomRings():
+        n = len(ring)
+        in_ring_atoms.update(ring)
+        if n <= MAX_RING_SIZE:
+            for start_atom in range(n):
+                ring_permutation = ring[start_atom:] + ring[:start_atom]
+                atom_extensions.setdefault(f"ring{n}", []).append(tuple(ring_permutation))
+                atom_extensions[f"ring{n}"].append(tuple(reversed(ring_permutation)))
+            in_ringN_atoms.setdefault(n, set()).update(ring)
+    if in_ring_atoms:
+        atom_extensions["in_ring"] = sorted(in_ring_atoms)
+    for n, atoms in in_ringN_atoms.items():
+        atom_extensions[f"in_ring{n}"] = sorted(atoms)
+    return atom_extensions
+
+
+def get_steroid_positions(mol: Chem.Mol) -> dict[str, list]:
+    # Steroid nucleus positions (steroid_1 … steroid_17)
+    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)
+    return atom_extensions
+
+
+def get_numerical_facts(mol: Chem.Mol) -> dict[str, list]:
+    """Molecular weight and ring size as numerical values.
+
+    Expresses "this molecule has weight ..." and "this molecule has a ring of size ..."
+    as molecule-integer value relations.
+    """
+    atom_extensions: dict[str, list] = {}
+    atom_extensions["mol_weight"] = [round(Descriptors.MolWt(mol))]
+    for ring in mol.GetRingInfo().AtomRings():
+        atom_extensions.setdefault("ring_size", []).append(len(ring))
+    return atom_extensions

chebi_utils/sample_filters.py

@@ -0,0 +1,76 @@
+# functionality for selecting specific sample subsets from the ChEBI dataset
+
+import networkx as nx
+
+from chebi_utils.obo_extractor import get_hierarchy_subgraph
+
+
+def get_closest_negatives(
+    samples: list[str], chebi_graph: nx.DiGraph, target_id: str, min_samples=25, max_samples=None
+) -> set[str]:
+    # from the list of samples, find those that are not subclasses of the target_id, but close
+    # to it in the hierarchy.
+    # goal: reach min_samples, but continue collecting samples (until max_samples) if they are
+    # siblings.
+    hierarchy_graph = nx.transitive_closure_dag(get_hierarchy_subgraph(chebi_graph))
+    undirected_graph = get_hierarchy_subgraph(chebi_graph).to_undirected()
+    import queue
+
+    q = queue.Queue()
+    q.put(target_id)
+    visited = set()  # visit closest labels
+    selected = (
+        set()
+    )  # select samples that are subclasses of closest labels until we have enough samples
+    siblings = True
+    while not q.empty():
+        current = q.get()
+        for neighbor in undirected_graph.neighbors(current):
+            if neighbor not in visited:
+                visited.add(neighbor)
+                q.put(neighbor)
+                for neighbor_sub in hierarchy_graph.predecessors(neighbor):
+                    if str(neighbor_sub) in samples:
+                        selected.add(str(neighbor_sub))
+                    if (max_samples and len(selected) >= max_samples) or (
+                        len(selected) >= min_samples and not siblings
+                    ):
+                        return selected
+        if len(selected) >= min_samples:
+            break
+        siblings = False
+
+    return selected
+
+
+def get_direct_neighbors(
+    samples: list[str],
+    chebi_graph: nx.DiGraph,
+    target_id: str,
+) -> tuple[list[str], list[str]]:
+    """
+    Filter samples and sort into two groups:
+    positive: sample is a descendant of the target_id
+    negative: sample is not a descendant of the target_id, but a "direct neighbor" -> a
+        descendant of all direct parents of the target_id.
+
+    Returns:
+        pos_ids: list of positive validation molecule IDs
+        neg_ids: list of negative validation molecule IDs
+                 (empty when target has no siblings)
+    """
+    hierarchy_graph = nx.transitive_closure_dag(get_hierarchy_subgraph(chebi_graph))
+    pos_ids = [str(d) for d in hierarchy_graph.predecessors(target_id) if str(d) in samples]
+
+    sample_space_by_parent = dict()
+    for parent in chebi_graph.successors(target_id):
+        sample_space_by_parent[parent] = set()
+        for desc in hierarchy_graph.predecessors(parent):
+            s = str(desc)
+            if s in samples:
+                sample_space_by_parent[parent].add(s)
+    if len(sample_space_by_parent) == 0:
+        return pos_ids, []
+    sample_space = set.intersection(*sample_space_by_parent.values())
+    neg_ids = list(sample_space - set(pos_ids))
+    return pos_ids, neg_ids

chebi_utils/splitter.py

@@ -44,7 +44,7 @@ def create_multilabel_splits(
     Returns
     -------
     dict
-        Dictionary with keys ``'train'``, ``'val'``, ``'test'``, each
+        Dictionary with keys ``'train'``, ``'validation'``, ``'test'``, each
         containing a DataFrame.
 
     Raises
@@ -104,6 +104,6 @@ def create_multilabel_splits(
 
     return {
         "train": df_train.reset_index(drop=True),
-        "val": df_val.reset_index(drop=True),
+        "validation": df_val.reset_index(drop=True),
         "test": df_test.reset_index(drop=True),
     }

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "chebi-utils"
-version = "0.2.1"
+version = "0.3"
 description = "Common processing functionality for the ChEBI ontology"
 readme = "README.md"
 license = { file = "LICENSE" }

tests/test_splitter.py

@@ -51,7 +51,7 @@ def singlelabel_df():
 class TestCreateMultilabelSplits:
     def test_returns_three_splits(self, multilabel_df):
         splits = create_multilabel_splits(multilabel_df)
-        assert set(splits.keys()) == {"train", "val", "test"}
+        assert set(splits.keys()) == {"train", "validation", "test"}
 
     def test_sizes_sum_to_total(self, multilabel_df):
         splits = create_multilabel_splits(multilabel_df)
@@ -60,7 +60,7 @@ def test_sizes_sum_to_total(self, multilabel_df):
     def test_no_overlap(self, multilabel_df):
         splits = create_multilabel_splits(multilabel_df)
         train_ids = set(splits["train"]["chebi_id"])
-        val_ids = set(splits["val"]["chebi_id"])
+        val_ids = set(splits["validation"]["chebi_id"])
         test_ids = set(splits["test"]["chebi_id"])
         assert train_ids.isdisjoint(val_ids)
         assert train_ids.isdisjoint(test_ids)
@@ -70,7 +70,7 @@ def test_all_rows_covered(self, multilabel_df):
         splits = create_multilabel_splits(multilabel_df)
         all_ids = (
             set(splits["train"]["chebi_id"])
-            | set(splits["val"]["chebi_id"])
+            | set(splits["validation"]["chebi_id"])
             | set(splits["test"]["chebi_id"])
         )
         assert all_ids == set(multilabel_df["chebi_id"])
@@ -96,7 +96,7 @@ def test_approximate_split_sizes(self, multilabel_df):
         )
         n = len(multilabel_df)
         assert abs(len(splits["test"]) - int(n * 0.1)) <= 2
-        assert abs(len(splits["val"]) - int(n * 0.1)) <= 2
+        assert abs(len(splits["validation"]) - int(n * 0.1)) <= 2
 
     def test_custom_label_start_col(self, multilabel_df):
         # Drop the 'mol' column so labels start at index 1
@@ -117,7 +117,7 @@ def test_singlelabel_path(self, singlelabel_df):
         splits = create_multilabel_splits(singlelabel_df)
         assert sum(len(v) for v in splits.values()) == len(singlelabel_df)
         train_ids = set(splits["train"]["chebi_id"])
-        val_ids = set(splits["val"]["chebi_id"])
+        val_ids = set(splits["validation"]["chebi_id"])
         test_ids = set(splits["test"]["chebi_id"])
         assert train_ids.isdisjoint(val_ids)
         assert train_ids.isdisjoint(test_ids)