Active Learning for Drug Discovery Campaigns

Traditional HTS screens millions of compounds in a single batch; ML-guided active learning instead runs iterative cycles of: predict → rank → select → test → retrain. Each new batch of experimental results refines the model, progressively focussing the search on productive chemical regions. Drug discovery timelines and costs have motivated several industry deployments (Merck, BenevolentAI, Recursion) demonstrating that 2–3 active-learning cycles can recover >80% of all potent hits from a virtual library while testing <5% of the space.

Acquisition Functions for Molecular Optimisation

import numpy as np
from scipy.stats import norm


def upper_confidence_bound(
    mu: np.ndarray,
    sigma: np.ndarray,
    kappa: float = 2.0,
) -> np.ndarray:
    """
    UCB acquisition: score = μ + κ·σ.

    Higher κ → more exploration. Typical range: κ ∈ [1, 5].
    """
    return mu + kappa * sigma


def expected_improvement(
    mu: np.ndarray,
    sigma: np.ndarray,
    y_best: float,
    xi: float = 0.01,
) -> np.ndarray:
    """
    Expected Improvement (EI) acquisition function.

    EI(x) = (μ−y*−ξ)·Φ(Z) + σ·φ(Z), where Z=(μ−y*−ξ)/σ.
    Higher xi → more exploration vs exploitation.

    Args:
        mu: Predicted mean for each candidate
        sigma: Predicted std for each candidate
        y_best: Best observed value so far (e.g. best pIC50)
        xi: Jitter term to control exploration

    Returns:
        EI scores (non-negative) — higher means more promising to test.
    """
    improvement = mu - y_best - xi
    Z = np.where(sigma > 1e-8, improvement / sigma, 0.0)
    ei = improvement * norm.cdf(Z) + sigma * norm.pdf(Z)
    return np.maximum(ei, 0.0)


def probability_of_improvement(
    mu: np.ndarray,
    sigma: np.ndarray,
    y_best: float,
    xi: float = 0.01,
) -> np.ndarray:
    """PI(x) = Φ((μ − y* − ξ) / σ). Simple but tends to over-exploit."""
    Z = np.where(sigma > 1e-8, (mu - y_best - xi) / sigma, 0.0)
    return norm.cdf(Z)

Batch Selection with Diversity Constraint

import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem, DataStructs


def select_diverse_batch(
    candidate_smiles: list[str],
    acquisition_scores: np.ndarray,
    batch_size: int = 48,
    diversity_threshold: float = 0.4,
    top_k_pool: int = 500,
) -> list[int]:
    """
    Select a batch of candidates balancing acquisition score and chemical diversity.

    Strategy: Pre-filter to top-k by acquisition score, then greedily pick
    the highest-scoring compound, then iteratively add compounds that are
    sufficiently dissimilar (Tanimoto < 1 − diversity_threshold) from all
    already-selected compounds.

    Args:
        candidate_smiles: SMILES of all unscreened candidates
        acquisition_scores: EI/UCB scores aligned with candidate_smiles
        batch_size: Target number of compounds to select
        diversity_threshold: Minimum diversity (1 - max Tanimoto) to any selected compound
        top_k_pool: Restrict greedy search to top-K by acquisition score first

    Returns:
        Indices of selected compounds into candidate_smiles list.
    """
    sorted_idx = np.argsort(acquisition_scores)[::-1]
    pool_idx = sorted_idx[:top_k_pool]

    fps = []
    valid_pool = []
    for i in pool_idx:
        mol = Chem.MolFromSmiles(candidate_smiles[i])
        if mol is not None:
            fps.append(AllChem.GetMorganFingerprintAsBitVect(mol, 2, 1024))
            valid_pool.append(int(i))

    selected = [valid_pool[0]]
    selected_fps = [fps[0]]

    for i, fp in zip(valid_pool[1:], fps[1:]):
        if len(selected) >= batch_size:
            break
        sims = DataStructs.BulkTanimotoSimilarity(fp, selected_fps)
        if max(sims) < (1.0 - diversity_threshold):
            selected.append(i)
            selected_fps.append(fp)

    return selected


def simulate_active_learning_campaign(
    all_smiles: list[str],
    true_labels: np.ndarray,
    n_initial: int = 96,
    n_rounds: int = 5,
    batch_size: int = 96,
    acquisition: str = "ei",
    random_seed: int = 42,
) -> dict:
    """
    Simulate an active learning drug discovery campaign using RF with variance estimation.

    Tracks hit recovery: fraction of all true positives found vs compounds tested.
    """
    from rdkit import Chem
    from rdkit.Chem import AllChem
    from sklearn.ensemble import RandomForestClassifier
    import numpy as np

    rng = np.random.default_rng(random_seed)
    n = len(all_smiles)

    # Featurise
    fps = []
    valid_idx = []
    for i, smi in enumerate(all_smiles):
        mol = Chem.MolFromSmiles(smi)
        if mol:
            fps.append(list(AllChem.GetMorganFingerprintAsBitVect(mol, 2, 1024)))
            valid_idx.append(i)
    X = np.array(fps)
    y = true_labels[valid_idx]
    N = len(valid_idx)

    # Initial random subset
    initial_idx = rng.choice(N, n_initial, replace=False).tolist()
    tested = set(initial_idx)
    history = {"n_tested": [], "hit_recovery": [], "round": []}

    total_hits = int(y.sum())

    for round_i in range(n_rounds + 1):
        tested_l = list(tested)
        untested_l = [i for i in range(N) if i not in tested]

        X_train, y_train = X[tested_l], y[tested_l]
        clf = RandomForestClassifier(n_estimators=200, random_state=random_seed)
        clf.fit(X_train, y_train)

        hits_found = int(y[tested_l].sum())
        history["n_tested"].append(len(tested))
        history["hit_recovery"].append(round(hits_found / max(total_hits, 1), 4))
        history["round"].append(round_i)

        if round_i == n_rounds or not untested_l:
            break

        X_pool = X[untested_l]
        probas = np.array([t.predict_proba(X_pool)[:, 1] for t in clf.estimators_])
        mu = probas.mean(axis=0)
        sigma = probas.std(axis=0)
        y_best = float(y_train.max())

        if acquisition == "ei":
            acq_scores = expected_improvement(mu, sigma, y_best=y_best)
        else:
            acq_scores = upper_confidence_bound(mu, sigma)

        pool_smiles = [all_smiles[valid_idx[i]] for i in untested_l]
        selected_rel = select_diverse_batch(pool_smiles, acq_scores, batch_size=batch_size)
        for rel_i in selected_rel:
            tested.add(untested_l[rel_i])

    return history

Active Learning vs Random Baseline

Key Takeaways

• Expected Improvement with diversity-constrained batch selection outperforms pure greedy exploitation in drug discovery because productive chemical space is often multi-modal — novel scaffolds emerge from exploration steps that greedy strategies miss.
• Uncertainty quantification is essential for acquisition functions. Ensembles of neural networks or random forests with tree-level variance, Gaussian processes, or MC-Dropout all provide reasonable uncertainty estimates for molecular inputs.
• The initial training set composition (n₀ ~96–384 random compounds) has a large impact on convergence speed. Maximally diverse initial sets (Butina-clustered random selection) provide better coverage than random sampling alone.
• Active learning ROI is highest when the assay is expensive or slow (enzymatic/cellular). For cheap HTS, the overhead of model training and synthesis routing may not be worthwhile unless the virtual screening space is very large (>10⁶ compounds).

Try it live

Explore a related model hands-on: open Molecule Screening Engine — Similarity Search Live and change its parameters while it runs. Everything stays in your browser.