{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {"name": "ipython", "version": 3},
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "version": "3.10.0"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-01",
   "metadata": {},
   "source": [
    "# avGFP Fluorescence Optimisation with ALSEBO\n",
    "\n",
    "This tutorial walks through a **complete ALSEBO run** end-to-end:\n",
    "from a FASTA file of VAE-generated sequences to a recommended set of\n",
    "high-fluorescence avGFP variants, using only ~65 simulated experiments.\n",
    "\n",
    "---\n",
    "\n",
    "## Biological context\n",
    "\n",
    "**avGFP** (*Aequorea victoria* Green Fluorescent Protein) is one of the most\n",
    "widely used reporter proteins in cell biology.\n",
    "Engineering brighter variants is a classic protein engineering benchmark.\n",
    "\n",
    "We use the deep mutational scanning dataset from:\n",
    "\n",
    "> Sarkisyan *et al.* (2016). Local fitness landscape of the green fluorescent\n",
    "> protein. *Nature*, **533**, 397–401.\n",
    "> https://doi.org/10.1038/nature17995\n",
    "\n",
    "A pre-trained **Support Vector Regression (SVR)** model (`best_svr_model.pkl`)\n",
    "trained on this dataset acts as our fitness oracle — standing in for wet-lab\n",
    "fluorescence measurements.\n",
    "\n",
    "---\n",
    "\n",
    "## Files in this tutorial\n",
    "\n",
    "| File | Description |\n",
    "|---|---|\n",
    "| `generated_seqs.fasta` | VAE-generated avGFP variant sequences (candidate pool) |\n",
    "| `avgfp_jhmmer.fasta` | Multiple sequence alignment used for DCA featurisation |\n",
    "| `best_svr_model.pkl` | Pre-trained SVR fitness oracle |\n",
    "\n",
    "## Pipeline\n",
    "\n",
    "```\n",
    "generated_seqs.fasta  +  avgfp_jhmmer.fasta\n",
    "           │\n",
    "           ▼\n",
    "  Step 1: Featurise sequence space  (DCA → seq_space.csv)\n",
    "           │\n",
    "           ▼\n",
    "  Step 2: Explore fitness landscape  (t-SNE + SVR oracle)\n",
    "           │\n",
    "           ▼\n",
    "  Step 3: Sample initial training set  (k-means diversity, n=15)\n",
    "           │\n",
    "           ▼\n",
    "  Step 4: Simulate experiments  (SVR oracle)\n",
    "           │\n",
    "           ▼\n",
    "  Step 5: Bayesian Optimisation loop  (GPR + UCB, 10 rounds × 5 sequences)\n",
    "           │\n",
    "           ▼\n",
    "  Step 6: Analyse results  (convergence + trajectory)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-02",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-03",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import joblib\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "from alsebo.seq_space import generate_seq_space\n",
    "from alsebo.optimizer import (\n",
    "    read_seq_files,\n",
    "    gpr,\n",
    "    seq_space_prediction,\n",
    "    get_next_seq_bo,\n",
    "    save_next_batch_results,\n",
    ")\n",
    "from alsebo.training_space import (\n",
    "    sample_initial_training_sequnces,\n",
    "    generate_sequence_training_file,\n",
    ")\n",
    "\n",
    "plt.rcParams.update({\n",
    "    \"font.size\": 11,\n",
    "    \"axes.linewidth\": 1.2,\n",
    "    \"axes.labelweight\": \"bold\",\n",
    "    \"figure.dpi\": 120,\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-04",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ── Configuration ──────────────────────────────────────────────────────────\n",
    "EXP_DIR     = \"./\"                   # directory containing all data files\n",
    "MSA_FILE    = \"avgfp_jhmmer.fasta\"   # MSA for DCA featurisation\n",
    "GEN_SEQS    = \"generated_seqs.fasta\" # VAE-generated candidate sequences\n",
    "ORACLE_PATH = \"best_svr_model.pkl\"\n",
    "\n",
    "INIT_BATCH  = 15    # sequences in the initial training set\n",
    "BO_ROUNDS   = 10    # number of Bayesian optimisation rounds\n",
    "BATCH_SIZE  = 5     # sequences recommended per round\n",
    "BETA        = 2.0   # UCB exploration-exploitation trade-off\n",
    "\n",
    "obj_config = {\n",
    "    \"names\":      [\"fitness\"],\n",
    "    \"directions\": [\"max\"],\n",
    "}\n",
    "# ───────────────────────────────────────────────────────────────────────────\n",
    "\n",
    "# Remove leftovers from a previous run so we start clean\n",
    "for fname in [\"seq_exp_data.csv\", \"training_seqs.csv\", \"seq_space.csv\"]:\n",
    "    fpath = os.path.join(EXP_DIR, fname)\n",
    "    if os.path.exists(fpath):\n",
    "        os.remove(fpath)\n",
    "\n",
    "oracle = joblib.load(ORACLE_PATH)\n",
    "print(\"Oracle loaded. Ready to start.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-05",
   "metadata": {},
   "source": [
    "## Step 1 — Featurise the Sequence Space\n",
    "\n",
    "We first convert each candidate sequence in `generated_seqs.fasta` into a\n",
    "numerical feature vector using **Direct Coupling Analysis (DCA)**.\n",
    "\n",
    "DCA fits a maximum-entropy statistical model to the MSA and computes\n",
    "per-position features that capture both single-site preferences and\n",
    "residue–residue co-evolutionary couplings.\n",
    "\n",
    "The result is written to `seq_space.csv` — one row per sequence.\n",
    "\n",
    "> **Note:** This step fits a DCA model on the MSA and is the slowest part\n",
    "> of the pipeline (~1–3 min depending on MSA depth). It only needs to run once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-06",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_seq_space(\n",
    "    exp_dir=EXP_DIR,\n",
    "    msa_fname=MSA_FILE,\n",
    "    gen_seq_fasta_fname=GEN_SEQS,\n",
    "    featuarization_method=\"DCA\",\n",
    ")\n",
    "\n",
    "seq_df = pd.read_csv(f\"{EXP_DIR}seq_space.csv\")\n",
    "print(f\"Sequence space: {len(seq_df):,} sequences × {seq_df.shape[1]-1} DCA features\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-07",
   "metadata": {},
   "source": [
    "## Step 2 — Explore the Fitness Landscape\n",
    "\n",
    "We use the SVR oracle to predict the fluorescence fitness of every sequence\n",
    "in the candidate pool, then project the high-dimensional DCA feature space\n",
    "down to 2-D with **t-SNE** to visualise the landscape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-08",
   "metadata": {},
   "outputs": [],
   "source": [
    "features    = seq_df.drop(\"seq_id\", axis=1)\n",
    "all_seqs    = seq_df[\"seq_id\"].tolist()\n",
    "fitness_all = oracle.predict(features.values)\n",
    "\n",
    "print(\"Running t-SNE projection...\")\n",
    "tsne_2d = TSNE(n_components=2, random_state=42).fit_transform(features)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 5))\n",
    "sc = ax.scatter(tsne_2d[:, 0], tsne_2d[:, 1],\n",
    "                c=fitness_all, cmap=\"magma\", alpha=0.75, s=8)\n",
    "plt.colorbar(sc, ax=ax, label=\"Predicted fluorescence fitness\")\n",
    "ax.set_xlabel(\"t-SNE 1\")\n",
    "ax.set_ylabel(\"t-SNE 2\")\n",
    "ax.set_title(\"avGFP Sequence Space — Predicted Fitness Landscape\")\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(f\"Fitness range : {fitness_all.min():.2f}  →  {fitness_all.max():.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-09",
   "metadata": {},
   "source": [
    "## Step 3 — Sample a Diverse Initial Training Set\n",
    "\n",
    "Before running BO we need a small starting set of measured sequences.\n",
    "A random selection risks clustering in one region of the landscape.\n",
    "\n",
    "ALSEBO avoids this with a **t-SNE + k-means** strategy:\n",
    "1. Project feature space to 2-D\n",
    "2. Partition into *k* = `INIT_BATCH` clusters\n",
    "3. Pick the sequence closest to each centroid\n",
    "\n",
    "This ensures the 15 initial sequences span the full landscape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-10",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_initial_training_sequnces(EXP_DIR, training_seq_size=INIT_BATCH, manipold=\"TSNE\")\n",
    "\n",
    "training_df = pd.read_csv(f\"{EXP_DIR}training_seqs.csv\")\n",
    "train_idx   = [all_seqs.index(s) for s in training_df[\"seq_id\"] if s in all_seqs]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 5))\n",
    "sc = ax.scatter(tsne_2d[:, 0], tsne_2d[:, 1],\n",
    "                c=fitness_all, cmap=\"magma\", alpha=0.45, s=8, label=\"All sequences\")\n",
    "ax.scatter(tsne_2d[train_idx, 0], tsne_2d[train_idx, 1],\n",
    "           c=\"cyan\", edgecolors=\"black\", s=70, zorder=5,\n",
    "           label=f\"Initial batch (n={INIT_BATCH})\")\n",
    "plt.colorbar(sc, ax=ax, label=\"Predicted fluorescence fitness\")\n",
    "ax.set_xlabel(\"t-SNE 1\")\n",
    "ax.set_ylabel(\"t-SNE 2\")\n",
    "ax.set_title(\"Diverse Initial Training Set\")\n",
    "ax.legend(frameon=False)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-11",
   "metadata": {},
   "source": [
    "## Step 4 — Simulate Initial Experiments\n",
    "\n",
    "In a real campaign you would now synthesise and measure the 15 selected\n",
    "sequences in the lab. Here the SVR oracle simulates those measurements\n",
    "instantly. The results are written to `seq_exp_data.csv` — the training\n",
    "log that ALSEBO appends to after every BO round."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-12",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_features = training_df.drop(\"seq_id\", axis=1)\n",
    "init_fitness   = oracle.predict(train_features.values)\n",
    "\n",
    "# generate_sequence_training_file expects list[list[float]]\n",
    "obj_values = [[float(v)] for v in init_fitness]\n",
    "generate_sequence_training_file(EXP_DIR, obj_config, obj_values)\n",
    "\n",
    "print(\"Initial training data written to seq_exp_data.csv\")\n",
    "print(f\"  Sequences measured : {INIT_BATCH}\")\n",
    "print(f\"  Mean fitness       : {np.mean(init_fitness):.3f}\")\n",
    "print(f\"  Best fitness       : {np.max(init_fitness):.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-13",
   "metadata": {},
   "source": [
    "## Step 5 — Bayesian Optimisation Loop\n",
    "\n",
    "Each round ALSEBO:\n",
    "\n",
    "1. Fits a **Gaussian Process** surrogate on all sequences measured so far\n",
    "2. Predicts posterior **mean** and **uncertainty** across every untested candidate\n",
    "3. Scores each candidate with **UCB**: `score = mean + β × std`\n",
    "4. Selects the top-5 and evaluates them with the oracle\n",
    "5. Appends results and repeats\n",
    "\n",
    "After 10 rounds we will have evaluated\n",
    "15 + 10 × 5 = **65 sequences** — just ~3% of the candidate pool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-14",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = []  # stores (round, mean_fitness, max_fitness)\n",
    "\n",
    "for round_idx in range(BO_ROUNDS):\n",
    "\n",
    "    # Load current training data + remaining candidates\n",
    "    x_train, y_train, x_space, seq_ids_remaining = read_seq_files(EXP_DIR, obj_config)\n",
    "\n",
    "    # Fit GPR surrogate — one model per objective\n",
    "    models = gpr(x_train, y_train)\n",
    "\n",
    "    # Predict mean + uncertainty over the full candidate space\n",
    "    preds = seq_space_prediction(models, x_space)\n",
    "\n",
    "    # UCB acquisition → select top-k candidates\n",
    "    next_seqs, scores, idx = get_next_seq_bo(\n",
    "        seq_ids_remaining, obj_config, preds,\n",
    "        top_k=BATCH_SIZE, strategy=\"UCB\", beta=BETA,\n",
    "    )\n",
    "\n",
    "    # Oracle evaluation (simulates wet-lab measurement)\n",
    "    new_fitness = oracle.predict(x_space.iloc[idx].values)\n",
    "\n",
    "    # Append results to seq_exp_data.csv\n",
    "    save_next_batch_results(\n",
    "        EXP_DIR, next_seqs, idx, x_space, obj_config,\n",
    "        obj_values=new_fitness.tolist(),\n",
    "    )\n",
    "\n",
    "    history.append({\"round\": round_idx + 1,\n",
    "                    \"mean\":  float(np.mean(new_fitness)),\n",
    "                    \"max\":   float(np.max(new_fitness))})\n",
    "\n",
    "    best_so_far = max(h[\"max\"] for h in history)\n",
    "    print(f\"Round {round_idx+1:2d}  |  batch max = {np.max(new_fitness):.3f}  \"\n",
    "          f\"|  best so far = {best_so_far:.3f}\")\n",
    "\n",
    "print(\"\\nOptimisation complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-15",
   "metadata": {},
   "source": [
    "## Step 6 — Results\n",
    "\n",
    "### Convergence\n",
    "\n",
    "Mean and best fitness per round, compared against the oracle maximum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-16",
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_df = pd.read_csv(f\"{EXP_DIR}seq_exp_data.csv\")\n",
    "\n",
    "rounds_labels = [\"Init\"] + [f\"R{h['round']}\" for h in history]\n",
    "mean_by_round = [exp_df[\"fitness\"][:INIT_BATCH].mean()] + [h[\"mean\"] for h in history]\n",
    "max_by_round  = [exp_df[\"fitness\"][:INIT_BATCH].max()]  + [h[\"max\"]  for h in history]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "ax.plot(rounds_labels, mean_by_round, \"-o\", label=\"Mean fitness\", lw=2)\n",
    "ax.plot(rounds_labels, max_by_round,  \"-o\", label=\"Best fitness\",  lw=2)\n",
    "ax.axhline(fitness_all.max(), color=\"grey\", linestyle=\"--\",\n",
    "           label=f\"Oracle max ({fitness_all.max():.2f})\")\n",
    "ax.set_xlabel(\"Optimisation round\")\n",
    "ax.set_ylabel(\"Predicted fluorescence fitness\")\n",
    "ax.set_title(\"ALSEBO Convergence on avGFP\")\n",
    "ax.legend(frameon=False)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-17",
   "metadata": {},
   "source": [
    "### Optimisation Trajectory\n",
    "\n",
    "Best sequence per round overlaid on the fitness landscape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cell-18",
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_seqs = exp_df[\"seq_id\"].tolist()\n",
    "\n",
    "opt_tsne = []\n",
    "for i, h in enumerate(history):\n",
    "    start     = INIT_BATCH + i * BATCH_SIZE\n",
    "    batch_fit = exp_df[\"fitness\"].iloc[start : start + BATCH_SIZE].values\n",
    "    best_seq  = exp_seqs[start + int(np.argmax(batch_fit))]\n",
    "    if best_seq in all_seqs:\n",
    "        opt_tsne.append(tsne_2d[all_seqs.index(best_seq)])\n",
    "\n",
    "opt_tsne = np.array(opt_tsne)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 5))\n",
    "sc = ax.scatter(tsne_2d[:, 0], tsne_2d[:, 1],\n",
    "                c=fitness_all, cmap=\"magma\", alpha=0.45, s=8)\n",
    "ax.scatter(tsne_2d[train_idx, 0], tsne_2d[train_idx, 1],\n",
    "           c=\"cyan\", edgecolors=\"black\", s=50, zorder=4, label=\"Initial training\")\n",
    "\n",
    "if len(opt_tsne) > 1:\n",
    "    ax.plot(opt_tsne[:, 0], opt_tsne[:, 1], \"w-\", lw=2, zorder=5)\n",
    "if len(opt_tsne) > 0:\n",
    "    ax.scatter(opt_tsne[:, 0], opt_tsne[:, 1],\n",
    "               c=\"white\", edgecolors=\"black\", s=50, zorder=6, label=\"Best per round\")\n",
    "    ax.scatter(opt_tsne[-1, 0], opt_tsne[-1, 1],\n",
    "               c=\"red\", edgecolors=\"black\", s=90, zorder=7, label=\"Final best\")\n",
    "\n",
    "plt.colorbar(sc, ax=ax, label=\"Predicted fluorescence fitness\")\n",
    "ax.set_xlabel(\"t-SNE 1\")\n",
    "ax.set_ylabel(\"t-SNE 2\")\n",
    "ax.set_title(\"Optimisation Trajectory on the Fitness Landscape\")\n",
    "ax.legend(frameon=False, fontsize=9)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-19",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "| | |\n",
    "|---|---|\n",
    "| Initial training set | 15 sequences (k-means diversity sampling) |\n",
    "| BO rounds | 10 |\n",
    "| Batch size per round | 5 |\n",
    "| **Total sequences evaluated** | **65 out of the full candidate pool (~3%)** |\n",
    "\n",
    "ALSEBO navigated the avGFP fitness landscape using only a small fraction of\n",
    "the candidate pool, guided by the GPR surrogate and UCB acquisition function.\n",
    "\n",
    "---\n",
    "\n",
    "### Things to try\n",
    "\n",
    "- **`manipold=\"PCA\"`** — swap t-SNE for PCA in the initial sampling step (faster, linear)\n",
    "- **`beta=0.5`** — more exploitative; converges faster but may miss the global optimum\n",
    "- **`beta=5.0`** — more exploratory; better for rugged or multi-modal landscapes\n",
    "- **ESM features** — set `featuarization_method=\"ESM\"` in `generate_seq_space()` for protein language model embeddings\n",
    "- **Multi-objective** — add a second objective (e.g. thermostability) to `obj_config` and provide two values per sequence in `obj_values`"
   ]
  }
 ]
}