{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dca40d3f",
   "metadata": {},
   "source": [
    "# gb1 dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "43acb29b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-11-11T22:23:27.743043Z",
     "start_time": "2021-11-11T22:23:25.875662Z"
    }
   },
   "outputs": [],
   "source": [
    "# Standard imports\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Special imports\n",
    "import mavenn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5336886a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-11-11T17:26:47.608641Z",
     "start_time": "2021-11-11T17:26:47.392567Z"
    }
   },
   "source": [
    "## Summary\n",
    "\n",
    "The DMS dataset from Olson et al., 2014. The authors used an RNA display selection experiment to assay the binding of over half a million protein GB1 variants to IgG. These variants included all 1-point and 2-point mutations within the 55 residue GB1 sequence. Only the 2-point variants are included in this dataset.\n",
    "\n",
    "**Name:** ``'gb1'``\n",
    "\n",
    "**Reference**: Olson C, Wu N, Sun R. A comprehensive biophysical description of pairwise epistasis throughout an entire protein domain. [Curr Biol 24(22):2643-2651 (2014).](https://pubmed.ncbi.nlm.nih.gov/25455030/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3072cf25",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-11-11T22:23:28.280522Z",
     "start_time": "2021-11-11T22:23:27.744157Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>set</th>\n",
       "      <th>dist</th>\n",
       "      <th>input_ct</th>\n",
       "      <th>selected_ct</th>\n",
       "      <th>y</th>\n",
       "      <th>x</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>173</td>\n",
       "      <td>33</td>\n",
       "      <td>-3.145154</td>\n",
       "      <td>AAKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>18</td>\n",
       "      <td>8</td>\n",
       "      <td>-1.867676</td>\n",
       "      <td>ACKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>2</td>\n",
       "      <td>-5.270800</td>\n",
       "      <td>ADKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>72</td>\n",
       "      <td>1</td>\n",
       "      <td>-5.979498</td>\n",
       "      <td>AEKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>69</td>\n",
       "      <td>168</td>\n",
       "      <td>0.481923</td>\n",
       "      <td>AFKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>530732</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>462</td>\n",
       "      <td>139</td>\n",
       "      <td>-2.515259</td>\n",
       "      <td>QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>530733</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>317</td>\n",
       "      <td>84</td>\n",
       "      <td>-2.693165</td>\n",
       "      <td>QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>530734</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>335</td>\n",
       "      <td>77</td>\n",
       "      <td>-2.896589</td>\n",
       "      <td>QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>530735</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>148</td>\n",
       "      <td>28</td>\n",
       "      <td>-3.150861</td>\n",
       "      <td>QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>530736</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>95</td>\n",
       "      <td>16</td>\n",
       "      <td>-3.287173</td>\n",
       "      <td>QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>530737 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "             set  dist  input_ct  selected_ct         y  \\\n",
       "0       training     2       173           33 -3.145154   \n",
       "1       training     2        18            8 -1.867676   \n",
       "2       training     2        66            2 -5.270800   \n",
       "3       training     2        72            1 -5.979498   \n",
       "4       training     2        69          168  0.481923   \n",
       "...          ...   ...       ...          ...       ...   \n",
       "530732  training     2       462          139 -2.515259   \n",
       "530733  training     2       317           84 -2.693165   \n",
       "530734  training     2       335           77 -2.896589   \n",
       "530735  training     2       148           28 -3.150861   \n",
       "530736  training     2        95           16 -3.287173   \n",
       "\n",
       "                                                        x  \n",
       "0       AAKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "1       ACKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "2       ADKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "3       AEKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "4       AFKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "...                                                   ...  \n",
       "530732  QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "530733  QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "530734  QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "530735  QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "530736  QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "\n",
       "[530737 rows x 6 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mavenn.load_example_dataset('gb1')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fe018cb",
   "metadata": {},
   "source": [
    "## Preprocessing\n",
    "\n",
    "First we load the double-mutation dataset published by Olson et al. (2021).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ecccf7df",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-11-11T17:34:28.129865Z",
     "start_time": "2021-11-11T17:34:28.124430Z"
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Mut1 WT amino acid</th>\n",
       "      <th>Mut1 Position</th>\n",
       "      <th>Mut1 Mutation</th>\n",
       "      <th>Mut2 WT amino acid</th>\n",
       "      <th>Mut2 Position</th>\n",
       "      <th>Mut2 Mutation</th>\n",
       "      <th>Input Count</th>\n",
       "      <th>Selection Count</th>\n",
       "      <th>Mut1 Fitness</th>\n",
       "      <th>Mut2 Fitness</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Q</td>\n",
       "      <td>2</td>\n",
       "      <td>A</td>\n",
       "      <td>Y</td>\n",
       "      <td>3</td>\n",
       "      <td>A</td>\n",
       "      <td>173</td>\n",
       "      <td>33</td>\n",
       "      <td>1.518</td>\n",
       "      <td>0.579</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Q</td>\n",
       "      <td>2</td>\n",
       "      <td>A</td>\n",
       "      <td>Y</td>\n",
       "      <td>3</td>\n",
       "      <td>C</td>\n",
       "      <td>18</td>\n",
       "      <td>8</td>\n",
       "      <td>1.518</td>\n",
       "      <td>0.616</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Q</td>\n",
       "      <td>2</td>\n",
       "      <td>A</td>\n",
       "      <td>Y</td>\n",
       "      <td>3</td>\n",
       "      <td>D</td>\n",
       "      <td>66</td>\n",
       "      <td>2</td>\n",
       "      <td>1.518</td>\n",
       "      <td>0.010</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Q</td>\n",
       "      <td>2</td>\n",
       "      <td>A</td>\n",
       "      <td>Y</td>\n",
       "      <td>3</td>\n",
       "      <td>E</td>\n",
       "      <td>72</td>\n",
       "      <td>1</td>\n",
       "      <td>1.518</td>\n",
       "      <td>0.009</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Q</td>\n",
       "      <td>2</td>\n",
       "      <td>A</td>\n",
       "      <td>Y</td>\n",
       "      <td>3</td>\n",
       "      <td>F</td>\n",
       "      <td>69</td>\n",
       "      <td>168</td>\n",
       "      <td>1.518</td>\n",
       "      <td>1.054</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>535912</th>\n",
       "      <td>E</td>\n",
       "      <td>56</td>\n",
       "      <td>Y</td>\n",
       "      <td>T</td>\n",
       "      <td>55</td>\n",
       "      <td>R</td>\n",
       "      <td>462</td>\n",
       "      <td>139</td>\n",
       "      <td>0.190</td>\n",
       "      <td>0.941</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>535913</th>\n",
       "      <td>E</td>\n",
       "      <td>56</td>\n",
       "      <td>Y</td>\n",
       "      <td>T</td>\n",
       "      <td>55</td>\n",
       "      <td>S</td>\n",
       "      <td>317</td>\n",
       "      <td>84</td>\n",
       "      <td>0.190</td>\n",
       "      <td>0.840</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>535914</th>\n",
       "      <td>E</td>\n",
       "      <td>56</td>\n",
       "      <td>Y</td>\n",
       "      <td>T</td>\n",
       "      <td>55</td>\n",
       "      <td>V</td>\n",
       "      <td>335</td>\n",
       "      <td>77</td>\n",
       "      <td>0.190</td>\n",
       "      <td>0.669</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>535915</th>\n",
       "      <td>E</td>\n",
       "      <td>56</td>\n",
       "      <td>Y</td>\n",
       "      <td>T</td>\n",
       "      <td>55</td>\n",
       "      <td>W</td>\n",
       "      <td>148</td>\n",
       "      <td>28</td>\n",
       "      <td>0.190</td>\n",
       "      <td>0.798</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>535916</th>\n",
       "      <td>E</td>\n",
       "      <td>56</td>\n",
       "      <td>Y</td>\n",
       "      <td>T</td>\n",
       "      <td>55</td>\n",
       "      <td>Y</td>\n",
       "      <td>95</td>\n",
       "      <td>16</td>\n",
       "      <td>0.190</td>\n",
       "      <td>0.663</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>535917 rows × 10 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       Mut1 WT amino acid  Mut1 Position Mut1 Mutation Mut2 WT amino acid  \\\n",
       "0                       Q              2             A                  Y   \n",
       "1                       Q              2             A                  Y   \n",
       "2                       Q              2             A                  Y   \n",
       "3                       Q              2             A                  Y   \n",
       "4                       Q              2             A                  Y   \n",
       "...                   ...            ...           ...                ...   \n",
       "535912                  E             56             Y                  T   \n",
       "535913                  E             56             Y                  T   \n",
       "535914                  E             56             Y                  T   \n",
       "535915                  E             56             Y                  T   \n",
       "535916                  E             56             Y                  T   \n",
       "\n",
       "        Mut2 Position Mut2 Mutation  Input Count  Selection Count  \\\n",
       "0                   3             A          173               33   \n",
       "1                   3             C           18                8   \n",
       "2                   3             D           66                2   \n",
       "3                   3             E           72                1   \n",
       "4                   3             F           69              168   \n",
       "...               ...           ...          ...              ...   \n",
       "535912             55             R          462              139   \n",
       "535913             55             S          317               84   \n",
       "535914             55             V          335               77   \n",
       "535915             55             W          148               28   \n",
       "535916             55             Y           95               16   \n",
       "\n",
       "        Mut1 Fitness  Mut2 Fitness  \n",
       "0              1.518         0.579  \n",
       "1              1.518         0.616  \n",
       "2              1.518         0.010  \n",
       "3              1.518         0.009  \n",
       "4              1.518         1.054  \n",
       "...              ...           ...  \n",
       "535912         0.190         0.941  \n",
       "535913         0.190         0.840  \n",
       "535914         0.190         0.669  \n",
       "535915         0.190         0.798  \n",
       "535916         0.190         0.663  \n",
       "\n",
       "[535917 rows x 10 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Dataset is is at this URL:\n",
    "# url = 'https://ars.els-cdn.com/content/image/1-s2.0-S0960982214012688-mmc2.xlsx'\n",
    "\n",
    "# We have downloaded this Excel file and reformatted it into a more parseable form\n",
    "raw_data_file = '../../mavenn/examples/datasets/raw/gb1_raw.xlsx'\n",
    "\n",
    "# Load data (takes a while)\n",
    "double_mut_df = pd.read_excel(raw_data_file, sheet_name='double_mutants')\n",
    "double_mut_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e865c5e0",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next we reconstruct the wild-type GB1 sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c7dc8e80",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WT sequence: QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE\n"
     ]
    }
   ],
   "source": [
    "# Get unique WT pos-aa associations, sorted by position\n",
    "wt_1_df = double_mut_df[['Mut1 Position', 'Mut1 WT amino acid']].copy()\n",
    "wt_1_df.columns = ['pos','aa']\n",
    "wt_2_df = double_mut_df[['Mut2 Position', 'Mut2 WT amino acid']].copy()\n",
    "wt_2_df.columns = ['pos','aa']\n",
    "wt_seq_df = pd.concat([wt_1_df, wt_2_df], axis=0).drop_duplicates().sort_values(by='pos').reset_index(drop=True)\n",
    "\n",
    "# Confirm that each position occurs at most once\n",
    "assert np.all(wt_seq_df['pos'].value_counts()==1)\n",
    "\n",
    "# Confirm that the set of unique positions is correct\n",
    "L = len(wt_seq_df)\n",
    "assert set(wt_seq_df['pos'].values) == set(range(2,L+2))\n",
    "\n",
    "# Compute wt_seq and confirm its identity\n",
    "wt_seq = ''.join(wt_seq_df['aa'])\n",
    "known_wt_seq = 'QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE'\n",
    "assert wt_seq == known_wt_seq\n",
    "\n",
    "# Print final wt sequence\n",
    "print(f'WT sequence: {wt_seq}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adac1d2d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-11-11T22:23:57.928708Z",
     "start_time": "2021-11-11T22:23:57.115529Z"
    },
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next we convert the list of mutations to an array `x` of variant sequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7fc92598",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['AAKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE',\n",
       "       'ACKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE',\n",
       "       'ADKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE', ...,\n",
       "       'QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVVY',\n",
       "       'QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVWY',\n",
       "       'QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVYY'],\n",
       "      dtype='<U55')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Introduce double mutations into WT sequence and append to a growing list\n",
    "pos1s = double_mut_df['Mut1 Position'].values-2\n",
    "pos2s = double_mut_df['Mut2 Position'].values-2\n",
    "aa1s = double_mut_df['Mut1 Mutation'].values\n",
    "aa2s = double_mut_df['Mut2 Mutation'].values\n",
    "x_list = []\n",
    "for pos1, aa1, pos2, aa2 in zip(pos1s, aa1s, pos2s, aa2s):\n",
    "    mut_seq_list = list(wt_seq)\n",
    "    mut_seq_list[pos1] = aa1\n",
    "    mut_seq_list[pos2] = aa2\n",
    "    mut_seq = ''.join(mut_seq_list)\n",
    "    x_list.append(mut_seq)\n",
    "\n",
    "# Convert list to an np.array and preview it\n",
    "x = np.array(x_list)\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5d1622a",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We then compute the log2 enrichment `y` for each variant sequence in `x`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a8763c9e",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-3.14515399, -1.86767585, -5.27080003, ..., -2.89658854,\n",
       "       -3.15086086, -3.287173  ])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Extract input and output counts\n",
    "in_ct = double_mut_df['Input Count'].values\n",
    "out_ct = double_mut_df['Selection Count'].values\n",
    "\n",
    "# Read in wt data (on separate sheet) and compute baseline for y\n",
    "wt_df = pd.read_excel(raw_data_file, sheet_name='wild_type')\n",
    "wt_in_ct = wt_df['Input Count'][0]\n",
    "wt_out_ct = wt_df['Selection Count'][0]\n",
    "\n",
    "# Compute log2 enrichment values relative to WT\n",
    "y = np.log2((out_ct+1)/(in_ct+1)) - np.log2((wt_out_ct+1)/(wt_in_ct+1))\n",
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee76586e",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next we assign each sequence to the training, validation, or test set\n",
    "using a split of 90%:5%:5%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "09d2fef3",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['training', 'training', 'training', ..., 'training', 'training',\n",
       "       'training'], dtype='<U10')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Assign to training, validation, or test set\n",
    "np.random.seed(0)\n",
    "sets = np.random.choice(a=['training', 'validation', 'test'],\n",
    "                        p=[.90,.05,.05],\n",
    "                        size=len(x))\n",
    "sets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "191f00c6",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Finally we assemble all relevant information into a dataframe and save."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d18625d5",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>set</th>\n",
       "      <th>dist</th>\n",
       "      <th>input_ct</th>\n",
       "      <th>selected_ct</th>\n",
       "      <th>y</th>\n",
       "      <th>x</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>173</td>\n",
       "      <td>33</td>\n",
       "      <td>-3.145154</td>\n",
       "      <td>AAKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>18</td>\n",
       "      <td>8</td>\n",
       "      <td>-1.867676</td>\n",
       "      <td>ACKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>66</td>\n",
       "      <td>2</td>\n",
       "      <td>-5.270800</td>\n",
       "      <td>ADKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>72</td>\n",
       "      <td>1</td>\n",
       "      <td>-5.979498</td>\n",
       "      <td>AEKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>69</td>\n",
       "      <td>168</td>\n",
       "      <td>0.481923</td>\n",
       "      <td>AFKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>530732</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>462</td>\n",
       "      <td>139</td>\n",
       "      <td>-2.515259</td>\n",
       "      <td>QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>530733</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>317</td>\n",
       "      <td>84</td>\n",
       "      <td>-2.693165</td>\n",
       "      <td>QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>530734</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>335</td>\n",
       "      <td>77</td>\n",
       "      <td>-2.896589</td>\n",
       "      <td>QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>530735</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>148</td>\n",
       "      <td>28</td>\n",
       "      <td>-3.150861</td>\n",
       "      <td>QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>530736</th>\n",
       "      <td>training</td>\n",
       "      <td>2</td>\n",
       "      <td>95</td>\n",
       "      <td>16</td>\n",
       "      <td>-3.287173</td>\n",
       "      <td>QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>530737 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "             set  dist  input_ct  selected_ct         y  \\\n",
       "0       training     2       173           33 -3.145154   \n",
       "1       training     2        18            8 -1.867676   \n",
       "2       training     2        66            2 -5.270800   \n",
       "3       training     2        72            1 -5.979498   \n",
       "4       training     2        69          168  0.481923   \n",
       "...          ...   ...       ...          ...       ...   \n",
       "530732  training     2       462          139 -2.515259   \n",
       "530733  training     2       317           84 -2.693165   \n",
       "530734  training     2       335           77 -2.896589   \n",
       "530735  training     2       148           28 -3.150861   \n",
       "530736  training     2        95           16 -3.287173   \n",
       "\n",
       "                                                        x  \n",
       "0       AAKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "1       ACKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "2       ADKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "3       AEKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "4       AFKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "...                                                   ...  \n",
       "530732  QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "530733  QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "530734  QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "530735  QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "530736  QYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDD...  \n",
       "\n",
       "[530737 rows x 6 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Assemble into dataframe\n",
    "final_df = pd.DataFrame({'set':sets, 'dist':2, 'input_ct':in_ct, 'selected_ct':out_ct, 'y':y, 'x':x})\n",
    "\n",
    "# Keep only sequences with input_ct >= 10\n",
    "final_df = final_df[final_df['input_ct']>=10].reset_index(drop=True)\n",
    "\n",
    "# Save to file (uncomment to execute)\n",
    "# final_df.to_csv('gb1_data.csv.gz', index=False, compression='gzip')\n",
    "\n",
    "# Preview dataframe\n",
    "final_df"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.2"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}