normalize-filter-aggregate-barcodes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Titeseq barcode aggregation\n",
    "\n",
    "This notebook reads in the barcode-level counts, barcode-variant mapping, and sequencing runs total cells, so that we can prep Titeseq data for fitting by:\n",
    "\n",
    "1. Merging in the variants substitution annotations to barcodes.\n",
    "2. Normalize the counts by the total cells in each sequencing run.\n",
    "3. Aggregate the counts by the variant substitution annotations.\n",
    "4. Filter out low count barcodes and variants.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:26:35.711410Z",
     "iopub.status.busy": "2024-11-08T18:26:35.711186Z",
     "iopub.status.idle": "2024-11-08T18:26:36.542328Z",
     "shell.execute_reply": "2024-11-08T18:26:36.541858Z"
    }
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "import yaml\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from IPython.display import display, HTML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ignore warnings that clutter output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:26:36.544636Z",
     "iopub.status.busy": "2024-11-08T18:26:36.544388Z",
     "iopub.status.idle": "2024-11-08T18:26:36.546849Z",
     "shell.execute_reply": "2024-11-08T18:26:36.546464Z"
    }
   },
   "outputs": [],
   "source": [
    "warnings.simplefilter('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters for notebook\n",
    "Read the configuration file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:26:36.548503Z",
     "iopub.status.busy": "2024-11-08T18:26:36.548267Z",
     "iopub.status.idle": "2024-11-08T18:26:36.554737Z",
     "shell.execute_reply": "2024-11-08T18:26:36.554374Z"
    }
   },
   "outputs": [],
   "source": [
    "with open('config.yaml') as f:\n",
    "    config = yaml.safe_load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make output directory if needed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:26:36.556436Z",
     "iopub.status.busy": "2024-11-08T18:26:36.556249Z",
     "iopub.status.idle": "2024-11-08T18:26:36.558576Z",
     "shell.execute_reply": "2024-11-08T18:26:36.558264Z"
    }
   },
   "outputs": [],
   "source": [
    "os.makedirs(config['counts_dir'], exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:26:36.560261Z",
     "iopub.status.busy": "2024-11-08T18:26:36.559963Z",
     "iopub.status.idle": "2024-11-08T18:26:36.583509Z",
     "shell.execute_reply": "2024-11-08T18:26:36.583169Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th>sample_type</th>\n",
       "      <th>sort_bin</th>\n",
       "      <th>concentration</th>\n",
       "      <th>date</th>\n",
       "      <th>number_cells</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>library</th>\n",
       "      <th>sample</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"4\" valign=\"top\">lib1</th>\n",
       "      <th>SortSeq_bin1</th>\n",
       "      <td>SortSeq</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>210621</td>\n",
       "      <td>210800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>SortSeq_bin2</th>\n",
       "      <td>SortSeq</td>\n",
       "      <td>2</td>\n",
       "      <td>NaN</td>\n",
       "      <td>210621</td>\n",
       "      <td>1984000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>SortSeq_bin3</th>\n",
       "      <td>SortSeq</td>\n",
       "      <td>3</td>\n",
       "      <td>NaN</td>\n",
       "      <td>210621</td>\n",
       "      <td>2940000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>SortSeq_bin4</th>\n",
       "      <td>SortSeq</td>\n",
       "      <td>4</td>\n",
       "      <td>NaN</td>\n",
       "      <td>210621</td>\n",
       "      <td>3575000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>lib2</th>\n",
       "      <th>SortSeq_bin1</th>\n",
       "      <td>SortSeq</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>210621</td>\n",
       "      <td>275500</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(80, 5)\n"
     ]
    }
   ],
   "source": [
    "barcode_runs = pd.read_csv(config['barcode_runs']).drop(columns=[\"R1\"])\n",
    "barcode_runs.query(\"sample.str.startswith('TiteSeq') | sample.str.startswith('SortSeq')\", inplace=True)\n",
    "barcode_runs.set_index([\"library\", \"sample\"], inplace=True)\n",
    "display(HTML(barcode_runs.head().to_html(index=True)))\n",
    "print(barcode_runs.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:26:36.608592Z",
     "iopub.status.busy": "2024-11-08T18:26:36.608191Z",
     "iopub.status.idle": "2024-11-08T18:26:48.782928Z",
     "shell.execute_reply": "2024-11-08T18:26:48.782490Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>barcode</th>\n",
       "      <th>count</th>\n",
       "      <th>library</th>\n",
       "      <th>sample</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>AGCATACCCTTAACAA</td>\n",
       "      <td>26343</td>\n",
       "      <td>lib1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>TGACGCCTTATCCTCC</td>\n",
       "      <td>20015</td>\n",
       "      <td>lib1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>TGCGATGGTACGTCAA</td>\n",
       "      <td>15678</td>\n",
       "      <td>lib1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>AACTACACGGATAGGT</td>\n",
       "      <td>14906</td>\n",
       "      <td>lib1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CATAATGAATGTGCAA</td>\n",
       "      <td>14667</td>\n",
       "      <td>lib1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(7697160, 4)\n"
     ]
    }
   ],
   "source": [
    "variant_counts = pd.read_csv(config['variant_counts_file'])\n",
    "variant_counts.query(\"sample.str.startswith('TiteSeq') | sample.str.startswith('SortSeq')\", inplace=True)\n",
    "display(HTML(variant_counts.head().to_html(index=True)))\n",
    "print(variant_counts.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:26:48.784568Z",
     "iopub.status.busy": "2024-11-08T18:26:48.784322Z",
     "iopub.status.idle": "2024-11-08T18:26:48.926474Z",
     "shell.execute_reply": "2024-11-08T18:26:48.926103Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>target</th>\n",
       "      <th>library</th>\n",
       "      <th>barcode</th>\n",
       "      <th>variant_call_support</th>\n",
       "      <th>codon_substitutions</th>\n",
       "      <th>aa_substitutions</th>\n",
       "      <th>n_codon_substitutions</th>\n",
       "      <th>n_aa_substitutions</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>CGG_naive</td>\n",
       "      <td>lib1</td>\n",
       "      <td>AAAAAAAAAACACCGG</td>\n",
       "      <td>6</td>\n",
       "      <td>GGC119GGT TTA200ACT</td>\n",
       "      <td>L200T</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>CGG_naive</td>\n",
       "      <td>lib1</td>\n",
       "      <td>AAAAAAAAAACATGAG</td>\n",
       "      <td>1</td>\n",
       "      <td>CAG16TGG</td>\n",
       "      <td>Q16W</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CGG_naive</td>\n",
       "      <td>lib1</td>\n",
       "      <td>AAAAAAAAAAGCGACG</td>\n",
       "      <td>1</td>\n",
       "      <td>GTG156CAT</td>\n",
       "      <td>V156H</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CGG_naive</td>\n",
       "      <td>lib1</td>\n",
       "      <td>AAAAAAAAAAGGAAAG</td>\n",
       "      <td>6</td>\n",
       "      <td>GTG110GGT</td>\n",
       "      <td>V110G</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CGG_naive</td>\n",
       "      <td>lib1</td>\n",
       "      <td>AAAAAAAAAATATAGA</td>\n",
       "      <td>1</td>\n",
       "      <td>TAC47CCA</td>\n",
       "      <td>Y47P</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(192429, 8)\n"
     ]
    }
   ],
   "source": [
    "codon_variant_table = pd.read_csv(config['codon_variant_table_file'])\n",
    "display(HTML(codon_variant_table.head().to_html(index=True)))\n",
    "print(codon_variant_table.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define concentrations and bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:26:48.928139Z",
     "iopub.status.busy": "2024-11-08T18:26:48.927936Z",
     "iopub.status.idle": "2024-11-08T18:26:48.931161Z",
     "shell.execute_reply": "2024-11-08T18:26:48.930859Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.e-06, 1.e-07, 1.e-08, 1.e-09, 1.e-10, 1.e-11, 1.e-12, 1.e-13,\n",
       "       0.e+00])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "concs = np.array(config['concentrations']['CGG']).astype(float)\n",
    "concs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Combine tables to make a single barcode level frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:26:48.932629Z",
     "iopub.status.busy": "2024-11-08T18:26:48.932393Z",
     "iopub.status.idle": "2024-11-08T18:26:51.453827Z",
     "shell.execute_reply": "2024-11-08T18:26:51.453449Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>barcode</th>\n",
       "      <th>count</th>\n",
       "      <th>library</th>\n",
       "      <th>sample</th>\n",
       "      <th>target</th>\n",
       "      <th>variant_call_support</th>\n",
       "      <th>codon_substitutions</th>\n",
       "      <th>aa_substitutions</th>\n",
       "      <th>n_codon_substitutions</th>\n",
       "      <th>n_aa_substitutions</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>AGCATACCCTTAACAA</td>\n",
       "      <td>26343</td>\n",
       "      <td>lib1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>CGG_naive</td>\n",
       "      <td>65</td>\n",
       "      <td>CGT38ATT GTA148GAA</td>\n",
       "      <td>R38I V148E</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>TGACGCCTTATCCTCC</td>\n",
       "      <td>20015</td>\n",
       "      <td>lib1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>CGG_naive</td>\n",
       "      <td>23</td>\n",
       "      <td>GAT100GTT CAG217TAG</td>\n",
       "      <td>D100V Q217*</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>TGCGATGGTACGTCAA</td>\n",
       "      <td>15678</td>\n",
       "      <td>lib1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>CGG_naive</td>\n",
       "      <td>26</td>\n",
       "      <td>GAC72TGT</td>\n",
       "      <td>D72C</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>AACTACACGGATAGGT</td>\n",
       "      <td>14906</td>\n",
       "      <td>lib1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>CGG_naive</td>\n",
       "      <td>31</td>\n",
       "      <td>CAA133AAA TTT137TAG</td>\n",
       "      <td>Q133K F137*</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CATAATGAATGTGCAA</td>\n",
       "      <td>14667</td>\n",
       "      <td>lib1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>CGG_naive</td>\n",
       "      <td>35</td>\n",
       "      <td>TAT94ATT CAG217TGG GAG232TAG</td>\n",
       "      <td>Y94I Q217W E232*</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(7697160, 10)\n"
     ]
    }
   ],
   "source": [
    "df_barcodes = variant_counts.merge(codon_variant_table, on=(\"barcode\", \"library\"), how=\"left\")\n",
    "display(HTML(df_barcodes.head().to_html(index=True)))\n",
    "print(df_barcodes.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add concentrations by parsing the sample name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:26:51.455508Z",
     "iopub.status.busy": "2024-11-08T18:26:51.455355Z",
     "iopub.status.idle": "2024-11-08T18:27:02.509547Z",
     "shell.execute_reply": "2024-11-08T18:27:02.509118Z"
    }
   },
   "outputs": [],
   "source": [
    "# parse the sample name to get the antigen concentration and bin\n",
    "df_barcodes[\"antigen_concentration\"] = np.nan\n",
    "df_barcodes.loc[df_barcodes['sample'].str.contains('TiteSeq'), \"antigen_concentration\"] = concs[\n",
    "    df_barcodes.query(f\"sample.str.contains('TiteSeq')\")[\"sample\"].str.extract(r\"TiteSeq_(\\d+)\").astype(int) - 1\n",
    "]\n",
    "df_barcodes[\"bin\"] = df_barcodes[\"sample\"].str[-1].astype(int)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filter synonymous variants with silent mutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:27:02.511934Z",
     "iopub.status.busy": "2024-11-08T18:27:02.511747Z",
     "iopub.status.idle": "2024-11-08T18:27:08.390127Z",
     "shell.execute_reply": "2024-11-08T18:27:08.389744Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 27000 synonymous variants with codon substitutions, dropping them.\n"
     ]
    },
    {
     "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>sample</th>\n",
       "      <th>library</th>\n",
       "      <th>variant</th>\n",
       "      <th>n_aa_substitutions</th>\n",
       "      <th>barcode</th>\n",
       "      <th>antigen_concentration</th>\n",
       "      <th>bin</th>\n",
       "      <th>read_count</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>42492</th>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>AAAAACATCAGTTGGT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>17</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>29591</th>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>AAAACACTATCTAGGA</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>36</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>59773</th>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>AAAATTCAAAATTATC</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>60801</th>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>AACAAAAGTGTATGTT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>62532</th>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>AAGTTATGAATACCCT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</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",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7697113</th>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>lib2</td>\n",
       "      <td>WT</td>\n",
       "      <td>0</td>\n",
       "      <td>TTTTTCATGTATATGC</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7697148</th>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>lib2</td>\n",
       "      <td>WT</td>\n",
       "      <td>0</td>\n",
       "      <td>TTTTTTAAAGTTCATA</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7697151</th>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>lib2</td>\n",
       "      <td>WT</td>\n",
       "      <td>0</td>\n",
       "      <td>TTTTTTACCTTTACCT</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7697153</th>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>lib2</td>\n",
       "      <td>WT</td>\n",
       "      <td>0</td>\n",
       "      <td>TTTTTTAGAAGCGAAG</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7697157</th>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>lib2</td>\n",
       "      <td>WT</td>\n",
       "      <td>0</td>\n",
       "      <td>TTTTTTGCTGTGTATC</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>7670160 rows × 8 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                  sample library variant  n_aa_substitutions  \\\n",
       "42492       SortSeq_bin1    lib1   A104C                   1   \n",
       "29591       SortSeq_bin1    lib1   A104C                   1   \n",
       "59773       SortSeq_bin1    lib1   A104C                   1   \n",
       "60801       SortSeq_bin1    lib1   A104C                   1   \n",
       "62532       SortSeq_bin1    lib1   A104C                   1   \n",
       "...                  ...     ...     ...                 ...   \n",
       "7697113  TiteSeq_09_bin4    lib2      WT                   0   \n",
       "7697148  TiteSeq_09_bin4    lib2      WT                   0   \n",
       "7697151  TiteSeq_09_bin4    lib2      WT                   0   \n",
       "7697153  TiteSeq_09_bin4    lib2      WT                   0   \n",
       "7697157  TiteSeq_09_bin4    lib2      WT                   0   \n",
       "\n",
       "                  barcode  antigen_concentration  bin  read_count  \n",
       "42492    AAAAACATCAGTTGGT                    NaN    1          17  \n",
       "29591    AAAACACTATCTAGGA                    NaN    1          36  \n",
       "59773    AAAATTCAAAATTATC                    NaN    1           0  \n",
       "60801    AACAAAAGTGTATGTT                    NaN    1           0  \n",
       "62532    AAGTTATGAATACCCT                    NaN    1           0  \n",
       "...                   ...                    ...  ...         ...  \n",
       "7697113  TTTTTCATGTATATGC                    0.0    4           0  \n",
       "7697148  TTTTTTAAAGTTCATA                    0.0    4           0  \n",
       "7697151  TTTTTTACCTTTACCT                    0.0    4           0  \n",
       "7697153  TTTTTTAGAAGCGAAG                    0.0    4           0  \n",
       "7697157  TTTTTTGCTGTGTATC                    0.0    4           0  \n",
       "\n",
       "[7670160 rows x 8 columns]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "synonymous = df_barcodes.query(\"n_aa_substitutions == 0 & n_codon_substitutions > 0\")\n",
    "print(f\"There are {synonymous.shape[0]} synonymous variants with codon substitutions, dropping them.\")\n",
    "df_barcodes.query(\"~(n_aa_substitutions == 0 & n_codon_substitutions > 0)\", inplace=True)\n",
    "\n",
    "# drop columns that are not needed\n",
    "df_barcodes.drop(columns=[\"codon_substitutions\", \"n_codon_substitutions\", \"target\", \"variant_call_support\"], inplace=True)\n",
    "df_barcodes.rename(columns={\"aa_substitutions\": \"variant\",\n",
    "                            \"count\": \"read_count\"},\n",
    "                   inplace=True)\n",
    "df_barcodes = df_barcodes.loc[:, [\"sample\", \"library\", \"variant\", \"n_aa_substitutions\", \"barcode\", \"antigen_concentration\", \"bin\", \"read_count\"]]\n",
    "df_barcodes.sort_values(by=list(df_barcodes.columns), inplace=True)\n",
    "df_barcodes.variant = df_barcodes.variant.fillna(\"WT\")\n",
    "\n",
    "df_barcodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use total cell counts and total read counts in each concentration and bin to estimate the number of cells with each barcode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:27:08.391972Z",
     "iopub.status.busy": "2024-11-08T18:27:08.391796Z",
     "iopub.status.idle": "2024-11-08T18:27:12.157509Z",
     "shell.execute_reply": "2024-11-08T18:27:12.157090Z"
    }
   },
   "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>sample</th>\n",
       "      <th>library</th>\n",
       "      <th>variant</th>\n",
       "      <th>n_aa_substitutions</th>\n",
       "      <th>barcode</th>\n",
       "      <th>antigen_concentration</th>\n",
       "      <th>bin</th>\n",
       "      <th>read_count</th>\n",
       "      <th>estimated_cell_count</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>AAAAACATCAGTTGGT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>17</td>\n",
       "      <td>0.610535</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>AAAACACTATCTAGGA</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>36</td>\n",
       "      <td>1.292898</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>AAAATTCAAAATTATC</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>AACAAAAGTGTATGTT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>AAGTTATGAATACCCT</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</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",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7670155</th>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>lib2</td>\n",
       "      <td>WT</td>\n",
       "      <td>0</td>\n",
       "      <td>TTTTTCATGTATATGC</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7670156</th>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>lib2</td>\n",
       "      <td>WT</td>\n",
       "      <td>0</td>\n",
       "      <td>TTTTTTAAAGTTCATA</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7670157</th>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>lib2</td>\n",
       "      <td>WT</td>\n",
       "      <td>0</td>\n",
       "      <td>TTTTTTACCTTTACCT</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7670158</th>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>lib2</td>\n",
       "      <td>WT</td>\n",
       "      <td>0</td>\n",
       "      <td>TTTTTTAGAAGCGAAG</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7670159</th>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>lib2</td>\n",
       "      <td>WT</td>\n",
       "      <td>0</td>\n",
       "      <td>TTTTTTGCTGTGTATC</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>7670160 rows × 9 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                  sample library variant  n_aa_substitutions  \\\n",
       "0           SortSeq_bin1    lib1   A104C                   1   \n",
       "1           SortSeq_bin1    lib1   A104C                   1   \n",
       "2           SortSeq_bin1    lib1   A104C                   1   \n",
       "3           SortSeq_bin1    lib1   A104C                   1   \n",
       "4           SortSeq_bin1    lib1   A104C                   1   \n",
       "...                  ...     ...     ...                 ...   \n",
       "7670155  TiteSeq_09_bin4    lib2      WT                   0   \n",
       "7670156  TiteSeq_09_bin4    lib2      WT                   0   \n",
       "7670157  TiteSeq_09_bin4    lib2      WT                   0   \n",
       "7670158  TiteSeq_09_bin4    lib2      WT                   0   \n",
       "7670159  TiteSeq_09_bin4    lib2      WT                   0   \n",
       "\n",
       "                  barcode  antigen_concentration  bin  read_count  \\\n",
       "0        AAAAACATCAGTTGGT                    NaN    1          17   \n",
       "1        AAAACACTATCTAGGA                    NaN    1          36   \n",
       "2        AAAATTCAAAATTATC                    NaN    1           0   \n",
       "3        AACAAAAGTGTATGTT                    NaN    1           0   \n",
       "4        AAGTTATGAATACCCT                    NaN    1           0   \n",
       "...                   ...                    ...  ...         ...   \n",
       "7670155  TTTTTCATGTATATGC                    0.0    4           0   \n",
       "7670156  TTTTTTAAAGTTCATA                    0.0    4           0   \n",
       "7670157  TTTTTTACCTTTACCT                    0.0    4           0   \n",
       "7670158  TTTTTTAGAAGCGAAG                    0.0    4           0   \n",
       "7670159  TTTTTTGCTGTGTATC                    0.0    4           0   \n",
       "\n",
       "         estimated_cell_count  \n",
       "0                    0.610535  \n",
       "1                    1.292898  \n",
       "2                    0.000000  \n",
       "3                    0.000000  \n",
       "4                    0.000000  \n",
       "...                       ...  \n",
       "7670155              0.000000  \n",
       "7670156              0.000000  \n",
       "7670157              0.000000  \n",
       "7670158              0.000000  \n",
       "7670159              0.000000  \n",
       "\n",
       "[7670160 rows x 9 columns]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def normalize_read_count(df):\n",
    "    library = df.library.iloc[0]\n",
    "    sample = df[\"sample\"].iloc[0]\n",
    "    total_reads = df.read_count.sum()\n",
    "    total_cells = barcode_runs.number_cells[(library, sample)]\n",
    "    df[\"estimated_cell_count\"] = total_cells * df.read_count / total_reads\n",
    "    return df\n",
    "\n",
    "df_barcodes = df_barcodes.groupby([\"library\", \"sample\"]).apply(normalize_read_count).reset_index(drop=True)\n",
    "df_barcodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Barcode aggregation\n",
    "\n",
    "The current $K_D$ estimation procedure does a separate estimate for each barcode, then computes the median of the estimated $\\log K_D$ across barcodes for each variant.\n",
    "We should instead estimate a single $K_D$ parameter for each variants. We will do this by aggregating read counts from all barcodes for a given variant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:27:12.159306Z",
     "iopub.status.busy": "2024-11-08T18:27:12.159135Z",
     "iopub.status.idle": "2024-11-08T18:27:15.049499Z",
     "shell.execute_reply": "2024-11-08T18:27:15.049035Z"
    }
   },
   "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>library</th>\n",
       "      <th>variant</th>\n",
       "      <th>n_aa_substitutions</th>\n",
       "      <th>antigen_concentration</th>\n",
       "      <th>bin</th>\n",
       "      <th>sample</th>\n",
       "      <th>read_count</th>\n",
       "      <th>estimated_cell_count</th>\n",
       "      <th>barcode</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>1</td>\n",
       "      <td>TiteSeq_09_bin1</td>\n",
       "      <td>2663</td>\n",
       "      <td>1250.687875</td>\n",
       "      <td>25</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>2</td>\n",
       "      <td>TiteSeq_09_bin2</td>\n",
       "      <td>25</td>\n",
       "      <td>24.273527</td>\n",
       "      <td>25</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>3</td>\n",
       "      <td>TiteSeq_09_bin3</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>25</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>0.000000e+00</td>\n",
       "      <td>4</td>\n",
       "      <td>TiteSeq_09_bin4</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>25</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>lib1</td>\n",
       "      <td>A104C</td>\n",
       "      <td>1</td>\n",
       "      <td>1.000000e-13</td>\n",
       "      <td>1</td>\n",
       "      <td>TiteSeq_08_bin1</td>\n",
       "      <td>578</td>\n",
       "      <td>1494.790836</td>\n",
       "      <td>25</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",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>912835</th>\n",
       "      <td>lib2</td>\n",
       "      <td>Y94W R145M</td>\n",
       "      <td>2</td>\n",
       "      <td>1.000000e-06</td>\n",
       "      <td>4</td>\n",
       "      <td>TiteSeq_01_bin4</td>\n",
       "      <td>16</td>\n",
       "      <td>12.947315</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>912836</th>\n",
       "      <td>lib2</td>\n",
       "      <td>Y94W R145M</td>\n",
       "      <td>2</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>SortSeq_bin1</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>912837</th>\n",
       "      <td>lib2</td>\n",
       "      <td>Y94W R145M</td>\n",
       "      <td>2</td>\n",
       "      <td>NaN</td>\n",
       "      <td>2</td>\n",
       "      <td>SortSeq_bin2</td>\n",
       "      <td>14</td>\n",
       "      <td>2.535604</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>912838</th>\n",
       "      <td>lib2</td>\n",
       "      <td>Y94W R145M</td>\n",
       "      <td>2</td>\n",
       "      <td>NaN</td>\n",
       "      <td>3</td>\n",
       "      <td>SortSeq_bin3</td>\n",
       "      <td>57</td>\n",
       "      <td>10.063316</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>912839</th>\n",
       "      <td>lib2</td>\n",
       "      <td>Y94W R145M</td>\n",
       "      <td>2</td>\n",
       "      <td>NaN</td>\n",
       "      <td>4</td>\n",
       "      <td>SortSeq_bin4</td>\n",
       "      <td>112</td>\n",
       "      <td>37.721971</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>912840 rows × 9 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       library     variant  n_aa_substitutions  antigen_concentration  bin  \\\n",
       "0         lib1       A104C                   1           0.000000e+00    1   \n",
       "1         lib1       A104C                   1           0.000000e+00    2   \n",
       "2         lib1       A104C                   1           0.000000e+00    3   \n",
       "3         lib1       A104C                   1           0.000000e+00    4   \n",
       "4         lib1       A104C                   1           1.000000e-13    1   \n",
       "...        ...         ...                 ...                    ...  ...   \n",
       "912835    lib2  Y94W R145M                   2           1.000000e-06    4   \n",
       "912836    lib2  Y94W R145M                   2                    NaN    1   \n",
       "912837    lib2  Y94W R145M                   2                    NaN    2   \n",
       "912838    lib2  Y94W R145M                   2                    NaN    3   \n",
       "912839    lib2  Y94W R145M                   2                    NaN    4   \n",
       "\n",
       "                 sample  read_count  estimated_cell_count  barcode  \n",
       "0       TiteSeq_09_bin1        2663           1250.687875       25  \n",
       "1       TiteSeq_09_bin2          25             24.273527       25  \n",
       "2       TiteSeq_09_bin3           0              0.000000       25  \n",
       "3       TiteSeq_09_bin4           0              0.000000       25  \n",
       "4       TiteSeq_08_bin1         578           1494.790836       25  \n",
       "...                 ...         ...                   ...      ...  \n",
       "912835  TiteSeq_01_bin4          16             12.947315        1  \n",
       "912836     SortSeq_bin1           0              0.000000        1  \n",
       "912837     SortSeq_bin2          14              2.535604        1  \n",
       "912838     SortSeq_bin3          57             10.063316        1  \n",
       "912839     SortSeq_bin4         112             37.721971        1  \n",
       "\n",
       "[912840 rows x 9 columns]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_variants = (\n",
    "    df_barcodes\n",
    "    .groupby(\n",
    "        [\"library\", \"variant\", \"n_aa_substitutions\", \"antigen_concentration\", \"bin\", \"sample\"],\n",
    "        dropna=False\n",
    "    )\n",
    "    .agg(\n",
    "        {\n",
    "            \"read_count\": \"sum\",\n",
    "            \"estimated_cell_count\": \"sum\",\n",
    "            \"barcode\": \"count\"\n",
    "        }\n",
    "    )\n",
    "    .reset_index()\n",
    ")\n",
    "df_variants.sort_values(by=list(df_variants.columns), inplace=True)\n",
    "df_variants"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filter missing concentrations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We filter missing concentrations for each barcode in the barcode data, and for each variant in the variant data. Missing means the sum of reads is zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:27:15.051330Z",
     "iopub.status.busy": "2024-11-08T18:27:15.051162Z",
     "iopub.status.idle": "2024-11-08T18:27:15.053323Z",
     "shell.execute_reply": "2024-11-08T18:27:15.052941Z"
    }
   },
   "outputs": [],
   "source": [
    "def conc_filter_fn(df):\n",
    "    return all(df.groupby(\"antigen_concentration\", dropna=False).read_count.sum() > 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:27:15.054836Z",
     "iopub.status.busy": "2024-11-08T18:27:15.054684Z",
     "iopub.status.idle": "2024-11-08T18:27:21.626561Z",
     "shell.execute_reply": "2024-11-08T18:27:21.626114Z"
    }
   },
   "outputs": [],
   "source": [
    "# make a 'sample_type' column, such that we can group sample type variants accross all bins when filtering for concentrations\n",
    "df_variants[\"sample_type\"] = df_variants[\"sample\"].str.extract(r\"^(TiteSeq|SortSeq)\")\n",
    "df_barcodes[\"sample_type\"] = df_barcodes[\"sample\"].str.extract(r\"^(TiteSeq|SortSeq)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:27:21.628759Z",
     "iopub.status.busy": "2024-11-08T18:27:21.628462Z",
     "iopub.status.idle": "2024-11-08T18:29:20.715430Z",
     "shell.execute_reply": "2024-11-08T18:29:20.714921Z"
    }
   },
   "outputs": [],
   "source": [
    "df_barcodes = (\n",
    "    df_barcodes\n",
    "    .groupby(\n",
    "        [\"sample_type\", \"library\", \"variant\", \"n_aa_substitutions\", \"barcode\"]\n",
    "    )\n",
    "    .filter(conc_filter_fn)\n",
    "    .drop(columns=[\"sample_type\"])\n",
    ")\n",
    "df_variants = (\n",
    "    df_variants\n",
    "    .groupby(\n",
    "        [\"sample_type\", \"library\", \"variant\", \"n_aa_substitutions\"]\n",
    "    )\n",
    "    .filter(conc_filter_fn)\n",
    "    .drop(columns=[\"sample_type\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:29:20.717436Z",
     "iopub.status.busy": "2024-11-08T18:29:20.717289Z",
     "iopub.status.idle": "2024-11-08T18:29:20.906619Z",
     "shell.execute_reply": "2024-11-08T18:29:20.906162Z"
    }
   },
   "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></th>\n",
       "      <th>count</th>\n",
       "      <th>mean</th>\n",
       "      <th>std</th>\n",
       "      <th>min</th>\n",
       "      <th>25%</th>\n",
       "      <th>50%</th>\n",
       "      <th>75%</th>\n",
       "      <th>max</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sample</th>\n",
       "      <th>library</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">SortSeq_bin1</th>\n",
       "      <th>lib1</th>\n",
       "      <td>10938.0</td>\n",
       "      <td>8.470104</td>\n",
       "      <td>105.405655</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>15.0</td>\n",
       "      <td>10986.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>lib2</th>\n",
       "      <td>11686.0</td>\n",
       "      <td>8.464060</td>\n",
       "      <td>114.128714</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>15.0</td>\n",
       "      <td>12299.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">SortSeq_bin2</th>\n",
       "      <th>lib1</th>\n",
       "      <td>10938.0</td>\n",
       "      <td>8.470104</td>\n",
       "      <td>105.405655</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>15.0</td>\n",
       "      <td>10986.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>lib2</th>\n",
       "      <td>11686.0</td>\n",
       "      <td>8.464060</td>\n",
       "      <td>114.128714</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>15.0</td>\n",
       "      <td>12299.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>SortSeq_bin3</th>\n",
       "      <th>lib1</th>\n",
       "      <td>10938.0</td>\n",
       "      <td>8.470104</td>\n",
       "      <td>105.405655</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>15.0</td>\n",
       "      <td>10986.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\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",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>TiteSeq_09_bin2</th>\n",
       "      <th>lib2</th>\n",
       "      <td>10200.0</td>\n",
       "      <td>9.548922</td>\n",
       "      <td>122.122509</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>17.0</td>\n",
       "      <td>12299.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">TiteSeq_09_bin3</th>\n",
       "      <th>lib1</th>\n",
       "      <td>9437.0</td>\n",
       "      <td>9.654551</td>\n",
       "      <td>113.434854</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>10986.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>lib2</th>\n",
       "      <td>10200.0</td>\n",
       "      <td>9.548922</td>\n",
       "      <td>122.122509</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>17.0</td>\n",
       "      <td>12299.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th rowspan=\"2\" valign=\"top\">TiteSeq_09_bin4</th>\n",
       "      <th>lib1</th>\n",
       "      <td>9437.0</td>\n",
       "      <td>9.654551</td>\n",
       "      <td>113.434854</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>10986.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>lib2</th>\n",
       "      <td>10200.0</td>\n",
       "      <td>9.548922</td>\n",
       "      <td>122.122509</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>17.0</td>\n",
       "      <td>12299.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>80 rows × 8 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                           count      mean         std  min  25%  50%   75%  \\\n",
       "sample          library                                                       \n",
       "SortSeq_bin1    lib1     10938.0  8.470104  105.405655  1.0  1.0  1.0  15.0   \n",
       "                lib2     11686.0  8.464060  114.128714  1.0  1.0  1.0  15.0   \n",
       "SortSeq_bin2    lib1     10938.0  8.470104  105.405655  1.0  1.0  1.0  15.0   \n",
       "                lib2     11686.0  8.464060  114.128714  1.0  1.0  1.0  15.0   \n",
       "SortSeq_bin3    lib1     10938.0  8.470104  105.405655  1.0  1.0  1.0  15.0   \n",
       "...                          ...       ...         ...  ...  ...  ...   ...   \n",
       "TiteSeq_09_bin2 lib2     10200.0  9.548922  122.122509  1.0  1.0  1.0  17.0   \n",
       "TiteSeq_09_bin3 lib1      9437.0  9.654551  113.434854  1.0  1.0  1.0  16.0   \n",
       "                lib2     10200.0  9.548922  122.122509  1.0  1.0  1.0  17.0   \n",
       "TiteSeq_09_bin4 lib1      9437.0  9.654551  113.434854  1.0  1.0  1.0  16.0   \n",
       "                lib2     10200.0  9.548922  122.122509  1.0  1.0  1.0  17.0   \n",
       "\n",
       "                             max  \n",
       "sample          library           \n",
       "SortSeq_bin1    lib1     10986.0  \n",
       "                lib2     12299.0  \n",
       "SortSeq_bin2    lib1     10986.0  \n",
       "                lib2     12299.0  \n",
       "SortSeq_bin3    lib1     10986.0  \n",
       "...                          ...  \n",
       "TiteSeq_09_bin2 lib2     12299.0  \n",
       "TiteSeq_09_bin3 lib1     10986.0  \n",
       "                lib2     12299.0  \n",
       "TiteSeq_09_bin4 lib1     10986.0  \n",
       "                lib2     12299.0  \n",
       "\n",
       "[80 rows x 8 columns]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# plot the distribution of 'barcodes' values across the various each of the sample groups\n",
    "df_variants.groupby([\"sample\", \"library\"]).barcode.describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-08T18:29:20.908361Z",
     "iopub.status.busy": "2024-11-08T18:29:20.908228Z",
     "iopub.status.idle": "2024-11-08T18:29:37.192343Z",
     "shell.execute_reply": "2024-11-08T18:29:37.191820Z"
    }
   },
   "outputs": [],
   "source": [
    "df_barcodes.to_csv(config['prepped_barcode_counts_file'], index=False)\n",
    "df_variants.to_csv(config['prepped_variant_counts_file'], index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "replay",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}