diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
new file mode 100644
index 0000000..0a27fdb
--- /dev/null
+++ b/.github/workflows/docs.yml
@@ -0,0 +1,90 @@
+name: Deploy Documentation
+
+# Publishes versioned docs to the gh-pages branch with mike.
+# - a GitHub release -> deploy that version and move the `latest` alias
+# - a push to main -> deploy the floating `dev` version (unreleased docs)
+# - manual dispatch -> deploy an arbitrary version/alias (safe for testing)
+#
+# GitHub Pages for this repo serves "from a branch" (gh-pages); mike commits and
+# pushes the built site there, and Pages rebuilds automatically. This is NOT the
+# actions/deploy-pages OIDC flow, so the only permission needed is contents:write.
+on:
+ release:
+ types: [published]
+ push:
+ branches: [main]
+ paths:
+ - "docs/**"
+ - "mkdocs.yml"
+ - "themap/**"
+ - "pyproject.toml"
+ - ".github/workflows/docs.yml"
+ workflow_dispatch:
+ inputs:
+ version:
+ description: "Version identifier to deploy (e.g. 0.5.0 or test)"
+ required: true
+ default: "test"
+ alias:
+ description: "Alias to point at this version (leave blank for none)"
+ required: false
+ default: ""
+
+# Serialize deploys: they all mutate the single gh-pages branch, and concurrent
+# pushes would race and reject. Do not cancel in-flight deploys (a half-finished
+# mike push could leave gh-pages inconsistent) — queue them instead.
+concurrency:
+ group: docs-deploy
+ cancel-in-progress: false
+
+jobs:
+ deploy:
+ runs-on: ubuntu-latest
+ permissions:
+ # mike pushes the built site to the gh-pages branch.
+ contents: write
+ steps:
+ - uses: actions/checkout@v4
+ with:
+ # mike needs full history and the gh-pages branch; a shallow clone
+ # breaks versioned deploys.
+ fetch-depth: 0
+
+ - name: Set up Python
+ uses: actions/setup-python@v5
+ with:
+ python-version: "3.10"
+
+ - name: Install uv
+ uses: astral-sh/setup-uv@v4
+ with:
+ enable-cache: true
+ cache-dependency-glob: "pyproject.toml"
+
+ - name: Install dependencies
+ run: uv pip install --system -e ".[docs]"
+
+ - name: Configure git identity
+ run: |
+ git config user.name "github-actions[bot]"
+ git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
+
+ - name: Deploy released version
+ if: github.event_name == 'release'
+ run: |
+ VERSION="${GITHUB_REF_NAME#v}"
+ mike deploy --push --update-aliases "$VERSION" latest
+ mike set-default --push latest
+
+ - name: Deploy dev docs from main
+ if: github.event_name == 'push'
+ run: mike deploy --push dev
+
+ - name: Deploy manual dispatch
+ if: github.event_name == 'workflow_dispatch'
+ run: |
+ if [ -n "${{ inputs.alias }}" ]; then
+ mike deploy --push --update-aliases "${{ inputs.version }}" "${{ inputs.alias }}"
+ else
+ mike deploy --push "${{ inputs.version }}"
+ fi
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
deleted file mode 100644
index dfdfba5..0000000
--- a/.readthedocs.yaml
+++ /dev/null
@@ -1,20 +0,0 @@
-# .readthedocs.yaml
-# Read the Docs configuration file
-# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
-
-version: 2
-
-build:
- os: ubuntu-22.04
- tools:
- python: "3.10"
-
-mkdocs:
- configuration: mkdocs.yml
-
-python:
- install:
- - method: pip
- path: .
- extra_requirements:
- - docs
diff --git a/docs/assets/css/custom-splito.css b/docs/assets/css/themap.css
similarity index 84%
rename from docs/assets/css/custom-splito.css
rename to docs/assets/css/themap.css
index b8e43ed..a27b6b9 100644
--- a/docs/assets/css/custom-splito.css
+++ b/docs/assets/css/themap.css
@@ -14,11 +14,20 @@
}
[data-md-color-scheme="slate"] {
+ /* Lighter purple in dark mode so header text and links keep AA contrast. */
+ --md-primary-fg-color: #b39ddb;
+ --md-primary-fg-color--light: #d1c4e9;
+ --md-primary-fg-color--dark: #9575cd;
+ --md-accent-fg-color: #b388ff;
--md-code-bg-color: #1e1e2e;
--md-default-bg-color: #0d1117;
}
/* ===== Typography ===== */
+.md-typeset {
+ line-height: 1.7;
+}
+
.md-typeset h1 {
font-weight: 700;
letter-spacing: -0.02em;
@@ -27,9 +36,9 @@
.md-typeset h2 {
font-weight: 600;
- margin-top: 2rem;
- padding-bottom: 0.5rem;
- border-bottom: 2px solid var(--md-primary-fg-color--light);
+ margin-top: 2.5rem;
+ padding-bottom: 0.3rem;
+ border-bottom: 1px solid var(--md-default-fg-color--lightest);
}
.md-typeset h3 {
@@ -67,11 +76,11 @@
/* ===== Navigation ===== */
.md-nav__title {
font-weight: 600;
- font-size: 0.75rem;
+ font-size: 0.8rem;
}
.md-nav__link {
- font-size: 0.8rem;
+ font-size: 0.85rem;
}
.md-nav__item {
@@ -85,7 +94,7 @@
/* Right-side table of contents */
.md-nav--secondary .md-nav__link {
- font-size: 0.75rem;
+ font-size: 0.8rem;
}
/* Tab navigation styling */
@@ -105,7 +114,7 @@
overflow: hidden;
box-shadow: 0 1px 4px rgba(0, 0, 0, 0.08);
border: 1px solid rgba(0, 0, 0, 0.08);
- font-size: 0.85rem;
+ font-size: 0.9rem;
}
.md-typeset table:not([class]) th {
@@ -116,6 +125,12 @@
text-align: left;
}
+/* In dark mode the primary is a light purple; use a darker header tone so the
+ white header text keeps contrast. */
+[data-md-color-scheme="slate"] .md-typeset table:not([class]) th {
+ background: var(--md-primary-fg-color--dark);
+}
+
.md-typeset table:not([class]) td {
padding: 0.65rem 1rem;
border-bottom: 1px solid rgba(0, 0, 0, 0.05);
@@ -165,17 +180,17 @@
/* ===== Footer ===== */
.md-footer {
- background: linear-gradient(135deg, var(--themap-gradient-start), var(--themap-gradient-end));
+ background-color: var(--md-primary-fg-color--dark);
}
.md-footer-meta {
- background: rgba(0, 0, 0, 0.15);
+ background: rgba(0, 0, 0, 0.2);
}
/* ===== Header ===== */
.md-header {
- background: linear-gradient(90deg, var(--themap-gradient-start), var(--themap-gradient-end));
- box-shadow: 0 2px 8px rgba(0, 0, 0, 0.15);
+ background-color: var(--md-primary-fg-color);
+ box-shadow: 0 2px 6px rgba(0, 0, 0, 0.12);
}
/* ===== Search ===== */
@@ -196,8 +211,8 @@
}
.md-typeset a.md-button--primary {
- background: linear-gradient(135deg, var(--themap-gradient-start), var(--themap-gradient-end));
- border: none;
+ background: var(--md-primary-fg-color);
+ border-color: var(--md-primary-fg-color);
}
/* ===== API Documentation ===== */
@@ -222,7 +237,7 @@
/* ===== Content Width ===== */
.md-content__inner {
- max-width: 900px;
+ max-width: 54rem;
}
/* ===== Version Selector ===== */
diff --git a/docs/assets/images/banner.png b/docs/assets/images/banner.png
new file mode 100644
index 0000000..2b1e76c
Binary files /dev/null and b/docs/assets/images/banner.png differ
diff --git a/docs/index.md b/docs/index.md
index 228906a..3754f3b 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -1,5 +1,14 @@
# THEMAP: Task Hardness Estimation for Molecular Activity Prediction
+[](https://doi.org/10.1021/acs.jcim.4c00160)
+[](https://www.python.org/downloads/)
+[](https://opensource.org/licenses/MIT)
+[](https://pypi.org/project/themap/)
+
+
+ 
+
+
**THEMAP** is a Python library for computing distances between chemical datasets and estimating task hardness for bioactivity prediction. It helps researchers identify similar tasks for transfer learning and quantify prediction difficulty.
@@ -41,18 +50,23 @@
## Installation
```bash
-# Clone and install (uv-based)
-git clone https://github.com/HFooladi/THEMAP.git
-cd THEMAP
-source install.sh
+# From PyPI (recommended)
+pip install themap
+
+# With all optional dependencies (ML, protein, OTDD)
+pip install "themap[all]"
+```
-# Or install with pip
-pip install -e .
+To install from source for development:
-# With all optional dependencies
-pip install -e ".[all]"
+```bash
+git clone https://github.com/HFooladi/THEMAP.git
+cd THEMAP
+source install.sh # uv-based; creates .venv and installs dev extras
```
+See [Getting Started](user-guide/getting-started.md) for the optional dependency groups.
+
## Quick Examples
### Command Line
diff --git a/docs/tutorials/Basics.ipynb b/docs/tutorials/Basics.ipynb
deleted file mode 100644
index 5dc8329..0000000
--- a/docs/tutorials/Basics.ipynb
+++ /dev/null
@@ -1,511 +0,0 @@
-{
- "cells": [
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "# Importing Required Libraries\n",
- "\n",
- "First, we are importing the libraries and modules that are required for running this notebook."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# import general packages\n",
- "import os\n",
- "import sys\n",
- "\n",
- "import pandas as pd\n",
- "import torch\n",
- "from dpu_utils.utils.richpath import RichPath\n",
- "from tqdm.notebook import tqdm\n",
- "\n",
- "# Setting up local details:\n",
- "# This should be the location of the checkout of the THEMAP repository:\n",
- "repo_path = os.path.dirname(os.path.abspath(\"\"))\n",
- "CHECKOUT_PATH = repo_path\n",
- "DATASET_PATH = os.path.join(repo_path, \"datasets\")\n",
- "\n",
- "os.chdir(CHECKOUT_PATH)\n",
- "sys.path.insert(0, CHECKOUT_PATH)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "from third_party.otdd.otdd.pytorch.datasets import MolDataset, load_molecule_data\n",
- "from third_party.otdd.otdd.pytorch.distance import DatasetDistance"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# import visualization packages\n",
- "%matplotlib inline\n",
- "\n",
- "import ipywidgets as widgets\n",
- "import matplotlib\n",
- "import matplotlib.pyplot as plt\n",
- "\n",
- "from themap.data import MoleculeDataset, ProteinDataset\n",
- "\n",
- "light_color = plt.get_cmap(\"plasma\").colors[170]\n",
- "dark_color = \"black\"\n",
- "\n",
- "matplotlib.rcParams.update(\n",
- " {\n",
- " \"pgf.texsystem\": \"pdflatex\",\n",
- " \"font.family\": \"serif\",\n",
- " \"font.serif\": \"Computer Modern Roman\",\n",
- " \"font.size\": 20,\n",
- " \"text.usetex\": True,\n",
- " \"pgf.rcfonts\": False,\n",
- " }\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "# Create source and target datasets (Data)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "source_dataset_path = RichPath.create(os.path.join(DATASET_PATH, \"train\", \"CHEMBL1023359.jsonl.gz\"))\n",
- "target_dataset_path = RichPath.create(os.path.join(DATASET_PATH, \"test\", \"CHEMBL2219358.jsonl.gz\"))\n",
- "\n",
- "source_dataset = MoleculeDataset.load_from_file(source_dataset_path)\n",
- "target_dataset = MoleculeDataset.load_from_file(target_dataset_path)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "Molecule_Feaurizer = widgets.Dropdown(\n",
- " options=[\"gin_supervised_infomax\", \"gin_supervised_masking\", \"gin_supervised_edgepred\"],\n",
- " value=\"gin_supervised_infomax\",\n",
- " description=\"Molecule Featurizer:\",\n",
- " disabled=False,\n",
- ")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "Molecule_Feaurizer"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "## compute and load the embeddings\n",
- "molecule_feaurizer = Molecule_Feaurizer.value\n",
- "source_features = source_dataset.get_features(molecule_feaurizer)\n",
- "target_features = target_dataset.get_features(molecule_feaurizer)\n",
- "assert source_features.shape[1] == target_features.shape[1]"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "source_dataset_otdd = MolDataset(source_dataset)\n",
- "target_dataset_otdd = MolDataset(target_dataset)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "source_dataset_loader = load_molecule_data(source_dataset)\n",
- "target_dataset_loader = load_molecule_data(target_dataset)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Instantiate distance\n",
- "dist = DatasetDistance(\n",
- " source_dataset_loader,\n",
- " target_dataset_loader,\n",
- " inner_ot_method=\"exact\",\n",
- " debiased_loss=True,\n",
- " p=2,\n",
- " entreg=1e-1,\n",
- " device=\"cuda\" if torch.cuda.is_available() else \"cpu\",\n",
- ")\n",
- "\n",
- "d = dist.distance(maxsamples=1000)\n",
- "print(f\"OTDD(src,tgt)={d}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "## Calculate chemcial distance between target datasets with all the source datasets"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "import glob\n",
- "\n",
- "source_datasets_path = glob.glob(os.path.join(DATASET_PATH, \"train\", \"CHEMBL*\"))\n",
- "target_datasets_path = glob.glob(os.path.join(DATASET_PATH, \"test\", \"CHEMBL*\"))\n",
- "chem_distances = {}\n",
- "for target_path in tqdm(target_datasets_path):\n",
- " chem_distance = {}\n",
- " target_dataset_path = RichPath.create(target_path)\n",
- " target_dataset = MoleculeDataset.load_from_file(target_dataset_path)\n",
- " target_features = target_dataset.get_features(molecule_feaurizer)\n",
- " target_dataset_otdd = MolDataset(target_dataset)\n",
- " target_dataset_loader = load_molecule_data(target_dataset)\n",
- " for source_path in source_datasets_path:\n",
- " source_dataset_path = RichPath.create(source_path)\n",
- " source_dataset = MoleculeDataset.load_from_file(source_dataset_path)\n",
- " source_features = source_dataset.get_features(molecule_feaurizer)\n",
- " source_dataset_otdd = MolDataset(source_dataset)\n",
- " source_dataset_loader = load_molecule_data(source_dataset)\n",
- "\n",
- " dist = DatasetDistance(\n",
- " source_dataset_loader,\n",
- " target_dataset_loader,\n",
- " inner_ot_method=\"exact\",\n",
- " debiased_loss=True,\n",
- " p=2,\n",
- " entreg=1e-1,\n",
- " device=\"cuda\" if torch.cuda.is_available() else \"cpu\",\n",
- " )\n",
- "\n",
- " d = dist.distance(maxsamples=1000)\n",
- " print(f\"OTDD({source_dataset.task_id},{target_dataset.task_id})= {d}\")\n",
- " chem_distance[source_dataset.task_id] = d.cpu().item()\n",
- " chem_distances[target_dataset.task_id] = chem_distance"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "## Choose your target from chem_distances.keys()\n",
- "your_tasks = \"CHEMBL2219236\"\n",
- "chem_dist = chem_distances[your_tasks]\n",
- "fig = plt.figure(figsize=(12, 5))\n",
- "plt.bar(chem_dist.keys(), chem_dist.values())\n",
- "plt.xlabel(\"Source datasets\")\n",
- "plt.ylabel(\"OTDD\")\n",
- "plt.title(f\"OTDD between source datasets and target {your_tasks}\")\n",
- "plt.xticks(rotation=90)\n",
- "plt.grid(axis=\"y\", linestyle=\"--\", alpha=1.0)\n",
- "plt.show()"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "chem_distances.keys()"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "## Calculate protein distance between target datasets with all the source datasets"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "Protein_Feaurizer = widgets.Dropdown(\n",
- " options=[\n",
- " \"esm2_t6_8M_UR50D\",\n",
- " \"esm2_t12_35M_UR50D\",\n",
- " \"esm2_t30_150M_UR50D\",\n",
- " \"esm2_t33_650M_UR50D\",\n",
- " \"esm2_t36_3B_UR50D\",\n",
- " ],\n",
- " value=\"esm2_t33_650M_UR50D\",\n",
- " description=\"Protein Featurizer:\",\n",
- " disabled=False,\n",
- ")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "Protein_Feaurizer"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "source_protein = ProteinDataset.load_from_file(\"datasets/train/train_proteins.fasta\")\n",
- "target_protein = ProteinDataset.load_from_file(\"datasets/test/test_proteins.fasta\")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "protein_featurizer = Protein_Feaurizer.value\n",
- "source_protein_features = source_protein.get_features(protein_featurizer)\n",
- "target_protein_features = target_protein.get_features(protein_featurizer)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "from scipy.spatial.distance import cdist\n",
- "\n",
- "dist = cdist(source_protein.features, target_protein.features)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "prot_distances = {}\n",
- "for i, target_prot in enumerate(target_protein.task_id):\n",
- " prot_distance = {}\n",
- " for j, source_prot in enumerate(source_protein.task_id):\n",
- " prot_distance[source_prot] = dist[j, i]\n",
- " prot_distances[target_prot] = prot_distance"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "## Choose your target from chem_distances.keys()\n",
- "your_tasks = \"CHEMBL2219236\"\n",
- "prot_dist = prot_distances[your_tasks]\n",
- "fig = plt.figure(figsize=(12, 5))\n",
- "plt.bar(prot_dist.keys(), prot_dist.values())\n",
- "plt.xlabel(\"Source datasets\")\n",
- "plt.ylabel(\"Protein Distance\")\n",
- "plt.title(f\"Protein Distance between source datasets and target {your_tasks}\")\n",
- "plt.xticks(rotation=90)\n",
- "plt.grid(axis=\"y\", linestyle=\"--\", alpha=1.0)\n",
- "plt.show()"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "## Combine Two Distances"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "Now, we can answer to the following questions:\n",
- "- Given a target task, what is the closest source task in terms of chemical and protein distances?\n",
- "- Given the source tasks, which target task is hardest to transfer to in terms of chemical and protein space?\n",
- "- Given a target task and source tasks, how to pick the k nearset source tasks fo transfer learning?\n",
- "\n",
- "So, let's answer to this questions in the following sections."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "chem_df = pd.DataFrame.from_dict(chem_distances)\n",
- "prot_df = pd.DataFrame.from_dict(prot_distances)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "## Given a target task, what is the closest source task in terms of chemical and protein distances?\n",
- "your_task = \"CHEMBL2219236\"\n",
- "\n",
- "chem_distance = chem_df[your_task]\n",
- "prot_distance = prot_df[your_task]\n",
- "\n",
- "normalized_chem_distance = (chem_distance - chem_distance.min()) / (chem_distance.max() - chem_distance.min())\n",
- "normalized_prot_distance = (prot_distance - prot_distance.min()) / (prot_distance.max() - prot_distance.min())\n",
- "normalized_prot_distance = normalized_prot_distance.reindex(normalized_chem_distance.index)\n",
- "normalized_comb_distance = (normalized_chem_distance + normalized_prot_distance) / 2\n",
- "\n",
- "\n",
- "print(f\"Closest source task in terms of chemical distance: {chem_distance.idxmin()}\")\n",
- "print(f\"Closest source task in terms of protein distance: {prot_distance.idxmin()}\")\n",
- "print(\n",
- " f\"Closest source task in terms of combination of chemical and protein distance: {normalized_comb_distance.idxmin()}\"\n",
- ")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "## Given the source tasks, which target task is hardest to transfer to in terms of chemical and protein space?\n",
- "## Answering this question require to define hardness, which we consider here as the average of k-nearest source tasks.\n",
- "k = 3\n",
- "\n",
- "target_tasks = chem_df.columns\n",
- "source_tasks = chem_df.index\n",
- "\n",
- "hardness_all = {}\n",
- "for target_task in target_tasks:\n",
- " hardness = {}\n",
- " chem_distance = chem_df[target_task]\n",
- " prot_distance = prot_df[target_task]\n",
- "\n",
- " chem_distance = chem_distance.sort_values()\n",
- " prot_distance = prot_distance.sort_values()\n",
- "\n",
- " hardness[\"EXT_CHEM\"] = chem_distance[:k].sum() / k\n",
- " hardness[\"EXT_PROT\"] = prot_distance[:k].sum() / k\n",
- "\n",
- " hardness_all[target_task] = hardness\n",
- "\n",
- "hardness_df = pd.DataFrame.from_dict(hardness_all).T\n",
- "hardness_df[\"all\"] = (hardness_df[\"EXT_CHEM\"] + hardness_df[\"EXT_PROT\"]) / 2\n",
- "\n",
- "print(f\"Easiest target task in terms of chemical distance: {hardness_df['EXT_CHEM'].idxmin()}\")\n",
- "print(f\"Easiest target task in terms of protein distance: {hardness_df['EXT_PROT'].idxmin()}\")\n",
- "print(\n",
- " f\"Easiest target task in terms of combination of chemical and protein distance: {hardness_df['all'].idxmin()}\"\n",
- ")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "plt.figure(figsize=(12, 5))\n",
- "plt.bar(hardness_df.index, hardness_df[\"all\"])\n",
- "plt.xlabel(\"Target datasets\")\n",
- "plt.ylabel(\"Hardness\")\n",
- "plt.title(\"Hardness of target datasets\")\n",
- "plt.xticks(rotation=90)\n",
- "plt.grid(axis=\"y\", linestyle=\"--\", alpha=1.0)\n",
- "plt.show()"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "## Given a target task and source tasks, how to pick the k nearset source tasks fo transfer learning?\n",
- "## Given the source tasks, which target task is hardest to transfer to in terms of chemical and protein space?\n",
- "## Answering this question require to define hardness, which we consider here as the average of k-nearest source tasks.\n",
- "k = 3\n",
- "\n",
- "target_tasks = chem_df.columns\n",
- "source_tasks = chem_df.index\n",
- "\n",
- "closest_tasks = {}\n",
- "for target_task in target_tasks:\n",
- " closest = {}\n",
- " chem_distance = chem_df[target_task]\n",
- " prot_distance = prot_df[target_task]\n",
- "\n",
- " chem_distance = chem_distance.sort_values()\n",
- " prot_distance = prot_distance.sort_values()\n",
- "\n",
- " closest[\"EXT_CHEM\"] = chem_distance[:k].index.to_list()\n",
- " closest[\"EXT_PROT\"] = prot_distance[:k].index.to_list()\n",
- "\n",
- " closest_tasks[target_task] = closest\n",
- "\n",
- "closest_df = pd.DataFrame.from_dict(closest_tasks).T"
- ]
- }
- ],
- "metadata": {
- "kernelspec": {
- "display_name": "pytorch",
- "language": "python",
- "name": "python3"
- },
- "language_info": {
- "codemirror_mode": {
- "name": "ipython",
- "version": 3
- },
- "file_extension": ".py",
- "mimetype": "text/x-python",
- "name": "python",
- "nbconvert_exporter": "python",
- "pygments_lexer": "ipython3",
- "version": "3.10.12"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/docs/tutorials/colab-notebooks.md b/docs/tutorials/colab-notebooks.md
new file mode 100644
index 0000000..82ac185
--- /dev/null
+++ b/docs/tutorials/colab-notebooks.md
@@ -0,0 +1,20 @@
+# Run on Colab
+
+Want to see what THEMAP does before installing anything? Run these notebooks
+directly in your browser on Google Colab — no local setup needed.
+
+| Notebook | What it covers | Runtime | Open |
+| --- | --- | --- | --- |
+| **5-minute quick tour** | Install THEMAP, download 5 sample ChEMBL datasets, compute a 2×3 distance matrix with `quick_distance`, and visualise it as a heatmap. | CPU | [](https://colab.research.google.com/github/HFooladi/THEMAP/blob/main/notebooks/colab/01_quick_tour.ipynb) |
+| **API deep dive** | Walk through the building blocks: `DatasetLoader`, `MoleculeFeaturizer`, `DatasetDistance`, YAML pipelines, featurizer/metric comparison, and a PCA task landscape over 13 datasets. | CPU | [](https://colab.research.google.com/github/HFooladi/THEMAP/blob/main/notebooks/colab/02_api_deep_dive.ipynb) |
+| **OTDD deep dive** | Run OTDD on the same ChEMBL assays with three featurizers (ECFP, `desc2D`, ChemBERTa). Compare OTDD against Euclidean/Cosine, and see why OTDD's Gaussian inner approximation rewards continuous representations and struggles with binary fingerprints. | GPU (T4) | [](https://colab.research.google.com/github/HFooladi/THEMAP/blob/main/notebooks/colab/03_otdd_deep_dive.ipynb) |
+
+All three notebooks live in
+[`notebooks/colab/`](https://github.com/HFooladi/THEMAP/tree/main/notebooks/colab).
+The first two run end-to-end on a free Colab CPU runtime; the OTDD notebook needs
+a GPU (a free T4 works) because OTDD's Wasserstein computation is GPU-bound at any
+practical scale.
+
+!!! tip
+ On Colab, select **Runtime → Change runtime type → T4 GPU** before running the
+ OTDD notebook.
diff --git a/docs/tutorials/working-with-tasks.md b/docs/tutorials/working-with-tasks.md
index c1b8e5f..e6b93e7 100644
--- a/docs/tutorials/working-with-tasks.md
+++ b/docs/tutorials/working-with-tasks.md
@@ -1,15 +1,23 @@
# Working with Tasks
-This tutorial covers THEMAP's unified task system that integrates molecular data, protein data, and metadata for comprehensive analysis.
+This tutorial covers THEMAP's unified task system that integrates molecular data,
+protein data, and other metadata for comprehensive analysis.
+
+!!! note
+ The code on this page is written against the current THEMAP API. The examples
+ are illustrative — running them requires a real dataset directory laid out as
+ shown below.
## Understanding the Task System
-The Task system in THEMAP provides a unified interface for working with multi-modal bioactivity prediction data:
+The Task system in THEMAP provides a unified interface for working with
+multi-modal bioactivity prediction data:
-- **Tasks**: Individual prediction problems (e.g., CHEMBL1023359)
-- **Data modalities**: Molecules, proteins, metadata
-- **Splits**: Train/test/validation divisions
-- **Caching**: Efficient feature storage and retrieval
+- **Tasks**: individual prediction problems (e.g. `CHEMBL1023359`)
+- **Data modalities**: molecules (`molecule_dataset`) and metadata such as protein
+ embeddings (`metadata_datasets`)
+- **Folds**: train / validation / test divisions, addressed with the `DataFold` enum
+- **Caching**: efficient feature storage and retrieval
## Loading Tasks
@@ -25,7 +33,7 @@ tasks = Tasks.from_directory(
load_molecules=True,
load_proteins=True,
load_metadata=True,
- cache_dir="cache/"
+ cache_dir="cache/",
)
print(f"Loaded {len(tasks)} tasks")
@@ -33,7 +41,7 @@ print(f"Loaded {len(tasks)} tasks")
### Task Configuration File
-The `sample_tasks_list.json` file defines which tasks belong to which split:
+The `sample_tasks_list.json` file defines which tasks belong to which fold:
```json
{
@@ -49,8 +57,8 @@ The `sample_tasks_list.json` file defines which tasks belong to which split:
datasets/
├── train/
│ ├── CHEMBL1023359.jsonl.gz # Molecular data
-│ ├── CHEMBL1023359.fasta # Protein sequences
-│ └── CHEMBL1023359_metadata.json
+│ ├── CHEMBL1023359.fasta # Protein sequence
+│ └── ...
├── test/
│ └── ...
├── valid/
@@ -60,134 +68,157 @@ datasets/
## Exploring Tasks
-### Basic Task Information
+Fold-aware methods take a `DataFold` value (`TRAIN`, `VALIDATION`, or `TEST`), not
+a string. There is no single "all folds" accessor, so concatenate the three folds
+when you need every task ID.
```python
-# Get task statistics
+from themap.data.enums import DataFold
+
+# Per-fold counts
print(f"Total tasks: {len(tasks)}")
-print(f"Train tasks: {tasks.get_num_fold_tasks('TRAIN')}")
-print(f"Test tasks: {tasks.get_num_fold_tasks('TEST')}")
-print(f"Valid tasks: {tasks.get_num_fold_tasks('VALID')}")
+print(f"Train tasks: {tasks.get_num_fold_tasks(DataFold.TRAIN)}")
+print(f"Test tasks: {tasks.get_num_fold_tasks(DataFold.TEST)}")
+print(f"Valid tasks: {tasks.get_num_fold_tasks(DataFold.VALIDATION)}")
-# Get task IDs
-train_task_ids = tasks.get_task_ids(fold='TRAIN')
+# Task IDs for a single fold
+train_task_ids = tasks.get_task_ids(DataFold.TRAIN)
print(f"Training task IDs: {train_task_ids}")
-# Access individual tasks
-task = tasks.get_task("CHEMBL1023359")
-print(f"Task {task.task_id} has {len(task.molecules)} molecules")
+
+def all_task_ids(tasks):
+ """Return task IDs across every fold."""
+ return (
+ tasks.get_task_ids(DataFold.TRAIN)
+ + tasks.get_task_ids(DataFold.VALIDATION)
+ + tasks.get_task_ids(DataFold.TEST)
+ )
+
+
+# Access an individual task
+task = tasks.get_task_by_id("CHEMBL1023359")
+print(f"Task {task.task_id} has {len(task.molecule_dataset)} molecules")
```
### Data Modality Access
+A `Task` exposes its molecules through `molecule_dataset` and any other modality
+(such as protein embeddings) through the `metadata_datasets` dictionary, where
+protein data lives under the `"protein"` key.
+
```python
-# Access different data types
-for task_id in tasks.get_task_ids()[:3]: # First 3 tasks
- task = tasks.get_task(task_id)
+for task_id in all_task_ids(tasks)[:3]: # first 3 tasks
+ task = tasks.get_task_by_id(task_id)
print(f"\nTask: {task_id}")
# Molecular data
- if hasattr(task, 'molecules') and task.molecules:
- print(f" Molecules: {len(task.molecules)}")
- print(f" Sample SMILES: {task.molecules[0].smiles}")
-
- # Protein data
- if hasattr(task, 'proteins') and task.proteins:
- print(f" Proteins: {len(task.proteins)}")
- print(f" Sample protein length: {len(task.proteins[0].sequence)}")
-
- # Metadata
- if hasattr(task, 'metadata') and task.metadata:
- print(f" Metadata keys: {list(task.metadata.keys())}")
+ if task.molecule_dataset is not None:
+ print(f" Molecules: {len(task.molecule_dataset)}")
+ print(f" Sample SMILES: {task.molecule_dataset.smiles[0]}")
+
+ # Metadata modalities (e.g. protein embeddings)
+ if task.metadata_datasets:
+ print(f" Metadata modalities: {list(task.metadata_datasets.keys())}")
+ if "protein" in task.metadata_datasets:
+ print(" Has protein metadata")
```
## Task-Based Distance Computation
### Unified Distance Calculation
+`TaskDistanceCalculator` orchestrates molecule, protein, and combined distances.
+`compute_all_distances` returns a dictionary keyed by `"molecules"`, `"protein"`,
+and `"combined"`; each value is a nested `{target_id: {source_id: distance}}` map.
+
```python
-from themap.distance import TaskDistance
+from themap.distance import TaskDistanceCalculator
-# Create task distance calculator
-task_distance = TaskDistance(
+# Create the calculator (dataset_method drives molecule distances,
+# metadata_method drives protein/metadata distances)
+task_distance = TaskDistanceCalculator(
tasks=tasks,
- molecule_method="cosine",
- protein_method="euclidean"
+ dataset_method="cosine",
+ metadata_method="euclidean",
)
-# Compute all distance types
+# Compute every distance type at once
all_distances = task_distance.compute_all_distances(
- combination_strategy="weighted_average",
- molecule_weight=0.7,
- protein_weight=0.3
+ molecule_featurizer="ecfp",
+ combination="weighted_average",
+ weights={"molecules": 0.7, "protein": 0.3},
)
print("Distance types computed:")
-print(f" Molecule distances: {len(all_distances['molecule'])} tasks")
-print(f" Protein distances: {len(all_distances['protein'])} tasks")
-print(f" Combined distances: {len(all_distances['combined'])} tasks")
+print(f" Molecule distances: {len(all_distances['molecules'])} target tasks")
+print(f" Protein distances: {len(all_distances['protein'])} target tasks")
+print(f" Combined distances: {len(all_distances['combined'])} target tasks")
```
### Specific Distance Types
```python
-# Compute only molecular distances
-molecule_distances = task_distance.get_molecule_distances()
+# Molecule distances only
+molecule_distances = task_distance.compute_molecule_distance(
+ molecule_featurizer="ecfp",
+)
-# Compute only protein distances
-protein_distances = task_distance.get_protein_distances()
+# Protein distances only
+protein_distances = task_distance.compute_protein_distance(
+ protein_featurizer="esm2_t33_650M_UR50D",
+)
-# Compute combined distances with custom weights
-combined_distances = task_distance.get_combined_distances(
- molecule_weight=0.6,
- protein_weight=0.4
+# Combined distances with custom weights
+combined_distances = task_distance.compute_combined_distance(
+ molecule_featurizer="ecfp",
+ weights={"molecules": 0.6, "protein": 0.4},
+ combination="weighted_average",
)
```
## Working with Folds
-### Train-Test Analysis
+### Train–Test Analysis
+
+Each distance matrix is indexed as `distances[target_id][source_id]`. With the
+default `source_fold=DataFold.TRAIN`, train tasks are the sources (inner keys) and
+the evaluated tasks are the targets (outer keys).
```python
-def analyze_train_test_distances(tasks, distance_type="molecule"):
- """Analyze distances between train and test tasks."""
+import pandas as pd
+from themap.data.enums import DataFold
- # Get task IDs by fold
- train_ids = set(tasks.get_task_ids(fold='TRAIN'))
- test_ids = set(tasks.get_task_ids(fold='TEST'))
- # Compute distances
- task_distance = TaskDistance(tasks=tasks)
- distances = task_distance.compute_all_distances()[distance_type]
+def analyze_train_test_distances(tasks, distance_type="molecules"):
+ """Collect distances between train (source) and test (target) tasks."""
+ train_ids = set(tasks.get_task_ids(DataFold.TRAIN))
+ test_ids = set(tasks.get_task_ids(DataFold.TEST))
- # Extract train-test distances
- train_test_distances = []
+ task_distance = TaskDistanceCalculator(tasks=tasks)
+ distances = task_distance.compute_all_distances()[distance_type]
+ pairs = []
for test_id in test_ids:
- if test_id in distances:
- for train_id in train_ids:
- if train_id in distances[test_id]:
- train_test_distances.append({
- 'test_task': test_id,
- 'train_task': train_id,
- 'distance': distances[test_id][train_id]
- })
-
- return train_test_distances
-
-# Analyze train-test relationships
-tt_distances = analyze_train_test_distances(tasks, "molecule")
+ if test_id not in distances:
+ continue
+ for train_id in train_ids:
+ if train_id in distances[test_id]:
+ pairs.append({
+ "test_task": test_id,
+ "train_task": train_id,
+ "distance": distances[test_id][train_id],
+ })
+ return pairs
+
+
+tt_distances = analyze_train_test_distances(tasks, "molecules")
print(f"Found {len(tt_distances)} train-test pairs")
-# Find closest training tasks for each test task
-import pandas as pd
-
+# Closest training tasks for each test task
df = pd.DataFrame(tt_distances)
-for test_task in df['test_task'].unique():
- task_distances = df[df['test_task'] == test_task]
- closest = task_distances.nsmallest(3, 'distance')
-
+for test_task in df["test_task"].unique():
+ closest = df[df["test_task"] == test_task].nsmallest(3, "distance")
print(f"\nTest task {test_task} - closest training tasks:")
for _, row in closest.iterrows():
print(f" {row['train_task']}: {row['distance']:.3f}")
@@ -196,31 +227,30 @@ for test_task in df['test_task'].unique():
### Cross-Validation Support
```python
-def create_cv_folds(tasks, n_folds=5):
- """Create cross-validation folds from tasks."""
+import random
+from themap.data.enums import DataFold
- all_task_ids = tasks.get_task_ids()
- random.shuffle(all_task_ids) # Randomize
- fold_size = len(all_task_ids) // n_folds
+def create_cv_folds(tasks, n_folds=5):
+ """Create cross-validation folds from all task IDs."""
+ task_ids = (
+ tasks.get_task_ids(DataFold.TRAIN)
+ + tasks.get_task_ids(DataFold.VALIDATION)
+ + tasks.get_task_ids(DataFold.TEST)
+ )
+ random.shuffle(task_ids)
+
+ fold_size = len(task_ids) // n_folds
folds = []
-
for i in range(n_folds):
- start_idx = i * fold_size
- end_idx = start_idx + fold_size if i < n_folds - 1 else len(all_task_ids)
-
- test_tasks = all_task_ids[start_idx:end_idx]
- train_tasks = [tid for tid in all_task_ids if tid not in test_tasks]
-
- folds.append({
- 'fold': i,
- 'train': train_tasks,
- 'test': test_tasks
- })
-
+ start = i * fold_size
+ end = start + fold_size if i < n_folds - 1 else len(task_ids)
+ test_tasks = task_ids[start:end]
+ train_tasks = [tid for tid in task_ids if tid not in test_tasks]
+ folds.append({"fold": i, "train": train_tasks, "test": test_tasks})
return folds
-# Create CV folds
+
cv_folds = create_cv_folds(tasks, n_folds=5)
print(f"Created {len(cv_folds)} CV folds")
```
@@ -229,31 +259,31 @@ print(f"Created {len(cv_folds)} CV folds")
### Feature Caching
+Pass a `cache_dir` and THEMAP caches computed features there, reusing them on
+subsequent loads.
+
```python
-# Tasks automatically cache computed features
tasks_with_cache = Tasks.from_directory(
directory="datasets/",
task_list_file="datasets/sample_tasks_list.json",
load_molecules=True,
load_proteins=True,
- cache_dir="cache/", # Features will be cached here
- force_reload=False # Use cached features if available
+ cache_dir="cache/", # features are cached here
)
-# Check cache status
print(f"Cache directory: {tasks_with_cache.cache_dir}")
```
### Selective Loading
```python
-# Load only specific data types for better performance
+# Load only what you need for faster processing
molecule_only_tasks = Tasks.from_directory(
directory="datasets/",
task_list_file="datasets/sample_tasks_list.json",
load_molecules=True,
- load_proteins=False, # Skip protein loading
- load_metadata=False # Skip metadata loading
+ load_proteins=False, # skip protein loading
+ load_metadata=False, # skip other metadata
)
print("Loaded molecular data only for faster processing")
@@ -266,49 +296,37 @@ print("Loaded molecular data only for faster processing")
```python
def filter_tasks_by_size(tasks, min_molecules=10, max_molecules=1000):
"""Filter tasks by number of molecules."""
+ filtered = []
+ for task_id in all_task_ids(tasks):
+ task = tasks.get_task_by_id(task_id)
+ if task.molecule_dataset is None:
+ continue
+ n_molecules = len(task.molecule_dataset)
+ if min_molecules <= n_molecules <= max_molecules:
+ filtered.append(task_id)
+ return filtered
- filtered_task_ids = []
-
- for task_id in tasks.get_task_ids():
- task = tasks.get_task(task_id)
- if hasattr(task, 'molecules') and task.molecules:
- n_molecules = len(task.molecules)
- if min_molecules <= n_molecules <= max_molecules:
- filtered_task_ids.append(task_id)
- return filtered_task_ids
-
-# Filter tasks
good_tasks = filter_tasks_by_size(tasks, min_molecules=20, max_molecules=500)
-print(f"Found {len(good_tasks)} tasks with appropriate size")
+print(f"Found {len(good_tasks)} tasks with an appropriate size")
```
-### Task Metadata Analysis
+### Inspecting Modalities
```python
-def analyze_task_metadata(tasks):
- """Analyze metadata across all tasks."""
-
- metadata_summary = {}
-
- for task_id in tasks.get_task_ids():
- task = tasks.get_task(task_id)
-
- if hasattr(task, 'metadata') and task.metadata:
- for key, value in task.metadata.items():
- if key not in metadata_summary:
- metadata_summary[key] = []
- metadata_summary[key].append(value)
-
- # Print summary
- for key, values in metadata_summary.items():
- unique_values = len(set(str(v) for v in values))
- print(f"Metadata '{key}': {len(values)} tasks, {unique_values} unique values")
-
- return metadata_summary
-
-# Analyze metadata
-metadata = analyze_task_metadata(tasks)
+def summarize_modalities(tasks):
+ """Count how many tasks carry each metadata modality."""
+ counts = {}
+ for task_id in all_task_ids(tasks):
+ task = tasks.get_task_by_id(task_id)
+ for modality in (task.metadata_datasets or {}):
+ counts[modality] = counts.get(modality, 0) + 1
+ for modality, n in counts.items():
+ print(f"Modality '{modality}': present in {n} tasks")
+ return counts
+
+
+summarize_modalities(tasks)
```
## Integration with External Tools
@@ -316,46 +334,25 @@ metadata = analyze_task_metadata(tasks)
### Export for Analysis
```python
-def export_task_summary(tasks, output_file="task_summary.csv"):
- """Export task information to CSV for external analysis."""
-
- import pandas as pd
-
- summary_data = []
-
- for task_id in tasks.get_task_ids():
- task = tasks.get_task(task_id)
-
- row = {'task_id': task_id}
-
- # Add molecule info
- if hasattr(task, 'molecules') and task.molecules:
- row['n_molecules'] = len(task.molecules)
- else:
- row['n_molecules'] = 0
-
- # Add protein info
- if hasattr(task, 'proteins') and task.proteins:
- row['n_proteins'] = len(task.proteins)
- row['avg_protein_length'] = sum(len(p.sequence) for p in task.proteins) / len(task.proteins)
- else:
- row['n_proteins'] = 0
- row['avg_protein_length'] = 0
-
- # Add metadata
- if hasattr(task, 'metadata') and task.metadata:
- for key, value in task.metadata.items():
- row[f'metadata_{key}'] = value
+import pandas as pd
- summary_data.append(row)
- df = pd.DataFrame(summary_data)
+def export_task_summary(tasks, output_file="task_summary.csv"):
+ """Export per-task information to CSV for external analysis."""
+ rows = []
+ for task_id in all_task_ids(tasks):
+ task = tasks.get_task_by_id(task_id)
+ rows.append({
+ "task_id": task_id,
+ "n_molecules": len(task.molecule_dataset) if task.molecule_dataset else 0,
+ "modalities": ",".join((task.metadata_datasets or {}).keys()),
+ })
+ df = pd.DataFrame(rows)
df.to_csv(output_file, index=False)
print(f"Task summary exported to {output_file}")
-
return df
-# Export summary
+
task_df = export_task_summary(tasks)
```
@@ -363,75 +360,34 @@ task_df = export_task_summary(tasks)
### Performance Tips
-1. **Use caching**: Always specify a cache directory
-2. **Load selectively**: Only load needed data types
-3. **Filter early**: Remove unsuitable tasks before distance computation
-4. **Batch processing**: Process tasks in groups for large datasets
+1. **Use caching**: always specify a `cache_dir`.
+2. **Load selectively**: only load the modalities you need.
+3. **Filter early**: remove unsuitable tasks before distance computation.
+4. **Batch processing**: process tasks in groups for large datasets.
### Memory Management
-```python
-# For large task sets, process in batches
-def process_tasks_in_batches(tasks, batch_size=10):
- """Process tasks in smaller batches to manage memory."""
-
- task_ids = tasks.get_task_ids()
-
- for i in range(0, len(task_ids), batch_size):
- batch_ids = task_ids[i:i + batch_size]
-
- # Create subset tasks
- batch_tasks = Tasks()
- for task_id in batch_ids:
- batch_tasks.add_task(tasks.get_task(task_id))
-
- # Process batch
- yield batch_tasks
-
-# Example usage
-for batch_tasks in process_tasks_in_batches(tasks, batch_size=5):
- print(f"Processing batch with {len(batch_tasks)} tasks")
- # Perform distance computation on batch
-```
-
-## Troubleshooting
-
-### Common Issues
-
-1. **Missing files**: Ensure all task files exist in expected directories
-2. **Cache conflicts**: Use `force_reload=True` to refresh cached features
-3. **Memory errors**: Reduce batch size or load fewer data modalities
-4. **Inconsistent data**: Validate that all tasks have required data types
-
-### Validation
+For large task sets, build smaller `Tasks` collections with the constructor
+(there is no incremental `add_task`; pass fold lists instead).
```python
-def validate_tasks(tasks):
- """Validate task data consistency."""
-
- issues = []
+from themap.data.enums import DataFold
- for task_id in tasks.get_task_ids():
- task = tasks.get_task(task_id)
- # Check for required data
- if not hasattr(task, 'molecules') or not task.molecules:
- issues.append(f"Task {task_id}: No molecular data")
-
- if not hasattr(task, 'proteins') or not task.proteins:
- issues.append(f"Task {task_id}: No protein data")
-
- if issues:
- print("Validation issues found:")
- for issue in issues:
- print(f" - {issue}")
- else:
- print("All tasks validated successfully")
+def process_train_tasks_in_batches(tasks, batch_size=10):
+ """Yield smaller Tasks collections over the train fold."""
+ train_ids = tasks.get_task_ids(DataFold.TRAIN)
+ for i in range(0, len(train_ids), batch_size):
+ batch_ids = train_ids[i:i + batch_size]
+ batch_tasks = Tasks(
+ train_tasks=[tasks.get_task_by_id(tid) for tid in batch_ids],
+ )
+ yield batch_tasks
- return issues
-# Validate loaded tasks
-validation_issues = validate_tasks(tasks)
+for batch_tasks in process_train_tasks_in_batches(tasks, batch_size=5):
+ print(f"Processing batch with {len(batch_tasks)} tasks")
+ # Perform distance computation on the batch
```
## Next Steps
diff --git a/docs/user-guide/getting-started.md b/docs/user-guide/getting-started.md
index 404f2fc..6c1aa62 100644
--- a/docs/user-guide/getting-started.md
+++ b/docs/user-guide/getting-started.md
@@ -8,32 +8,37 @@ This guide covers installation and three ways to compute distances between molec
- Python 3.10 or higher
-### Quick Install
+### Install from PyPI (recommended)
```bash
-git clone https://github.com/HFooladi/THEMAP.git
-cd THEMAP
-source install.sh # creates .venv with uv
+pip install themap # core
+pip install "themap[all]" # everything (ML, protein, OTDD)
```
-To reactivate later:
+Install only the extras you need:
```bash
-source .venv/bin/activate
+pip install "themap[ml]" # molecular analysis
+pip install "themap[protein]" # protein analysis (ESM2)
+pip install "themap[otdd]" # OTDD distance computation
```
-### Optional Dependencies
+### Install from source (development)
+
+```bash
+git clone https://github.com/HFooladi/THEMAP.git
+cd THEMAP
+source install.sh # creates .venv with uv, installs dev + test extras
+```
-THEMAP has optional dependency groups for different functionality:
+To reactivate the environment later:
```bash
-pip install -e ".[ml]" # molecular analysis
-pip install -e ".[protein]" # protein analysis (ESM2)
-pip install -e ".[otdd]" # OTDD distance computation
-pip install -e ".[all]" # everything
-pip install -e ".[dev,test]" # development
+source .venv/bin/activate
```
+The editable equivalent of the extras above is `pip install -e ".[all]"` (or `".[ml]"`, `".[dev,test]"`, etc.).
+
### Verify Installation
```python
diff --git a/docs/user-guide/reproducing-fsmol.md b/docs/user-guide/reproducing-fsmol.md
new file mode 100644
index 0000000..2ddf26f
--- /dev/null
+++ b/docs/user-guide/reproducing-fsmol.md
@@ -0,0 +1,78 @@
+# Reproducing FS-Mol Experiments
+
+The companion data for our paper *"Quantifying the hardness of bioactivity
+prediction tasks for transfer learning"* (J. Chem. Inf. Model. 64(10), 4031–4046,
+2024) is published on [Zenodo (record 10605093)](https://zenodo.org/records/10605093).
+It contains pre-computed OTDD distance matrices across multiple molecular
+featurizers, ESM-2 protein embeddings, internal chemical hardness measures, and
+ProtoNet evaluation summaries on the FS-Mol benchmark — everything needed to
+reproduce the figures and tables without re-running the expensive embedding
+pipelines.
+
+## 1. Install dependencies
+
+```bash
+source install.sh # creates .venv and installs themap[all,dev,test]
+```
+
+The reproduction notebooks rely on the optional `ml` extras (torch, ESM, etc.);
+the all-in-one install above covers them.
+
+## 2. Download the dataset (~16 GB)
+
+You need ~35 GB of free disk space (16 GB zip + ~16 GB extracted). The script
+downloads with resume support, verifies the MD5 checksum, extracts into
+`datasets/fsmol_hardness/`, and removes the zip when done.
+
+```bash
+make download-fsmol
+# or, equivalently:
+python scripts/download_fsmol_data.py
+```
+
+Useful flags: `--keep-zip` (don't delete the archive after extraction), `--force`
+(re-download), `--no-verify` (skip MD5 — only if you've already verified
+out-of-band), `--dest DIR` (custom location).
+
+After it completes (~31 GB extracted) you should see:
+
+```
+datasets/fsmol_hardness/
+├── ext_chem/ # OTDD distance matrices per molecular featurizer
+├── ext_prot/ # ESM-2 protein-distance matrices (t6_8M ... t36_3B)
+├── int_chem/{train,test}/ # Internal chemical hardness (RF baselines)
+├── embeddings/ # Per-task molecular embeddings used to compute the OTDDs
+├── FSMol_Eval_ProtoNet/summary/ # ProtoNet performance per support-set size (16/32/64/128)
+└── FSMol_Eval_randomForest/summary/ # Random-forest baseline performance summaries
+```
+
+The reproduction notebooks read from `ext_chem/`, `ext_prot/`, `int_chem/`, and
+`FSMol_Eval_ProtoNet/`; the other two directories are provided so you can rebuild
+the OTDD matrices from raw embeddings if desired.
+
+??? note "Manual download (no Python)"
+
+ ```bash
+ mkdir -p datasets/fsmol_hardness
+ cd datasets
+ wget -c https://zenodo.org/records/10605093/files/fsmol_hardness.zip
+ echo "10644660a53d8d106b6883cb53eb1f3b fsmol_hardness.zip" | md5sum -c -
+ unzip fsmol_hardness.zip -d fsmol_hardness/
+ ```
+
+## 3. Run the reproduction notebooks
+
+```bash
+cd notebooks
+jupyter lab # or: jupyter notebook
+```
+
+| Notebook | What it reproduces |
+| --- | --- |
+| [`external_chemical_hardness.ipynb`](https://github.com/HFooladi/THEMAP/blob/main/notebooks/external_chemical_hardness.ipynb) | External chemical-space hardness: correlation between k-nearest source-task OTDD distance and ProtoNet performance, across molecular featurizers (GIN, UniMol, ChemBERTa/Roberta-Zinc, desc2D). |
+| [`external_protein_hardness.ipynb`](https://github.com/HFooladi/THEMAP/blob/main/notebooks/external_protein_hardness.ipynb) | External protein-space hardness: correlation between target/source protein-embedding distance and performance, across ESM-2 model sizes (t6_8M → t36_3B). |
+| [`task_hardness.ipynb`](https://github.com/HFooladi/THEMAP/blob/main/notebooks/task_hardness.ipynb) | Combined task-hardness score (external chemical + external protein + internal chemical) and its correlation with ProtoNet performance at support-set sizes 16/32/64/128. |
+
+Notebook paths are resolved relative to the `notebooks/` directory, so launch
+Jupyter from there. Outputs are auto-stripped on commit by the pre-commit hook
+(`nbstripout`).
diff --git a/mkdocs.yml b/mkdocs.yml
index bbfb791..9dd2e02 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -19,9 +19,11 @@ nav:
- User Guide:
- Distance Computation: user-guide/distance-computation.md
- Command Line Interface: user-guide/cli.md
+ - Reproducing FS-Mol: user-guide/reproducing-fsmol.md
- Tutorials:
- Overview: tutorials/index.md
+ - Run on Colab: tutorials/colab-notebooks.md
- Basic Distance Computation: tutorials/basic-distance-computation.md
- Working with Tasks: tutorials/working-with-tasks.md
- Performance Optimization: tutorials/performance-optimization.md
@@ -57,6 +59,7 @@ theme:
features:
- navigation.tabs
- navigation.tabs.sticky
+ - navigation.sections
- navigation.path
- navigation.top
- navigation.footer
@@ -64,6 +67,7 @@ theme:
- toc.follow
- content.code.copy
- content.code.annotate
+ - content.tabs.link
- search.suggest
- search.highlight
@@ -71,7 +75,7 @@ theme:
logo: assets/images/logo-white.svg
extra_css:
- - assets/css/custom-splito.css
+ - assets/css/themap.css
extra_javascript:
- https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js
@@ -122,12 +126,19 @@ plugins:
- sys.path.append("themap")
options:
show_root_heading: yes
+ show_root_full_path: false
heading_level: 3
show_source: false
group_by_category: true
members_order: source
separate_signature: true
show_signature_annotations: true
+ signature_crossrefs: true
+ merge_init_into_class: true
+ show_symbol_type_heading: true
+ show_symbol_type_toc: true
+ docstring_section_style: table
+ filters: ["!^_"]
line_length: 80
docstring_style: google
@@ -149,6 +160,6 @@ extra:
link: https://github.com/HFooladi/THEMAP
- icon: fontawesome/brands/python
link: https://pypi.org/project/themap/
- analytics:
- provider: google
- property: G-XXXXXXXXXX
+ # analytics:
+ # provider: google
+ # property: G-XXXXXXXXXX # replace with a real Measurement ID to enable