ALineMol Splitter Quick-Start¶

A guided tour of molecular dataset splitters in ALineMol, running entirely in the browser via Google Colab.

What you'll see. Train/test splits are the foundation of any honest ML evaluation. In molecular property prediction, how you split your data determines whether you're measuring in-distribution memorisation or genuine out-of-distribution (OOD) generalisation. ALineMol ships a unified API for over a dozen splitting strategies — scaffold-based, clustering-based, property-based, diversity-based. This notebook applies seven of them to the same dataset, then visualises and quantifies how the resulting partitions differ.

No local install required. Everything below runs on a free Colab CPU runtime in a few minutes. If you like what you see, the last section points you at the full repo.

1. Install ALineMol¶

The default install pulls in only the splitter API — no PyTorch or DGL — so the install is small (~1 minute) and works cleanly on Colab.

In [ ]:

!pip install -q git+https://github.com/HFooladi/ALineMol.git

!pip install -q git+https://github.com/HFooladi/ALineMol.git

If Colab prompts you to restart the runtime after install, restart and continue from the next cell.

2. Load BBBP¶

We use BBBP (Blood-Brain Barrier Penetration) — ~2,000 small molecules with binary labels for whether they cross the blood-brain barrier. We pull the CSV straight from the MoleculeNet mirror so this notebook has no auxiliary dataset-loader dependency.

In [ ]:

import pandas as pd
from rdkit import Chem
from rdkit import RDLogger

RDLogger.DisableLog("rdApp.*")  # silence rdkit parsing warnings

URL = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/BBBP.csv"
raw = pd.read_csv(URL)

# Filter rows whose SMILES RDKit refuses to parse (BBBP contains a few salts
# and broken strings); keep the largest fragment for any mixture so scaffold
# splitters can compute Bemis-Murcko frameworks reliably.
def _clean(smi):
    mol = Chem.MolFromSmiles(smi)
    if mol is None:
        return None
    frags = Chem.GetMolFrags(mol, asMols=True)
    if not frags:
        return None
    largest = max(frags, key=lambda m: m.GetNumHeavyAtoms())
    return Chem.MolToSmiles(largest)

raw["clean_smiles"] = raw["smiles"].map(_clean)
data = raw.dropna(subset=["clean_smiles"]).reset_index(drop=True)

smiles = data["clean_smiles"].tolist()
labels = data["p_np"].tolist()

print(f"Dataset: BBBP ({len(smiles)} molecules, {len(raw) - len(data)} dropped during parse)")
print(f"Class balance: {sum(labels)}/{len(labels)} positive "
      f"({100 * sum(labels) / len(labels):.1f}%)")
print(f"Example SMILES: {smiles[0]}")

import pandas as pd from rdkit import Chem from rdkit import RDLogger RDLogger.DisableLog("rdApp.*") # silence rdkit parsing warnings URL = "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/BBBP.csv" raw = pd.read_csv(URL) # Filter rows whose SMILES RDKit refuses to parse (BBBP contains a few salts # and broken strings); keep the largest fragment for any mixture so scaffold # splitters can compute Bemis-Murcko frameworks reliably. def _clean(smi): mol = Chem.MolFromSmiles(smi) if mol is None: return None frags = Chem.GetMolFrags(mol, asMols=True) if not frags: return None largest = max(frags, key=lambda m: m.GetNumHeavyAtoms()) return Chem.MolToSmiles(largest) raw["clean_smiles"] = raw["smiles"].map(_clean) data = raw.dropna(subset=["clean_smiles"]).reset_index(drop=True) smiles = data["clean_smiles"].tolist() labels = data["p_np"].tolist() print(f"Dataset: BBBP ({len(smiles)} molecules, {len(raw) - len(data)} dropped during parse)") print(f"Class balance: {sum(labels)}/{len(labels)} positive " f"({100 * sum(labels) / len(labels):.1f}%)") print(f"Example SMILES: {smiles[0]}")

3. Your first splitter¶

Every splitter in ALineMol is created via the get_splitter(name, **kwargs) factory and produces train/test index pairs via .split(smiles).

We start with the classic Bemis-Murcko scaffold split: molecules sharing a scaffold are kept together, so the test set sees only scaffolds the model never trained on.

In [ ]:

from alinemol.splitters import get_splitter, get_splitter_names

print("Available splitters:")
print(sorted(get_splitter_names()))

from alinemol.splitters import get_splitter, get_splitter_names print("Available splitters:") print(sorted(get_splitter_names()))

In [ ]:

splitter = get_splitter("scaffold", n_splits=1, test_size=0.2)
train_idx, test_idx = next(splitter.split(smiles))

print(f"Train: {len(train_idx)} molecules")
print(f"Test:  {len(test_idx)} molecules")
print()
print("First 3 train SMILES:")
for i in train_idx[:3]:
    print(f"  {smiles[i]}")
print()
print("First 3 test SMILES:")
for i in test_idx[:3]:
    print(f"  {smiles[i]}")

splitter = get_splitter("scaffold", n_splits=1, test_size=0.2) train_idx, test_idx = next(splitter.split(smiles)) print(f"Train: {len(train_idx)} molecules") print(f"Test: {len(test_idx)} molecules") print() print("First 3 train SMILES:") for i in train_idx[:3]: print(f" {smiles[i]}") print() print("First 3 test SMILES:") for i in test_idx[:3]: print(f" {smiles[i]}")

4. What does that split look like?¶

SplitAnalyzer computes Tanimoto similarities, scaffold overlap, property shifts, and more. We instantiate it once on the full SMILES list and reuse it throughout the notebook (fingerprints and scaffolds are cached lazily).

In [ ]:

from alinemol.splitters import SplitAnalyzer

analyzer = SplitAnalyzer(smiles)
report = analyzer.analyze_split(train_idx, test_idx, splitter_name="scaffold")

print(f"Train size:           {report.size_metrics.train_size}")
print(f"Test size:            {report.size_metrics.test_size}")
print(f"Mean max-Tanimoto:    {report.similarity_metrics.mean_sim:.3f}")
print(f"Median max-Tanimoto:  {report.similarity_metrics.median_sim:.3f}")
print(f"Scaffold overlap:     {report.scaffold_metrics.scaffold_overlap_percentage:.1f}%")

from alinemol.splitters import SplitAnalyzer analyzer = SplitAnalyzer(smiles) report = analyzer.analyze_split(train_idx, test_idx, splitter_name="scaffold") print(f"Train size: {report.size_metrics.train_size}") print(f"Test size: {report.size_metrics.test_size}") print(f"Mean max-Tanimoto: {report.similarity_metrics.mean_sim:.3f}") print(f"Median max-Tanimoto: {report.similarity_metrics.median_sim:.3f}") print(f"Scaffold overlap: {report.scaffold_metrics.scaffold_overlap_percentage:.1f}%")

How to read the similarity number. For each test molecule, we measure its maximum Tanimoto similarity to any training molecule. A low mean means the test set is structurally far from train — a harder, more honest OOD evaluation. A high mean (close to 1.0) means most test molecules have a near-twin in train.

Let's project the split into 2D with UMAP to see the partition geometry.

In [ ]:

from alinemol.splitters.analyzer_plots import plot_chemical_space
import matplotlib.pyplot as plt

fig = plot_chemical_space(analyzer, train_idx, test_idx, method="umap")
plt.show()

from alinemol.splitters.analyzer_plots import plot_chemical_space import matplotlib.pyplot as plt fig = plot_chemical_space(analyzer, train_idx, test_idx, method="umap") plt.show()

5. Compare seven splitters side-by-side¶

We pick one canonical splitter from each conceptual family:

Family	Splitter	What it does
Baseline	random	Uniform shuffle (in-distribution)
Structure	scaffold	Group by Bemis-Murcko scaffold
Clustering	kmeans	k-Means on Morgan fingerprints
Clustering	butina	Taylor-Butina clustering on fingerprints
Clustering	umap	UMAP + agglomerative clustering
Property	molecular_weight	Heavy molecules in test
Diversity	perimeter	Test set on the edge of chemical space

For the visual comparison we pull a single split from each (fast). We'll do a proper multi-split statistical comparison in section 7.

In [ ]:

splitter_names = [
    "random",
    "scaffold",
    "kmeans",
    "butina",
    "umap",
    "molecular_weight",
    "perimeter",
]

splits = {}
for name in splitter_names:
    splitter = get_splitter(name, n_splits=1, test_size=0.2)
    splits[name] = next(splitter.split(smiles))
    train_n, test_n = len(splits[name][0]), len(splits[name][1])
    print(f"  {name:20s}  train={train_n:5d}  test={test_n:4d}")

splitter_names = [ "random", "scaffold", "kmeans", "butina", "umap", "molecular_weight", "perimeter", ] splits = {} for name in splitter_names: splitter = get_splitter(name, n_splits=1, test_size=0.2) splits[name] = next(splitter.split(smiles)) train_n, test_n = len(splits[name][0]), len(splits[name][1]) print(f" {name:20s} train={train_n:5d} test={test_n:4d}")

6. Chemical-space grid¶

We compute a single UMAP embedding of all molecules and use it as a shared backdrop for every panel — that way the visual differences between panels reflect only the splitter, not a re-randomised projection.

In [ ]:

import numpy as np
import umap

# Cached on the analyzer instance — first access is the only one that pays.
fps = analyzer.fingerprints

reducer = umap.UMAP(n_components=2, random_state=42)
embedding = reducer.fit_transform(fps)
print(f"UMAP embedding shape: {embedding.shape}")

import numpy as np import umap # Cached on the analyzer instance — first access is the only one that pays. fps = analyzer.fingerprints reducer = umap.UMAP(n_components=2, random_state=42) embedding = reducer.fit_transform(fps) print(f"UMAP embedding shape: {embedding.shape}")

In [ ]:

ncols = 4
nrows = (len(splitter_names) + ncols - 1) // ncols
fig, axes = plt.subplots(nrows, ncols, figsize=(4.5 * ncols, 4 * nrows))
axes = axes.flatten()

for i, name in enumerate(splitter_names):
    train_idx_i, test_idx_i = splits[name]
    ax = axes[i]
    ax.scatter(embedding[train_idx_i, 0], embedding[train_idx_i, 1],
               c="#1f77b4", alpha=0.4, s=10,
               label=f"Train ({len(train_idx_i)})")
    ax.scatter(embedding[test_idx_i, 0], embedding[test_idx_i, 1],
               c="#ff7f0e", alpha=0.8, s=12,
               label=f"Test ({len(test_idx_i)})")
    ax.set_title(name, fontsize=11)
    ax.set_xlabel("UMAP 1")
    ax.set_ylabel("UMAP 2")
    ax.legend(loc="best", fontsize=8)

for j in range(len(splitter_names), len(axes)):
    axes[j].set_visible(False)

plt.suptitle("Same dataset, same UMAP projection, different splitters", y=1.02, fontsize=13)
plt.tight_layout()
plt.show()

ncols = 4 nrows = (len(splitter_names) + ncols - 1) // ncols fig, axes = plt.subplots(nrows, ncols, figsize=(4.5 * ncols, 4 * nrows)) axes = axes.flatten() for i, name in enumerate(splitter_names): train_idx_i, test_idx_i = splits[name] ax = axes[i] ax.scatter(embedding[train_idx_i, 0], embedding[train_idx_i, 1], c="#1f77b4", alpha=0.4, s=10, label=f"Train ({len(train_idx_i)})") ax.scatter(embedding[test_idx_i, 0], embedding[test_idx_i, 1], c="#ff7f0e", alpha=0.8, s=12, label=f"Test ({len(test_idx_i)})") ax.set_title(name, fontsize=11) ax.set_xlabel("UMAP 1") ax.set_ylabel("UMAP 2") ax.legend(loc="best", fontsize=8) for j in range(len(splitter_names), len(axes)): axes[j].set_visible(False) plt.suptitle("Same dataset, same UMAP projection, different splitters", y=1.02, fontsize=13) plt.tight_layout() plt.show()

Notice the contrast: random scatters the orange test points uniformly through the blue cloud; scaffold, umap, and perimeter collect them into compact regions or along the periphery; molecular_weight sorts by a property axis instead of clustering.

7. How structurally similar is test to train?¶

The most intuitive OOD-difficulty diagnostic: for each test molecule, find its max Tanimoto similarity to any training molecule, then look at the distribution. random should sit near 1.0 (every test mol has near-twins in train); structural and clustering splitters should shift to the left.

In [ ]:

import seaborn as sns
import pandas as pd


def max_tanimoto_to_train(train_fps: np.ndarray, test_fps: np.ndarray) -> np.ndarray:
    train = train_fps.astype(bool)
    test = test_fps.astype(bool)
    intersect = test.astype(np.uint16) @ train.astype(np.uint16).T
    train_count = train.sum(axis=1)
    test_count = test.sum(axis=1)
    union = train_count[None, :] + test_count[:, None] - intersect
    with np.errstate(divide="ignore", invalid="ignore"):
        sim = np.where(union > 0, intersect / union, 0.0)
    return sim.max(axis=1)


records = []
for name in splitter_names:
    train_idx_i, test_idx_i = splits[name]
    max_sims = max_tanimoto_to_train(fps[train_idx_i], fps[test_idx_i])
    records.extend({"splitter": name, "max_tanimoto_to_train": s} for s in max_sims)

sim_df = pd.DataFrame(records)

import seaborn as sns import pandas as pd def max_tanimoto_to_train(train_fps: np.ndarray, test_fps: np.ndarray) -> np.ndarray: train = train_fps.astype(bool) test = test_fps.astype(bool) intersect = test.astype(np.uint16) @ train.astype(np.uint16).T train_count = train.sum(axis=1) test_count = test.sum(axis=1) union = train_count[None, :] + test_count[:, None] - intersect with np.errstate(divide="ignore", invalid="ignore"): sim = np.where(union > 0, intersect / union, 0.0) return sim.max(axis=1) records = [] for name in splitter_names: train_idx_i, test_idx_i = splits[name] max_sims = max_tanimoto_to_train(fps[train_idx_i], fps[test_idx_i]) records.extend({"splitter": name, "max_tanimoto_to_train": s} for s in max_sims) sim_df = pd.DataFrame(records)

In [ ]:

plt.figure(figsize=(10, 6))
palette = sns.color_palette("tab10", n_colors=len(splitter_names))
for name, color in zip(splitter_names, palette):
    sub = sim_df.loc[sim_df["splitter"] == name, "max_tanimoto_to_train"]
    sns.kdeplot(sub, label=name, color=color, linewidth=2, clip=(0, 1))

plt.xlabel("Max Tanimoto similarity of each test molecule to the training set")
plt.ylabel("Density")
plt.title("Lower curves on the left = harder OOD test set")
plt.legend()
plt.tight_layout()
plt.show()

plt.figure(figsize=(10, 6)) palette = sns.color_palette("tab10", n_colors=len(splitter_names)) for name, color in zip(splitter_names, palette): sub = sim_df.loc[sim_df["splitter"] == name, "max_tanimoto_to_train"] sns.kdeplot(sub, label=name, color=color, linewidth=2, clip=(0, 1)) plt.xlabel("Max Tanimoto similarity of each test molecule to the training set") plt.ylabel("Density") plt.title("Lower curves on the left = harder OOD test set") plt.legend() plt.tight_layout() plt.show()

8. Quantitative comparison across multiple splits¶

Single splits can be noisy. compare_splitters runs each strategy multiple times and aggregates the metrics. We use n_splits=3 here to keep the notebook quick; bump it to 5 or 10 for production analysis.

Columns explained:

• sim_mean_sim_mean — mean (over splits) of the mean test-to-train Tanimoto. Lower = harder OOD.
• scaffold_..._overlap_..._mean — fraction of test scaffolds also present in train.
• prop_mw_ks_stat_mean / prop_logp_ks_stat_mean — Kolmogorov-Smirnov statistic comparing train vs test distributions of MW and LogP. Higher = larger property shift.

In [ ]:

comparison = analyzer.compare_splitters(
    splitter_names=splitter_names,
    n_splits=3,
    test_size=0.2,
)

display_cols = [
    "splitter_name",
    "sim_mean_sim_mean",
    "sim_mean_sim_std",
    "scaffold_scaffold_overlap_percentage_mean",
    "prop_mw_ks_stat_mean",
    "prop_logp_ks_stat_mean",
]
display_cols = [c for c in display_cols if c in comparison.columns]
comparison[display_cols].round(3)

comparison = analyzer.compare_splitters( splitter_names=splitter_names, n_splits=3, test_size=0.2, ) display_cols = [ "splitter_name", "sim_mean_sim_mean", "sim_mean_sim_std", "scaffold_scaffold_overlap_percentage_mean", "prop_mw_ks_stat_mean", "prop_logp_ks_stat_mean", ] display_cols = [c for c in display_cols if c in comparison.columns] comparison[display_cols].round(3)

9. Radar chart: splitters at a glance¶

A single visual that summarises the comparison table. Each axis is normalised to [0, 1] across splitters.

In [ ]:

from alinemol.splitters.analyzer_plots import plot_splitter_comparison_radar

fig = plot_splitter_comparison_radar(comparison)
plt.show()

from alinemol.splitters.analyzer_plots import plot_splitter_comparison_radar fig = plot_splitter_comparison_radar(comparison) plt.show()

10. Property splits do what they say on the tin¶

Property-based splitters move the test distribution along an explicit axis (here, molecular weight). Compare random (no shift) against molecular_weight (heavy molecules in test).

In [ ]:

from alinemol.splitters.analyzer_plots import plot_property_distributions

print("Random split — train and test should overlap:")
fig = plot_property_distributions(
    analyzer, splits["random"][0], splits["random"][1],
    properties=["MW", "LogP"],
)
plt.show()

print("\nMolecular weight split — test should sit to the right of train on MW:")
fig = plot_property_distributions(
    analyzer, splits["molecular_weight"][0], splits["molecular_weight"][1],
    properties=["MW", "LogP"],
)
plt.show()

from alinemol.splitters.analyzer_plots import plot_property_distributions print("Random split — train and test should overlap:") fig = plot_property_distributions( analyzer, splits["random"][0], splits["random"][1], properties=["MW", "LogP"], ) plt.show() print("\nMolecular weight split — test should sit to the right of train on MW:") fig = plot_property_distributions( analyzer, splits["molecular_weight"][0], splits["molecular_weight"][1], properties=["MW", "LogP"], ) plt.show()

11. What else is available¶

ALineMol ships more splitters than we demoed here. Quick tour of the rest:

• scaffold_generic — like scaffold but with atom-generic scaffolds (looser grouping).
• scaffold_kmeans — hybrid: extract scaffolds, then k-means cluster them.
• molecular_weight_reverse — light molecules in test instead of heavy.
• molecular_logp — same idea, on LogP.
• max_dissimilarity — greedy maximum-diversity test selection.
• lo and hi — similarity-based splits for lead optimisation / hit identification (require continuous activity values, not just SMILES).
• datasail — minimises information leakage; requires the external datasail package.

Run get_splitter_names() for the full list, and list_splitters() to see the class behind each name.

12. Next steps¶

• Install locally:

  git clone https://github.com/HFooladi/ALineMol.git
  cd ALineMol && ./install.sh
  

• Run splitters from the command line: python scripts/splitting.py --help
• Source and docs: https://github.com/HFooladi/ALineMol
• Underlying paper / methodology: see the repository README.

Found a bug or want to add a splitter? PRs welcome.