forestgump 0.1.0__py3-none-any.whl

forestgump/__init__.py

@@ -0,0 +1,14 @@
+"""forestgump: did your model beat a random forest? Life is like a box of baselines."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("forestgump")
+except PackageNotFoundError:  # running from a source tree
+    __version__ = "0+unknown"
+
+from .core import run
+from .io import InputError, load
+from .report import Report
+
+__all__ = ["InputError", "Report", "__version__", "load", "run"]

forestgump/baselines.py

@@ -0,0 +1,180 @@
+"""The baselines, their small grids, and grid selection on training rows only."""
+
+from __future__ import annotations
+
+import itertools
+from dataclasses import dataclass, field
+
+import numpy as np
+from sklearn.dummy import DummyClassifier, DummyRegressor
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+from sklearn.model_selection import KFold, StratifiedKFold
+from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
+
+from .featurize import Fingerprints
+from .stats import LOWER_IS_BETTER, PRIMARY, metrics
+
+N_TREES = 500
+METHODS = ("mean", "knn", "forest")
+METHODS_WITH_CONTEXT = ("mean", "context", "knn", "forest")
+FOREST_SELECTION = ("oob", "inner-cv")
+
+GRIDS: dict[str, list[dict]] = {
+    "mean": [{}],
+    "context": [{}],
+    "knn": [
+        {"radius": r, "n_neighbors": k, "weights": w}
+        for r, k, w in itertools.product((2, 3), (1, 3, 5, 10), ("uniform", "distance"))
+    ],
+    "forest": [
+        {"radius": r, "n_estimators": N_TREES, "max_features": mf, "min_samples_leaf": leaf}
+        for r, mf, leaf in itertools.product((2, 3), ("sqrt", 0.3, 1.0), (1, 3))
+    ],
+}
+
+
+class ContextMean:
+    """Predict the mean training target of the test row's context group, or the training mean for a group that has
+    no training rows. For a binary target that mean is the group's frequency of the larger class."""
+
+    def fit(self, codes: np.ndarray, y: np.ndarray):
+        self.fallback_ = float(np.mean(y))
+        self.means_ = {g: float(np.mean(y[codes == g])) for g in np.unique(codes)}
+        return self
+
+    def predict(self, codes: np.ndarray) -> np.ndarray:
+        return np.array([self.means_.get(g, self.fallback_) for g in codes], dtype=float)
+
+    def predict_proba(self, codes: np.ndarray) -> np.ndarray:
+        p = self.predict(codes)
+        return np.column_stack([1 - p, p])
+
+
+def _model(method: str, task: str, params: dict, n_train: int, seed: int, n_jobs: int):
+    binary = task == "binary"
+    if method == "context":
+        return ContextMean()
+    if method == "mean":
+        return DummyClassifier(strategy="prior") if binary else DummyRegressor(strategy="mean")
+    if method == "knn":
+        cls = KNeighborsClassifier if binary else KNeighborsRegressor
+        return cls(n_neighbors=min(params["n_neighbors"], n_train), weights=params["weights"], metric="precomputed")
+    cls = RandomForestClassifier if binary else RandomForestRegressor
+    return cls(
+        n_estimators=params["n_estimators"],
+        max_features=params["max_features"],
+        min_samples_leaf=params["min_samples_leaf"],
+        random_state=seed,
+        n_jobs=n_jobs,
+    )
+
+
+def tanimoto_distance(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    """1 - Tanimoto similarity between the rows of two bit matrices (scipy's Jaccard distance; 0 if both empty)."""
+    a, b = a.astype(np.float64), b.astype(np.float64)
+    shared = a @ b.T
+    union = a.sum(1)[:, None] + b.sum(1)[None, :] - shared
+    with np.errstate(invalid="ignore", divide="ignore"):
+        distance = (union - shared) / union  # integer counts, so this rounds exactly as scipy's jaccard does
+    distance[union == 0] = 0.0
+    return distance
+
+
+def fingerprints_for(fps: Fingerprints, params: dict) -> np.ndarray | None:
+    """The bit matrix a config uses (all rows), or None for the mean baseline."""
+    return fps(params["radius"]) if "radius" in params else None
+
+
+def _features(method: str, X: np.ndarray | None, train_idx: np.ndarray, test_idx: np.ndarray):
+    """(X to fit on, X to predict). kNN gets precomputed Tanimoto distances to the training rows."""
+    if X is None:
+        # The mean baseline ignores X; one zero column keeps sklearn's input checks happy.
+        return np.zeros((len(train_idx), 1)), np.zeros((len(test_idx), 1))
+    if method == "knn":
+        return tanimoto_distance(X[train_idx], X[train_idx]), tanimoto_distance(X[test_idx], X[train_idx])
+    return X[train_idx], X[test_idx]
+
+
+def predict(model, task: str, X: np.ndarray) -> np.ndarray:
+    return model.predict_proba(X)[:, 1] if task == "binary" else model.predict(X)
+
+
+def fit_predict(method, task, params, X, y, train_idx, test_idx, seed, n_jobs):
+    """Fit one configuration on rows train_idx of X and predict rows test_idx. Returns (predictions, fitted model).
+
+    X is the config's fingerprint matrix over all rows (`fingerprints_for`), the per-row context group codes for the
+    context mean, or None for the mean baseline.
+    """
+    model = _model(method, task, params, len(train_idx), seed, n_jobs)
+    X_fit, X_pred = _features(method, X, train_idx, test_idx)
+    model.fit(X_fit, y[train_idx])
+    return predict(model, task, X_pred), model
+
+
+def _inner_cv_score(method, task, params, X, y, train_idx, seed, n_jobs=1) -> float:
+    """Mean primary metric over inner CV folds of the training rows."""
+    scores = []
+    for a, b in _inner_splits(task, y[train_idx], seed).split(train_idx, y[train_idx]):
+        p, _ = fit_predict(method, task, params, X, y, train_idx[a], train_idx[b], seed, n_jobs)
+        scores.append(metrics(task, y[train_idx[b]], p)[PRIMARY[task]])
+    return float(np.mean(scores))
+
+
+def _inner_splits(task: str, y_train: np.ndarray, seed: int):
+    n_splits = 5 if len(y_train) >= 50 else 3
+    if task == "binary":
+        n_splits = min(n_splits, int(np.bincount(y_train.astype(int)).min()))
+        return StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
+    return KFold(n_splits=n_splits, shuffle=True, random_state=seed)
+
+
+def _fit_forest_with_oob(task, params, X, y, train_idx, seed, n_jobs):
+    """Fit one forest on the training rows; score it on its out-of-bag predictions. Returns (score, model)."""
+    model = _model("forest", task, params, len(train_idx), seed, n_jobs).set_params(oob_score=True)
+    model.fit(X[train_idx], y[train_idx])
+    oob = model.oob_decision_function_[:, 1] if task == "binary" else model.oob_prediction_
+    seen = np.isfinite(oob)  # a row in every bootstrap sample has no out-of-bag prediction
+    return metrics(task, y[train_idx][seen], oob[seen])[PRIMARY[task]], model
+
+
+def _better(primary: str, a: float, b: float) -> bool:
+    return a < b if primary in LOWER_IS_BETTER else a > b
+
+
+@dataclass
+class Selection:
+    """The grid winner for one method on one cell, with every config's score on the training rows."""
+
+    params: dict
+    score: float
+    grid_scores: list[float] = field(default_factory=list)
+    model: object = None  # the forest winner, already fitted on every training row
+
+
+def select(method, task, fps, y, train_idx, seed, n_jobs, forest_selection="oob") -> Selection:
+    """Pick a method's config using only the rows in train_idx.
+
+    By default the forest fits each of its configs once on all training rows and is judged by its out-of-bag
+    score, which costs nothing extra; the winning forest is kept, so nothing is refit. With
+    forest_selection="inner-cv" it is judged by inner CV like kNN instead, then refit: k forests per config plus
+    one, about 5x the forest time. kNN fits take milliseconds, so it always uses inner CV. Ties go to the earlier
+    config in grid order, so selection is deterministic.
+    """
+    grid = GRIDS[method]
+    primary = PRIMARY[task]
+    if len(grid) == 1:
+        return Selection(grid[0], np.nan, [np.nan])
+    best, scores = None, []
+    for params in grid:
+        X = fingerprints_for(fps, params)
+        if method == "forest" and forest_selection == "oob":
+            score, model = _fit_forest_with_oob(task, params, X, y, train_idx, seed, n_jobs)
+        elif method == "forest":
+            score, model = _inner_cv_score(method, task, params, X, y, train_idx, seed, n_jobs), None
+        else:
+            score, model = _inner_cv_score(method, task, params, X, y, train_idx, seed), None
+        scores.append(score)
+        if best is None or _better(primary, score, best.score):
+            best = Selection(params, score, model=model)
+    best.grid_scores = scores
+    return best

forestgump/cli.py

@@ -0,0 +1,109 @@
+"""`forestgump run data.csv`: Run, Forest, run."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import warnings
+
+import pandas as pd
+
+from . import __version__, jokes
+from .core import INNER_CV_WARNING, run
+from .demo import MODELS, demo_frame
+from .io import TASKS, InputError
+
+EXIT_FOREST_WON = 1
+EXIT_BAD_INPUT = 2
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="forestgump", description=f"Did your model beat a random forest? {jokes.TAGLINE}"
+    )
+    parser.add_argument("--version", action="version", version=f"forestgump {__version__}. {jokes.TAGLINE}")
+    sub = parser.add_subparsers(dest="command", required=True)
+    r = sub.add_parser(
+        "run",
+        help="fit the baselines on your folds and judge your model",
+        description="One CSV row per (molecule, repeat). The test set of (repeat r, fold f) is its rows; "
+        "the training set is every other row of repeat r. --pred is your model's prediction for the row "
+        "(for binary targets, the probability of the larger class label).",
+    )
+    r.add_argument("csv")
+    _shared(r)
+    r.add_argument("--smiles", default="smiles", help="SMILES column (default: smiles)")
+    r.add_argument("--target", default="target", help="label column (default: target)")
+    r.add_argument("--fold", default="fold", help="test-fold column (default: fold)")
+    r.add_argument("--pred", default="pred", help="your model's out-of-fold prediction column (default: pred)")
+    r.add_argument("--repeat", default=None, help="repeat column for repeated CV (default: one repeat)")
+    r.add_argument("--task", choices=TASKS, default=None, help="default: binary if the target has 2 values")
+    r.add_argument(
+        "--context",
+        nargs="+",
+        default=None,
+        metavar="COL",
+        help="column(s) defining a context group; adds the context mean, which predicts each group's training mean",
+    )
+    d = sub.add_parser("demo", help="judge a pretend model on 105 synthetic molecules labelled with Crippen logP")
+    d.add_argument("--model", choices=MODELS, default="noise", help="noise loses, oracle wins (default: noise)")
+    d.add_argument("--repeats", type=int, default=1, help="repeats of 5-fold CV (default: 1)")
+    _shared(d)
+    return parser
+
+
+def _shared(p: argparse.ArgumentParser) -> None:
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--n-jobs", type=int, default=-1, help="sklearn n_jobs (default: all cores)")
+    p.add_argument("--alpha", type=float, default=0.05)
+    p.add_argument(
+        "--brag", default=jokes.DEFAULT_BRAG, help=f'how your model was trained (default: "{jokes.DEFAULT_BRAG}")'
+    )
+    p.add_argument("--seriously", action="store_true", help="plain output, no jokes")
+    p.add_argument("--json", metavar="PATH", help="write every number to a JSON file")
+    p.add_argument(
+        "--forest-inner-cv",
+        action="store_true",
+        help="pick the forest's hyperparameters by inner CV instead of out-of-bag score (about 5x slower)",
+    )
+    p.add_argument(
+        "--fail-if-forest-wins",
+        action="store_true",
+        help=f"exit {EXIT_FOREST_WON} if the forest won, so a CI pipeline fails when your model loses",
+    )
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _parser().parse_args(argv)
+    if not args.seriously:
+        print("Planting the forest...", file=sys.stderr, flush=True)
+    if args.command == "demo":
+        frame = demo_frame(args.model, n_repeats=args.repeats, seed=args.seed)
+        columns = {"repeat": "repeat"}
+    else:
+        frame = pd.read_csv(args.csv)
+        columns = {k: getattr(args, k) for k in ("smiles", "target", "fold", "pred", "repeat", "task", "context")}
+    if args.forest_inner_cv:
+        print(f"forestgump: {INNER_CV_WARNING}", file=sys.stderr, flush=True)
+    try:
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", message="forest_selection='inner-cv'")  # already printed above
+            report = run(
+                frame,
+                seed=args.seed,
+                n_jobs=args.n_jobs,
+                alpha=args.alpha,
+                forest_selection="inner-cv" if args.forest_inner_cv else "oob",
+                **columns,
+            )
+    except InputError as exc:
+        print(f"forestgump: {exc}", file=sys.stderr)
+        return EXIT_BAD_INPUT
+    print(report.render(seriously=args.seriously, brag=args.brag))
+    if args.json:
+        report.to_json(args.json)
+    return EXIT_FOREST_WON if args.fail_if_forest_wins and report.verdict == "lose" else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

forestgump/core.py

@@ -0,0 +1,108 @@
+"""Run every baseline on every cell and compare it with the user's model."""
+
+from __future__ import annotations
+
+import time
+import warnings
+
+import numpy as np
+import pandas as pd
+
+from .baselines import FOREST_SELECTION, METHODS, METHODS_WITH_CONTEXT, fingerprints_for, fit_predict, predict, select
+from .featurize import Fingerprints
+from .io import load
+from .report import Comparison, Report
+from .stats import MIN_CELLS, PRIMARY, corrected_t, holm, improvement, metrics
+
+INNER_CV_WARNING = (
+    "forest_selection='inner-cv' fits every forest config once per inner fold, then refits the winner: "
+    "about 5x slower than the default out-of-bag selection. The forest will take its time."
+)
+
+
+def run(
+    df: pd.DataFrame,
+    *,
+    smiles: str = "smiles",
+    target: str = "target",
+    fold: str = "fold",
+    pred: str = "pred",
+    repeat: str | None = None,
+    task: str | None = None,
+    context: str | list[str] | None = None,
+    seed: int = 0,
+    n_jobs: int = -1,
+    alpha: float = 0.05,
+    forest_selection: str = "oob",
+) -> Report:
+    """Fit the mean, kNN and forest baselines on the user's folds and test the user's model against each.
+
+    See `forestgump.io.load` for the input contract. Hyperparameters are chosen per cell from that cell's training
+    rows only; test rows are only ever predicted. kNN always uses inner CV. The forest uses its out-of-bag score
+    by default (one fit per config); forest_selection="inner-cv" uses inner CV instead, at about 5x the time.
+    Naming `context` column(s) adds a fourth baseline, the context mean: each test row gets the mean training
+    target of the rows that share its context values, or the training mean if no training row does.
+    """
+    if forest_selection not in FOREST_SELECTION:
+        raise ValueError(f"forest_selection must be one of {FOREST_SELECTION}, got {forest_selection!r}")
+    if forest_selection == "inner-cv":
+        warnings.warn(INNER_CV_WARNING, stacklevel=2)
+    ds = load(df, smiles=smiles, target=target, fold=fold, pred=pred, repeat=repeat, task=task, context=context)
+    fps = Fingerprints(ds.mols, ds.mol_index)
+    methods = METHODS if ds.context is None else METHODS_WITH_CONTEXT
+    rows, selections, predictions = [], [], []
+    forest_seconds, n_leaves = 0.0, 0
+    for i, cell in enumerate(ds.cells):
+        base = {"cell": i, "repeat": cell.repeat, "fold": cell.fold, "n_train": len(cell.train_idx)}
+        base["n_test"] = len(cell.test_idx)
+        rows.append(base | {"method": "user"} | metrics(ds.task, ds.y[cell.test_idx], ds.pred[cell.test_idx]))
+        for method in methods:
+            start = time.perf_counter()
+            sel = select(method, ds.task, fps, ds.y, cell.train_idx, seed, n_jobs, forest_selection)
+            if sel.model is None:
+                X = ds.context if method == "context" else fingerprints_for(fps, sel.params)
+                p, model = fit_predict(
+                    method, ds.task, sel.params, X, ds.y, cell.train_idx, cell.test_idx, seed, n_jobs
+                )
+            else:
+                model = sel.model
+                p = predict(model, ds.task, fingerprints_for(fps, sel.params)[cell.test_idx])
+            if method == "forest":
+                forest_seconds += time.perf_counter() - start
+                n_leaves = sum(int(t.get_n_leaves()) for t in model.estimators_)
+            rows.append(base | {"method": method} | metrics(ds.task, ds.y[cell.test_idx], p))
+            predictions.append(pd.DataFrame({"cell": i, "row": cell.test_idx, "method": method, "pred": p}))
+            selection = {"cell": i, "method": method, "params": sel.params, "selection_score": sel.score}
+            if method == "context":
+                seen = np.isin(ds.context[cell.test_idx], ds.context[cell.train_idx])
+                selection["n_test_unseen_context"] = int((~seen).sum())
+            selections.append(selection)
+
+    table = pd.DataFrame(rows)
+    primary = PRIMARY[ds.task]
+    wide = table.pivot(index="cell", columns="method", values=primary)
+    comparisons = []
+    if len(ds.cells) >= MIN_CELLS:
+        tests = [
+            corrected_t(improvement(primary, wide["user"].to_numpy(), wide[m].to_numpy()), ds.test_train_ratio, alpha)
+            for m in methods
+        ]
+        for m, t, p_adj in zip(methods, tests, holm([t.p for t in tests]), strict=True):
+            comparisons.append(Comparison(method=m, test=t, p_holm=p_adj))
+    return Report(
+        task=ds.task,
+        primary=primary,
+        table=table,
+        selections=selections,
+        predictions=pd.concat(predictions, ignore_index=True),
+        comparisons=comparisons,
+        n_cells=len(ds.cells),
+        test_train_ratio=ds.test_train_ratio,
+        forest_seconds=forest_seconds,
+        forest_leaves=n_leaves,
+        alpha=alpha,
+        seed=seed,
+        classes=ds.classes,
+        forest_selection=forest_selection,
+        context=ds.context_columns,
+    )

forestgump/demo.py

@@ -0,0 +1,42 @@
+"""A synthetic demo: para-disubstituted benzenes labelled with RDKit's Crippen logP.
+
+Two pretend models are on offer. "noise" predicts at random, so the forest beats it. "oracle" predicts the true
+label plus a little noise, so it beats the forest.
+"""
+
+from __future__ import annotations
+
+import itertools
+
+import numpy as np
+import pandas as pd
+from rdkit import Chem
+from rdkit.Chem import Crippen
+
+SUBSTITUENTS = ("[H]", "C", "CC", "CCC", "O", "N", "F", "Cl", "Br", "C(=O)O", "C#N", "OC", "C(F)(F)F", "[N+](=O)[O-]")
+MODELS = ("noise", "oracle")
+
+
+def molecules() -> list[str]:
+    """Unique canonical SMILES of every para-disubstituted benzene over SUBSTITUENTS (105 molecules)."""
+    out = set()
+    for a, b in itertools.combinations_with_replacement(SUBSTITUENTS, 2):
+        mol = Chem.MolFromSmiles(f"c1({a})ccc({b})cc1")
+        out.add(Chem.MolToSmiles(Chem.RemoveHs(mol)))
+    return sorted(out)
+
+
+def demo_frame(model: str = "noise", n_folds: int = 5, n_repeats: int = 1, seed: int = 0) -> pd.DataFrame:
+    """Columns smiles, target, repeat, fold, pred: ready for `forestgump.run(..., repeat="repeat")`."""
+    if model not in MODELS:
+        raise ValueError(f"model must be one of {MODELS}")
+    rng = np.random.default_rng(seed)
+    smiles = molecules()
+    y = np.array([Crippen.MolLogP(Chem.MolFromSmiles(s)) for s in smiles])
+    frames = []
+    for r in range(n_repeats):
+        fold = rng.permutation(np.arange(len(smiles)) % n_folds)
+        noise = rng.normal(0, 0.1 if model == "oracle" else 1.5, len(smiles))
+        pred = y + noise if model == "oracle" else y.mean() + noise
+        frames.append(pd.DataFrame({"smiles": smiles, "target": y, "repeat": r, "fold": fold, "pred": pred}))
+    return pd.concat(frames, ignore_index=True)

forestgump/featurize.py

@@ -0,0 +1,29 @@
+"""Morgan fingerprints: a message-passing network with zero trainable parameters."""
+
+from __future__ import annotations
+
+import numpy as np
+from rdkit.Chem import rdFingerprintGenerator
+
+FP_SIZE = 2048
+
+
+def morgan_bits(mols, radius: int, fp_size: int = FP_SIZE) -> np.ndarray:
+    """Return an (n_mols, fp_size) boolean matrix of Morgan bits."""
+    gen = rdFingerprintGenerator.GetMorganGenerator(radius=radius, fpSize=fp_size)
+    return np.stack([gen.GetFingerprintAsNumPy(m) for m in mols]).astype(bool)
+
+
+class Fingerprints:
+    """Per-row Morgan bits, computed once per radius over the unique molecules."""
+
+    def __init__(self, mols, mol_index: np.ndarray, fp_size: int = FP_SIZE):
+        self._mols = mols
+        self._mol_index = mol_index
+        self._fp_size = fp_size
+        self._cache: dict[int, np.ndarray] = {}
+
+    def __call__(self, radius: int) -> np.ndarray:
+        if radius not in self._cache:
+            self._cache[radius] = morgan_bits(self._mols, radius, self._fp_size)[self._mol_index]
+        return self._cache[radius]

forestgump/io.py

@@ -0,0 +1,173 @@
+"""Load the user's molecules, folds and predictions, and refuse anything that would make the comparison lie."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+from rdkit import Chem, RDLogger
+
+TASKS = ("regression", "binary")
+MIN_TRAIN_ROWS = 6
+
+
+class InputError(ValueError):
+    """The input cannot be compared fairly. The message says which rows and why."""
+
+
+@dataclass(frozen=True)
+class Cell:
+    """One (repeat, fold) test set: the paired unit of every comparison."""
+
+    repeat: object
+    fold: object
+    train_idx: np.ndarray
+    test_idx: np.ndarray
+
+
+@dataclass(frozen=True)
+class Dataset:
+    """Validated input. Row order follows the input frame; `mol_index` maps each row to `mols`."""
+
+    task: str
+    y: np.ndarray
+    pred: np.ndarray
+    mols: tuple
+    mol_index: np.ndarray
+    cells: tuple[Cell, ...]
+    classes: tuple | None = None
+    context: np.ndarray | None = None  # per-row group code over the context columns, if any were given
+    context_columns: tuple[str, ...] = ()
+
+    @property
+    def test_train_ratio(self) -> float:
+        """Mean n_test / n_train over cells, for the Nadeau-Bengio correction."""
+        return float(np.mean([len(c.test_idx) / len(c.train_idx) for c in self.cells]))
+
+
+def _rows(mask: np.ndarray, limit: int = 10) -> str:
+    """Name offending rows as 0-based positions in the input frame."""
+    rows = np.flatnonzero(mask)
+    shown = ", ".join(str(r) for r in rows[:limit])
+    return shown + (f", ... ({len(rows)} rows)" if len(rows) > limit else "")
+
+
+def _parse_smiles(smiles: np.ndarray) -> tuple[tuple, np.ndarray]:
+    unique, inverse = np.unique(smiles.astype(str), return_inverse=True)
+    RDLogger.DisableLog("rdApp.*")
+    try:
+        mols = [Chem.MolFromSmiles(s) for s in unique]
+    finally:
+        RDLogger.EnableLog("rdApp.*")
+    bad = np.array([m is None for m in mols])
+    if bad.any():
+        raise InputError(f"RDKit could not parse SMILES in rows {_rows(bad[inverse])}: {list(unique[bad][:5])}")
+    return tuple(mols), inverse
+
+
+def _resolve_task(target: pd.Series, task: str | None) -> tuple[str, np.ndarray, tuple | None]:
+    n_unique = target.nunique()
+    if task is None:
+        task = "binary" if n_unique == 2 else "regression"
+    if task not in TASKS:
+        raise InputError(f"task must be one of {TASKS}, got {task!r}")
+    if task == "binary":
+        if n_unique != 2:
+            raise InputError(f"task='binary' needs exactly 2 distinct target values, found {n_unique}")
+        classes = tuple(sorted(target.unique()))
+        return task, (target.to_numpy() == classes[1]).astype(float), classes
+    try:
+        y = pd.to_numeric(target, errors="raise").to_numpy(dtype=float)
+    except (ValueError, TypeError) as exc:
+        raise InputError(f"task='regression' needs a numeric target: {exc}") from None
+    return task, y, None
+
+
+def _build_cells(fold: pd.Series, repeat: pd.Series) -> tuple[Cell, ...]:
+    cells = []
+    for r in sorted(repeat.unique(), key=str):
+        in_repeat = (repeat == r).to_numpy()
+        folds = sorted(fold[in_repeat].unique(), key=str)
+        if len(folds) < 2:
+            raise InputError(f"repeat {r!r} has a single fold ({folds[0]!r}), so it has no training set")
+        for f in folds:
+            test = in_repeat & (fold == f).to_numpy()
+            cells.append(Cell(r, f, np.flatnonzero(in_repeat & ~test), np.flatnonzero(test)))
+    return tuple(cells)
+
+
+def _check_cells(cells: tuple[Cell, ...], task: str, y: np.ndarray) -> None:
+    for c in cells:
+        where = f"repeat {c.repeat!r}, fold {c.fold!r}"
+        if len(c.train_idx) < MIN_TRAIN_ROWS:
+            raise InputError(
+                f"{where} has {len(c.train_idx)} training rows; the forest needs at least {MIN_TRAIN_ROWS} "
+                "(it is a forest, not a shrub)"
+            )
+        if task == "binary":
+            train_counts = np.bincount(y[c.train_idx].astype(int), minlength=2)
+            if train_counts.min() < 2:
+                raise InputError(f"{where}: the training set needs at least 2 rows of each class, has {train_counts}")
+            if len(np.unique(y[c.test_idx])) < 2:
+                raise InputError(f"{where}: the test set has one class, so ROC-AUC is undefined. Use stratified folds.")
+
+
+def load(
+    df: pd.DataFrame,
+    *,
+    smiles: str = "smiles",
+    target: str = "target",
+    fold: str = "fold",
+    pred: str = "pred",
+    repeat: str | None = None,
+    task: str | None = None,
+    context: str | Sequence[str] | None = None,
+) -> Dataset:
+    """Validate a frame with one row per (molecule, repeat) and return the paired test cells.
+
+    The test set of cell (r, f) is the rows with ``repeat == r`` and ``fold == f``; its training set is every
+    other row of repeat r. `pred` holds the user's model's prediction for the row when it was in the test set:
+    a value for regression, the probability of the larger class label for binary. `context` names the column(s)
+    whose combined values define a row's context group for the context-mean baseline.
+    """
+    context_columns = (context,) if isinstance(context, str) else tuple(context or ())
+    columns = [smiles, target, fold, pred] + ([repeat] if repeat else []) + list(context_columns)
+    missing = [c for c in columns if c not in df.columns]
+    if missing:
+        raise InputError(f"missing columns {missing}; the frame has {list(df.columns)}")
+    df = df.reset_index(drop=True)
+    for c in columns:
+        if df[c].isna().any():
+            raise InputError(f"column {c!r} is empty in rows {_rows(df[c].isna().to_numpy())}")
+
+    task, y, classes = _resolve_task(df[target], task)
+    try:
+        p = pd.to_numeric(df[pred], errors="raise").to_numpy(dtype=float)
+    except (ValueError, TypeError) as exc:
+        raise InputError(f"column {pred!r} must be numeric: {exc}") from None
+    if task == "binary":
+        outside = (p < 0) | (p > 1)
+        if outside.any():
+            raise InputError(
+                f"for a binary target, {pred!r} must be the probability of class {classes[1]!r}; "
+                f"rows {_rows(outside)} are outside [0, 1]"
+            )
+
+    mols, mol_index = _parse_smiles(df[smiles].to_numpy())
+    repeats = df[repeat] if repeat else pd.Series(0, index=df.index)
+    cells = _build_cells(df[fold], repeats)
+    _check_cells(cells, task, y)
+    codes = df.groupby(list(context_columns), sort=False).ngroup().to_numpy() if context_columns else None
+    return Dataset(
+        task=task,
+        y=y,
+        pred=p,
+        mols=mols,
+        mol_index=mol_index,
+        cells=cells,
+        classes=classes,
+        context=codes,
+        context_columns=context_columns,
+    )