2-data-split-tcr.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train / test data split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook performs two important tasks:\n",
    "1. Generates negative training examples by randomly mispairing TCRs and epitopes from the positive training examples collated in the previous notebook.\n",
    "2. Splits the positive and negative training examples into train and test subsets, ensuring randomisation and mutual exclusivity by TCR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set a random seed for reproducibility\n",
    "seed = 3\n",
    "\n",
    "# Read in the collated dataframe from the previous notebook as the set of positive training examples\n",
    "path = 'data/input/collated/collated.csv'\n",
    "positives = pd.read_csv(path, sep=',', dtype=str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Drop the 'source' column, which is not needed at this stage\n",
    "positives.drop('source', axis=1, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "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>cdr3a</th>\n",
       "      <th>cdr3b</th>\n",
       "      <th>epitope</th>\n",
       "      <th>hla</th>\n",
       "      <th>v_a</th>\n",
       "      <th>j_a</th>\n",
       "      <th>v_b</th>\n",
       "      <th>j_b</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>IVVRSSNTGKLI</td>\n",
       "      <td>ASSQDRDTQY</td>\n",
       "      <td>VMAPRTLIL</td>\n",
       "      <td>HLA-E*01:01</td>\n",
       "      <td>TRAV26-1*01</td>\n",
       "      <td>TRAJ37*01</td>\n",
       "      <td>TRBV14*01</td>\n",
       "      <td>TRBJ2-3*01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>IVVRSSNTGKLI</td>\n",
       "      <td>ASSQDRDTQY</td>\n",
       "      <td>VMAPRTLIL</td>\n",
       "      <td>HLA-E*01:03</td>\n",
       "      <td>TRAV26-1*01</td>\n",
       "      <td>TRAJ37*01</td>\n",
       "      <td>TRBV14*01</td>\n",
       "      <td>TRBJ2-3*01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>AVRPLLDGTYIPT</td>\n",
       "      <td>ASSYLGNTGELF</td>\n",
       "      <td>SLLMWITQC</td>\n",
       "      <td>HLA-A*02:01</td>\n",
       "      <td>TRAV21*01</td>\n",
       "      <td>TRAJ6*01</td>\n",
       "      <td>TRBV6-5*01</td>\n",
       "      <td>TRBJ2-2*01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>IVWGGYQKVT</td>\n",
       "      <td>ASRYRDDSYNEQF</td>\n",
       "      <td>EENLLDFVRF</td>\n",
       "      <td>HLA-B*44:02</td>\n",
       "      <td>TRAV26-1*01</td>\n",
       "      <td>TRAJ13*01</td>\n",
       "      <td>TRBV7-9*01</td>\n",
       "      <td>TRBJ2-1*01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>AVTTDSWGKLQ</td>\n",
       "      <td>ASRPGLAGGRPEQY</td>\n",
       "      <td>LLFGYPVFV</td>\n",
       "      <td>HLA-A*02:01</td>\n",
       "      <td>TRAV12-2*01</td>\n",
       "      <td>TRAJ24*02</td>\n",
       "      <td>TRBV6-5*01</td>\n",
       "      <td>TRBJ2-7*01</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>24032</th>\n",
       "      <td>CARSTDSWGKLQF</td>\n",
       "      <td>CASTPEGSYNEQFF</td>\n",
       "      <td>LLWNGPMAV</td>\n",
       "      <td>HLA-A*02:01</td>\n",
       "      <td>TRAV24*01</td>\n",
       "      <td>TRAJ24*01</td>\n",
       "      <td>TRBV7-6*01</td>\n",
       "      <td>TRBJ2-1*01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24033</th>\n",
       "      <td>CWSPFGNEKLTF</td>\n",
       "      <td>CASSLGQGSYEQYF</td>\n",
       "      <td>GTSGSPIVNR</td>\n",
       "      <td>HLA-A*11:01</td>\n",
       "      <td>TRAV8-2*01</td>\n",
       "      <td>TRAJ48*01</td>\n",
       "      <td>TRBV11-2*01</td>\n",
       "      <td>TRBJ2-7*01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24034</th>\n",
       "      <td>CWSPFGNEKLTF</td>\n",
       "      <td>CASSLGQGSYEQYF</td>\n",
       "      <td>GTSGSPIINR</td>\n",
       "      <td>HLA-A*11:01</td>\n",
       "      <td>TRAV8-2*01</td>\n",
       "      <td>TRAJ48*01</td>\n",
       "      <td>TRBV11-2*01</td>\n",
       "      <td>TRBJ2-7*01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24035</th>\n",
       "      <td>SSGNQFYF</td>\n",
       "      <td>CASSQQTGTIGGYTF</td>\n",
       "      <td>NLVPMVATV</td>\n",
       "      <td>HLA-A*02</td>\n",
       "      <td>TRAV21*01</td>\n",
       "      <td>TRAJ49*01</td>\n",
       "      <td>TRBV6-5*01</td>\n",
       "      <td>TRBJ1-2*01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24036</th>\n",
       "      <td>CAGLTTDSWGKFQF</td>\n",
       "      <td>CASSVDVAGGYNEQFF</td>\n",
       "      <td>NLVPMVATV</td>\n",
       "      <td>HLA-A*02</td>\n",
       "      <td>TRAV25*01</td>\n",
       "      <td>TRAJ24*01</td>\n",
       "      <td>TRBV7-3*01</td>\n",
       "      <td>TRBJ2-1*01</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>24037 rows × 8 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                cdr3a             cdr3b     epitope          hla          v_a  \\\n",
       "0        IVVRSSNTGKLI        ASSQDRDTQY   VMAPRTLIL  HLA-E*01:01  TRAV26-1*01   \n",
       "1        IVVRSSNTGKLI        ASSQDRDTQY   VMAPRTLIL  HLA-E*01:03  TRAV26-1*01   \n",
       "2       AVRPLLDGTYIPT      ASSYLGNTGELF   SLLMWITQC  HLA-A*02:01    TRAV21*01   \n",
       "3          IVWGGYQKVT     ASRYRDDSYNEQF  EENLLDFVRF  HLA-B*44:02  TRAV26-1*01   \n",
       "4         AVTTDSWGKLQ    ASRPGLAGGRPEQY   LLFGYPVFV  HLA-A*02:01  TRAV12-2*01   \n",
       "...               ...               ...         ...          ...          ...   \n",
       "24032   CARSTDSWGKLQF    CASTPEGSYNEQFF   LLWNGPMAV  HLA-A*02:01    TRAV24*01   \n",
       "24033    CWSPFGNEKLTF    CASSLGQGSYEQYF  GTSGSPIVNR  HLA-A*11:01   TRAV8-2*01   \n",
       "24034    CWSPFGNEKLTF    CASSLGQGSYEQYF  GTSGSPIINR  HLA-A*11:01   TRAV8-2*01   \n",
       "24035        SSGNQFYF   CASSQQTGTIGGYTF   NLVPMVATV     HLA-A*02    TRAV21*01   \n",
       "24036  CAGLTTDSWGKFQF  CASSVDVAGGYNEQFF   NLVPMVATV     HLA-A*02    TRAV25*01   \n",
       "\n",
       "             j_a          v_b         j_b  \n",
       "0      TRAJ37*01    TRBV14*01  TRBJ2-3*01  \n",
       "1      TRAJ37*01    TRBV14*01  TRBJ2-3*01  \n",
       "2       TRAJ6*01   TRBV6-5*01  TRBJ2-2*01  \n",
       "3      TRAJ13*01   TRBV7-9*01  TRBJ2-1*01  \n",
       "4      TRAJ24*02   TRBV6-5*01  TRBJ2-7*01  \n",
       "...          ...          ...         ...  \n",
       "24032  TRAJ24*01   TRBV7-6*01  TRBJ2-1*01  \n",
       "24033  TRAJ48*01  TRBV11-2*01  TRBJ2-7*01  \n",
       "24034  TRAJ48*01  TRBV11-2*01  TRBJ2-7*01  \n",
       "24035  TRAJ49*01   TRBV6-5*01  TRBJ1-2*01  \n",
       "24036  TRAJ24*01   TRBV7-3*01  TRBJ2-1*01  \n",
       "\n",
       "[24037 rows x 8 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "positives"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the negative set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect a list of TCRs from the positive training examples\n",
    "tcr = positives.drop('epitope', axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "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>epitope</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>VMAPRTLIL</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SLLMWITQC</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>EENLLDFVRF</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>LLFGYPVFV</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>MLWGYLQYV</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>206</th>\n",
       "      <td>LPEPLPQGQLGAY</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>207</th>\n",
       "      <td>LPEPLPQGQLTGY</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>208</th>\n",
       "      <td>EFFWDANDIY</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>209</th>\n",
       "      <td>CPSQEPMSIYVY</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>210</th>\n",
       "      <td>VAANIVLTV</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>211 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "           epitope\n",
       "0        VMAPRTLIL\n",
       "1        SLLMWITQC\n",
       "2       EENLLDFVRF\n",
       "3        LLFGYPVFV\n",
       "4        MLWGYLQYV\n",
       "..             ...\n",
       "206  LPEPLPQGQLGAY\n",
       "207  LPEPLPQGQLTGY\n",
       "208     EFFWDANDIY\n",
       "209   CPSQEPMSIYVY\n",
       "210      VAANIVLTV\n",
       "\n",
       "[211 rows x 1 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Collect a list of epitopes from the positive training examples\n",
    "negepitopes = pd.DataFrame(positives['epitope'].unique(), columns=['epitope'])\n",
    "negepitopes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Determine the number of times to reuse each epitope to achieve a balanced dataset\n",
    "n = round(len(tcr)/len(negepitopes))\n",
    "\n",
    "negatives = pd.DataFrame(columns=positives.columns)\n",
    "\n",
    "# Randomly mispair the TCRs and epitopes to generate the negative training examples\n",
    "for i in range(n):\n",
    "    batch = tcr.sample(negepitopes.size, random_state=seed+i)\n",
    "    batch.reset_index(inplace=True, drop=True)\n",
    "    batch['epitope'] = negepitopes['epitope']\n",
    "    \n",
    "    negatives = negatives.append(batch, ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove any negative examples generated that already exist in the positive set\n",
    "intersection = pd.merge(positives, negatives, how='inner', on=['cdr3a', 'cdr3b', 'epitope'], suffixes=['', '_del'])\n",
    "intersection.drop(intersection.columns[intersection.columns.str.contains('_del')], axis=1, inplace=True)\n",
    "negatives = pd.concat([negatives, intersection, intersection]).drop_duplicates(keep=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add labels\n",
    "positives['y'] = 1.\n",
    "negatives['y'] = 0.\n",
    "\n",
    "# Union the negatives and positives\n",
    "data = positives.append(negatives, ignore_index=True)\n",
    "\n",
    "# Shuffle the data\n",
    "data = data.sample(frac=1, random_state=seed).reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "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>cdr3a</th>\n",
       "      <th>cdr3b</th>\n",
       "      <th>epitope</th>\n",
       "      <th>hla</th>\n",
       "      <th>v_a</th>\n",
       "      <th>j_a</th>\n",
       "      <th>v_b</th>\n",
       "      <th>j_b</th>\n",
       "      <th>y</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>CAVSPYNFNKFYF</td>\n",
       "      <td>CASSLAGFQETQYF</td>\n",
       "      <td>FLRGRAYGL</td>\n",
       "      <td>HLA-A*03:01</td>\n",
       "      <td>TRAV20*01</td>\n",
       "      <td>TRAJ21*01</td>\n",
       "      <td>TRBV11-2*01</td>\n",
       "      <td>TRBJ2-5*01</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>CAGPAGARLMF</td>\n",
       "      <td>CASSLYVNNGYTF</td>\n",
       "      <td>AVFDRKSDAK</td>\n",
       "      <td>HLA-A*11:01</td>\n",
       "      <td>TRAV35*01</td>\n",
       "      <td>TRAJ31*01</td>\n",
       "      <td>TRBV7-9*01</td>\n",
       "      <td>TRBJ1-2*01</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>CAVSDHSNNARLMF</td>\n",
       "      <td>CASSLGGTNEQFF</td>\n",
       "      <td>NLVPMVASV</td>\n",
       "      <td>HLA-B*08:01</td>\n",
       "      <td>TRAV8-4*01</td>\n",
       "      <td>TRAJ31*01</td>\n",
       "      <td>TRBV4-1*01</td>\n",
       "      <td>TRBJ2-1*01</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>CVVRSGGYNKLIF</td>\n",
       "      <td>CAITGENTGELFF</td>\n",
       "      <td>KLGGALQAK</td>\n",
       "      <td>HLA-A*03:01</td>\n",
       "      <td>TRAV10*01</td>\n",
       "      <td>TRAJ4*01</td>\n",
       "      <td>TRBV10-3*01</td>\n",
       "      <td>TRBJ2-2*01</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>CAMREFYNTDKLIF</td>\n",
       "      <td>CASSQDDHGSNEQYF</td>\n",
       "      <td>RAKFKQLL</td>\n",
       "      <td>HLA-B*08:01</td>\n",
       "      <td>TRAV14/DV4*01</td>\n",
       "      <td>TRAJ34*01</td>\n",
       "      <td>TRBV4-3*01</td>\n",
       "      <td>TRBJ2-7*01</td>\n",
       "      <td>1.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",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47804</th>\n",
       "      <td>CAATRYNNNDMRF</td>\n",
       "      <td>CASTREGAGWSAQHF</td>\n",
       "      <td>YLSNIIPAL</td>\n",
       "      <td>HLA-A*03:01</td>\n",
       "      <td>TRAV13-1*01</td>\n",
       "      <td>TRAJ43*01</td>\n",
       "      <td>TRBV12-3*01</td>\n",
       "      <td>TRBJ1-5*01</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47805</th>\n",
       "      <td>CAVKGTGGFKTIF</td>\n",
       "      <td>CASRSQGALYSNQPQHF</td>\n",
       "      <td>HSKKKCDEI</td>\n",
       "      <td>HLA-B*08:01</td>\n",
       "      <td>TRAV1-1*01</td>\n",
       "      <td>TRAJ9*01</td>\n",
       "      <td>TRBV28*01</td>\n",
       "      <td>TRBJ1-5*01</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47806</th>\n",
       "      <td>CAVTPRNQFYF</td>\n",
       "      <td>CASSQDRVGAGANVLTF</td>\n",
       "      <td>KLGGALQAK</td>\n",
       "      <td>HLA-A*03:01</td>\n",
       "      <td>TRAV8-6*01</td>\n",
       "      <td>TRAJ49*01</td>\n",
       "      <td>TRBV4-1*01</td>\n",
       "      <td>TRBJ2-6*01</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47807</th>\n",
       "      <td>CVVNEKEGNMLTF</td>\n",
       "      <td>CSVGEGLAYEQYF</td>\n",
       "      <td>LLWNGPMAV</td>\n",
       "      <td>HLA-A*02:01</td>\n",
       "      <td>TRAV12-1</td>\n",
       "      <td>TRAJ39</td>\n",
       "      <td>TRBV29-01</td>\n",
       "      <td>TRBJ2-7</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>47808</th>\n",
       "      <td>CAVDTGGGADGLTF</td>\n",
       "      <td>CARGQDAANTGELFF</td>\n",
       "      <td>ELAGIGILTV</td>\n",
       "      <td>HLA-A*02:01</td>\n",
       "      <td>TRAV12-2*01</td>\n",
       "      <td>TRAJ45*01</td>\n",
       "      <td>TRBV2*01</td>\n",
       "      <td>TRBJ2-2*01</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>47809 rows × 9 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                cdr3a              cdr3b     epitope          hla  \\\n",
       "0       CAVSPYNFNKFYF     CASSLAGFQETQYF   FLRGRAYGL  HLA-A*03:01   \n",
       "1         CAGPAGARLMF      CASSLYVNNGYTF  AVFDRKSDAK  HLA-A*11:01   \n",
       "2      CAVSDHSNNARLMF      CASSLGGTNEQFF   NLVPMVASV  HLA-B*08:01   \n",
       "3       CVVRSGGYNKLIF      CAITGENTGELFF   KLGGALQAK  HLA-A*03:01   \n",
       "4      CAMREFYNTDKLIF    CASSQDDHGSNEQYF    RAKFKQLL  HLA-B*08:01   \n",
       "...               ...                ...         ...          ...   \n",
       "47804   CAATRYNNNDMRF    CASTREGAGWSAQHF   YLSNIIPAL  HLA-A*03:01   \n",
       "47805   CAVKGTGGFKTIF  CASRSQGALYSNQPQHF   HSKKKCDEI  HLA-B*08:01   \n",
       "47806     CAVTPRNQFYF  CASSQDRVGAGANVLTF   KLGGALQAK  HLA-A*03:01   \n",
       "47807   CVVNEKEGNMLTF      CSVGEGLAYEQYF   LLWNGPMAV  HLA-A*02:01   \n",
       "47808  CAVDTGGGADGLTF    CARGQDAANTGELFF  ELAGIGILTV  HLA-A*02:01   \n",
       "\n",
       "                 v_a        j_a          v_b         j_b    y  \n",
       "0          TRAV20*01  TRAJ21*01  TRBV11-2*01  TRBJ2-5*01  0.0  \n",
       "1          TRAV35*01  TRAJ31*01   TRBV7-9*01  TRBJ1-2*01  1.0  \n",
       "2         TRAV8-4*01  TRAJ31*01   TRBV4-1*01  TRBJ2-1*01  0.0  \n",
       "3          TRAV10*01   TRAJ4*01  TRBV10-3*01  TRBJ2-2*01  1.0  \n",
       "4      TRAV14/DV4*01  TRAJ34*01   TRBV4-3*01  TRBJ2-7*01  1.0  \n",
       "...              ...        ...          ...         ...  ...  \n",
       "47804    TRAV13-1*01  TRAJ43*01  TRBV12-3*01  TRBJ1-5*01  0.0  \n",
       "47805     TRAV1-1*01   TRAJ9*01    TRBV28*01  TRBJ1-5*01  0.0  \n",
       "47806     TRAV8-6*01  TRAJ49*01   TRBV4-1*01  TRBJ2-6*01  1.0  \n",
       "47807       TRAV12-1     TRAJ39    TRBV29-01     TRBJ2-7  1.0  \n",
       "47808    TRAV12-2*01  TRAJ45*01     TRBV2*01  TRBJ2-2*01  1.0  \n",
       "\n",
       "[47809 rows x 9 columns]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the train and test data ensuring that they do not include any of the same TCRs\n",
    "tcr.drop_duplicates(inplace=True)\n",
    "train_tcrs, test_tcrs = train_test_split(tcr, test_size=0.10, random_state=seed)\n",
    "\n",
    "train = pd.merge(train_tcrs, data, how='inner')\n",
    "test = pd.merge(test_tcrs, data, how='inner')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Output the train and test sets\n",
    "path = 'data/input/collated/train.csv'\n",
    "train.to_csv(path, index=False)\n",
    "\n",
    "path = 'data/input/collated/test.csv'\n",
    "test.to_csv(path, index=False)"
   ]
  }
 ],
 "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}