update documentation

README.md

@@ -1,6 +1,12 @@
 # python-chebi-utils
 
-Common processing functionality for the ChEBI ontology — download versioned data files, build an ontology graph, extract molecules, assemble labeled datasets, and generate stratified train/val/test splits.
+Common processing functionality for the ChEBI ontology — download versioned data files, build an ontology graph, extract molecules, assemble labeled datasets, generate stratified train/validation/test splits, extract first-order-logic molecular properties, and select hierarchy-aware sample subsets.
+
+> **⚠️ Breaking change in v0.3**
+>
+> `create_multilabel_splits` now returns the validation split under the key
+> `"validation"` instead of `"val"`. Update any code that reads `splits["val"]`
+> to use `splits["validation"]`.
 
 ## Installation
 
@@ -86,7 +92,7 @@ from chebi_utils import create_multilabel_splits
 
 splits = create_multilabel_splits(dataset, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1)
 train_df = splits["train"]
-val_df   = splits["val"]
+val_df   = splits["validation"]   # renamed from "val" in v0.3
 test_df  = splits["test"]
 ```
 
@@ -96,6 +102,51 @@ present, `MultilabelStratifiedShuffleSplit` from the
 `iterative-stratification` package is used; for a single label column,
 `StratifiedShuffleSplit` from scikit-learn is used.
 
+### Extract molecular properties as first-order-logic facts
+
+```python
+from chebi_utils.extract_properties import mol_to_fol_atoms, get_numerical_facts
+
+atom_facts, mol_facts = mol_to_fol_atoms(mol, with_rings=True, with_steroids=True)
+# atom_facts — dict[str, list] of predicates over atom indices:
+#   unary  (e.g. "c", "charge_p", "has_2_hs", "cip_code_R", "in_ring6", "steroid_3")
+#          → list[int] of atom indices
+#   binary (e.g. "has_bond_to", "bSINGLE", "ring6") → list[tuple[int, ...]]
+# mol_facts — set[str] of molecule-level predicates that hold for the whole
+#             molecule (e.g. "net_charge_neutral", "aromatic")
+
+numerical_facts = get_numerical_facts(mol)
+# {"mol_weight": [<rounded MolWt>], "ring_size": [<size per ring>, …]}
+```
+
+Turns an RDKit `Mol` into a symbolic model suitable for building FOL structures
+for reasoning tasks. Facts cover per-atom element, formal charge, hydrogen
+counts, and CIP chirality; symmetric bond and bond-stereo relations; ring
+membership up to `MAX_RING_SIZE` (8); and steroid-nucleus positions
+(`steroid_1` … `steroid_17`) matched against the gonane core via IUPAC steroid
+numbering. Ring and steroid extraction can be toggled with `with_rings` and
+`with_steroids`.
+
+### Select hierarchy-aware sample subsets
+
+```python
+from chebi_utils.sample_filters import get_closest_negatives, get_direct_neighbors
+
+# Nearest negatives: samples that are NOT subclasses of the target but close to
+# it in the ontology, expanding outward until min_samples (up to max_samples) is met.
+negatives = get_closest_negatives(
+    samples, graph, target_id="15841", min_samples=25, max_samples=None
+)
+
+# Split samples into positives (descendants of the target) and "direct neighbor"
+# negatives (descendants of ALL direct parents of the target, but not the target).
+pos_ids, neg_ids = get_direct_neighbors(samples, graph, target_id="15841")
+```
+
+Useful for constructing balanced positive/negative sets for a given ChEBI class
+by leveraging the `is_a` hierarchy. `samples` is a list of ChEBI IDs (as
+strings) and `graph` is a graph from `build_chebi_graph`.
+
 ## Running Tests
 
 ```bash

chebi_utils/extract_properties.py

@@ -1,5 +1,8 @@
-# extract basic (and not so basic) properties from molecules. This is used to construct
-# FOL structures for reasoning tasks on molecules.
+"""Extract basic (and not so basic) properties from molecules.
+
+These are used to construct first-order-logic (FOL) structures for reasoning
+tasks on molecules.
+"""
 
 import logging
 
@@ -31,12 +34,28 @@ def mol_to_fol_atoms(
 ) -> 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``).
+    Aggregates atom, bond, ring, and steroid predicates into a single atom-level
+    extension mapping plus the set of molecule-level predicates.
+
+    Parameters
+    ----------
+    mol : Chem.Mol
+        The molecule to convert.
+    with_rings : bool
+        If ``True`` (default), include ring predicates from :func:`get_rings`.
+    with_steroids : bool
+        If ``True`` (default), include steroid-nucleus position predicates from
+        :func:`get_steroid_positions`.
+
+    Returns
+    -------
+    tuple[dict[str, list], set[str]]
+        ``(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] = {}
 
@@ -55,6 +74,25 @@ def mol_to_fol_atoms(
 
 
 def get_atom_properties(mol: Chem.Mol) -> dict[str, list]:
+    """Extract per-atom unary predicates.
+
+    For each atom, emits predicates for its element symbol, formal charge
+    (both a sign predicate ``charge_p``/``charge_n``/``charge0`` and an exact
+    predicate ``charge{n}``), total hydrogen count (exact ``has_{n}_hs`` and
+    cumulative ``has_at_least_{n}_hs``), and CIP chirality label
+    (``cip_code_R``/``cip_code_S``) when assigned.
+
+    Parameters
+    ----------
+    mol : Chem.Mol
+        The molecule to inspect.
+
+    Returns
+    -------
+    dict[str, list]
+        Mapping from predicate name to the ``list[int]`` of atom indices for
+        which the predicate holds.
+    """
     try:
         Chem.rdCIPLabeler.AssignCIPLabels(mol)
     except Exception as e:
@@ -65,7 +103,6 @@ def get_atom_properties(mol: Chem.Mol) -> dict[str, list]:
 
     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()
@@ -93,7 +130,24 @@ def get_atom_properties(mol: Chem.Mol) -> dict[str, list]:
 
 
 def get_bond_properties(mol: Chem.Mol) -> dict[str, list]:
-    # Bond predicates (symmetric)
+    """Extract symmetric per-bond binary predicates.
+
+    For each bond, emits a bond-type predicate (e.g. ``bSINGLE``), a generic
+    ``has_bond_to`` predicate, and a bond-stereo predicate (e.g.
+    ``bSTEREOE``) when the bond carries stereochemistry. Every predicate is
+    stored symmetrically, i.e. both ``(left, right)`` and ``(right, left)``.
+
+    Parameters
+    ----------
+    mol : Chem.Mol
+        The molecule to inspect.
+
+    Returns
+    -------
+    dict[str, list]
+        Mapping from predicate name to a ``list[tuple[int, int]]`` of ordered
+        atom-index pairs for which the predicate holds.
+    """
     atom_extensions: dict[str, list] = {}
     for bond in mol.GetBonds():
         left = bond.GetBeginAtomIdx()
@@ -110,7 +164,23 @@ def get_bond_properties(mol: Chem.Mol) -> dict[str, list]:
 
 
 def get_molecule_level_properties(mol: Chem.Mol) -> set[str]:
-    # Molecule-level (global) properties (either true or false for the whole molecule)
+    """Extract global (molecule-level) predicates.
+
+    Emits exactly one net-charge predicate
+    (``net_charge_positive``/``net_charge_negative``/``net_charge_neutral``) and
+    one aromaticity predicate (``aromatic`` if the molecule has at least one
+    aromatic atom, otherwise ``aliphatic``).
+
+    Parameters
+    ----------
+    mol : Chem.Mol
+        The molecule to inspect.
+
+    Returns
+    -------
+    set[str]
+        The set of molecule-level predicate names that hold for this molecule.
+    """
     mol_extensions: set[str] = set()
     net_charge = Chem.GetFormalCharge(mol)
     if net_charge > 0:
@@ -119,7 +189,6 @@ def get_molecule_level_properties(mol: Chem.Mol) -> set[str]:
         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:
@@ -128,11 +197,28 @@ def get_molecule_level_properties(mol: Chem.Mol) -> set[str]:
 
 
 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.
+    """Extract ring-membership predicates.
+
+    Emits three kinds of predicate:
+
+    - ``ring{N}(A1, …, AN)`` – ``A1…AN`` form an N-membered ring, stored for all
+      rotations in both directions. 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.
+
+    Parameters
+    ----------
+    mol : Chem.Mol
+        The molecule to inspect.
+
+    Returns
+    -------
+    dict[str, list]
+        Mapping from predicate name to its extension: ``ring{N}`` maps to a
+        ``list[tuple[int, ...]]`` of atom-index tuples; ``in_ring`` and
+        ``in_ring{N}`` map to a sorted ``list[int]`` of atom indices.
+    """
     atom_extensions: dict[str, list] = {}
     in_ring_atoms: set[int] = set()
     in_ringN_atoms: dict[int, set[int]] = {}
@@ -153,7 +239,23 @@ def get_rings(mol: Chem.Mol) -> dict[str, list]:
 
 
 def get_steroid_positions(mol: Chem.Mol) -> dict[str, list]:
-    # Steroid nucleus positions (steroid_1 … steroid_17)
+    """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.
+
+    Parameters
+    ----------
+    mol : Chem.Mol
+        The molecule to inspect.
+
+    Returns
+    -------
+    dict[str, list]
+        Mapping from ``steroid_{position}`` to the ``list[int]`` of matched atom
+        indices. Empty when the molecule has no steroid nucleus.
+    """
     atom_extensions: dict[str, list] = {}
     steroid_match = mol.GetSubstructMatch(_GONANE_PATTERN, useChirality=False)
     if steroid_match:
@@ -167,8 +269,20 @@ def get_steroid_positions(mol: Chem.Mol) -> dict[str, list]:
 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.
+    Expresses "this molecule has weight ..." and "this molecule has a ring of
+    size ..." as molecule-integer value relations.
+
+    Parameters
+    ----------
+    mol : Chem.Mol
+        The molecule to inspect.
+
+    Returns
+    -------
+    dict[str, list]
+        ``{"mol_weight": [<rounded MolWt>], "ring_size": [<size per ring>, …]}``.
+        The ``ring_size`` key is present only when the molecule has at least one
+        ring.
     """
     atom_extensions: dict[str, list] = {}
     atom_extensions["mol_weight"] = [round(Descriptors.MolWt(mol))]

chebi_utils/obo_extractor.py

@@ -141,8 +141,22 @@ def build_chebi_graph(filepath: str | Path, top_class: str | None = "23367") ->
 
 
 def get_hierarchy_subgraph(chebi_graph: nx.DiGraph) -> nx.DiGraph:
-    """Subgraph of ChEBI including only edges corresponding to hierarchical relations (is_a).
-    Also removes nodes that are not connected by any is_a edges to other nodes."""
+    """Extract the ``is_a`` hierarchy of a ChEBI graph as a subgraph.
+
+    Keeps only edges whose ``relation`` attribute is ``is_a``, and drops nodes
+    that are not connected to any other node by an ``is_a`` edge.
+
+    Parameters
+    ----------
+    chebi_graph : nx.DiGraph
+        Full ChEBI ontology graph from :func:`build_chebi_graph`.
+
+    Returns
+    -------
+    nx.DiGraph
+        Subgraph containing only ``is_a`` edges (child → parent) and the nodes
+        they connect.
+    """
     return nx.DiGraph(
         chebi_graph.edge_subgraph(
             (u, v) for u, v, d in chebi_graph.edges(data=True) if d.get("relation") == "is_a"