{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d78f13cc",
   "metadata": {},
   "source": [
    "# HIPPIE machine-learning tutorial\n",
    "\n",
    "Companion notebook to steps 5 and 6 of the HIPPIE Machine Learning page. The page explains each step; this notebook only runs it.\n",
    "\n",
    "Set `DATA` in the *Load the data* cell to your unzipped split folder."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ca0f1d0",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "With a GPU, install `torch` first as shown on pytorch.org so you get a CUDA build."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aae6e810",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install pandas numpy scikit-learn biopython torch transformers tqdm ipywidgets matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c789f28a",
   "metadata": {},
   "source": [
    "## Load the data and check the split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a2f6a41",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "DATA = (\n",
    "    \"hippie_split\"  # the unzipped ZIP from step 3; for your own run: \"results/hippie\"\n",
    ")\n",
    "splits = {\n",
    "    name: pd.read_csv(f\"{DATA}/{name}.csv\")\n",
    "    for name in [\"train\", \"val\", \"test_balanced\", \"test_realistic\"]\n",
    "}\n",
    "# columns: protein1, protein2, label (1 = interacts, 0 = negative), score\n",
    "\n",
    "# sanity check 1: class balance -- about 50 % positives, except test_realistic (about 9 %)\n",
    "for name, df in splits.items():\n",
    "    print(f\"{name:15s} {len(df):7d} pairs   positives: {df.label.mean():.1%}\")\n",
    "\n",
    "\n",
    "# sanity check 2: held-out proteins that also occur in training -- should be close to 0 %\n",
    "def proteins(df):\n",
    "    return set(df.protein1) | set(df.protein2)\n",
    "\n",
    "\n",
    "train_prots = proteins(splits[\"train\"])\n",
    "for name in [\"val\", \"test_balanced\", \"test_realistic\"]:\n",
    "    prots = proteins(splits[name])\n",
    "    print(\n",
    "        f\"{name:15s} {len(prots & train_prots) / len(prots):.1%} of proteins also in train\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1631adc",
   "metadata": {},
   "source": [
    "## A — Cheap baseline (scikit-learn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b4f989b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "from Bio import SeqIO\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "# FASTA headers are UniProt accessions; the ZIP has sequences.fasta at its root\n",
    "# (if your headers look like sp|P12345|NAME, the split below keeps just the accession)\n",
    "fasta = f\"{DATA}/sequences.fasta\"\n",
    "if not os.path.exists(fasta):  # for your own run\n",
    "    fasta = f\"{DATA}/data/sequences.fasta\"\n",
    "seqs = {\n",
    "    r.id.split(\"|\")[1] if \"|\" in r.id else r.id: str(r.seq)\n",
    "    for r in SeqIO.parse(fasta, \"fasta\")\n",
    "}\n",
    "\n",
    "\n",
    "def usable(df):  # drop pairs whose sequence is missing from the FASTA\n",
    "    keep = df.protein1.isin(seqs.keys()) & df.protein2.isin(seqs.keys())\n",
    "    if (~keep).any():\n",
    "        print(f\"dropping {(~keep).sum()} pairs without a sequence\")\n",
    "    return df[keep]\n",
    "\n",
    "\n",
    "splits = {name: usable(df) for name, df in splits.items()}\n",
    "train, val = splits[\"train\"], splits[\"val\"]\n",
    "\n",
    "AA = \"ACDEFGHIKLMNPQRSTVWY\"\n",
    "\n",
    "\n",
    "def composition(seq):\n",
    "    counts = np.array([seq.count(a) for a in AA], dtype=float)\n",
    "    return np.append(counts / len(seq), np.log(len(seq)))\n",
    "\n",
    "\n",
    "comp = {p: composition(s) for p, s in seqs.items()}\n",
    "\n",
    "\n",
    "# symmetric pair features: (A, B) and (B, A) give the same vector\n",
    "def pair_features(vec, df):\n",
    "    return np.array(\n",
    "        [\n",
    "            np.concatenate([vec[a] + vec[b], np.abs(vec[a] - vec[b])])\n",
    "            for a, b in df[[\"protein1\", \"protein2\"]].itertuples(index=False)\n",
    "        ]\n",
    "    )\n",
    "\n",
    "\n",
    "base_clf = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000))\n",
    "base_clf.fit(pair_features(comp, train), train[\"label\"])\n",
    "print(\n",
    "    \"val AUROC baseline:\",\n",
    "    roc_auc_score(val[\"label\"], base_clf.predict_proba(pair_features(comp, val))[:, 1]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09329585",
   "metadata": {},
   "source": [
    "## B — Sequence-based model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8956a55c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import AutoModel, AutoTokenizer\n",
    "\n",
    "# smallest ESM-2 (8M parameters, 320-d). Bigger and usually better:\n",
    "# facebook/esm2_t12_35M_UR50D, facebook/esm2_t33_650M_UR50D (GPU recommended)\n",
    "MODEL = \"facebook/esm2_t6_8M_UR50D\"\n",
    "cache = f\"embeddings_{MODEL.split('/')[-1]}_{DATA.strip('/').replace('/', '-')}.npz\"\n",
    "\n",
    "if os.path.exists(cache):\n",
    "    emb = dict(np.load(cache))\n",
    "else:\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "    print(\"embedding on\", device)\n",
    "    tokenizer = AutoTokenizer.from_pretrained(MODEL)\n",
    "    model = AutoModel.from_pretrained(MODEL, add_pooling_layer=False).eval().to(device)\n",
    "\n",
    "    # sort by length so each batch pads as little as possible\n",
    "    needed = sorted(\n",
    "        set().union(*(proteins(df) for df in splits.values())),\n",
    "        key=lambda p: len(seqs[p]),\n",
    "    )\n",
    "    emb, BATCH = {}, 32  # lower BATCH if you run out of (GPU) memory\n",
    "    for i in tqdm(range(0, len(needed), BATCH), desc=\"ESM-2\", unit=\"batch\"):\n",
    "        batch = needed[i : i + BATCH]\n",
    "        chunk = [seqs[p][:1022] for p in batch]  # ESM-2 context limit\n",
    "        tokens = tokenizer(chunk, return_tensors=\"pt\", padding=True).to(device)\n",
    "        with torch.no_grad():\n",
    "            hidden = model(**tokens).last_hidden_state\n",
    "        for j, (p, s) in enumerate(zip(batch, chunk)):\n",
    "            # mean over the residues only: skip the start token and the padding\n",
    "            emb[p] = hidden[j, 1 : len(s) + 1].mean(0).cpu().numpy()\n",
    "    np.savez(cache, **emb)\n",
    "\n",
    "\n",
    "# a random forest prints no progress, so grow it in chunks of trees behind a progress bar\n",
    "# (same forest as one fit; raise step if you have many more cores than 25)\n",
    "def fit_forest(X, y, desc, n_trees=300, step=25):\n",
    "    rf = RandomForestClassifier(\n",
    "        n_estimators=0, warm_start=True, n_jobs=-1, random_state=0\n",
    "    )\n",
    "    for _ in tqdm(\n",
    "        range(0, n_trees, step), desc=f\"{desc} ({n_trees} trees)\", unit=\"step\"\n",
    "    ):\n",
    "        rf.n_estimators += step\n",
    "        rf.fit(X, y)\n",
    "    return rf\n",
    "\n",
    "\n",
    "X_train, X_val = pair_features(emb, train), pair_features(emb, val)\n",
    "seq_clf = fit_forest(X_train, train[\"label\"], \"sequence model\")\n",
    "\n",
    "# tune hyperparameters on the validation split (never on the test sets)\n",
    "print(\n",
    "    \"val AUROC sequence:\",\n",
    "    roc_auc_score(val[\"label\"], seq_clf.predict_proba(X_val)[:, 1]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7477183b",
   "metadata": {},
   "source": [
    "## Control: shuffled labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d264476",
   "metadata": {},
   "outputs": [],
   "source": [
    "shuffled = train[\"label\"].sample(frac=1, random_state=0).to_numpy()\n",
    "ctrl = fit_forest(X_train, shuffled, \"shuffled control\")\n",
    "print(\n",
    "    \"val AUROC shuffled labels:\",\n",
    "    roc_auc_score(val[\"label\"], ctrl.predict_proba(X_val)[:, 1]),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fec8f343",
   "metadata": {},
   "source": [
    "## Step 6: evaluate on both test sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dcb93fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import average_precision_score\n",
    "\n",
    "models = {  # name -> function that scores a table of pairs\n",
    "    \"baseline\": lambda df: base_clf.predict_proba(pair_features(comp, df))[:, 1],\n",
    "    \"sequence\": lambda df: seq_clf.predict_proba(pair_features(emb, df))[:, 1],\n",
    "    \"shuffled\": lambda df: ctrl.predict_proba(pair_features(emb, df))[:, 1],\n",
    "}\n",
    "test_sets = [\"test_balanced\", \"test_realistic\"]\n",
    "\n",
    "rows = []\n",
    "for test_name in test_sets:\n",
    "    test = splits[test_name]\n",
    "    y = test[\"label\"].to_numpy()\n",
    "    for name, score in models.items():\n",
    "        proba = score(test)\n",
    "        rows.append(\n",
    "            {\n",
    "                \"test set\": test_name,\n",
    "                \"model\": name,\n",
    "                \"AUROC\": roc_auc_score(y, proba),\n",
    "                \"AUPRC\": average_precision_score(y, proba),\n",
    "            }\n",
    "        )\n",
    "results = pd.DataFrame(rows)\n",
    "print(results.round(3).to_string(index=False))\n",
    "\n",
    "# overview plot: one panel per metric, one bar per model\n",
    "colors = [\"#2a78d6\", \"#eb6834\", \"#1baf7a\"]  # baseline, sequence, shuffled\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4), sharey=True)\n",
    "for ax, metric in zip(axes, [\"AUROC\", \"AUPRC\"]):\n",
    "    table = results.pivot(index=\"test set\", columns=\"model\", values=metric).loc[\n",
    "        test_sets, list(models)\n",
    "    ]\n",
    "    table.plot.bar(\n",
    "        ax=ax,\n",
    "        rot=0,\n",
    "        width=0.75,\n",
    "        color=colors,\n",
    "        edgecolor=\"white\",\n",
    "        linewidth=2,\n",
    "        legend=False,\n",
    "    )\n",
    "    for x, test_name in enumerate(test_sets):  # random-guess level per test set\n",
    "        chance = 0.5 if metric == \"AUROC\" else splits[test_name][\"label\"].mean()\n",
    "        ax.hlines(\n",
    "            chance,\n",
    "            x - 0.4,\n",
    "            x + 0.4,\n",
    "            colors=\"#555555\",\n",
    "            linestyles=\"--\",\n",
    "            linewidth=1.5,\n",
    "            label=\"random guess\" if x == 0 else None,\n",
    "        )\n",
    "    ax.set(title=metric, xlabel=\"\", ylim=(0, 1))\n",
    "    ax.grid(axis=\"y\", alpha=0.3)\n",
    "    ax.set_axisbelow(True)\n",
    "    ax.spines[[\"top\", \"right\"]].set_visible(False)\n",
    "axes[0].set_ylabel(\"score\")\n",
    "fig.legend(\n",
    "    *axes[1].get_legend_handles_labels(), loc=\"upper center\", ncol=4, frameon=False\n",
    ")\n",
    "fig.tight_layout(rect=(0, 0, 1, 0.9))\n",
    "fig.savefig(\"evaluation.png\", dpi=150)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}