Add create_multilabel_splits with iterstrat/sklearn for multilabel stratified splits

Co-authored-by: sfluegel05 <43573433+sfluegel05@users.noreply.github.com>

Author: Copilot

chebi_utils/__init__.py

@@ -1,12 +1,13 @@
 from chebi_utils.downloader import download_chebi_obo, download_chebi_sdf
 from chebi_utils.obo_extractor import build_chebi_graph
 from chebi_utils.sdf_extractor import extract_molecules
-from chebi_utils.splitter import create_splits
+from chebi_utils.splitter import create_multilabel_splits, create_splits
 
 __all__ = [
     "download_chebi_obo",
     "download_chebi_sdf",
     "build_chebi_graph",
     "extract_molecules",
     "create_splits",
+    "create_multilabel_splits",
 ]

chebi_utils/splitter.py

@@ -6,6 +6,116 @@
 import pandas as pd
 
 
+def create_multilabel_splits(
+    df: pd.DataFrame,
+    labels_col: str,
+    train_ratio: float = 0.8,
+    val_ratio: float = 0.1,
+    test_ratio: float = 0.1,
+    seed: int | None = 42,
+) -> dict[str, pd.DataFrame]:
+    """Create stratified train/validation/test splits for multilabel DataFrames.
+
+    Automatically detects whether the dataset is multilabel (each entry has
+    more than one label) or single-label, and applies the appropriate
+    stratification strategy:
+
+    - Multilabel: uses ``MultilabelStratifiedShuffleSplit`` from the
+      ``iterative-stratification`` package.
+    - Single-label: uses ``StratifiedShuffleSplit`` from ``scikit-learn``.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        Input data to split.  Must contain a column ``labels_col`` whose
+        values are sequences of labels (e.g. lists of strings or ints).
+    labels_col : str
+        Name of the column that contains the label sequences.
+    train_ratio : float
+        Fraction of data for training (default 0.8).
+    val_ratio : float
+        Fraction of data for validation (default 0.1).
+    test_ratio : float
+        Fraction of data for testing (default 0.1).
+    seed : int or None
+        Random seed for reproducibility.
+
+    Returns
+    -------
+    dict
+        Dictionary with keys ``'train'``, ``'val'``, ``'test'``, each
+        containing a DataFrame.
+
+    Raises
+    ------
+    ValueError
+        If the ratios do not sum to 1, any ratio is outside ``[0, 1]``, or
+        ``labels_col`` is not found in *df*.
+    """
+    if abs(train_ratio + val_ratio + test_ratio - 1.0) > 1e-6:
+        raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")
+    if any(r < 0 or r > 1 for r in [train_ratio, val_ratio, test_ratio]):
+        raise ValueError("All ratios must be between 0 and 1")
+    if labels_col not in df.columns:
+        raise ValueError(f"Column '{labels_col}' not found in DataFrame")
+
+    from iterstrat.ml_stratifiers import MultilabelStratifiedShuffleSplit
+    from sklearn.model_selection import StratifiedShuffleSplit
+    from sklearn.preprocessing import MultiLabelBinarizer
+
+    labels_list: list[list] = df[labels_col].tolist()
+    is_multilabel = any(len(lbl) > 1 for lbl in labels_list)
+
+    df_reset = df.reset_index(drop=True)
+
+    if is_multilabel:
+        mlb = MultiLabelBinarizer()
+        labels_matrix = mlb.fit_transform(labels_list)
+    else:
+        labels_matrix = [lbl[0] for lbl in labels_list]
+
+    # ── Step 1: carve out the test set ──────────────────────────────────────
+    if is_multilabel:
+        test_splitter = MultilabelStratifiedShuffleSplit(
+            n_splits=1, test_size=test_ratio, random_state=seed
+        )
+        train_val_idx, test_idx = next(test_splitter.split(labels_matrix, labels_matrix))
+    else:
+        test_splitter = StratifiedShuffleSplit(n_splits=1, test_size=test_ratio, random_state=seed)
+        train_val_idx, test_idx = next(test_splitter.split(labels_matrix, labels_matrix))
+
+    df_test = df_reset.iloc[test_idx]
+    df_trainval = df_reset.iloc[train_val_idx]
+
+    # ── Step 2: split train/val from the remaining data ─────────────────────
+    labels_trainval = (
+        labels_matrix[train_val_idx]
+        if is_multilabel
+        else [labels_matrix[i] for i in train_val_idx]
+    )
+    val_ratio_adjusted = val_ratio / (1.0 - test_ratio)
+
+    if is_multilabel:
+        val_splitter = MultilabelStratifiedShuffleSplit(
+            n_splits=1, test_size=val_ratio_adjusted, random_state=seed
+        )
+        train_idx_inner, val_idx_inner = next(val_splitter.split(labels_trainval, labels_trainval))
+    else:
+        val_splitter = StratifiedShuffleSplit(
+            n_splits=1, test_size=val_ratio_adjusted, random_state=seed
+        )
+        train_idx_inner, val_idx_inner = next(val_splitter.split(labels_trainval, labels_trainval))
+
+    df_train = df_trainval.iloc[train_idx_inner]
+    df_val = df_trainval.iloc[val_idx_inner]
+
+    return {
+        "train": df_train.reset_index(drop=True),
+        "val": df_val.reset_index(drop=True),
+        "test": df_test.reset_index(drop=True),
+    }
+
+
 def create_splits(
     df: pd.DataFrame,
     train_ratio: float = 0.8,

pyproject.toml

@@ -11,10 +11,12 @@ license = { file = "LICENSE" }
 requires-python = ">=3.10"
 dependencies = [
     "fastobo>=0.14",
+    "iterative-stratification>=0.1.9",
     "networkx>=3.0",
     "numpy>=1.24",
     "pandas>=2.0",
     "rdkit>=2022.09",
+    "scikit-learn>=1.0",
     "chembl_structure_pipeline>=1.2.4",
 ]
 

tests/test_splitter.py

@@ -5,7 +5,7 @@
 import pandas as pd
 import pytest
 
-from chebi_utils.splitter import create_splits
+from chebi_utils.splitter import create_multilabel_splits, create_splits
 
 
 @pytest.fixture
@@ -103,3 +103,89 @@ def test_stratified_reproducible(self, sample_df):
         splits1 = create_splits(sample_df, stratify_col="category", seed=42)
         splits2 = create_splits(sample_df, stratify_col="category", seed=42)
         pd.testing.assert_frame_equal(splits1["train"], splits2["train"])
+
+
+@pytest.fixture
+def multilabel_df():
+    """DataFrame with multilabel 'labels' column (200 rows)."""
+    all_labels = [["A"], ["B"], ["C"], ["A", "B"], ["A", "C"], ["B", "C"]]
+    labels = [all_labels[i % len(all_labels)] for i in range(200)]
+    return pd.DataFrame(
+        {
+            "id": [f"CHEBI:{i}" for i in range(200)],
+            "labels": labels,
+        }
+    )
+
+
+@pytest.fixture
+def singlelabel_df():
+    """DataFrame with single-label 'labels' column."""
+    return pd.DataFrame(
+        {
+            "id": [f"CHEBI:{i}" for i in range(200)],
+            "labels": [["A"] if i % 2 == 0 else ["B"] for i in range(200)],
+        }
+    )
+
+
+class TestCreateMultilabelSplits:
+    def test_returns_three_splits(self, multilabel_df):
+        splits = create_multilabel_splits(multilabel_df, labels_col="labels")
+        assert set(splits.keys()) == {"train", "val", "test"}
+
+    def test_sizes_sum_to_total(self, multilabel_df):
+        splits = create_multilabel_splits(multilabel_df, labels_col="labels")
+        assert sum(len(v) for v in splits.values()) == len(multilabel_df)
+
+    def test_no_overlap(self, multilabel_df):
+        splits = create_multilabel_splits(multilabel_df, labels_col="labels")
+        train_ids = set(splits["train"]["id"])
+        val_ids = set(splits["val"]["id"])
+        test_ids = set(splits["test"]["id"])
+        assert train_ids.isdisjoint(val_ids)
+        assert train_ids.isdisjoint(test_ids)
+        assert val_ids.isdisjoint(test_ids)
+
+    def test_all_rows_covered(self, multilabel_df):
+        splits = create_multilabel_splits(multilabel_df, labels_col="labels")
+        all_ids = set(splits["train"]["id"]) | set(splits["val"]["id"]) | set(splits["test"]["id"])
+        assert all_ids == set(multilabel_df["id"])
+
+    def test_reproducible_with_same_seed(self, multilabel_df):
+        splits1 = create_multilabel_splits(multilabel_df, labels_col="labels", seed=7)
+        splits2 = create_multilabel_splits(multilabel_df, labels_col="labels", seed=7)
+        pd.testing.assert_frame_equal(splits1["train"], splits2["train"])
+
+    def test_different_seeds_give_different_splits(self, multilabel_df):
+        splits1 = create_multilabel_splits(multilabel_df, labels_col="labels", seed=1)
+        splits2 = create_multilabel_splits(multilabel_df, labels_col="labels", seed=2)
+        assert not splits1["train"]["id"].equals(splits2["train"]["id"])
+
+    def test_approximate_split_sizes(self, multilabel_df):
+        splits = create_multilabel_splits(
+            multilabel_df, labels_col="labels", train_ratio=0.8, val_ratio=0.1, test_ratio=0.1
+        )
+        n = len(multilabel_df)
+        assert abs(len(splits["test"]) - int(n * 0.1)) <= 2
+        assert abs(len(splits["val"]) - int(n * 0.1)) <= 2
+
+    def test_invalid_ratios_raise_error(self, multilabel_df):
+        with pytest.raises(ValueError, match="must equal 1.0"):
+            create_multilabel_splits(
+                multilabel_df, labels_col="labels", train_ratio=0.5, val_ratio=0.3, test_ratio=0.3
+            )
+
+    def test_missing_labels_col_raises_error(self, multilabel_df):
+        with pytest.raises(ValueError, match="not found in DataFrame"):
+            create_multilabel_splits(multilabel_df, labels_col="nonexistent")
+
+    def test_singlelabel_path(self, singlelabel_df):
+        """Single-label lists should use StratifiedShuffleSplit without error."""
+        splits = create_multilabel_splits(singlelabel_df, labels_col="labels")
+        assert sum(len(v) for v in splits.values()) == len(singlelabel_df)
+        train_ids = set(splits["train"]["id"])
+        val_ids = set(splits["val"]["id"])
+        test_ids = set(splits["test"]["id"])
+        assert train_ids.isdisjoint(val_ids)
+        assert train_ids.isdisjoint(test_ids)