cnn_protein_landscapes.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "cnn_protein_landscapes.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "jeV1ZCdJY4ST",
        "790Onzvqaa34",
        "2-beFIcol5c1",
        "shyZZvBsfbu1",
        "71BaMdNtlldo",
        "zVLely0SaeWP",
        "ZTbAjWrznCb5",
        "w09WErJpajjj",
        "LaWQ1KqljHN0",
        "RPvSKHf07tpj",
        "oLFZxzG2Kyoh",
        "hSlcNUW_KiHM",
        "v24_ZV7xKwHo",
        "AqZebxjQ6nim",
        "VASVo0goqzjJ",
        "WRc7urS3Mjfi",
        "3RgnuxcXMzCc",
        "2iUyp91eeY8h",
        "hwjV0A6eec3N",
        "0Hyo9e7Ij8MX",
        "sZ1K84k0yktc",
        "8aVROiHnzg4o",
        "bH0uaMhr0s1j",
        "0yO_4_B00LR8",
        "ZU1mXY0Z0V2w",
        "uHUniDwNW3pk",
        "R6OUT7DP2-ES",
        "MJ9I_Qr8nVsX",
        "tvun0LAR5ssO",
        "4nO7x4BiQmbZ",
        "PPl-2dUoQoe5",
        "S3Rks14tDwki",
        "n3PDOMMrGYsN",
        "40duS8d7IY9Z",
        "FMJwaVULIfIL",
        "7RuI0qSWFPNJ",
        "y6_6ZvKxE9ms",
        "KrG-n92mE_sn",
        "qmeob2F5O-5E",
        "BdjNebjPkGJh",
        "N4toUE5VfiIy",
        "ZTHUdtlogmOm"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/amirshane/protein-embedding-retrieval/blob/master/cnn_protein_landscapes.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hjxDefAj-IJj"
      },
      "source": [
        "CNN models are powerful tools for learning protein fitness landscapes. The code in this notebook applies them to three tasks: Fluorescence, Stability, and Variant Activity."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jeV1ZCdJY4ST"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NHdX2789av8o"
      },
      "source": [
        "# Install TORCH for TAPE and install JAX, JAXLIB, FLAX.\n",
        "!pip install --upgrade -q torch==1.6.0 jax==0.1.75 jaxlib==0.1.52 flax==0.2.0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "whGeiP0lav_L"
      },
      "source": [
        "import requests\n",
        "import os\n",
        "if 'TPU_DRIVER_MODE' not in globals():\n",
        "  url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver0.1-dev20200320'\n",
        "  resp = requests.post(url)\n",
        "  TPU_DRIVER_MODE = 1\n",
        "\n",
        "# Connect Jax XLA backend to TPU\n",
        "from jax.config import config\n",
        "config.enable_omnistaging()\n",
        "config.FLAGS.jax_xla_backend = \"tpu_driver\"\n",
        "config.FLAGS.jax_backend_target = \"grpc://\" + os.environ['COLAB_TPU_ADDR']\n",
        "print(config.FLAGS.jax_backend_target)\n",
        "\n",
        "# Specifying TPU configuration\n",
        "# Results are dependent on configuration since TPU computations are non-deterministic\n",
        "# Reported results are from \"tpu_driver_nightly\" configuration\n",
        "# See https://cloud.google.com/tpu/docs/version-switching and\n",
        "# https://github.com/google/jax/issues/4408\n",
        "from tensorflow.python.tpu.client.client import Client\n",
        "c = Client()\n",
        "c.configure_tpu_version('tpu_driver0.1-dev20200320', restart_type='ifNeeded')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O7_4yHfXY2dJ"
      },
      "source": [
        "import os\n",
        "import sys\n",
        "\n",
        "import jax.numpy as jnp\n",
        "\n",
        "import flax.nn as nn\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "import pandas as pd\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import scipy.stats\n",
        "\n",
        "import sklearn\n",
        "from sklearn.linear_model import Ridge\n",
        "import sklearn.metrics as metrics\n",
        "from sklearn.decomposition import PCA\n",
        "from sklearn.manifold import TSNE\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "import pprint\n",
        "\n",
        "import collections\n",
        "\n",
        "import copy"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "790Onzvqaa34"
      },
      "source": [
        "## clone and pip install tape"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DXOhWTS6_0Bd"
      },
      "source": [
        "!git clone https://github.com/songlab-cal/tape\n",
        "!pip install -q -r tape/requirements.txt\n",
        "os.chdir('tape/')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_wn-x9k3FJce"
      },
      "source": [
        "import tape\n",
        "from tape.datasets import LMDBDataset"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WQiEmyEZy2Wk"
      },
      "source": [
        "os.chdir('../')"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2-beFIcol5c1"
      },
      "source": [
        "## clone protein-embedding-retrieval repo and import"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QCyBGRHplvKD"
      },
      "source": [
        "# !git clone --recurse-submodules https://github.com/googleinterns/protein-embedding-retrieval.git\n",
        "!git clone --recurse-submodules https://github.com/amirshane/protein-embedding-retrieval.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9f0doVeYlxui"
      },
      "source": [
        "os.chdir('protein-embedding-retrieval')"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EJp_Mj8Yo7Aq"
      },
      "source": [
        "pip install -e ."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k_D33oz3l1TX"
      },
      "source": [
        "sys.path.insert(1, 'google_research/')"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X2e9kejFl1xh"
      },
      "source": [
        "from google_research.protein_lm import domains"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "shyZZvBsfbu1"
      },
      "source": [
        "# Utils"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ivex2BL8R5sR"
      },
      "source": [
        "from contextual_lenses.contextual_lenses import max_pool, linear_max_pool\n",
        "\n",
        "def flatten(x, padding_mask=None):\n",
        "  \"\"\"Apply padding and flatten over sequence length axis.\"\"\"\n",
        "\n",
        "  if padding_mask is not None:\n",
        "    x = x * padding_mask\n",
        "\n",
        "  rep = x.reshape(x.shape[0], x.shape[1]*x.shape[2])\n",
        "\n",
        "  return rep"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DzAkxImZSHgG"
      },
      "source": [
        "from contextual_lenses.encoders import one_hot_encoder, cnn_one_hot_encoder\n",
        "\n",
        "def flattened_one_hot_encoder(batch_inds, num_categories):\n",
        "  \"\"\"Flattens padded one-hot encoding from jax.nn.\"\"\"\n",
        "  \n",
        "  padding_mask = jnp.expand_dims(jnp.where(batch_inds < num_categories-1, 1, 0), axis=2)\n",
        "  \n",
        "  one_hots = one_hot_encoder(batch_inds, num_categories)\n",
        "  flattened_one_hots = flatten(one_hots, padding_mask)\n",
        "\n",
        "  return flattened_one_hots"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wYcOhZWySDwy"
      },
      "source": [
        "from contextual_lenses.loss_fns import mse_loss\n",
        "from contextual_lenses.train_utils import create_data_iterator, \\\n",
        "create_optimizer, create_representation_model, train\n",
        "from contextual_lenses.pfam_utils import compute_embeddings"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mFE04KffSnv8"
      },
      "source": [
        "def get_num_params(model):\n",
        "  \"\"\"Computes number of parameters in flax model.\"\"\"\n",
        "\n",
        "  # Code source: https://stackoverflow.com/questions/6027558/flatten-nested-dictionaries-compressing-keys\n",
        "  def dict_flatten(d, parent_key='', sep='_'):\n",
        "    items = []\n",
        "    for k, v in d.items():\n",
        "        new_key = parent_key + sep + k if parent_key else k\n",
        "        if isinstance(v, collections.MutableMapping):\n",
        "            items.extend(dict_flatten(v, new_key, sep=sep).items())\n",
        "        else:\n",
        "            items.append((new_key, v))\n",
        "    return dict(items)\n",
        "\n",
        "  params = model.params\n",
        "  params = dict_flatten(params)\n",
        "\n",
        "  num_params = 0\n",
        "  for layer in params.keys():\n",
        "    num_params += np.prod(params[layer].shape)\n",
        "  \n",
        "  return num_params"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "71BaMdNtlldo"
      },
      "source": [
        "# Fluorescence"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NMNX-Kkn-a7x"
      },
      "source": [
        "The data originally comes from [Sarkisyan et. al.](https://www.nature.com/articles/nature17995) and was formatted by [TAPE](https://github.com/songlab-cal/tape)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zVLely0SaeWP"
      },
      "source": [
        "## Open data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P2ZuTl_2FZdV"
      },
      "source": [
        "!wget http://s3.amazonaws.com/proteindata/data_pytorch/fluorescence.tar.gz"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f_vEoTDwFb61"
      },
      "source": [
        "!tar xzf fluorescence.tar.gz"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ksuiR_xyHg1V"
      },
      "source": [
        "def gfp_dataset_to_df(in_name):\n",
        "  dataset = LMDBDataset(in_name)\n",
        "  df = pd.DataFrame(list(dataset)[:])\n",
        "  df['log_fluorescence'] = df.log_fluorescence.apply(lambda x: x[0])\n",
        "  return df"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZTbAjWrznCb5"
      },
      "source": [
        "### Padding and one-hot encoding"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DxUGrV1u-41g"
      },
      "source": [
        "GFP_SEQ_LEN = 237"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pnFYMapDMTwQ"
      },
      "source": [
        "GFP_AMINO_ACID_VOCABULARY = [\n",
        "    'A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R',\n",
        "    'S', 'T', 'V', 'W', 'Y', '-'\n",
        "]"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mY3Rquilsfju"
      },
      "source": [
        "GFP_PROTEIN_DOMAIN = domains.VariableLengthDiscreteDomain(\n",
        "    vocab=domains.ProteinVocab(include_anomalous_amino_acids=False,\n",
        "                               include_eos=True,\n",
        "                               include_pad=True),\n",
        "    length=GFP_SEQ_LEN)"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XOs4w9LbmpwB"
      },
      "source": [
        "def gfp_seq_to_inds(seq):\n",
        "  \"\"\"Encode GFP amino acid sequence.\"\"\"\n",
        "\n",
        "  return GFP_PROTEIN_DOMAIN.encode([seq])[0]"
      ],
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w09WErJpajjj"
      },
      "source": [
        "### Add one-hots and batch data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0-6Nx1dD42-P"
      },
      "source": [
        "def create_gfp_df(test=False):\n",
        "  \"\"\"Processes GFP data into a featurized dataframe.\"\"\"\n",
        "  \n",
        "  if test:\n",
        "    gfp_df = gfp_dataset_to_df('fluorescence/fluorescence_test.lmdb')\n",
        "  else:\n",
        "    gfp_df = gfp_dataset_to_df('fluorescence/fluorescence_train.lmdb')\n",
        "  \n",
        "  gfp_df['one_hot_inds'] = gfp_df.primary.apply(lambda x: gfp_seq_to_inds(x[:GFP_SEQ_LEN]))\n",
        "\n",
        "  return gfp_df\n",
        "\n",
        "\n",
        "def create_gfp_batches(batch_size, epochs=1, test=False, buffer_size=None, \n",
        "                       seed=0, drop_remainder=False):\n",
        "  \"\"\"Creates iterable object of GFP batches.\"\"\"\n",
        "  \n",
        "  if test:\n",
        "    buffer_size = 1\n",
        "  \n",
        "  gfp_df = create_gfp_df(test=test)\n",
        "    \n",
        "  fluorescences = gfp_df['log_fluorescence'].values\n",
        "\n",
        "  gfp_batches = create_data_iterator(df=gfp_df, input_col='one_hot_inds', \n",
        "                                     output_col='log_fluorescence', \n",
        "                                     batch_size=batch_size, epochs=epochs, \n",
        "                                     buffer_size=buffer_size, seed=seed, \n",
        "                                     drop_remainder=drop_remainder)\n",
        "\n",
        "  return gfp_batches, fluorescences"
      ],
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LaWQ1KqljHN0"
      },
      "source": [
        "## Model evaluation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n1BfAaTkI_wn"
      },
      "source": [
        "def gfp_evaluate(predict_fn, title, batch_size=256, \n",
        "                 test_data=None, pred_fluorescences=None,\n",
        "                 clip_min=-999999, clip_max=999999):\n",
        "  \"\"\"Computes predicted fluorescences and measures performance in MSE and spearman correlation.\"\"\"\n",
        "  \n",
        "  test_batches, test_fluorescences = create_gfp_batches(batch_size=batch_size, \n",
        "                                                        test=True, buffer_size=1)\n",
        "  \n",
        "  if test_data is not None:\n",
        "    test_batches = test_data\n",
        "\n",
        "  if pred_fluorescences is None:\n",
        "    pred_fluorescences = []\n",
        "    for batch in iter(test_batches):\n",
        "      X, Y = batch\n",
        "      preds = predict_fn(X)\n",
        "      for pred in preds:\n",
        "        pred_fluorescences.append(pred[0])\n",
        "  \n",
        "  pred_fluorescences = np.array(pred_fluorescences)\n",
        "  pred_fluorescences = np.clip(pred_fluorescences, clip_min, clip_max)\n",
        "    \n",
        "  spearmanr = scipy.stats.spearmanr(test_fluorescences, pred_fluorescences).correlation\n",
        "  mse = sklearn.metrics.mean_squared_error(test_fluorescences, pred_fluorescences)\n",
        "  plt.scatter(test_fluorescences, pred_fluorescences, s=1, alpha=0.5)\n",
        "  plt.xlabel('True LogFluorescence')\n",
        "  plt.ylabel('Predicted LogFluorescence')\n",
        "  plt.title(title)\n",
        "  plt.show()\n",
        "\n",
        "  bright_inds = np.where(test_fluorescences > 2.5)\n",
        "  bright_test_fluorescences = test_fluorescences[bright_inds]\n",
        "  bright_pred_fluorescences = pred_fluorescences[bright_inds]\n",
        "  bright_spearmanr = scipy.stats.spearmanr(bright_test_fluorescences, bright_pred_fluorescences).correlation\n",
        "  bright_mse = sklearn.metrics.mean_squared_error(bright_test_fluorescences, bright_pred_fluorescences)\n",
        "  plt.scatter(bright_test_fluorescences, bright_pred_fluorescences, s=1, alpha=0.5)\n",
        "  plt.xlabel('True LogFluorescence')\n",
        "  plt.ylabel('Predicted LogFluorescence')\n",
        "  bright_title = title + ' (Bright)'\n",
        "  plt.title(bright_title)\n",
        "  plt.show()\n",
        "\n",
        "  dark_inds = np.where(test_fluorescences < 2.5)\n",
        "  dark_test_fluorescences = test_fluorescences[dark_inds]\n",
        "  dark_pred_fluorescences = pred_fluorescences[dark_inds]\n",
        "  dark_spearmanr = scipy.stats.spearmanr(dark_test_fluorescences, dark_pred_fluorescences).correlation\n",
        "  dark_mse = sklearn.metrics.mean_squared_error(dark_test_fluorescences, dark_pred_fluorescences)\n",
        "  plt.scatter(dark_test_fluorescences, dark_pred_fluorescences, s=1, alpha=0.5)\n",
        "  plt.xlabel('True LogFluorescence')\n",
        "  plt.ylabel('Predicted LogFluorescence')\n",
        "  dark_title = title + ' (Dark)'\n",
        "  plt.title(dark_title)\n",
        "  plt.show()\n",
        "\n",
        "  results = {\n",
        "      'title': title,\n",
        "      'spearmanr': round(spearmanr, 3),\n",
        "      'mse': round(mse, 3),\n",
        "      'bright_spearmanr': round(bright_spearmanr, 3),\n",
        "      'bright_mse': round(bright_mse, 3),\n",
        "      'dark_spearmanr': round(dark_spearmanr, 3),\n",
        "      'dark_mse': round(dark_mse, 3),\n",
        "  }\n",
        "\n",
        "  pprint.pprint(results)\n",
        "\n",
        "  return results, pred_fluorescences"
      ],
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RPvSKHf07tpj"
      },
      "source": [
        "##  Experiments"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b8sTon3kxlhq"
      },
      "source": [
        "gfp_train_df = create_gfp_df()\n",
        "train_fluorescences = gfp_train_df['log_fluorescence']\n",
        "gfp_train_one_hot_inds = np.array([x for x in gfp_train_df['one_hot_inds'].values])\n",
        "gfp_train_one_hots = flattened_one_hot_encoder(gfp_train_one_hot_inds, num_categories=len(GFP_AMINO_ACID_VOCABULARY))"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ue8RuUYXxmOO"
      },
      "source": [
        "gfp_test_df = create_gfp_df(test=True)\n",
        "gfp_test_one_hot_inds = np.array([x for x in gfp_test_df['one_hot_inds'].values])\n",
        "gfp_test_one_hots = flattened_one_hot_encoder(gfp_test_one_hot_inds, num_categories=len(GFP_AMINO_ACID_VOCABULARY))"
      ],
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oLFZxzG2Kyoh"
      },
      "source": [
        "### Linear regression"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PUoKzVwBfhGt"
      },
      "source": [
        "gfp_linear_model = Ridge()"
      ],
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_-hH-HF9vcHs",
        "outputId": "3bcf994d-1586-4b2d-c9f0-7f39271e311b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "gfp_linear_model.fit(X=gfp_train_one_hots, y=train_fluorescences)"
      ],
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Ridge(alpha=1.0, copy_X=True, fit_intercept=True, max_iter=None,\n",
              "      normalize=False, random_state=None, solver='auto', tol=0.001)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N6uuTzwSf6Gb"
      },
      "source": [
        "linear_model_pred_fluorescences = gfp_linear_model.predict(gfp_test_one_hots)"
      ],
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7BE2J_YIEgwL"
      },
      "source": [
        "gfp_linear_model_results, linear_model_pred_fluorescences = \\\n",
        "gfp_evaluate(predict_fn=None, \n",
        "             title='Linear Regression',\n",
        "             pred_fluorescences=linear_model_pred_fluorescences,\n",
        "             clip_min=min(train_fluorescences),\n",
        "             clip_max=max(train_fluorescences))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YnmwmIkP_Q6j"
      },
      "source": [
        "print('Number of Parameters for Fluorescence Linear Regression: ' + str(len(gfp_linear_model.coef_)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hSlcNUW_KiHM"
      },
      "source": [
        "### CNN + MaxPool"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yWqTzayWKkIF"
      },
      "source": [
        "epochs = 50\n",
        "gfp_train_batches, train_fluorescences = create_gfp_batches(batch_size=256, epochs=epochs)\n",
        "\n",
        "layers = ['CNN_0', 'Dense_1']                                 \n",
        "learning_rate = [1e-3, 5e-6]\n",
        "weight_decay = [0.0, 0.05]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 1,\n",
        "    'n_features': [1024],\n",
        "    'n_kernel_sizes': [5],\n",
        "    'n_kernel_dilations': None\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "\n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "gfp_model = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                        encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                        reduce_fn=reduce_fn,\n",
        "                                        reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                        num_categories=len(GFP_AMINO_ACID_VOCABULARY),\n",
        "                                        output_features=1)\n",
        "\n",
        "gfp_optimizer = train(model=gfp_model,\n",
        "                      train_data=gfp_train_batches, \n",
        "                      loss_fn=mse_loss,\n",
        "                      loss_fn_kwargs=loss_fn_kwargs,\n",
        "                      learning_rate=learning_rate, \n",
        "                      weight_decay=weight_decay,\n",
        "                      layers=layers)\n",
        "\n",
        "gfp_results, pred_fluorescences = gfp_evaluate(predict_fn=gfp_optimizer.target,\n",
        "                                               title='CNN + MaxPool',\n",
        "                                               batch_size=256,\n",
        "                                               clip_min=min(train_fluorescences),\n",
        "                                               clip_max=max(train_fluorescences))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6no3UFPd9wVM"
      },
      "source": [
        "print('Number of Parameters for Fluorescence CNN + MaxPool: ' + str(get_num_params(gfp_optimizer.target)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v24_ZV7xKwHo"
      },
      "source": [
        "### CNN + LinearMaxPool"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3a9NTtWN0xJi"
      },
      "source": [
        "epochs = 50\n",
        "gfp_train_batches, train_fluorescences = create_gfp_batches(batch_size=256, epochs=epochs)\n",
        "\n",
        "layers = ['CNN_0', 'Dense_1', 'Dense_2']                                 \n",
        "learning_rate = [1e-3, 5e-5, 5e-6]\n",
        "weight_decay = [0.0, 0.05, 0.05]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 1,\n",
        "    'n_features': [1024],\n",
        "    'n_kernel_sizes': [5],\n",
        "    'n_kernel_dilations': None\n",
        "}\n",
        "reduce_fn = linear_max_pool\n",
        "reduce_fn_kwargs = {\n",
        "    'rep_size': 2048\n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "gfp_model_l = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                          encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                          reduce_fn=reduce_fn,\n",
        "                                          reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                          num_categories=len(GFP_AMINO_ACID_VOCABULARY),\n",
        "                                          output_features=1)\n",
        "\n",
        "gfp_optimizer_l = train(model=gfp_model_l,\n",
        "                        train_data=gfp_train_batches, \n",
        "                        loss_fn=mse_loss,\n",
        "                        loss_fn_kwargs=loss_fn_kwargs,\n",
        "                        learning_rate=learning_rate, \n",
        "                        weight_decay=weight_decay,\n",
        "                        layers=layers)\n",
        "\n",
        "gfp_results_l, pred_fluorescences_l = gfp_evaluate(predict_fn=gfp_optimizer_l.target,\n",
        "                                                   title='CNN + LinearMaxPool',\n",
        "                                                   batch_size=256,\n",
        "                                                   clip_min=min(train_fluorescences),\n",
        "                                                   clip_max=max(train_fluorescences))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lPBV7lXC-lkM"
      },
      "source": [
        "print('Number of Parameters for Fluorescence CNN + LinearMaxPool: ' + str(get_num_params(gfp_optimizer_l.target)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AqZebxjQ6nim"
      },
      "source": [
        "### TAPE (Best Performance)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "idUlyheU6q49"
      },
      "source": [
        "tape_fluorescence_results = {\n",
        "    'title': 'TAPE',\n",
        "    'spearmanr': 0.68,\n",
        "    'mse': 0.19,\n",
        "    'bright_spearmanr': 0.63,\n",
        "    'bright_mse': 0.07,\n",
        "    'dark_mse': 0.22,\n",
        "    'dark_spearmanr': 0.05,\n",
        "}"
      ],
      "execution_count": 36,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "noKdH97Z6nO-"
      },
      "source": [
        "pprint.pprint(tape_fluorescence_results)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VASVo0goqzjJ"
      },
      "source": [
        "## Visualization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TlbbAl_cid3M"
      },
      "source": [
        "def gfp_embeddings(encoder_fn, encoder_fn_kwargs, reduce_fn, reduce_fn_kwargs, optimizer):\n",
        "  \"\"\"Computes GFP embeddings from given optimizer.\"\"\"\n",
        "\n",
        "  gfp_encoding_model = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                                   encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                                   reduce_fn=reduce_fn,\n",
        "                                                   reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                                   num_categories=len(GFP_AMINO_ACID_VOCABULARY),\n",
        "                                                   output='embedding',\n",
        "                                                   output_features=1)\n",
        "  \n",
        "  trained_params = copy.deepcopy(optimizer.target.params)\n",
        "\n",
        "  gfp_encoding_optimizer = create_optimizer(gfp_encoding_model,\n",
        "                                            learning_rate=learning_rate,\n",
        "                                            weight_decay=weight_decay,\n",
        "                                            layers=layers)\n",
        "  \n",
        "  for layer in gfp_encoding_optimizer.target.params.keys():\n",
        "    gfp_encoding_optimizer.target.params[layer] = trained_params[layer]\n",
        "\n",
        "  gfp_train_batches, train_fluorescences = create_gfp_batches(batch_size=256, buffer_size=1)\n",
        "  gfp_train_embeddings = compute_embeddings(gfp_encoding_optimizer.target, gfp_train_batches)\n",
        "  \n",
        "  gfp_test_batches, test_fluorescences = create_gfp_batches(batch_size=256, test=True)\n",
        "  gfp_test_embeddings = compute_embeddings(gfp_encoding_optimizer.target, gfp_test_batches)\n",
        "\n",
        "  return gfp_train_embeddings, gfp_test_embeddings"
      ],
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MstpsYX9kE3A"
      },
      "source": [
        "# CNN + MaxPool Optimizer Embeddings\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 1,\n",
        "    'n_features': [1024],\n",
        "    'n_kernel_sizes': [5],\n",
        "    'n_kernel_dilations': None\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "gfp_train_embeddings, gfp_test_embeddings = gfp_embeddings(encoder_fn, \n",
        "                                                           encoder_fn_kwargs,\n",
        "                                                           reduce_fn,\n",
        "                                                           reduce_fn_kwargs, \n",
        "                                                           gfp_optimizer)"
      ],
      "execution_count": 39,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0q9TVJ4o2qUE"
      },
      "source": [
        "# Random CNN + MaxPool Optimizer Embeddings\n",
        "layers = ['CNN_0', 'Dense_1']                                 \n",
        "learning_rate = [1e-3, 5e-6]\n",
        "weight_decay = [0.0, 0.05]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 1,\n",
        "    'n_features': [1024],\n",
        "    'n_kernel_sizes': [5],\n",
        "    'n_kernel_dilations': None\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "\n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "gfp_model = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                        encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                        reduce_fn=reduce_fn,\n",
        "                                        reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                        num_categories=len(GFP_AMINO_ACID_VOCABULARY),\n",
        "                                        output_features=1)\n",
        "\n",
        "gfp_optimizer_random = create_optimizer(gfp_model, learning_rate, weight_decay, layers)\n",
        "\n",
        "gfp_train_embeddings_random, gfp_test_embeddings_random = gfp_embeddings(encoder_fn, \n",
        "                                                                         encoder_fn_kwargs,\n",
        "                                                                         reduce_fn,\n",
        "                                                                         reduce_fn_kwargs, \n",
        "                                                                         gfp_optimizer_random)"
      ],
      "execution_count": 40,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e-6PLUHajMc3"
      },
      "source": [
        "# CNN + LinearMaxPool Optimizer Embeddings\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 1,\n",
        "    'n_features': [1024],\n",
        "    'n_kernel_sizes': [5],\n",
        "    'n_kernel_dilations': None\n",
        "}\n",
        "reduce_fn = linear_max_pool\n",
        "reduce_fn_kwargs = {\n",
        "    'rep_size': 2048\n",
        "}\n",
        "gfp_train_embeddings_l, gfp_test_embeddings_l = gfp_embeddings(encoder_fn,\n",
        "                                                               encoder_fn_kwargs,\n",
        "                                                               reduce_fn,\n",
        "                                                               reduce_fn_kwargs,\n",
        "                                                               gfp_optimizer_l)"
      ],
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WRc7urS3Mjfi"
      },
      "source": [
        "### Train PCA"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aLIqTubRIfC9"
      },
      "source": [
        "def gfp_train_pca_plot(train_embeddings, model_name):\n",
        "  \"\"\"Applies and plots PCA on GFP train embeddings.\"\"\"\n",
        "  \n",
        "  pca = PCA(n_components=2).fit(train_embeddings)\n",
        "  gfp_train_embeddings_pca = pca.transform(train_embeddings)\n",
        "\n",
        "  X_train = [g[0] for g in gfp_train_embeddings_pca]\n",
        "  Y_train = [g[1] for g in gfp_train_embeddings_pca]\n",
        "\n",
        "  gfp_train_df = create_gfp_df()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_train, Y_train, c=gfp_train_df.log_fluorescence.values, s=1, alpha=0.5)\n",
        "  plt.title('PCA of Train GFP Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  plt.colorbar().set_label('LogFluorescence')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_train, Y_train, c=gfp_train_df.num_mutations.values, s=1, alpha=0.5)\n",
        "  plt.title('PCA of Train GFP Embeddings( ' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('Mutations (Edit Distance)')\n",
        "  cbar.set_ticks(range(4))\n",
        "  plt.show()\n",
        "\n",
        "  return pca"
      ],
      "execution_count": 49,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t1F0Oq9gMmt1"
      },
      "source": [
        "gfp_pca_random = gfp_train_pca_plot(gfp_train_embeddings_random, model_name='Random CNN + MaxPool')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wyM6IBpyMqnI"
      },
      "source": [
        "gfp_pca = gfp_train_pca_plot(gfp_train_embeddings, model_name='CNN + MaxPool')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QWxjthryMqp1"
      },
      "source": [
        "gfp_pca_l = gfp_train_pca_plot(gfp_train_embeddings_l, model_name='CNN + LinearMaxPool')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3RgnuxcXMzCc"
      },
      "source": [
        "### Test PCA"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OEP0DwO6IfGC"
      },
      "source": [
        "def gfp_test_pca_plot(test_embeddings, model_name, gfp_pca, max_mut=8):\n",
        "  \"\"\"Applies and plots PCA on GFP test embeddings.\"\"\"\n",
        "\n",
        "  gfp_test_embeddings_pca = gfp_pca.transform(test_embeddings)\n",
        "  \n",
        "  X_test = [g[0] for g in gfp_test_embeddings_pca]\n",
        "  Y_test = [g[1] for g in gfp_test_embeddings_pca]\n",
        "\n",
        "  gfp_test_df = create_gfp_df(test=True)\n",
        "  inds = [i for i in range(len(X_test)) if gfp_test_df.num_mutations.values[i]<=max_mut]\n",
        "\n",
        "  X_test_mut_cap = [X_test[i] for i in inds]\n",
        "  Y_test_mut_cap = [Y_test[i] for i in inds]\n",
        "  mutations_mut_cap = [gfp_test_df.num_mutations.values[i] for i in inds]\n",
        "  fluorescences_mut_cap = [gfp_test_df.log_fluorescence.values[i] for i in inds]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_test_mut_cap, Y_test_mut_cap, c=fluorescences_mut_cap, s=1, alpha=0.25)\n",
        "  plt.title('PCA of Test GFP Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('LogFluorescence')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_test_mut_cap, Y_test_mut_cap, c=mutations_mut_cap, s=1, alpha=0.25)\n",
        "  plt.title('PCA of Test GFP Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('Mutations (Edit Distance)')\n",
        "  cbar.set_ticks(range(4, max_mut+1))\n",
        "  plt.show()"
      ],
      "execution_count": 71,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L6K16W7UIfJJ"
      },
      "source": [
        "gfp_test_pca_plot(gfp_test_embeddings_random, model_name='Random CNN + MaxPool', gfp_pca=gfp_pca_random, max_mut=9)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NgUjVxC-JCsA"
      },
      "source": [
        "gfp_test_pca_plot(gfp_test_embeddings, model_name='CNN + MaxPool', gfp_pca=gfp_pca, max_mut=9)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4BI1OM-0JCvT"
      },
      "source": [
        "gfp_test_pca_plot(gfp_test_embeddings_l, model_name='CNN + LinearMaxPool', gfp_pca=gfp_pca_l, max_mut=9)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2iUyp91eeY8h"
      },
      "source": [
        "### Train t-SNE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QQtQysz2p5M7"
      },
      "source": [
        "def gfp_train_tsne_plot(train_embeddings, model_name, num_samples=2500):\n",
        "  \"\"\"Applies and plots t-SNE on GFP train embeddings.\"\"\"\n",
        "\n",
        "  gfp_train_df = create_gfp_df()\n",
        "\n",
        "  np.random.seed(0)\n",
        "  gfp_train_pairs = np.random.permutation(np.array([(train_embeddings[i], gfp_train_df.log_fluorescence.values[i], gfp_train_df.num_mutations.values[i]) for i in range(len(train_embeddings))]))\n",
        "  sub_gfp_train_pairs = gfp_train_pairs[:num_samples]\n",
        "\n",
        "  sub_gfp_train_embeddings = np.array([g[0] for g in sub_gfp_train_pairs])\n",
        "  sub_gfp_train_fluorescences = np.array([g[1] for g in sub_gfp_train_pairs])\n",
        "  sub_gfp_train_mutations = np.array([g[2] for g in sub_gfp_train_pairs])\n",
        "\n",
        "  gfp_train_embeddings_tsne = TSNE(n_components=2).fit_transform(sub_gfp_train_embeddings)\n",
        "\n",
        "  X_tsne_train = [g[0] for g in gfp_train_embeddings_tsne]\n",
        "  Y_tsne_train = [g[1] for g in gfp_train_embeddings_tsne]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_train, Y_tsne_train, c=sub_gfp_train_fluorescences, s=5, alpha=0.5)\n",
        "  plt.title('t-SNE of Train GFP Embeddings (' + model_name + ')')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('LogFluorescence')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_train, Y_tsne_train, c=sub_gfp_train_mutations, s=5, alpha=0.5)\n",
        "  plt.title('t-SNE of Train GFP Embeddings (' + model_name + ')')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('Mutations (Edit Distance)')\n",
        "  cbar.set_ticks(range(4))\n",
        "  plt.show()"
      ],
      "execution_count": 68,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "blu_cukp5AJ4"
      },
      "source": [
        "gfp_train_tsne_plot(gfp_train_embeddings_random, model_name='Random CNN + MaxPool', num_samples=2500)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V9vuSquYELqP"
      },
      "source": [
        "gfp_train_tsne_plot(gfp_train_embeddings, model_name='CNN + MaxPool', num_samples=2500)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wcUtG3-Vp5TF"
      },
      "source": [
        "gfp_train_tsne_plot(gfp_train_embeddings_l, model_name='CNN + LinearMaxPool', num_samples=2500)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hwjV0A6eec3N"
      },
      "source": [
        "### Test t-SNE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ybJv0pPIs5TR"
      },
      "source": [
        "def gfp_test_tsne_plot(test_embeddings, model_name, max_mut=8, num_samples=2500):\n",
        "  \"\"\"Applies and plots t-SNE on GFP test embeddings.\"\"\"\n",
        "\n",
        "  gfp_test_df = create_gfp_df(test=True)\n",
        "\n",
        "  np.random.seed(0)\n",
        "  gfp_test_pairs_mut_cap = np.random.permutation(np.array([(test_embeddings[i], gfp_test_df.log_fluorescence.values[i], gfp_test_df.num_mutations.values[i]) for i in range(len(test_embeddings)) if gfp_test_df.num_mutations.values[i]<=max_mut]))\n",
        "  sub_gfp_test_pairs_mut_cap = gfp_test_pairs_mut_cap[:num_samples]\n",
        "\n",
        "  sub_gfp_test_embeddings_mut_cap = np.array([g[0] for g in sub_gfp_test_pairs_mut_cap])\n",
        "  sub_gfp_test_fluorescences_mut_cap = np.array([g[1] for g in sub_gfp_test_pairs_mut_cap])\n",
        "  sub_gfp_test_mutations_mut_cap = np.array([g[2] for g in sub_gfp_test_pairs_mut_cap])\n",
        "\n",
        "  gfp_test_embeddings_tsne_mut_cap = TSNE(n_components=2).fit_transform(sub_gfp_test_embeddings_mut_cap)\n",
        "\n",
        "  X_tsne_test_mut_cap = [g[0] for g in gfp_test_embeddings_tsne_mut_cap]\n",
        "  Y_tsne_test_mut_cap = [g[1] for g in gfp_test_embeddings_tsne_mut_cap]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_test_mut_cap, Y_tsne_test_mut_cap, c=sub_gfp_test_fluorescences_mut_cap, s=5, alpha=0.5)\n",
        "  plt.title('t-SNE of Test GFP Embeddings (' + model_name + ')')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('LogFluorescence')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_test_mut_cap, Y_tsne_test_mut_cap, c=sub_gfp_test_mutations_mut_cap, s=5, alpha=0.5)\n",
        "  plt.title('t-SNE of Test GFP Embeddings (' + model_name + ')')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('Mutations (Edit Distance)')\n",
        "  cbar.set_ticks(range(4, max_mut+1))\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZX-UdRX95DR1"
      },
      "source": [
        "gfp_test_tsne_plot(gfp_test_embeddings_random, model_name='Random CNN + MaxPool', max_mut=9, num_samples=2500)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "un70trFzu6XI"
      },
      "source": [
        "gfp_test_tsne_plot(gfp_test_embeddings, model_name='CNN + MaxPool', max_mut=9, num_samples=2500)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8-kmF76bu9u8"
      },
      "source": [
        "gfp_test_tsne_plot(gfp_test_embeddings_l, model_name='CNN + LinearMaxPool', max_mut=9, num_samples=2500)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gXTQ0QMhnA5O"
      },
      "source": [
        "# Stability"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "frInyFZk-61d"
      },
      "source": [
        "The data originally comes [from Rocklin et. al.](https://science.sciencemag.org/content/357/6347/168) and was formatted by [TAPE](https://github.com/songlab-cal/tape)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Hyo9e7Ij8MX"
      },
      "source": [
        "## Open data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P9k6zeSNnBwN"
      },
      "source": [
        "!wget http://s3.amazonaws.com/proteindata/data_pytorch/stability.tar.gz"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m1QEG5WqkAiB"
      },
      "source": [
        "!tar xzf stability.tar.gz"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ExpKadm3kAtE"
      },
      "source": [
        "def stability_dataset_to_df(in_name):\n",
        "  dataset = LMDBDataset(in_name)\n",
        "  df = pd.DataFrame(list(dataset)[:])\n",
        "  df['stability'] = df.stability_score.apply(lambda x: x[0])\n",
        "  df['id_str'] = df.id.apply(lambda x: x.decode('utf-8'))\n",
        "  return df"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sZ1K84k0yktc"
      },
      "source": [
        "### Padding and one-hot encoding"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S4QvX0sQynUu"
      },
      "source": [
        "STABILITY_SEQ_LEN = 50"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NxmDlbJAMW4a"
      },
      "source": [
        "STABILITY_AMINO_ACID_VOCABULARY = [\n",
        "    'A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R',\n",
        "    'S', 'T', 'V', 'W', 'Y', '-'\n",
        "]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iVFsRaXyu7Ig"
      },
      "source": [
        "STABILITY_PROTEIN_DOMAIN = domains.VariableLengthDiscreteDomain(\n",
        "    vocab=domains.ProteinVocab(include_anomalous_amino_acids=False,\n",
        "                               include_eos=True,\n",
        "                               include_pad=True),\n",
        "    length=STABILITY_SEQ_LEN)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kW6mNWHdvGWB"
      },
      "source": [
        "def stability_seq_to_inds(seq):\n",
        "  \"\"\"Encode stability amino acid sequence.\"\"\"\n",
        "\n",
        "  return STABILITY_PROTEIN_DOMAIN.encode([seq])[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8aVROiHnzg4o"
      },
      "source": [
        "### Add one-hots and batch data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QSwWfUH_8aXN"
      },
      "source": [
        "stability_train_df = stability_dataset_to_df('stability/stability_train.lmdb')\n",
        "stability_test_df = stability_dataset_to_df('stability/stability_test.lmdb')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3DEQqBIlAAkE"
      },
      "source": [
        "parent_to_parent_stability = {}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S3icCCB78NB3"
      },
      "source": [
        "for parent in set(stability_train_df.parent.values): \n",
        "  stabilities = stability_train_df[stability_train_df['id_str']==parent.decode('utf-8') + '.pdb'].stability.values\n",
        "  if len(stabilities) == 0:\n",
        "    stabilities = stability_train_df[stability_train_df['id_str']==parent.decode('utf-8')].stability.values\n",
        "    if len(stabilities) == 0:\n",
        "      parent_to_parent_stability[parent] = None\n",
        "    else:\n",
        "      parent_to_parent_stability[parent] = stabilities[0]\n",
        "  else:\n",
        "    parent_to_parent_stability[parent] = stabilities[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QXzZbFdkBKzO"
      },
      "source": [
        "for parent in set(stability_test_df.parent.values): \n",
        "  stabilities = stability_test_df[stability_test_df['id_str']==parent.decode('utf-8') + '.pdb'].stability.values\n",
        "  if len(stabilities) == 0:\n",
        "    stabilities = stability_test_df[stability_test_df['id_str']==parent.decode('utf-8')].stability.values\n",
        "    if len(stabilities) == 0:\n",
        "      parent_to_parent_stability[parent] = None\n",
        "    else:\n",
        "      parent_to_parent_stability[parent] = stabilities[0]\n",
        "  else:\n",
        "    parent_to_parent_stability[parent] = stabilities[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zF4w7F6n8BHO"
      },
      "source": [
        "topology_to_ind = {'HHH': 0, 'HEEH': 1, 'EHEE': 3, 'EEHEE': 4}\n",
        "def topology_to_index(top):\n",
        "  \"\"\"Returns categorical variable corresponding to a topology.\"\"\"\n",
        "  \n",
        "  top = top.decode('utf-8')\n",
        "  if top in topology_to_ind.keys():\n",
        "    return topology_to_ind[top]\n",
        "  else:\n",
        "    return 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HND4bMpyzhSM"
      },
      "source": [
        "def create_stability_df(test=False):\n",
        "  \"\"\"Processes stability data into a featurized dataframe.\"\"\"\n",
        "  \n",
        "  if test:\n",
        "    stability_df = stability_dataset_to_df('stability/stability_test.lmdb')\n",
        "  else:\n",
        "    stability_df = stability_dataset_to_df('stability/stability_train.lmdb')\n",
        "  \n",
        "  stability_df['one_hot_inds'] = stability_df.primary.apply(lambda x: stability_seq_to_inds(x[:STABILITY_SEQ_LEN]))\n",
        "\n",
        "  stability_df['parent_stability'] = stability_df.parent.apply(lambda x: parent_to_parent_stability[x])\n",
        "\n",
        "  stability_df['topology_ind'] = stability_df.topology.apply(lambda x: topology_to_index(x))\n",
        "\n",
        "  return stability_df\n",
        "\n",
        "\n",
        "def create_stability_batches(batch_size, epochs=1, test=False, buffer_size=None,\n",
        "                             seed=0, drop_remainder=False):\n",
        "  \"\"\"Creates iterable object of Stability batches.\"\"\"\n",
        "  \n",
        "  if test:\n",
        "    buffer_size = 1\n",
        "  \n",
        "  stability_df = create_stability_df(test=test)\n",
        "    \n",
        "  stabilities = stability_df['stability'].values\n",
        "\n",
        "  stability_batches = create_data_iterator(df=stability_df, input_col='one_hot_inds',\n",
        "                                           output_col='stability', batch_size=batch_size,\n",
        "                                           epochs=epochs, buffer_size=buffer_size, \n",
        "                                           seed=seed, drop_remainder=drop_remainder)\n",
        "\n",
        "  return stability_batches, stabilities"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bH0uaMhr0s1j"
      },
      "source": [
        "## Model evaluation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9gqbT2oFE6JT"
      },
      "source": [
        "def stability_evaluate(predict_fn, title, batch_size=256, \n",
        "                       test_data=None, pred_stabilities=None,\n",
        "                       clip_min=-999999, clip_max=999999):\n",
        "  \"\"\"Computes predicted stabilities and measures performance in spearman correlation.\"\"\"\n",
        "  \n",
        "  test_batches, test_stabilities = create_stability_batches(batch_size=batch_size, test=True, buffer_size=1)\n",
        "  \n",
        "  if test_data is not None:\n",
        "    test_batches = test_data\n",
        "\n",
        "  if pred_stabilities is None:\n",
        "    pred_stabilities = []\n",
        "    for batch in iter(test_batches):\n",
        "      X, Y = batch\n",
        "      preds = predict_fn(X)\n",
        "      for pred in preds:\n",
        "        pred_stabilities.append(pred[0])\n",
        "  \n",
        "  pred_stabilities = np.array(pred_stabilities)\n",
        "  pred_stabilities = np.clip(pred_stabilities, clip_min, clip_max)\n",
        "    \n",
        "  spearmanr = scipy.stats.spearmanr(test_stabilities, pred_stabilities).correlation\n",
        "  plt.scatter(test_stabilities, pred_stabilities, s=1, alpha=0.5)\n",
        "  plt.xlabel('True Stability')\n",
        "  plt.ylabel('Predicted Stability')\n",
        "  plt.title(title)\n",
        "  plt.show()\n",
        "  \n",
        "  stability_test_df = create_stability_df(test=True)\n",
        "  stability_test_df['pred_stability'] = pred_stabilities\n",
        "  stability_test_df['topology_str'] = stability_test_df.topology.apply(lambda x: x.decode('utf-8'))\n",
        "\n",
        "  parent_to_pred_parent_stability = {}\n",
        "  for parent in set(stability_test_df.parent.values): \n",
        "    pred_parent_stabilities = stability_test_df[stability_test_df['id_str']==parent.decode('utf-8') + '.pdb'].pred_stability.values\n",
        "    if len(pred_parent_stabilities) == 0:\n",
        "      pred_parent_stabilities = stability_test_df[stability_test_df['id_str']==parent.decode('utf-8')].pred_stability.values\n",
        "      if len(pred_parent_stabilities) == 0:\n",
        "        parent_to_pred_parent_stability[parent] = None\n",
        "      else:\n",
        "        parent_to_pred_parent_stability[parent] = pred_parent_stabilities[0]\n",
        "    else:\n",
        "      parent_to_pred_parent_stability[parent] = pred_parent_stabilities[0]\n",
        "  \n",
        "  stability_test_df['pred_parent_stability'] = stability_test_df.parent.apply(lambda x: parent_to_pred_parent_stability[x])\n",
        "  \n",
        "  parent_stabilities = stability_test_df['parent_stability'].values\n",
        "  pred_parent_stabilities = stability_test_df['pred_parent_stability'].values\n",
        "\n",
        "  correct_direction = 0\n",
        "  for test_stability, pred_stability, parent_stability, pred_parent_stability in zip(test_stabilities, pred_stabilities, parent_stabilities, pred_parent_stabilities):\n",
        "    if parent_stability is not None:\n",
        "      if (test_stability >= parent_stability and pred_stability >= pred_parent_stability) or (test_stability <= parent_stability and pred_stability <= pred_parent_stability):\n",
        "        correct_direction += 1\n",
        "  accuracy = correct_direction / len(np.where(parent_stabilities!=None)[0])\n",
        "  \n",
        "  topologies = [('EEHEE', 'BBABB'), ('EHEE', 'BABB'), ('HEEH', 'ABBA'), ('HHH', 'AAA')]\n",
        "  topology_results = {}\n",
        "  \n",
        "  for topology_pair in topologies:\n",
        "    topology, topology_name = topology_pair\n",
        "    topology_test_df = stability_test_df[stability_test_df['topology_str']==topology]\n",
        "\n",
        "    topology_test_stabilities = topology_test_df['stability'].values\n",
        "    topology_pred_stabilities = topology_test_df['pred_stability'].values\n",
        "      \n",
        "    topology_spearmanr = scipy.stats.spearmanr(topology_test_stabilities, topology_pred_stabilities).correlation\n",
        "    plt.scatter(topology_test_stabilities, topology_pred_stabilities, s=1, alpha=0.5)\n",
        "    plt.xlabel('True Stability')\n",
        "    plt.ylabel('Predicted Stability')\n",
        "    plt.title(title + ' (' + topology_name + ')')\n",
        "    plt.show()\n",
        "    \n",
        "    topology_parent_stabilities = topology_test_df['parent_stability'].values\n",
        "    topology_pred_parent_stabilities = topology_test_df['pred_parent_stability'].values\n",
        "\n",
        "    topology_correct_direction = 0\n",
        "    for test_stability, pred_stability, parent_stability, pred_parent_stability in zip(topology_test_stabilities, topology_pred_stabilities, topology_parent_stabilities, topology_pred_parent_stabilities):\n",
        "      if parent_stability is not None:\n",
        "        if (test_stability >= parent_stability and pred_stability >= pred_parent_stability) or (test_stability <= parent_stability and pred_stability <= pred_parent_stability):\n",
        "          topology_correct_direction += 1\n",
        "    topology_accuracy = topology_correct_direction / len(np.where(topology_parent_stabilities!=None)[0])\n",
        "    \n",
        "    topology_results[topology_name] = (topology_spearmanr, topology_accuracy)\n",
        "\n",
        "  results = {\n",
        "      'title': title,\n",
        "      'spearmanr': round(spearmanr, 3),\n",
        "      'accuracy': round(accuracy, 3)\n",
        "  }\n",
        "  for topology_name in topology_results.keys():\n",
        "    results[topology_name + '_spearmanr'] = round(topology_results[topology_name][0], 3)\n",
        "    results[topology_name + '_accuracy'] = round(topology_results[topology_name][1], 3)\n",
        "  \n",
        "  pprint.pprint(results)\n",
        "\n",
        "  return results, pred_stabilities"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0yO_4_B00LR8"
      },
      "source": [
        "## Experiments"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_6fg0DSUyAJW"
      },
      "source": [
        "stability_train_df = create_stability_df()\n",
        "train_stabilities = stability_train_df['stability']\n",
        "stability_train_one_hot_inds = np.array([x for x in stability_train_df['one_hot_inds'].values])\n",
        "stability_train_one_hots = flattened_one_hot_encoder(stability_train_one_hot_inds, num_categories=len(STABILITY_AMINO_ACID_VOCABULARY))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LieskquwyCWp"
      },
      "source": [
        "stability_test_df = create_stability_df(test=True)\n",
        "stability_test_one_hot_inds = np.array([x for x in stability_test_df['one_hot_inds'].values])\n",
        "stability_test_one_hots = flattened_one_hot_encoder(stability_test_one_hot_inds, num_categories=len(STABILITY_AMINO_ACID_VOCABULARY))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZU1mXY0Z0V2w"
      },
      "source": [
        "### Linear regression"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d0wa1Tmm0YCB"
      },
      "source": [
        "stability_linear_model = Ridge()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f1vYTZYDyFx3"
      },
      "source": [
        "stability_linear_model.fit(X=stability_train_one_hots, y=train_stabilities)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6qlIVGFL0X8u"
      },
      "source": [
        "linear_model_pred_stabilities = stability_linear_model.predict(stability_test_one_hots)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OPI8Blhocfjk"
      },
      "source": [
        "stability_linear_model_results, linear_model_pred_stabilities = \\\n",
        "stability_evaluate(predict_fn=None, \n",
        "                   title='Linear Regression',\n",
        "                   pred_stabilities=linear_model_pred_stabilities,\n",
        "                   clip_min=min(train_stabilities),\n",
        "                   clip_max=max(train_stabilities))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ai4ZvsJv_vD8"
      },
      "source": [
        "print('Number of Parameters for Stability Linear Regression: ' + str(len(stability_linear_model.coef_)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uHUniDwNW3pk"
      },
      "source": [
        "### CNN + MaxPool"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qvCTeHusW3yZ"
      },
      "source": [
        "epochs = 5\n",
        "stability_train_batches, train_stabilities = create_stability_batches(batch_size=16, epochs=epochs)\n",
        "\n",
        "layers = ['CNN_0', 'Dense_1']                                 \n",
        "learning_rate = [5e-4, 5e-5]\n",
        "weight_decay = [0.025, 0.025]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 3,\n",
        "    'n_features': [1024, 1024, 1024],\n",
        "    'n_kernel_sizes': [5, 5, 5],\n",
        "    'n_kernel_dilations': None\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "stability_model = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                              encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                              reduce_fn=reduce_fn,\n",
        "                                              reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                              num_categories=len(STABILITY_AMINO_ACID_VOCABULARY),\n",
        "                                              output_features=1)\n",
        "\n",
        "stability_optimizer = train(model=stability_model,\n",
        "                            train_data=stability_train_batches, \n",
        "                            loss_fn=mse_loss,\n",
        "                            loss_fn_kwargs=loss_fn_kwargs,\n",
        "                            learning_rate=learning_rate, \n",
        "                            weight_decay=weight_decay,\n",
        "                            layers=layers)\n",
        "\n",
        "stability_results, pred_stabilities = stability_evaluate(predict_fn=stability_optimizer.target,\n",
        "                                                         title='CNN + MaxPool',\n",
        "                                                         batch_size=1024,\n",
        "                                                         clip_min=min(train_stabilities),\n",
        "                                                         clip_max=max(train_stabilities))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZJ_qVMwb-uC9"
      },
      "source": [
        "print('Number of Parameters for Stability CNN + MaxPool: ' + str(get_num_params(stability_optimizer.target)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-mlZGFNmdyFt"
      },
      "source": [
        "epochs = 5\n",
        "stability_train_batches, train_stabilities = create_stability_batches(batch_size=16, epochs=epochs)\n",
        "\n",
        "layers = ['CNN_0', 'Dense_1']                                 \n",
        "learning_rate = [5e-4, 5e-5]\n",
        "weight_decay = [0.025, 0.025]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 3,\n",
        "    'n_features': [1024, 1024, 1024],\n",
        "    'n_kernel_sizes': [5, 5, 5],\n",
        "    'n_kernel_dilations': [1, 2, 1]\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "stability_model_d = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                                encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                                reduce_fn=reduce_fn,\n",
        "                                                reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                                num_categories=len(STABILITY_AMINO_ACID_VOCABULARY),\n",
        "                                                output_features=1)\n",
        "\n",
        "stability_optimizer_d = train(model=stability_model_d,\n",
        "                              train_data=stability_train_batches, \n",
        "                              loss_fn=mse_loss,\n",
        "                              loss_fn_kwargs=loss_fn_kwargs,\n",
        "                              learning_rate=learning_rate, \n",
        "                              weight_decay=weight_decay,\n",
        "                              layers=layers)\n",
        "\n",
        "stability_results_d, pred_stabilities_d = stability_evaluate(predict_fn=stability_optimizer_d.target,\n",
        "                                                             title='Dilated CNN + MaxPool',\n",
        "                                                             batch_size=1024,\n",
        "                                                             clip_min=min(train_stabilities),\n",
        "                                                             clip_max=max(train_stabilities))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uxU5r_F7dyOw"
      },
      "source": [
        "print('Number of Parameters for Dilated Stability CNN + MaxPool: ' + str(get_num_params(stability_optimizer_d.target)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R6OUT7DP2-ES"
      },
      "source": [
        "### CNN + LinearMaxPool"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6JyyIntz9AHY"
      },
      "source": [
        "epochs = 10\n",
        "stability_train_batches, train_stabilities = create_stability_batches(batch_size=256, epochs=epochs)\n",
        "\n",
        "layers = ['CNN_0', 'Dense_1', 'Dense_2']  \n",
        "                              \n",
        "learning_rate = [1e-5, 5e-5, 5e-6]\n",
        "weight_decay = [0.0, 0.0, 0.0]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 3,\n",
        "    'n_features': [1024, 1024, 1024],\n",
        "    'n_kernel_sizes': [5, 5, 5],\n",
        "    'n_kernel_dilations': None\n",
        "}\n",
        "reduce_fn = linear_max_pool\n",
        "reduce_fn_kwargs = {\n",
        "    'rep_size': 2048\n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "stability_model_l = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                                encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                                reduce_fn=reduce_fn,\n",
        "                                                reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                                num_categories=len(STABILITY_AMINO_ACID_VOCABULARY),\n",
        "                                                output_features=1)\n",
        "\n",
        "stability_optimizer_l = train(model=stability_model_l,\n",
        "                              train_data=stability_train_batches, \n",
        "                              loss_fn=mse_loss,\n",
        "                              loss_fn_kwargs=loss_fn_kwargs,\n",
        "                              learning_rate=learning_rate, \n",
        "                              weight_decay=weight_decay,\n",
        "                              layers=layers)\n",
        "\n",
        "stability_results_l, pred_stabilities_l = stability_evaluate(predict_fn=stability_optimizer_l.target,\n",
        "                                                             title='CNN + LinearMaxPool',\n",
        "                                                             batch_size=1024,\n",
        "                                                             clip_min=min(train_stabilities),\n",
        "                                                             clip_max=max(train_stabilities))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yBwPuJv6-18H"
      },
      "source": [
        "print('Number of Parameters for Stability CNN + LinearMaxPool: ' + str(get_num_params(stability_optimizer_l.target)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MJ9I_Qr8nVsX"
      },
      "source": [
        "### Ensemble of CNNs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CrfvP-qEssci"
      },
      "source": [
        "ensemble_pred_stabilities = [(pred_stabilities[i]+pred_stabilities_d[i]+pred_stabilities_l[i])/3 for i in range(len(pred_stabilities))]\n",
        "ensemble_results, ensemble_pred_stabilities = stability_evaluate(predict_fn=None, \n",
        "                                                                 title='Ensemble of CNNs',\n",
        "                                                                 pred_stabilities=ensemble_pred_stabilities,\n",
        "                                                                 clip_min=min(train_stabilities),\n",
        "                                                                 clip_max=max(train_stabilities))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BNmobGqdECCz"
      },
      "source": [
        "print('Number of Parameters for Ensemble of CNNs: ' + str(get_num_params(stability_optimizer.target) + get_num_params(stability_optimizer_d.target) + get_num_params(stability_optimizer_l.target)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvun0LAR5ssO"
      },
      "source": [
        "### TAPE (Best Performance)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rArqT44j5sLE"
      },
      "source": [
        "tape_stability_results = {\n",
        "    'title': 'TAPE',\n",
        "    'spearmanr': 0.73,\n",
        "    'accuracy': 0.70,\n",
        "    'AAA_spearmanr': 0.72,\n",
        "    'AAA_accuracy': 0.70,\n",
        "    'ABBA_spearmanr': 0.48,\n",
        "    'ABBA_accuracy': 0.79,\n",
        "    'BABB_spearmanr': 0.68,\n",
        "    'BABB_accuracy': 0.71,\n",
        "    'BBABB_spearmanr': 0.67,\n",
        "    'BBABB_accuracy': 0.70\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w8WYARpX6aa8"
      },
      "source": [
        "pprint.pprint(tape_stability_results)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QcrMzPvxW40L"
      },
      "source": [
        "## Visualization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fY0YxgT9zw0L"
      },
      "source": [
        "def stability_embeddings(encoder_fn, encoder_fn_kwargs, reduce_fn, reduce_fn_kwargs, optimizer):\n",
        "  \"\"\"Computes stability embeddings from given optimizer.\"\"\"\n",
        "\n",
        "  stability_encoding_model = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                                         encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                                         reduce_fn=reduce_fn,\n",
        "                                                         reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                                         num_categories=len(STABILITY_AMINO_ACID_VOCABULARY),\n",
        "                                                         output='embedding',\n",
        "                                                         output_features=1)\n",
        "  \n",
        "  trained_params = copy.deepcopy(optimizer.target.params)\n",
        "\n",
        "  stability_encoding_optimizer = create_optimizer(stability_encoding_model,\n",
        "                                                  learning_rate=learning_rate,\n",
        "                                                  weight_decay=weight_decay,\n",
        "                                                  layers=layers)\n",
        "  \n",
        "  for layer in stability_encoding_optimizer.target.params.keys():\n",
        "    stability_encoding_optimizer.target.params[layer] = trained_params[layer]\n",
        "\n",
        "  stability_train_batches, train_stabilities = create_stability_batches(batch_size=256, buffer_size=1)\n",
        "  stability_train_embeddings = compute_embeddings(stability_encoding_optimizer.target, stability_train_batches)\n",
        "  \n",
        "  stability_test_batches, test_stabilities = create_stability_batches(batch_size=256, test=True)\n",
        "  stability_test_embeddings = compute_embeddings(stability_encoding_optimizer.target, stability_test_batches)\n",
        "\n",
        "  return stability_train_embeddings, stability_test_embeddings"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6ierIaqKT-x8"
      },
      "source": [
        "# CNN + MaxPool Optimizer Embeddings\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 1,\n",
        "    'n_features': [1024],\n",
        "    'n_kernel_sizes': [5],\n",
        "    'n_kernel_dilations': None\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "\n",
        "}\n",
        "stability_train_embeddings, stability_test_embeddings = stability_embeddings(encoder_fn,\n",
        "                                                                             encoder_fn_kwargs,\n",
        "                                                                             reduce_fn,\n",
        "                                                                             reduce_fn_kwargs,\n",
        "                                                                             stability_optimizer)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m3_OA_bh0C1w"
      },
      "source": [
        "# Random CNN + MaxPool Optimizer Embeddings\n",
        "layers = ['CNN_0', 'Dense_1']                                 \n",
        "learning_rate = [5e-4, 5e-5]\n",
        "weight_decay = [0.025, 0.025]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 3,\n",
        "    'n_features': [1024, 1024, 1024],\n",
        "    'n_kernel_sizes': [5, 5, 5],\n",
        "    'n_kernel_dilations': None\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "stability_model = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                              encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                              reduce_fn=reduce_fn,\n",
        "                                              reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                              num_categories=len(STABILITY_AMINO_ACID_VOCABULARY),\n",
        "                                              output_features=1)\n",
        "\n",
        "stability_optimizer_random = create_optimizer(stability_model, learning_rate, weight_decay, layers)\n",
        "\n",
        "stability_train_embeddings_random, stability_test_embeddings_random = stability_embeddings(encoder_fn,\n",
        "                                                                                           encoder_fn_kwargs,\n",
        "                                                                                           reduce_fn,\n",
        "                                                                                           reduce_fn_kwargs,\n",
        "                                                                                           stability_optimizer_random)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kfmSlEoBzw6c"
      },
      "source": [
        "# CNN + LinearMaxPool Optimizer Embeddings\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 1,\n",
        "    'n_features': [1024],\n",
        "    'n_kernel_sizes': [5]\n",
        "}\n",
        "reduce_fn = linear_max_pool\n",
        "reduce_fn_kwargs = {\n",
        "    'rep_size': 2048\n",
        "}\n",
        "stability_train_embeddings_l, stability_test_embeddings_l = stability_embeddings(encoder_fn,\n",
        "                                                                                 encoder_fn_kwargs,\n",
        "                                                                                 reduce_fn,\n",
        "                                                                                 reduce_fn_kwargs,\n",
        "                                                                                 stability_optimizer_l)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lWR2UXSpYKrD"
      },
      "source": [
        "# Ensemble Embeddings (Concatention of Above Embeddings)\n",
        "stability_train_embeddings_e = np.concatenate((stability_train_embeddings, stability_train_embeddings_l), axis=1)\n",
        "stability_test_embeddings_e = np.concatenate((stability_test_embeddings, stability_test_embeddings_l), axis=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eooK40Ut1Per"
      },
      "source": [
        "### Train PCA"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NTqKUZAsW64L"
      },
      "source": [
        "def stability_train_pca_plot(train_embeddings, model_name):\n",
        "  \"\"\"Applies and plots PCA on stability train embeddings.\"\"\"\n",
        "\n",
        "  stability_train_df = create_stability_df()\n",
        "\n",
        "  pca = PCA(n_components=2).fit(train_embeddings)\n",
        "  stability_train_embeddings_pca = pca.transform(train_embeddings)\n",
        "  \n",
        "  X_train = [s[0] for s in stability_train_embeddings_pca]\n",
        "  Y_train = [s[1] for s in stability_train_embeddings_pca]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_train, Y_train, c=stability_train_df.stability.values, s=1, alpha=0.25)\n",
        "  plt.title('PCA of Train Stability Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  plt.colorbar().set_label('Stability')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_train, Y_train, c=stability_train_df.parent_stability.values, s=1, alpha=0.5)\n",
        "  plt.title('PCA of Train Stability Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  plt.colorbar().set_label('Parent Stability')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_train, Y_train, c=stability_train_df.topology_ind.values, s=1, alpha=0.5)\n",
        "  plt.title('PCA of Train Stability Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  \n",
        "  topology_names = ['AAA', 'ABBA', 'Other', 'BABB', 'BBABB']\n",
        "  formatter = plt.FuncFormatter(lambda val, loc: topology_names[val])\n",
        "  cbar = plt.colorbar(ticks=range(5), format=formatter)\n",
        "  cbar.set_label('Topology')\n",
        "  plt.show()\n",
        "\n",
        "  return pca"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fPJItPWQ0ofT"
      },
      "source": [
        "stability_pca_random = stability_train_pca_plot(stability_train_embeddings_random, model_name='Random CNN + MaxPool')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fEtsiil8IQxP"
      },
      "source": [
        "stability_pca = stability_train_pca_plot(stability_train_embeddings, model_name='CNN + MaxPool')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OuItkHRpW6-o"
      },
      "source": [
        "stability_pca_l = stability_train_pca_plot(stability_train_embeddings_l, model_name='CNN + LinearMaxPool')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2pkjDpNDZFPO"
      },
      "source": [
        "stability_pca_e = stability_train_pca_plot(stability_train_embeddings_e, model_name='Ensemble of CNNs')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WLtjJpFsJTzT"
      },
      "source": [
        "### Test PCA"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WTPi-j-U6yaX"
      },
      "source": [
        "def stability_test_pca_plot(test_embeddings, model_name, stability_pca):\n",
        "  \"\"\"Applies and plots PCA on stability test embeddings.\"\"\"\n",
        "\n",
        "  stability_test_df = create_stability_df(test=True)\n",
        "\n",
        "  stability_test_embeddings_pca = stability_pca.transform(test_embeddings)\n",
        "  \n",
        "  X_test = [s[0] for s in stability_test_embeddings_pca]\n",
        "  Y_test = [s[1] for s in stability_test_embeddings_pca]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_test, Y_test, c=stability_test_df.stability.values, s=1, alpha=0.25)\n",
        "  plt.title('PCA of Test Stability Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  plt.colorbar().set_label('Stability')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_test, Y_test, c=stability_test_df.parent_stability.values, s=1, alpha=0.5)\n",
        "  plt.title('PCA of Test Stability Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  plt.colorbar().set_label('Parent Stability')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_test, Y_test, c=stability_test_df.topology_ind.values, s=1, alpha=0.5)\n",
        "  plt.title('PCA of Test Stability Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  topology_names = ['AAA', 'ABBA', 'Other', 'BABB', 'BBABB']\n",
        "  formatter = plt.FuncFormatter(lambda val, loc: topology_names[val])\n",
        "  cbar = plt.colorbar(ticks=range(5), format=formatter)\n",
        "  cbar.set_label('Topology')\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jWwDdABQ1B1q"
      },
      "source": [
        "stability_test_pca_plot(stability_test_embeddings_random, model_name='Random CNN + MaxPool', stability_pca=stability_pca_random)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bQTKe1q97T4S"
      },
      "source": [
        "stability_test_pca_plot(stability_test_embeddings, model_name='CNN + MaxPool', stability_pca=stability_pca)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0Ae7MvPw7sor"
      },
      "source": [
        "stability_test_pca_plot(stability_test_embeddings_l, model_name='CNN + LinearMaxPool', stability_pca=stability_pca_l)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "khp_PzevZQQO"
      },
      "source": [
        "stability_test_pca_plot(stability_test_embeddings_e, model_name='Ensemble of CNNs', stability_pca=stability_pca_e)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4nO7x4BiQmbZ"
      },
      "source": [
        "### Train t-SNE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OHfeUmXU7sr6"
      },
      "source": [
        "def stability_train_tsne_plot(train_embeddings, model_name, num_samples=5000):\n",
        "  \"\"\"Applies and plots t-SNE on stability train embeddings.\"\"\"\n",
        "\n",
        "  stability_train_df = create_stability_df()\n",
        "\n",
        "  np.random.seed(0)\n",
        "  stability_train_pairs = np.random.permutation(np.array([(train_embeddings[i], stability_train_df.stability.values[i], stability_train_df.parent_stability.values[i], stability_train_df.topology_ind.values[i]) for i in range(len(train_embeddings))]))\n",
        "  sub_stability_train_pairs = stability_train_pairs[:num_samples]\n",
        "\n",
        "  sub_stability_train_embeddings = np.array([s[0] for s in sub_stability_train_pairs])\n",
        "  sub_stability_train_stabilities = np.array([s[1] for s in sub_stability_train_pairs])\n",
        "  sub_stability_train_parent_stabilities = np.array([s[2] for s in sub_stability_train_pairs])\n",
        "  sub_stability_train_topologies = np.array([s[3] for s in sub_stability_train_pairs])\n",
        "\n",
        "  stability_train_embeddings_tsne = TSNE(n_components=2).fit_transform(sub_stability_train_embeddings)\n",
        "\n",
        "  X_tsne_train = [s[0] for s in stability_train_embeddings_tsne]\n",
        "  Y_tsne_train = [s[1] for s in stability_train_embeddings_tsne]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_train, Y_tsne_train, c=sub_stability_train_stabilities, s=5, alpha=0.5)\n",
        "  plt.title('t-SNE of Train Stability Embeddings (' + model_name + ')')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('Stability')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_train, Y_tsne_train, c=sub_stability_train_parent_stabilities, s=5, alpha=0.5)\n",
        "  plt.title('t-SNE of Train Stability Embeddings (' + model_name + ')')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('Parent Stability')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_train, Y_tsne_train, c=sub_stability_train_topologies, s=5, alpha=0.5)\n",
        "  plt.title('t-SNE of Train Stability Embeddings (' + model_name + ')')\n",
        "  topology_names = ['AAA', 'ABBA', 'Other', 'BABB', 'BBABB']\n",
        "  formatter = plt.FuncFormatter(lambda val, loc: topology_names[val])\n",
        "  cbar = plt.colorbar(ticks=range(5), format=formatter)\n",
        "  cbar.set_label('Topology')\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IWAMloJx1LAM"
      },
      "source": [
        "stability_train_tsne_plot(stability_train_embeddings_random, model_name='Random CNN + MaxPool', num_samples=5000)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fxVe0pYyk5Ub"
      },
      "source": [
        "stability_train_tsne_plot(stability_train_embeddings, model_name='CNN + MaxPool', num_samples=5000)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FOwzxPQFk5Xd"
      },
      "source": [
        "stability_train_tsne_plot(stability_train_embeddings_l, model_name='CNN + LinearMaxPool', num_samples=5000)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rUOGIHyFk5aV"
      },
      "source": [
        "stability_train_tsne_plot(stability_train_embeddings_e, model_name='Ensemble of CNNs', num_samples=5000)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PPl-2dUoQoe5"
      },
      "source": [
        "### Test t-SNE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4lff7OBT7swF"
      },
      "source": [
        "def stability_test_tsne_plot(test_embeddings, model_name, num_samples=5000):\n",
        "  \"\"\"Applies and plots t-SNE on stability test embeddings.\"\"\"\n",
        "\n",
        "  stability_test_df = create_stability_df(test=True)\n",
        "\n",
        "  np.random.seed(0)\n",
        "  stability_test_pairs = np.random.permutation(np.array([(test_embeddings[i], stability_test_df.stability.values[i], stability_test_df.parent_stability.values[i], stability_test_df.topology_ind.values[i]) for i in range(len(test_embeddings))]))\n",
        "  sub_stability_test_pairs = stability_test_pairs[:num_samples]\n",
        "\n",
        "  sub_stability_test_embeddings = np.array([s[0] for s in sub_stability_test_pairs])\n",
        "  sub_stability_test_stabilities = np.array([s[1] for s in sub_stability_test_pairs])\n",
        "  sub_stability_test_parent_stabilities = np.array([s[2] for s in sub_stability_test_pairs])\n",
        "  sub_stability_test_topologies = np.array([s[3] for s in sub_stability_test_pairs])\n",
        "\n",
        "  stability_test_embeddings_tsne = TSNE(n_components=2).fit_transform(sub_stability_test_embeddings)\n",
        "\n",
        "  X_tsne_test = [s[0] for s in stability_test_embeddings_tsne]\n",
        "  Y_tsne_test = [s[1] for s in stability_test_embeddings_tsne]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_test, Y_tsne_test, c=sub_stability_test_stabilities, s=5, alpha=0.5)\n",
        "  plt.title('t-SNE of Test Stability Embeddings (' + model_name + ')')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('Stability')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_test, Y_tsne_test, c=sub_stability_test_parent_stabilities, s=5, alpha=0.5)\n",
        "  plt.title('t-SNE of Test Stability Embeddings (' + model_name + ')')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('Parent Stability')\n",
        "  plt.show()\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_test, Y_tsne_test, c=sub_stability_test_topologies, s=5, alpha=0.5)\n",
        "  plt.title('t-SNE of Test Stability Embeddings (' + model_name + ')')\n",
        "  topology_names = ['AAA', 'ABBA', 'Other', 'BABB', 'BBABB']\n",
        "  formatter = plt.FuncFormatter(lambda val, loc: topology_names[val])\n",
        "  cbar = plt.colorbar(ticks=range(5), format=formatter)\n",
        "  cbar.set_label('Topology')\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VT9WxnaO1PRj"
      },
      "source": [
        "stability_test_tsne_plot(stability_test_embeddings_random, model_name='Random CNN + MaxPool', num_samples=5000)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bSc96O7Sirxo"
      },
      "source": [
        "stability_test_tsne_plot(stability_test_embeddings, model_name='CNN + MaxPool', num_samples=5000)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SX5N5s2JirvU"
      },
      "source": [
        "stability_test_tsne_plot(stability_test_embeddings_l, model_name='CNN + LinearMaxPool', num_samples=5000)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w5sD3iWpirtF"
      },
      "source": [
        "stability_test_tsne_plot(stability_test_embeddings_e, model_name='Ensemble of CNNs', num_samples=5000)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7sET4OzCBJrN"
      },
      "source": [
        "# Variant prediction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ETGWvVU1_OhH"
      },
      "source": [
        "The data originally comes from [Envision (Gray et. al.)](https://pubmed.ncbi.nlm.nih.gov/29226803/) and was formatted by [FAIR ESM](https://github.com/facebookresearch/esm)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S3Rks14tDwki"
      },
      "source": [
        "## Open data\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n3PDOMMrGYsN"
      },
      "source": [
        "### Install FAIR esm"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YNqH6CY5BNV6"
      },
      "source": [
        "!pip install git+https://github.com/facebookresearch/esm.git\n",
        "!curl -O https://dl.fbaipublicfiles.com/fair-esm/examples/P62593.fasta\n",
        "!pwd\n",
        "!ls"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s5N0NxJXEfj1"
      },
      "source": [
        "import esm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "40duS8d7IY9Z"
      },
      "source": [
        "### Padding and one-hot encoding"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7unmF3cjH8CU"
      },
      "source": [
        "VARIANT_SEQ_LEN = 286"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KOHbFN5RIFJU"
      },
      "source": [
        "VARIANT_AMINO_ACID_VOCABULARY = [\n",
        "    'A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R',\n",
        "    'S', 'T', 'V', 'W', 'Y', '-'\n",
        "]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FaqUdNynIFOw"
      },
      "source": [
        "VARIANT_PROTEIN_DOMAIN = domains.VariableLengthDiscreteDomain(\n",
        "    vocab=domains.ProteinVocab(include_anomalous_amino_acids=False,\n",
        "                               include_eos=True,\n",
        "                               include_pad=True),\n",
        "    length=VARIANT_SEQ_LEN)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-A9VMvz6IT5L"
      },
      "source": [
        "def variant_seq_to_inds(seq):\n",
        "  \"\"\"Encode variant amino acid sequence.\"\"\"\n",
        "\n",
        "  return VARIANT_PROTEIN_DOMAIN.encode([seq])[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FMJwaVULIfIL"
      },
      "source": [
        "### Add one-hots and batch data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rgLhnToDEsLd"
      },
      "source": [
        "def create_variant_df(test=False, train_size=0.8):\n",
        "  \"\"\"Processes variant data into a featurized dataframe.\"\"\"\n",
        "  \n",
        "  FASTA_PATH = \"./P62593.fasta\"\n",
        "\n",
        "  ys = []\n",
        "  Xs = []\n",
        "  for header, _seq in esm.data.read_fasta(FASTA_PATH):\n",
        "    Xs.append(_seq)\n",
        "    scaled_effect = header.split('|')[-1]\n",
        "    ys.append(float(scaled_effect))\n",
        "  \n",
        "  Xs_train, Xs_test, ys_train, ys_test = \\\n",
        "    train_test_split(Xs, ys, train_size=train_size, random_state=42)\n",
        "\n",
        "  if test:\n",
        "    variant_df = pd.DataFrame(np.column_stack([Xs_test, ys_test]),\n",
        "                              columns=['primary', 'activity'])    \n",
        "  else:\n",
        "    variant_df = pd.DataFrame(np.column_stack([Xs_train, ys_train]),\n",
        "                              columns=['primary', 'activity'])\n",
        "\n",
        "  variant_df['activity'] = variant_df.activity.apply(lambda x: float(x))\n",
        "  variant_df['one_hot_inds'] = variant_df.primary.apply(lambda x: variant_seq_to_inds(x[:VARIANT_SEQ_LEN]))\n",
        "\n",
        "  return variant_df\n",
        "\n",
        "\n",
        "def create_variant_batches(batch_size, epochs=1, test=False, buffer_size=None, \n",
        "                           seed=0, drop_remainder=False, train_size=0.8):\n",
        "  \"\"\"Creates iterable object of variant batches.\"\"\"\n",
        "  \n",
        "  if test:\n",
        "    buffer_size = 1\n",
        "\n",
        "  variant_df = create_variant_df(test=test, train_size=train_size)\n",
        "    \n",
        "  activities = variant_df['activity'].values\n",
        "\n",
        "  variant_batches = create_data_iterator(df=variant_df, input_col='one_hot_inds', \n",
        "                                         output_col='activity', \n",
        "                                         batch_size=batch_size, epochs=epochs, \n",
        "                                         buffer_size=buffer_size, seed=seed, \n",
        "                                         drop_remainder=drop_remainder)\n",
        "\n",
        "  return variant_batches, activities"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7RuI0qSWFPNJ"
      },
      "source": [
        "## Model evaluation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nLDVgNrYFRHi"
      },
      "source": [
        "def variant_evaluate(predict_fn, title, batch_size=256, test_data=None, \n",
        "                     pred_activities=None, clip_min=-999999, clip_max=999999,\n",
        "                     train_size=0.8):\n",
        "  \"\"\"Computes variant predictions and measures performance in spearman correlation.\"\"\"\n",
        "  \n",
        "  test_batches, test_activities = create_variant_batches(batch_size=batch_size, \n",
        "                                                         test=True, buffer_size=1,\n",
        "                                                         train_size=train_size)\n",
        "  \n",
        "  if test_data is not None:\n",
        "    test_batches = test_data\n",
        "\n",
        "  if pred_activities is None:\n",
        "    pred_activities = []\n",
        "    for batch in iter(test_batches):\n",
        "      X, Y = batch\n",
        "      preds = predict_fn(X)\n",
        "      for pred in preds:\n",
        "        pred_activities.append(pred[0])\n",
        "  \n",
        "  pred_activities = np.array(pred_activities)\n",
        "  pred_activities = np.clip(pred_activities, clip_min, clip_max)\n",
        "    \n",
        "  spearmanr = scipy.stats.spearmanr(test_activities, pred_activities).correlation\n",
        "  mse = sklearn.metrics.mean_squared_error(test_activities, pred_activities)\n",
        "  plt.scatter(test_activities, pred_activities, s=1, alpha=0.5)\n",
        "  plt.xlabel('True Activity')\n",
        "  plt.ylabel('Predicted Activity')\n",
        "  plt.title(title)\n",
        "  plt.show()\n",
        "\n",
        "  # Source: https://stackoverflow.com/questions/2130016/splitting-a-list-into-n-parts-of-approximately-equal-length\n",
        "  def list_split(a, n):\n",
        "    k, m = divmod(len(a), n)\n",
        "    return [a[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n)]\n",
        "\n",
        "  all_pred_activities = copy.deepcopy(pred_activities)\n",
        "\n",
        "  np.random.seed(0)\n",
        "  test_activities = np.random.permutation(test_activities)\n",
        "  np.random.seed(0)\n",
        "  pred_activities = np.random.permutation(pred_activities)\n",
        "\n",
        "  test_activities = list_split(test_activities, 5)\n",
        "  pred_activities = list_split(pred_activities, 5)\n",
        "\n",
        "  spearmanrs = []\n",
        "  for i in range(5):\n",
        "    spearmanrs.append(scipy.stats.spearmanr(test_activities[i], pred_activities[i]).correlation)\n",
        "  \n",
        "  mean_spearmanr = np.mean(spearmanrs)\n",
        "  std_spearmanr = np.std(spearmanrs)\n",
        "\n",
        "  results = {\n",
        "      'title': title,\n",
        "      'spearmanr': round(spearmanr, 3),\n",
        "      'mse': round(mse, 3),\n",
        "      'mean_spearmanr': round(mean_spearmanr, 3),\n",
        "      'std_spearmanr': round(std_spearmanr, 3)\n",
        "  }\n",
        "\n",
        "  pprint.pprint(results)\n",
        "\n",
        "  return results, all_pred_activities"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6_6ZvKxE9ms"
      },
      "source": [
        "## Experiments"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S9cMj3nzyml4"
      },
      "source": [
        "variant_train_df = create_variant_df()\n",
        "train_activities = variant_train_df['activity']\n",
        "variant_train_one_hot_inds = np.array([x for x in variant_train_df['one_hot_inds'].values])\n",
        "variant_train_one_hots = flattened_one_hot_encoder(variant_train_one_hot_inds, num_categories=len(VARIANT_AMINO_ACID_VOCABULARY))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tnX_eeqHympi"
      },
      "source": [
        "variant_test_df = create_variant_df(test=True)\n",
        "variant_test_one_hot_inds = np.array([x for x in variant_test_df['one_hot_inds'].values])\n",
        "variant_test_one_hots = flattened_one_hot_encoder(variant_test_one_hot_inds, num_categories=len(VARIANT_AMINO_ACID_VOCABULARY))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KrG-n92mE_sn"
      },
      "source": [
        "### Linear regression"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rIAt5L5GE-x8"
      },
      "source": [
        "variant_linear_model = Ridge()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W-HFfn-4ykOl"
      },
      "source": [
        "variant_linear_model.fit(X=variant_train_one_hots, y=train_activities)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x4I0VOqKMWvF"
      },
      "source": [
        "linear_model_pred_activities = variant_linear_model.predict(variant_test_one_hots)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DXMnzxNyIFn_"
      },
      "source": [
        "variant_linear_model_results, linear_model_pred_activities = \\\n",
        "variant_evaluate(predict_fn=None, \n",
        "                 title='Linear Regression',\n",
        "                 pred_activities=linear_model_pred_activities,\n",
        "                 clip_min=min(train_activities),\n",
        "                 clip_max=max(train_activities))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qgv9rUzvNSR5"
      },
      "source": [
        "print('Number of Parameters for Variant Linear Regression: ' + str(len(variant_linear_model.coef_)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NCIQ0kR9ITZv"
      },
      "source": [
        "def fit_variant_linear_model(train_size=0.8):\n",
        "  \"\"\"Fits and evaluates Ridge regression model for given train_size.\"\"\"\n",
        "\n",
        "  variant_train_df = create_variant_df(train_size=train_size)\n",
        "  train_activities = variant_train_df['activity']\n",
        "  variant_train_one_hot_inds = np.array([x for x in variant_train_df['one_hot_inds'].values])\n",
        "  variant_train_one_hots = flattened_one_hot_encoder(variant_train_one_hot_inds, num_categories=len(VARIANT_AMINO_ACID_VOCABULARY))\n",
        "\n",
        "  variant_test_df = create_variant_df(test=True, train_size=train_size)\n",
        "  variant_test_one_hot_inds = np.array([x for x in variant_test_df['one_hot_inds'].values])\n",
        "  variant_test_one_hots = flattened_one_hot_encoder(variant_test_one_hot_inds, num_categories=len(VARIANT_AMINO_ACID_VOCABULARY))\n",
        "\n",
        "  variant_linear_model = Ridge()\n",
        "\n",
        "  variant_linear_model.fit(X=variant_train_one_hots, y=train_activities)\n",
        "\n",
        "  linear_model_pred_activities = variant_linear_model.predict(variant_test_one_hots)\n",
        "\n",
        "  variant_linear_model_results, linear_model_pred_activities = \\\n",
        "  variant_evaluate(predict_fn=None,\n",
        "                   title='Linear Regression (train_size=' + str(train_size) + ')',\n",
        "                   pred_activities=linear_model_pred_activities,\n",
        "                   clip_min=min(train_activities),\n",
        "                   clip_max=max(train_activities),\n",
        "                   train_size=train_size)\n",
        "    \n",
        "  return variant_linear_model_results, linear_model_pred_activities"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9mMhoyC8IThG"
      },
      "source": [
        "for train_size in [0.01, 0.1, 0.3, 0.5, 0.8]:\n",
        "  variant_linear_model_results, linear_model_pred_activities = \\\n",
        "  fit_variant_linear_model(train_size)\n",
        "  print()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qmeob2F5O-5E"
      },
      "source": [
        "### CNN + MaxPool"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bsf8S8seSnw9"
      },
      "source": [
        "epochs = 500\n",
        "variant_train_batches, train_activities = create_variant_batches(batch_size=32, epochs=epochs)\n",
        "\n",
        "layers = ['CNN_0', 'Dense_1']                                 \n",
        "learning_rate = [1e-3, 1e-3]\n",
        "weight_decay = [0.1, 0.05]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 2,\n",
        "    'n_features': [1024, 512],\n",
        "    'n_kernel_sizes': [9, 7],\n",
        "    'n_kernel_dilations': None\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "\n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "variant_model = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                            encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                            reduce_fn=reduce_fn,\n",
        "                                            reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                            num_categories=len(VARIANT_AMINO_ACID_VOCABULARY),\n",
        "                                            output_features=1)\n",
        "\n",
        "variant_optimizer = train(model=variant_model,\n",
        "                          train_data=variant_train_batches, \n",
        "                          loss_fn=mse_loss,\n",
        "                          loss_fn_kwargs=loss_fn_kwargs,\n",
        "                          learning_rate=learning_rate, \n",
        "                          weight_decay=weight_decay,\n",
        "                          layers=layers)\n",
        "\n",
        "variant_results, pred_activities = variant_evaluate(predict_fn=variant_optimizer.target,\n",
        "                                                    title='CNN + MaxPool',\n",
        "                                                    batch_size=256,\n",
        "                                                    clip_min=min(train_activities),\n",
        "                                                    clip_max=max(train_activities))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qdBGiaXdbGBD"
      },
      "source": [
        "print('Number of Parameters for Variant CNN + MaxPool: ' + str(get_num_params(variant_optimizer.target)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tYRNi9X9Usit"
      },
      "source": [
        "epochs = 500\n",
        "variant_train_batches, train_activities = create_variant_batches(batch_size=32, epochs=epochs)\n",
        "\n",
        "layers = ['CNN_0', 'Dense_1']                                 \n",
        "learning_rate = [1e-3, 1e-3]\n",
        "weight_decay = [0.1, 0.05]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 2,\n",
        "    'n_features': [1024, 512],\n",
        "    'n_kernel_sizes': [9, 7],\n",
        "    'n_kernel_dilations': [2, 1]\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "\n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "variant_model_d = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                              encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                              reduce_fn=reduce_fn,\n",
        "                                              reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                              num_categories=len(VARIANT_AMINO_ACID_VOCABULARY),\n",
        "                                              output_features=1)\n",
        "\n",
        "variant_optimizer_d = train(model=variant_model_d,\n",
        "                            train_data=variant_train_batches, \n",
        "                            loss_fn=mse_loss,\n",
        "                            loss_fn_kwargs=loss_fn_kwargs,\n",
        "                            learning_rate=learning_rate, \n",
        "                            weight_decay=weight_decay,\n",
        "                            layers=layers)\n",
        "\n",
        "variant_results_d, pred_activities_d = variant_evaluate(predict_fn=variant_optimizer_d.target,\n",
        "                                                        title='Dilated CNN + MaxPool',\n",
        "                                                        batch_size=256,\n",
        "                                                        clip_min=min(train_activities),\n",
        "                                                        clip_max=max(train_activities))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "podjKh0ebCjE"
      },
      "source": [
        "print('Number of Parameters for Variant Dilated CNN + MaxPool: ' + str(get_num_params(variant_optimizer_d.target)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BdjNebjPkGJh"
      },
      "source": [
        "### Ensemble of CNNs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RLrEuiv8agYc"
      },
      "source": [
        "ensemble_pred_activities = [(pred_activities[i]+pred_activities_d[i])/2 for i in range(len(pred_activities))]\n",
        "ensemble_results, ensemble_pred_activities = variant_evaluate(predict_fn=None, \n",
        "                                                              title='Ensemble of CNNs',\n",
        "                                                              pred_activities=ensemble_pred_activities,\n",
        "                                                              clip_min=min(train_activities),\n",
        "                                                              clip_max=max(train_activities))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S0PHVIO57-3d"
      },
      "source": [
        "print('Number of Parameters for Variant Ensemble of CNNs: ' + str(get_num_params(variant_optimizer.target) + get_num_params(variant_optimizer_d.target)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HeYhSHRuj9Ie"
      },
      "source": [
        "def fit_variant_cnn_models(train_size=0.8):\n",
        "  \"\"\"Fits and evaluates CNN models for given train_size.\"\"\"\n",
        "  \n",
        "  # CNN + MaxPool\n",
        "  epochs = 500\n",
        "  variant_train_batches, train_activities = create_variant_batches(batch_size=32, epochs=epochs, train_size=train_size)\n",
        "\n",
        "  layers = ['CNN_0', 'Dense_1']                                 \n",
        "  learning_rate = [1e-3, 1e-3]\n",
        "  weight_decay = [0.1, 0.05]\n",
        "\n",
        "  encoder_fn = cnn_one_hot_encoder\n",
        "  encoder_fn_kwargs = {\n",
        "      'n_layers': 2,\n",
        "      'n_features': [1024, 512],\n",
        "      'n_kernel_sizes': [9, 7],\n",
        "      'n_kernel_dilations': None\n",
        "  }\n",
        "  reduce_fn = max_pool\n",
        "  reduce_fn_kwargs = {\n",
        "\n",
        "  }\n",
        "  loss_fn_kwargs = {\n",
        "      \n",
        "  }\n",
        "\n",
        "  variant_model = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                              encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                              reduce_fn=reduce_fn,\n",
        "                                              reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                              num_categories=len(VARIANT_AMINO_ACID_VOCABULARY),\n",
        "                                              output_features=1)\n",
        "\n",
        "  variant_optimizer = train(model=variant_model,\n",
        "                            train_data=variant_train_batches, \n",
        "                            loss_fn=mse_loss,\n",
        "                            loss_fn_kwargs=loss_fn_kwargs,\n",
        "                            learning_rate=learning_rate, \n",
        "                            weight_decay=weight_decay,\n",
        "                            layers=layers)\n",
        "\n",
        "  variant_results, pred_activities = variant_evaluate(predict_fn=variant_optimizer.target,\n",
        "                                                      title='CNN + MaxPool (train_size=' + str(train_size) + ')',\n",
        "                                                      batch_size=256,\n",
        "                                                      clip_min=min(train_activities),\n",
        "                                                      clip_max=max(train_activities),\n",
        "                                                      train_size=train_size)\n",
        "\n",
        "  # Dilated CNN + MaxPool\n",
        "  epochs = 500\n",
        "  variant_train_batches, train_activities = create_variant_batches(batch_size=32, epochs=epochs, train_size=train_size)\n",
        "\n",
        "  layers = ['CNN_0', 'Dense_1']                                 \n",
        "  learning_rate = [1e-3, 1e-3]\n",
        "  weight_decay = [0.1, 0.05]\n",
        "\n",
        "  encoder_fn = cnn_one_hot_encoder\n",
        "  encoder_fn_kwargs = {\n",
        "      'n_layers': 2,\n",
        "      'n_features': [1024, 512],\n",
        "      'n_kernel_sizes': [9, 7],\n",
        "      'n_kernel_dilations': [2, 1]\n",
        "  }\n",
        "  reduce_fn = max_pool\n",
        "  reduce_fn_kwargs = {\n",
        "\n",
        "  }\n",
        "  loss_fn_kwargs = {\n",
        "      \n",
        "  }\n",
        "\n",
        "  variant_model_d = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                                encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                                reduce_fn=reduce_fn,\n",
        "                                                reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                                num_categories=len(VARIANT_AMINO_ACID_VOCABULARY),\n",
        "                                                output_features=1)\n",
        "\n",
        "  variant_optimizer_d = train(model=variant_model_d,\n",
        "                              train_data=variant_train_batches, \n",
        "                              loss_fn=mse_loss,\n",
        "                              loss_fn_kwargs=loss_fn_kwargs,\n",
        "                              learning_rate=learning_rate, \n",
        "                              weight_decay=weight_decay,\n",
        "                              layers=layers)\n",
        "\n",
        "  variant_results_d, pred_activities_d = variant_evaluate(predict_fn=variant_optimizer_d.target,\n",
        "                                                          title='Dilated CNN + MaxPool (train_size=' + str(train_size) + ')',\n",
        "                                                          batch_size=256,\n",
        "                                                          clip_min=min(train_activities),\n",
        "                                                          clip_max=max(train_activities),\n",
        "                                                          train_size=train_size)\n",
        "  \n",
        "  # Ensemble of CNNs\n",
        "  ensemble_pred_activities = [(pred_activities[i]+pred_activities_d[i])/2 for i in range(len(pred_activities))]\n",
        "  ensemble_results, ensemble_pred_activities = variant_evaluate(predict_fn=None, \n",
        "                                                                title='Ensemble of CNNs (train_size=' + str(train_size) + ')',\n",
        "                                                                pred_activities=ensemble_pred_activities,\n",
        "                                                                clip_min=min(train_activities),\n",
        "                                                                clip_max=max(train_activities),\n",
        "                                                                train_size=train_size)\n",
        "  \n",
        "  return variant_results, pred_activities, variant_results_d, pred_activities_d, ensemble_results, ensemble_pred_activities"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IzvCheljjHZs"
      },
      "source": [
        "for train_size in [0.01, 0.1, 0.3, 0.5, 0.8]:\n",
        "  fit_variant_cnn_models(train_size)\n",
        "  print()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4yoeGQ3O-3lP"
      },
      "source": [
        "## Visualization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kRfNwDdo-7do"
      },
      "source": [
        "def variant_embeddings(encoder_fn, encoder_fn_kwargs, reduce_fn, reduce_fn_kwargs, optimizer, train_size=0.8):\n",
        "  \"\"\"Computes variant embeddings from given optimizer.\"\"\"\n",
        "\n",
        "  variant_encoding_model = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                                       encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                                       reduce_fn=reduce_fn,\n",
        "                                                       reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                                       num_categories=len(VARIANT_AMINO_ACID_VOCABULARY),\n",
        "                                                       output='embedding',\n",
        "                                                       output_features=1)\n",
        "  \n",
        "  trained_params = copy.deepcopy(optimizer.target.params)\n",
        "\n",
        "  variant_encoding_optimizer = create_optimizer(variant_encoding_model,\n",
        "                                                learning_rate=learning_rate,\n",
        "                                                weight_decay=weight_decay,\n",
        "                                                layers=layers)\n",
        "  \n",
        "  for layer in variant_encoding_optimizer.target.params.keys():\n",
        "    variant_encoding_optimizer.target.params[layer] = trained_params[layer]\n",
        "\n",
        "  variant_train_batches, train_variant= create_variant_batches(batch_size=256, buffer_size=1, train_size=train_size)\n",
        "  variant_train_embeddings = compute_embeddings(variant_encoding_optimizer.target, variant_train_batches)\n",
        "  \n",
        "  variant_test_batches, test_variant = create_variant_batches(batch_size=256, test=True, train_size=train_size)\n",
        "  variant_test_embeddings = compute_embeddings(variant_encoding_optimizer.target, variant_test_batches)\n",
        "\n",
        "  return variant_train_embeddings, variant_test_embeddings"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CrrBpBORYaKl"
      },
      "source": [
        "# CNN + MaxPool Optimizer Embeddings\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 2,\n",
        "    'n_features': [1024, 512],\n",
        "    'n_kernel_sizes': [9, 7],\n",
        "    'n_kernel_dilations': [1, 1]\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "\n",
        "}\n",
        "variant_train_embeddings, variant_test_embeddings = variant_embeddings(encoder_fn,\n",
        "                                                                       encoder_fn_kwargs,\n",
        "                                                                       reduce_fn,\n",
        "                                                                       reduce_fn_kwargs,\n",
        "                                                                       variant_optimizer)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MyewfqIPYaPB"
      },
      "source": [
        "# Dilated CNN + MaxPool Optimizer Embeddings\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 2,\n",
        "    'n_features': [1024, 512],\n",
        "    'n_kernel_sizes': [9, 7],\n",
        "    'n_kernel_dilations': [2, 1]\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "\n",
        "}\n",
        "variant_train_embeddings_d, variant_test_embeddings_d = variant_embeddings(encoder_fn,\n",
        "                                                                           encoder_fn_kwargs,\n",
        "                                                                           reduce_fn,\n",
        "                                                                           reduce_fn_kwargs,\n",
        "                                                                           variant_optimizer_d)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sZRFaSWIYu1K"
      },
      "source": [
        "# Ensemble Embeddings (Concatention of Above Embeddings)\n",
        "variant_train_embeddings_e = np.concatenate((variant_train_embeddings, variant_train_embeddings_d), axis=1)\n",
        "variant_test_embeddings_e = np.concatenate((variant_test_embeddings, variant_test_embeddings_d), axis=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tldb9jDCTxyH"
      },
      "source": [
        "# Random CNN + MaxPool Optimizer Embeddings\n",
        "layers = ['CNN_0', 'Dense_1']                                 \n",
        "learning_rate = [1e-3, 1e-3]\n",
        "weight_decay = [0.1, 0.05]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 2,\n",
        "    'n_features': [1024, 512],\n",
        "    'n_kernel_sizes': [9, 7],\n",
        "    'n_kernel_dilations': [1, 1]\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "variant_model = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                            encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                            reduce_fn=reduce_fn,\n",
        "                                            reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                            num_categories=len(VARIANT_AMINO_ACID_VOCABULARY),\n",
        "                                            output_features=1)\n",
        "\n",
        "variant_optimizer_random = create_optimizer(variant_model, learning_rate, weight_decay, layers)\n",
        "\n",
        "variant_train_embeddings_random, variant_test_embeddings_random = variant_embeddings(encoder_fn,\n",
        "                                                                                     encoder_fn_kwargs,\n",
        "                                                                                     reduce_fn,\n",
        "                                                                                     reduce_fn_kwargs,\n",
        "                                                                                     variant_optimizer_random)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XiSsXufSTx1N"
      },
      "source": [
        "# Random Dilated CNN + MaxPool Optimizer Embeddings\n",
        "layers = ['CNN_0', 'Dense_1']                                 \n",
        "learning_rate = [1e-3, 1e-3]\n",
        "weight_decay = [0.1, 0.05]\n",
        "\n",
        "encoder_fn = cnn_one_hot_encoder\n",
        "encoder_fn_kwargs = {\n",
        "    'n_layers': 2,\n",
        "    'n_features': [1024, 512],\n",
        "    'n_kernel_sizes': [9, 7],\n",
        "    'n_kernel_dilations': [2, 1]\n",
        "}\n",
        "reduce_fn = max_pool\n",
        "reduce_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "loss_fn_kwargs = {\n",
        "    \n",
        "}\n",
        "\n",
        "variant_model_d = create_representation_model(encoder_fn=encoder_fn,\n",
        "                                              encoder_fn_kwargs=encoder_fn_kwargs,\n",
        "                                              reduce_fn=reduce_fn,\n",
        "                                              reduce_fn_kwargs=reduce_fn_kwargs,\n",
        "                                              num_categories=len(VARIANT_AMINO_ACID_VOCABULARY),\n",
        "                                              output_features=1)\n",
        "\n",
        "variant_optimizer_d_random = create_optimizer(variant_model_d, learning_rate, weight_decay, layers)\n",
        "\n",
        "variant_train_embeddings_d_random, variant_test_embeddings_d_random = variant_embeddings(encoder_fn,\n",
        "                                                                                         encoder_fn_kwargs,\n",
        "                                                                                         reduce_fn,\n",
        "                                                                                         reduce_fn_kwargs,\n",
        "                                                                                         variant_optimizer_d_random)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e4kZrpf0Tx4D"
      },
      "source": [
        "# Random Ensemble Embeddings (Concatention of Above Embeddings)\n",
        "variant_train_embeddings_e_random = np.concatenate((variant_train_embeddings_random, variant_train_embeddings_d_random), axis=1)\n",
        "variant_test_embeddings_e_random = np.concatenate((variant_test_embeddings_random, variant_test_embeddings_d_random), axis=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_CAGO32OdH1e"
      },
      "source": [
        "### Train PCA"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ldz8J4qbTx-F"
      },
      "source": [
        "def variant_train_pca_plot(train_embeddings, model_name, train_size=0.8):\n",
        "  \"\"\"Applies and plots PCA on variant train embeddings.\"\"\"\n",
        "\n",
        "  variant_train_df = create_variant_df(train_size=train_size)\n",
        "\n",
        "  pca = PCA(n_components=2).fit(train_embeddings)\n",
        "  variant_train_embeddings_pca = pca.transform(train_embeddings)\n",
        "  \n",
        "  X_train = [v[0] for v in variant_train_embeddings_pca]\n",
        "  Y_train = [v[1] for v in variant_train_embeddings_pca]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_train, Y_train, c=variant_train_df.activity.values, s=10, alpha=0.5)\n",
        "  plt.title('PCA of Train Variant Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  plt.colorbar().set_label('Activity')\n",
        "  plt.show()\n",
        "\n",
        "  return pca"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O16mG1lURKfd"
      },
      "source": [
        "variant_pca_random = variant_train_pca_plot(variant_train_embeddings_e_random, model_name='Random Ensemble of CNNs')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HAuJUMbwd1NM"
      },
      "source": [
        "variant_pca = variant_train_pca_plot(variant_train_embeddings_e, model_name='Ensemble of CNNs')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9xMRjtp2fBDz"
      },
      "source": [
        "### Test PCA"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wAOQV693eG8t"
      },
      "source": [
        "def variant_test_pca_plot(test_embeddings, model_name, variant_pca, train_size=0.8):\n",
        "  \"\"\"Applies and plots PCA on variant train embeddings.\"\"\"\n",
        "\n",
        "  variant_test_df = create_variant_df(test=True, train_size=train_size)\n",
        "\n",
        "  variant_test_embeddings_pca = variant_pca.transform(test_embeddings)\n",
        "  \n",
        "  X_test = [v[0] for v in variant_test_embeddings_pca]\n",
        "  Y_test = [v[1] for v in variant_test_embeddings_pca]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_test, Y_test, c=variant_test_df.activity.values, s=10, alpha=0.5)\n",
        "  plt.title('PCA of Test Variant Embeddings (' + model_name + ')')\n",
        "  plt.xlabel('Principal Component 1')\n",
        "  plt.ylabel('Principal Component 2')\n",
        "  plt.colorbar().set_label('Activity')\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1lCtLh3NfUpJ"
      },
      "source": [
        "variant_test_pca_plot(variant_test_embeddings_e_random, model_name='Random Ensemble of CNNs', variant_pca=variant_pca_random)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f-qsXfGaRHdS"
      },
      "source": [
        "variant_test_pca_plot(variant_test_embeddings_e, model_name='Ensemble of CNNs', variant_pca=variant_pca)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N4toUE5VfiIy"
      },
      "source": [
        "### Train t-SNE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8oPbLEhbfYqS"
      },
      "source": [
        "def variant_train_tsne_plot(train_embeddings, model_name, train_size=0.8):\n",
        "  \"\"\"Applies and plots t-SNE on GFP train embeddings.\"\"\"\n",
        "\n",
        "  variant_train_df = create_variant_df(train_size=train_size)\n",
        "\n",
        "  variant_train_embeddings_tsne = TSNE(n_components=2).fit_transform(train_embeddings)\n",
        "\n",
        "  X_tsne_train = [v[0] for v in variant_train_embeddings_tsne]\n",
        "  Y_tsne_train = [v[1] for v in variant_train_embeddings_tsne]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_train, Y_tsne_train, c=variant_train_df.activity.values, s=10, alpha=0.5)\n",
        "  plt.title('t-SNE of Train GFP Embeddings (' + model_name + ')')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('Activity')\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CGS1PdpWgdRU"
      },
      "source": [
        "variant_train_tsne_plot(variant_train_embeddings_e, model_name='Ensemble of CNNs')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Yh7bQhnIgTLR"
      },
      "source": [
        "variant_train_tsne_plot(variant_train_embeddings_e_random, model_name='Random Ensemble of CNNs')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZTHUdtlogmOm"
      },
      "source": [
        "### Test t-SNE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BnZKAHJxgXzM"
      },
      "source": [
        "def variant_test_tsne_plot(test_embeddings, model_name, train_size=0.8):\n",
        "  \"\"\"Applies and plots t-SNE on GFP train embeddings.\"\"\"\n",
        "\n",
        "  variant_test_df = create_variant_df(test=True, train_size=train_size)\n",
        "\n",
        "  variant_test_embeddings_tsne = TSNE(n_components=2).fit_transform(test_embeddings)\n",
        "\n",
        "  X_tsne_test = [v[0] for v in variant_test_embeddings_tsne]\n",
        "  Y_tsne_test = [v[1] for v in variant_test_embeddings_tsne]\n",
        "\n",
        "  plt.figure(figsize=(8, 6))\n",
        "  plt.scatter(X_tsne_test, Y_tsne_test, c=variant_test_df.activity.values, s=10, alpha=0.5)\n",
        "  plt.title('t-SNE of Test GFP Embeddings (' + model_name + ')')\n",
        "  cbar = plt.colorbar()\n",
        "  cbar.set_label('Activity')\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8QYo8fGfg4Am"
      },
      "source": [
        "variant_test_tsne_plot(variant_test_embeddings_e, model_name='Ensemble of CNNs')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o-VPPO2Kg4Tj"
      },
      "source": [
        "variant_test_tsne_plot(variant_test_embeddings_e_random, model_name='Random Ensemble of CNNs')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}