Make DataSAIL scaling portfolio ready

Author: Kashish Mendiratta

.github/workflows/ci.yml

@@ -0,0 +1,25 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+permissions:
+  contents: read
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: pip
+      - name: Install test dependencies
+        run: python -m pip install --upgrade pip && python -m pip install -e ".[test]"
+      - name: Run tests
+        run: pytest -q
+      - name: Check code quality
+        run: ruff check src tests

README.md

@@ -1,5 +1,8 @@
 # Scaling DataSAIL to Large Datasets
 
+[![CI](https://github.com/KashishMendiratta/DataSAIL_scaling/actions/workflows/ci.yml/badge.svg)](https://github.com/KashishMendiratta/DataSAIL_scaling/actions/workflows/ci.yml)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue)](https://www.python.org/)
+
 Master's thesis project investigating how to scale [DataSAIL](https://doi.org/10.1038/s41467-025-58606-8) — a leakage-aware dataset splitting framework for molecular ML — to large datasets (10k–100k+ samples), where full pairwise similarity computation and ILP optimization become computationally impractical.
 
 The project proposes a hybrid approximation: run DataSAIL on a representative subset of the data, then extend the resulting splits to the remaining samples using a **balance-aware k-nearest-neighbour assignment** strategy.
@@ -17,11 +20,13 @@ Across BACE, BBBP, Tox21, and HIV (MoleculeNet benchmarks), the scaled approxima
 | BACE (~1.5k) | 3.1s | 5.8s | overhead dominates on small data |
 | BBBP (~2k) | 3.8s | 6.3s | overhead dominates on small data |
 | Tox21 (~8k) | 31.5s | 21.9s | 1.44x speedup |
-| HIV (~41k) | ~1800s (est.) | ~312s | full DataSAIL impractical at this scale |
+| HIV (~41k) | ~1800s (est.) | 298.0s | full DataSAIL impractical at this scale |
+
+Leakage fraction (test molecules with a near-duplicate in train, similarity > 0.7) dropped substantially vs. random splitting across all datasets — e.g. for BACE: 84.9% (random) → 7.3% (full DataSAIL) → 1.2–4.1% (scaled). See [`reports/MASTER REPORT.pdf`](reports/MASTER%20REPORT.pdf) for the full leakage, split-quality, and downstream ML performance analysis (Accuracy, F1, MCC, ROC-AUC, and PR-AUC using Random Forest and Logistic Regression classifiers).
 
-Leakage fraction (test molecules with a near-duplicate in train, similarity > 0.7) dropped substantially vs. random splitting across all datasets — e.g. for BACE: 84.9% (random) → 7.3% (full DataSAIL) → 1.2–4.1% (scaled). See `report/` for full leakage, split-quality, and downstream ML performance analysis (Accuracy, F1, MCC, ROC-AUC, PR-AUC on Random Forest / Logistic Regression classifiers).
+**Balanced kNN assignment was a necessary addition, not an optional refinement**: naive nearest-neighbour propagation caused runaway train-split dominance (up to ±14.5 percentage-point deviation from target ratios on BBBP), since unassigned molecules statistically favor neighbours in the largest existing partition. The balance-aware correction reduced the maximum deviation to approximately 0.5 percentage points or less in the final four-dataset evaluation.
 
-**Balanced kNN assignment was a necessary addition, not an optional refinement**: naive nearest-neighbour propagation caused runaway train-split dominance (up to ±14.5 percentage-point deviation from target ratios on BBBP), since unassigned molecules statistically favor neighbours in the largest existing partition. The balance-aware correction reduced max deviation to ≤0.5pp across all four datasets.
+![Runtime comparison between full and scaled DataSAIL](results/analysis/plots/runtime_comparison.png)
 
 ---
 
@@ -88,8 +93,10 @@ Data-SAIL_scaling/
 │   ├── experiments/<dataset>/    # per-dataset splits, runtime logs, split statistics
 │   └── analysis/                 # aggregated CSVs (leakage, DRE, KL divergence, ML results) + plots/
 │
-├── report/                       # thesis report (LaTeX/PDF) and seminar slides
-├── environment.yml
+├── reports/                      # complete project report
+├── tests/                        # deterministic unit tests for core algorithms
+├── environment.yml              # complete research environment
+├── pyproject.toml                # lightweight package and CI dependencies
 └── README.md
 ```
 
@@ -99,11 +106,22 @@ Data-SAIL_scaling/
 
 ## Setup
 
+For the complete research environment:
+
 ```bash
 conda env create -f environment.yml
 conda activate datasail-scale
 ```
 
+For the lightweight core package and its tests:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+python -m pip install -e ".[test]"
+pytest -q
+```
+
 ---
 
 ## Running the Pipeline
@@ -168,7 +186,7 @@ All outputs are written to `results/analysis/` (CSVs) and `results/analysis/plot
 2. **Run DataSAIL** on the subset only — similarity computation and ILP optimization scale quadratically with dataset size, so restricting this step to the subset is what makes the approach tractable.
 3. **Assign remaining molecules** to train/val/test using balanced kNN: for each unassigned molecule, retrieve k nearest neighbours (Jaccard similarity on ECFP/Morgan fingerprints, radius=2, 2048 bits via RDKit), score each candidate split by similarity × balance-correction factor, and assign to the highest-scoring split. The balance factor penalizes already-oversized splits and boosts underrepresented ones, which is what prevents the train-dominance failure seen under naive nearest-neighbour propagation.
 
-Full derivation, complexity analysis (O(k² + (n−k)·K) vs. O(n²) for full DataSAIL), and evaluation framework (leakage metrics, split-quality metrics, downstream ML performance) are in `report/`.
+Full derivation, complexity analysis (O(k² + (n−k)·K) vs. O(n²) for full DataSAIL), and evaluation framework (leakage metrics, split-quality metrics, downstream ML performance) are in [`reports/MASTER REPORT.pdf`](reports/MASTER%20REPORT.pdf).
 
 ---
 
@@ -187,4 +205,4 @@ From [MoleculeNet](https://doi.org/10.1039/c7sc02664a):
 
 ## Status
 
-Preliminary/exploratory results (single run, seed=42) — see `report/` conclusion and future work for planned extensions (larger datasets, dynamic hyperparameter tuning for k, evaluation on D-MPNN models).
+Research prototype with preliminary single-run results (`seed=42`). The central scaling result is promising, but it should not be interpreted as a statistically replicated benchmark. See the report's conclusion and future-work section for planned extensions, including repeated runs, larger datasets, dynamic tuning of `k`, and evaluation with D-MPNN models.

pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+requires = ["setuptools>=69"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "datasail-scaling"
+version = "0.1.0"
+description = "A balance-aware approximation for scaling leakage-conscious molecular dataset splitting."
+requires-python = ">=3.10"
+dependencies = [
+    "numpy>=1.26,<3",
+    "pandas>=2.1,<3",
+    "scikit-learn>=1.4,<2",
+]
+
+[project.optional-dependencies]
+test = ["pytest>=8,<9", "ruff>=0.12,<1"]
+experiments = [
+    "datasail",
+    "lightgbm",
+    "matplotlib",
+    "memory-profiler",
+    "pyyaml",
+    "rdkit",
+    "seaborn",
+    "tqdm",
+    "xgboost",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F"]

src/assignment/balanced_knn.py

@@ -4,21 +4,22 @@
 This addresses the majority class bias in naive kNN assignment by
 considering both similarity to neighbors AND current split sizes.
 """
+from collections import Counter
+from typing import Mapping
+
 import numpy as np
 import pandas as pd
 from sklearn.neighbors import NearestNeighbors
-from collections import Counter
-from typing import Tuple
 
 
 def balanced_knn_assign(
     train_fps: np.ndarray,
     train_splits: pd.Series,
     remaining_fps: np.ndarray,
     k: int = 5,
-    target_ratios: dict = None,
-    balance_weight: float = 1.0
-) -> Tuple[np.ndarray, np.ndarray]:
+    target_ratios: Mapping[str, float] | None = None,
+    balance_weight: float = 1.0,
+) -> tuple[np.ndarray, np.ndarray]:
     """
     Assign remaining samples using balanced kNN.
     
@@ -37,7 +38,21 @@ def balanced_knn_assign(
         Tuple of (assignments, confidence_scores)
     """
     if target_ratios is None:
-        target_ratios = {'train': 0.7, 'val': 0.2, 'test': 0.1}
+        target_ratios = {"train": 0.7, "val": 0.2, "test": 0.1}
+
+    if len(train_fps) != len(train_splits):
+        raise ValueError("train_fps and train_splits must have equal lengths")
+    if not 1 <= k <= len(train_fps):
+        raise ValueError("k must be between 1 and the number of reference samples")
+    if balance_weight < 0:
+        raise ValueError("balance_weight must be non-negative")
+    if not target_ratios or any(ratio <= 0 for ratio in target_ratios.values()):
+        raise ValueError("target_ratios must contain positive values")
+    if not np.isclose(sum(target_ratios.values()), 1.0):
+        raise ValueError("target_ratios must sum to 1")
+    unknown_splits = set(train_splits) - set(target_ratios)
+    if unknown_splits:
+        raise ValueError(f"Reference data contains unknown splits: {sorted(unknown_splits)}")
 
     # Fit kNN
     knn = NearestNeighbors(n_neighbors=k, metric='jaccard', n_jobs=-1)

src/assignment/knn_assignment.py

@@ -8,7 +8,7 @@
 import numpy as np
 import pandas as pd
 from sklearn.neighbors import NearestNeighbors
-from typing import Dict, Tuple
+from typing import Dict
 from collections import Counter
 
 
@@ -244,4 +244,4 @@ def evaluate_assignment_quality(
             )
         )
 
-    return metrics
+    return metrics

src/downsampling/strategies.py

@@ -6,7 +6,7 @@
 """
 import numpy as np
 import pandas as pd
-from typing import List, Tuple
+from typing import Tuple
 from sklearn.cluster import KMeans
 from sklearn.metrics.pairwise import euclidean_distances
 
@@ -24,15 +24,19 @@ def random_downsample(data: pd.DataFrame, ratio: float, seed: int = 42) -> Tuple
         Tuple of (sampled_data, remaining_data)
         Original indices are preserved!
     """
-    np.random.seed(seed)
+    if not 0 < ratio <= 1:
+        raise ValueError("ratio must be in the interval (0, 1]")
+
+    rng = np.random.default_rng(seed)
     n_samples = int(len(data) * ratio)
 
     # Get original indices
     all_indices = data.index.tolist()
 
     # Randomly select indices
-    sampled_indices = np.random.choice(all_indices, size=n_samples, replace=False)
-    remaining_indices = [idx for idx in all_indices if idx not in sampled_indices]
+    sampled_indices = rng.choice(all_indices, size=n_samples, replace=False)
+    sampled_index_set = set(sampled_indices)
+    remaining_indices = [idx for idx in all_indices if idx not in sampled_index_set]
 
     # Use loc to preserve original indices
     sampled_data = data.loc[sampled_indices].copy()
@@ -55,7 +59,12 @@ def stratified_downsample(data: pd.DataFrame, labels: pd.Series, ratio: float, s
         Tuple of (sampled_data, remaining_data)
         Original indices are preserved!
     """
-    np.random.seed(seed)
+    if not 0 < ratio <= 1:
+        raise ValueError("ratio must be in the interval (0, 1]")
+    if not data.index.equals(labels.index):
+        raise ValueError("data and labels must have matching indices")
+
+    rng = np.random.default_rng(seed)
 
     sampled_indices = []
     for label in labels.unique():
@@ -64,10 +73,11 @@ def stratified_downsample(data: pd.DataFrame, labels: pd.Series, ratio: float, s
         n_class_samples = int(len(class_indices) * ratio)
 
         # Sample from this class
-        class_sampled = np.random.choice(class_indices, size=n_class_samples, replace=False)
+        class_sampled = rng.choice(class_indices, size=n_class_samples, replace=False)
         sampled_indices.extend(class_sampled)
 
-    remaining_indices = [idx for idx in data.index if idx not in sampled_indices]
+    sampled_index_set = set(sampled_indices)
+    remaining_indices = [idx for idx in data.index if idx not in sampled_index_set]
 
     sampled_data = data.loc[sampled_indices].copy()
     remaining_data = data.loc[remaining_indices].copy()

tests/test_assignment.py

@@ -0,0 +1,59 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from src.assignment.balanced_knn import balanced_knn_assign
+from src.assignment.knn_assignment import evaluate_assignment_quality, knn_assign
+
+
+def test_knn_assigns_nearest_reference_split() -> None:
+    reference = np.array([[1, 0, 0], [0, 1, 0]], dtype=bool)
+    splits = pd.Series(["train", "test"])
+    queries = np.array([[1, 0, 0], [0, 1, 0]], dtype=bool)
+
+    assignments, confidence = knn_assign(reference, splits, queries, k=1)
+
+    assert assignments.tolist() == ["train", "test"]
+    assert confidence.tolist() == [1.0, 1.0]
+
+
+def test_balanced_assignment_respects_requested_splits() -> None:
+    reference = np.eye(6, dtype=bool)
+    splits = pd.Series(["train", "train", "train", "val", "test", "test"])
+    queries = np.array([[1, 1, 0, 0, 0, 0], [0, 0, 0, 1, 1, 0]], dtype=bool)
+
+    assignments, confidence = balanced_knn_assign(
+        reference,
+        splits,
+        queries,
+        k=2,
+        target_ratios={"train": 0.5, "val": 0.25, "test": 0.25},
+    )
+
+    assert set(assignments).issubset({"train", "val", "test"})
+    assert np.all((confidence >= 0) & (confidence <= 1))
+
+
+def test_balanced_assignment_rejects_invalid_configuration() -> None:
+    reference = np.eye(2, dtype=bool)
+    splits = pd.Series(["train", "test"])
+
+    with pytest.raises(ValueError, match="sum to 1"):
+        balanced_knn_assign(
+            reference,
+            splits,
+            reference,
+            k=1,
+            target_ratios={"train": 0.8, "test": 0.3},
+        )
+
+
+def test_assignment_quality_summary() -> None:
+    metrics = evaluate_assignment_quality(
+        np.array(["train", "train", "test"]),
+        np.array([0.9, 0.4, 1.0]),
+    )
+
+    assert metrics["total_assigned"] == 3
+    assert metrics["split_distribution"] == {"train": 2, "test": 1}
+    assert metrics["low_confidence_count"] == 1