train_and_evaluate_allelic_inclusion.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "754adf82-7e74-48ca-8f23-94ac6bca7ad7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from tqdm import tqdm\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "import sys\n",
    "import os\n",
    "sys.path.insert(1, os.getcwd() + '/src')\n",
    "from cnn_paired import CNN_Paired\n",
    "from utils import embed_genes, embed_seqs, ALPHABET, embed_sl, embed_dl, train_epoch_cell\n",
    "\n",
    "from scipy.stats import fisher_exact\n",
    "\n",
    "EPOCHS = 10\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "DATA_DIR = os.getcwd() + '/data/'\n",
    "MODEL_DIR = os.getcwd() + '/models/'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59e4dd99-6925-4dd9-8cb9-69d2b4c07377",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "## Load Jaffe et al. BCR repertoires"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0d7a85d0-7d5f-49dc-b312-c36b673a0b33",
   "metadata": {},
   "outputs": [],
   "source": [
    "sl = pd.read_csv(DATA_DIR + 'jaffe/naive_single_light.csv')\n",
    "dl = pd.read_csv(DATA_DIR + 'jaffe/combined_double_light.csv')\n",
    "\n",
    "### Restrict to B cells annotated naive and realistic CDR3 lengths\n",
    "### Longer CDR3H are very rare\n",
    "### Longer CDR3L are likely misannotated\n",
    "\n",
    "donor_celltype = pd.read_csv(DATA_DIR + 'jaffe/donor_and_class.csv').set_index('dataset')\n",
    "sl['donor'] = sl.source.apply(lambda x: donor_celltype.donor.loc[x])\n",
    "dl['donor'] = dl.source.apply(lambda x: donor_celltype.donor.loc[x])\n",
    "sl['celltype'] = sl.source.apply(lambda x: donor_celltype.flow_class.loc[x])\n",
    "dl['celltype'] = dl.source.apply(lambda x: donor_celltype.flow_class.loc[x])\n",
    "dl = dl[dl.celltype == 'naive']\n",
    "sl = sl[sl.celltype == 'naive']\n",
    "\n",
    "dl['length_L_1'] = dl.apply(lambda x: len(x.CDRL3_1), axis=1)\n",
    "dl['length_L_2'] = dl.apply(lambda x: len(x.CDRL3_2), axis=1)\n",
    "dl = dl[(dl.length_L_1 <= 20) & (dl.length_L_2 <= 20)]\n",
    "dl['length_H'] = dl.apply(lambda x: len(x.CDRH3), axis=1)\n",
    "dl = dl[dl.length_H <= 32]\n",
    "\n",
    "sl['length_L'] = sl.apply(lambda x: len(x.CDRL3), axis=1)\n",
    "sl = sl[sl.length_L <= 20]\n",
    "sl['length_H'] = sl.apply(lambda x: len(x.CDRH3), axis=1)\n",
    "sl = sl[sl.length_H <= 32]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "424ff6aa-dd80-4e29-a45e-4890ee4c4460",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Restrict to heavy and light V-genes present in more than 1 donor. \n",
    "### This is true for all J-genes in both heavy and light\n",
    "\n",
    "v_gene_donor_counts = sl.groupby('v_gene_L').apply(lambda x: pd.unique(x.donor).size)\n",
    "v_genes = v_gene_donor_counts[v_gene_donor_counts > 1].index.values\n",
    "v_genes_l = np.array([entry for entry in v_genes if entry[2] in ['K', 'L']])\n",
    "v_genes_l = {v_genes_l[i]:i for i in range(v_genes_l.size)}\n",
    "\n",
    "j_genes_l = np.array(['IGLJ2', 'IGKJ1', 'IGKJ2', 'IGLJ3', 'IGKJ4', 'IGLJ1', 'IGKJ3', 'IGKJ5', 'IGLJ7'])\n",
    "j_genes_l = {j_genes_l[i]:i for i in range(j_genes_l.size)}\n",
    "\n",
    "v_gene_donor_counts = sl.groupby('v_gene_H').apply(lambda x: pd.unique(x.donor).size)\n",
    "v_genes = v_gene_donor_counts[v_gene_donor_counts > 1].index.values\n",
    "v_genes_h = np.array([entry for entry in v_genes if entry[2] in ['H']])\n",
    "v_genes_h = {v_genes_h[i]:i for i in range(v_genes_h.size)}\n",
    "\n",
    "j_genes_h = np.array(['IGHJ1', 'IGHJ2', 'IGHJ3', 'IGHJ4', 'IGHJ5', 'IGHJ6'])\n",
    "j_genes_h = {j_genes_h[i]:i for i in range(j_genes_h.size)}\n",
    "\n",
    "sl = sl[sl.v_gene_H.isin(v_genes_h) & sl.v_gene_L.isin(v_genes_l)]\n",
    "dl = dl[dl.v_gene_H.isin(v_genes_h) & dl.v_gene_L_1.isin(v_genes_l) & dl.v_gene_L_2.isin(v_genes_l)]\n",
    "\n",
    "vj_genes = (v_genes_h, j_genes_h, v_genes_l, j_genes_l)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98a8300e-6394-4d73-8cb5-21250ab9d553",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "## Train model on light chains only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cda04748-f4f2-4e43-93e8-994b5f856851",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 1/10 [00:05<00:46,  5.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [0], Train loss: 0.6849095821380615, Test recall: 0.21439628482972137"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 2/10 [00:10<00:43,  5.43s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1], Train loss: 0.6663867235183716, Test recall: 0.26393188854489164"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███       | 3/10 [00:16<00:40,  5.72s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [2], Train loss: 0.6542438268661499, Test recall: 0.25402476780185756"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 4/10 [00:22<00:34,  5.69s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [3], Train loss: 0.653691828250885, Test recall: 0.2657894736842105"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|█████     | 5/10 [00:28<00:28,  5.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [4], Train loss: 0.6474936008453369, Test recall: 0.27987616099071205"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 6/10 [00:31<00:20,  5.06s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [5], Train loss: 0.6442214250564575, Test recall: 0.2885448916408669"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 70%|███████   | 7/10 [00:36<00:15,  5.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [6], Train loss: 0.6404878497123718, Test recall: 0.29272445820433435"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 8/10 [00:41<00:09,  4.97s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [7], Train loss: 0.641072690486908, Test recall: 0.3086687306501548"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|█████████ | 9/10 [00:46<00:05,  5.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [8], Train loss: 0.640450656414032, Test recall: 0.31222910216718264"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:52<00:00,  5.28s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [9], Train loss: 0.6372255682945251, Test recall: 0.295046439628483"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "### Here we train on 3 donors and hold out the 4th\n",
    "### In the paper we report results by averaging all four possible such models\n",
    "### Here we train only one.\n",
    "### Results will differ slightly, also due to replicate variance\n",
    "\n",
    "test_donor = 4\n",
    "\n",
    "dl_train = dl[dl.donor != test_donor]\n",
    "sl_train = sl[sl.donor != test_donor]\n",
    "\n",
    "dl_test = dl[dl.donor == test_donor]\n",
    "sl_test = sl[sl.donor == test_donor].sample(2 * dl_test.shape[0])\n",
    "sl_test = embed_sl(sl_test, device, vj_genes, drop_heavy=True)\n",
    "(dl_test_1, dl_test_2) = embed_dl(dl_test, device, vj_genes, drop_heavy=True)\n",
    "\n",
    "model = CNN_Paired(use_vj=True, num_v=sl_test[1].shape[1], num_j=sl_test[2].shape[1]).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=5e-3)\n",
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "num_samp = 2 * dl_train.shape[0]\n",
    "\n",
    "for epoch in tqdm(range(EPOCHS)):\n",
    "    sl_train_subs = sl_train.sample(num_samp).reset_index(drop=True)\n",
    "    dl_train = dl_train.sample(frac=1).reset_index(drop=True)\n",
    "    train_epoch_cell(model, device,\n",
    "                     (sl_train_subs, dl_train),\n",
    "                     (sl_test, dl_test_1, dl_test_2),\n",
    "                     optimizer, criterion, epoch, vj_genes,\n",
    "                    drop_heavy=True)\n",
    "    torch.save(model.state_dict(),\n",
    "               MODEL_DIR + 'cnn_dl_noH_{0}.pt'.format(test_donor))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbb398c7-511e-4816-a9d6-bc698e150063",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "## Load and assess on van der Wijst data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "169ccfc5-10bc-49d4-98dd-0e358f33b1ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### Data is prefiltered to only include naive B cells and vj genes in Jaffe data\n",
    "\n",
    "sl_ucsf = pd.read_csv(DATA_DIR + 'van_der_wijst/naive_single_light.csv')\n",
    "dl_ucsf = pd.read_csv(DATA_DIR + 'van_der_wijst/naive_double_light.csv')\n",
    "\n",
    "x_sl, v_sl, j_sl = embed_sl(sl_ucsf, device, vj_genes, drop_heavy=True)\n",
    "embed_dl_tuple = embed_dl(dl_ucsf, device, vj_genes, drop_heavy=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3e8e5915-105d-442c-be0a-f92c02034314",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "scores = model(x_sl, v_sl, j_sl).detach().cpu().numpy()\n",
    "sl_ucsf['scores'] = scores[:, 1] - scores[:, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4ecbb5fc-9d17-45d5-9fe3-f23a845ae5a7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i in range(2):\n",
    "    x, v, j = embed_dl_tuple[i]\n",
    "    scores = model(x, v, j).detach().cpu().numpy()\n",
    "    dl_ucsf['scores_{0}'.format(i+1)] = scores[:, 1] - scores[:, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e83178cc-032f-4596-b433-c037565eda4f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.3655536028119508"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### Fraction of DL cells with sequence worse than 95% of those in SL cells\n",
    "quantile = 0.05\n",
    "thres = sl_ucsf.scores.quantile(quantile)\n",
    "(dl_ucsf[['scores_1', 'scores_2']].min(axis=1) < thres).mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "543d7567-d8ea-4039-b185-3a280f7d7763",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "## Load and assess on Wardemann polyreactivity data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dcaa5012-b632-4d58-b426-b6cd4cbfee9f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data_poly = pd.read_csv(DATA_DIR + 'polyreactivity/wardemann_polyreactivity.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "07076adf-7412-4b20-976e-4b2611335159",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "x, v, j = embed_sl(data_poly, device, vj_genes, drop_heavy=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "355bfd89-c8bd-4a47-9798-d193702a5dc1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "scores = model(x, v, j).detach().cpu().numpy()\n",
    "data_poly['scores'] = scores[:, 1] - scores[:, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "44b0a021-87ed-4cb8-9139-f55d19d970d7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7004864055874283"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "roc_auc_score(data_poly.auto_count <= 1, data_poly.scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a12fff15-5c66-4245-8e66-f05d7ce3c171",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "## Load and assess on expression data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "62502ed3-b3c5-494b-a325-919a99bb3efb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data_expr = pd.read_csv(DATA_DIR + 'expression/expression.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "362f6ac7-99b7-418c-908a-b3f0d12938ea",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "x, v, j = embed_sl(data_expr, device, vj_genes, drop_heavy=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "dd26afe8-f0b9-408f-bcb7-b9286565f0c3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "scores = model(x, v, j).detach().cpu().numpy()\n",
    "data_expr['scores'] = scores[:, 1] - scores[:, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "60b1546a-8cca-492b-a43e-08228b2a4eda",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SignificanceResult(statistic=1.5304585665635886, pvalue=7.727560238765047e-05)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "quantile = 0.03\n",
    "thres = np.quantile(data_expr.scores[data_expr.igm_level == 'HI'], quantile)\n",
    "\n",
    "counts = np.zeros((2, 2))\n",
    "counts[1,0] = ((data_expr.igm_level == 'LO') & (data_expr.scores > thres)).sum()\n",
    "counts[1,1] = ((data_expr.igm_level == 'LO') & (data_expr.scores < thres)).sum()\n",
    "counts[0,0] = ((data_expr.igm_level == 'HI') & (data_expr.scores > thres)).sum()\n",
    "counts[0,1] = ((data_expr.igm_level == 'HI') & (data_expr.scores < thres)).sum()\n",
    "\n",
    "fisher_exact(counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55c7a942-c348-4b32-b70d-208ff1d493f1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "immuno_kernel",
   "language": "python",
   "name": "immuno_kernel"
  },
  "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}