Optimize.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run all the cells up to the Experiments Heading to load the algorithm and sample data\n",
    "# The subsequent cell generates the designs that are needed to run the experiments, which will\n",
    "# be dumped into the \"outputs\" folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "import gzip\n",
    "\n",
    "from copy import deepcopy\n",
    "from matplotlib import pyplot as plt\n",
    "import Levenshtein as lev\n",
    "\n",
    "import torch\n",
    "from torch.nn import parallel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Import Data</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./data.pkl\", 'rb') as fin:\n",
    "    mhc1_data, mhc2_data = pickle.load(fin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "779it [00:00, 1287.24it/s]\n",
      "1200it [00:01, 651.71it/s]\n",
      "440it [00:00, 3434.99it/s]\n",
      "779it [00:00, 1260.70it/s]\n",
      "1200it [00:02, 599.11it/s]\n",
      "440it [00:00, 787.40it/s]\n",
      "537it [00:00, 2761.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1043, 298]) torch.Size([1100, 234])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "920it [00:01, 799.18it/s] \n",
      "502it [00:00, 3279.87it/s]\n",
      "537it [00:00, 1616.17it/s]\n",
      "920it [00:01, 891.37it/s] \n",
      "502it [00:00, 1323.38it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3934, 281]) torch.Size([4195, 281])\n"
     ]
    }
   ],
   "source": [
    "# Script for processing data\n",
    "\n",
    "def makeProductDistribution(population, dataset):\n",
    "    _,alleles,_ = dataset\n",
    "    intmap = {}\n",
    "    for i,allele in enumerate(alleles):\n",
    "        intmap[allele] = i\n",
    "    default = len(alleles)\n",
    "    population = [(frozenset([intmap.get(x[1], default) for x in genotype]),p) for genotype, p in population]\n",
    "    diploidDistribution = {}\n",
    "    for i,(haplotype1,p1) in tqdm(enumerate(population), position = 0, leave = True):\n",
    "        if haplotype1 not in diploidDistribution:\n",
    "            diploidDistribution[haplotype1] = 0\n",
    "        diploidDistribution[haplotype1] += p1 ** 2\n",
    "        for (haplotype2,p2) in population[:i]:\n",
    "            diplotype = haplotype1.union(haplotype2)\n",
    "            if diplotype not in diploidDistribution:\n",
    "                diploidDistribution[diplotype] = 0\n",
    "            diploidDistribution[diplotype] += 2 * p1 * p2\n",
    "    return diploidDistribution\n",
    "    \n",
    "def makeDiploidDistribution(populations, dataset):\n",
    "    default = len(dataset[1])\n",
    "    distributions = [makeProductDistribution(population, dataset) for population in populations]\n",
    "    for distribution in distributions[1:]:\n",
    "        for genotype in distribution:\n",
    "            if genotype not in distributions[0]:\n",
    "                distributions[0][genotype] = distribution[genotype]\n",
    "            else:\n",
    "                distributions[0][genotype] += distribution[genotype]\n",
    "    distribution = distributions[0]\n",
    "    indexes = []\n",
    "    weights = []\n",
    "    for genotype in distribution:\n",
    "        weights.append( distribution[genotype]/3 )\n",
    "        genotype = tuple(genotype)\n",
    "        if len(genotype) != 6:\n",
    "            genotype = genotype + ( (default,) * (6-len(genotype)) )\n",
    "        indexes.append(genotype)\n",
    "    indexes = np.array(indexes, dtype = np.int16)\n",
    "    weights = np.array(weights)\n",
    "    return indexes, weights\n",
    "\n",
    "def trimDataset(data, filtered):\n",
    "    data, alleles, peptides = data\n",
    "    filtered = set(filtered)\n",
    "    indexes = []\n",
    "    peptides2 = []\n",
    "    for i, peptide in enumerate(peptides):\n",
    "        if peptide in filtered:\n",
    "            indexes.append(i)\n",
    "            peptides2.append(peptide)\n",
    "    data = data[indexes]\n",
    "    return data, alleles, peptides2\n",
    "    \n",
    "def reformatData(data):\n",
    "    credence_data, binary_data, peptides, population = data\n",
    "    \n",
    "    credence_data = trimDataset(credence_data, peptides)\n",
    "    credence_distribution = makeDiploidDistribution(population, credence_data)\n",
    "    credence_data = torch.tensor(1-credence_data[0], dtype = torch.float32), credence_data[2],\\\n",
    "        credence_distribution[0], torch.tensor(credence_distribution[1], dtype = torch.float32)\n",
    "    \n",
    "    binary_distribution = makeDiploidDistribution(population, binary_data)\n",
    "    binary_data = torch.tensor(1-binary_data[0], dtype = torch.float32), binary_data[2],\\\n",
    "        binary_distribution[0], torch.tensor(binary_distribution[1], dtype = torch.float32)\n",
    "    \n",
    "    print (credence_data[0].size(), binary_data[0].size())\n",
    "    return credence_data, binary_data\n",
    "\n",
    "inputs1_credences, inputs1_binarized = reformatData(mhc1_data)\n",
    "inputs2_credences, inputs2_binarized = reformatData(mhc2_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Greedy Algorithm</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility functions\n",
    "\n",
    "# f(x) = min(n, x)\n",
    "def getThresholdUtility(n):\n",
    "    return torch.tensor(np.arange(0,n+1,1),dtype = torch.float32)\n",
    "\n",
    "def getMarginalImprovement(utility):\n",
    "    return torch.cat( (utility[1:] - utility[:-1], torch.zeros(1)) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#candidates: [candidate, 1 - pMHC hit probability]\n",
    "#columnIndex: [diplotype, allele in diplotype]\n",
    "#columnWeights: [diplotype]\n",
    "#distributions: [dummy, diplotype, distribution]\n",
    "#marginalImprovement: [improvement (shifting from i to i+1, so last entry should be 0)]\n",
    "\n",
    "def evaluateCandidates(candidates, columnIndex, columnWeights, distributions, marginalImprovement):\n",
    "    probabilityOfHit = 1 - torch.prod(candidates[:, columnIndex], dim = 2).unsqueeze(2)\n",
    "    shiftedMass = distributions * probabilityOfHit\n",
    "    improvement = torch.sum( shiftedMass * marginalImprovement, dim = 2)\n",
    "    weightedImprovement = torch.sum( improvement * columnWeights, dim = 1)\n",
    "    return weightedImprovement\n",
    "\n",
    "def updateDistribution(newRow, columnIndex, distributions):\n",
    "    probabilityOfMiss = torch.prod(newRow[columnIndex], dim = 1).reshape(1, -1, 1)\n",
    "    shiftedMass = distributions * (1-probabilityOfMiss)\n",
    "    \n",
    "    convolution = distributions * probabilityOfMiss\n",
    "    convolution[:,:,1:] += shiftedMass[:,:,:-1]\n",
    "    convolution[:,:,-1] += shiftedMass[:,:,-1]\n",
    "    return convolution\n",
    "\n",
    "def evaluateDesign(candidates, seqs, columnIndex, columnWeights, design, utility):\n",
    "    distributions = torch.zeros( (1, len(columnIndex), len(utility)) )\n",
    "    distributions[:, :, 0] = 1\n",
    "    \n",
    "    seqToIndex = {}\n",
    "    for i, seq in enumerate(seqs):\n",
    "        seqToIndex[seq] = i\n",
    "        \n",
    "    for seq in design:\n",
    "        row = seqToIndex[seq]\n",
    "        distributions = updateDistribution(candidates[row], columnIndex, distributions)\n",
    "        \n",
    "    scores = torch.sum( distributions * utility.view(1,1,-1), dim = 2 ).reshape(-1)\n",
    "    return torch.sum(scores * columnWeights).numpy()\n",
    "\n",
    "class evaluateCandidatesModule(torch.nn.Module):\n",
    "    def __init__(self, columnIndex, columnWeights, marginalImprovement, device):\n",
    "        super(evaluateCandidatesModule, self).__init__()\n",
    "        self.device = device\n",
    "        self.columnIndex = columnIndex\n",
    "        self.columnWeights = columnWeights.cuda(self.device)\n",
    "        self.marginalImprovement = marginalImprovement.cuda(self.device)\n",
    "        \n",
    "    def updateDistributions(self, distributions):\n",
    "        self.distributions = distributions.cuda(self.device)\n",
    "        \n",
    "    def forward(self, candidates):\n",
    "        probabilityOfHit = 1 - torch.prod(candidates[:, self.columnIndex], dim = 2).unsqueeze(2)\n",
    "        shiftedMass = self.distributions * probabilityOfHit\n",
    "        improvement = torch.sum( shiftedMass * self.marginalImprovement, dim = 2)\n",
    "        weightedImprovement = torch.sum( improvement * self.columnWeights, dim = 1)\n",
    "        return weightedImprovement.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def greedySelectionMulticore(candidates,\n",
    "                             seqs,\n",
    "                             columnIndex,\n",
    "                             columnWeights,\n",
    "                             designSize,\n",
    "                             marginalImprovement,\n",
    "                             threshold,\n",
    "                             batchSize,\n",
    "                             devices):\n",
    "    \n",
    "    # Set up modules on different devices\n",
    "    modules = [evaluateCandidatesModule(columnIndex, columnWeights, marginalImprovement, device)\n",
    "               for device in devices]\n",
    "    \n",
    "    # Distribute the computation between devices\n",
    "    numRows = candidates.shape[0]\n",
    "    numVertical = (numRows//(len(devices) * batchSize))\n",
    "    sliceSize = (numRows//(len(devices) * numVertical)) + 1\n",
    "    slices = []\n",
    "    z = 0\n",
    "    for _ in range(numVertical):\n",
    "        singleSlice = []\n",
    "        if z*sliceSize >= numRows: break\n",
    "        for device in devices:\n",
    "            if z == numVertical * len(devices) - 1:\n",
    "                singleSlice.append( candidates[z*sliceSize:].cuda(device) )\n",
    "            else:\n",
    "                singleSlice.append( candidates[z*sliceSize:(z+1)*sliceSize].cuda(device) )\n",
    "            z += 1\n",
    "        slices.append(singleSlice)\n",
    "    \n",
    "    # Initialize selected set and score\n",
    "    selectedSet = []\n",
    "    score = 0\n",
    "    selectable = np.ones(numRows)\n",
    "    \n",
    "    # Initialize coverage distributions\n",
    "    distributions = torch.zeros( (1, len(columnIndex), len(marginalImprovement)) )\n",
    "    distributions[:, :, 0] = 1\n",
    "    \n",
    "    numberOfSlices = len(slices)\n",
    "    pbarDescription = \"Sequence added: None, Objective: 0.00000, Delta: 0.00000, Iteration: {}/{}\".format(\n",
    "        \"{}\", numberOfSlices)\n",
    "    with tqdm(range(designSize), position = 0, leave = True) as pbar:\n",
    "        for _ in pbar:\n",
    "            # Update distributions in modules\n",
    "            for module in modules:\n",
    "                module.updateDistributions(distributions)\n",
    "\n",
    "            # Compute marginal utilities\n",
    "            allImprovements = []\n",
    "            # We need to batch the following vector operations due to space limitations\n",
    "            for sliceIndex, singleSlice in enumerate(slices):\n",
    "                improvements = parallel.parallel_apply(modules, singleSlice)\n",
    "                allImprovements.append(torch.cat(improvements))\n",
    "                pbar.set_description(pbarDescription.format(sliceIndex+1))\n",
    "            allImprovements = torch.cat(allImprovements).numpy()\n",
    "\n",
    "            # Argmax\n",
    "            selection = np.argmax(allImprovements * selectable)\n",
    "\n",
    "            # Add best sequence\n",
    "            selectedSeq = seqs[selection]\n",
    "            selectedSet.append(selectedSeq)\n",
    "\n",
    "            # Update score\n",
    "            delta = allImprovements[selection]\n",
    "            score += delta\n",
    "\n",
    "            pbarDescription = \"Sequence added: {}, Objective: {:.5f}, Delta: {:.5f}, Iteration: {}/{}\".format(\n",
    "                selectedSeq, score, delta, \"{}\", numberOfSlices)\n",
    "\n",
    "            # Update distributions for next round\n",
    "            distributions = updateDistribution(candidates[selection], columnIndex, distributions)\n",
    "\n",
    "            # Remove invalid candidates from consideration\n",
    "            for i, seq in enumerate(seqs):\n",
    "                if lev.distance(seq, selectedSeq) <= threshold:\n",
    "                    selectable[i] = 0\n",
    "    \n",
    "    torch.cuda.empty_cache()\n",
    "    return selectedSet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 1, 2, 3, 4, 5, 6, 7]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Number of GPUs available\n",
    "list(range(torch.cuda.device_count()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Experiments</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The arguments to greedySelectionMulticore are as follows:\n",
    "# 1. A tuple consisting of credences for individual allele binding, a list of peptides, a list of diplotypes\n",
    "#   where each diplotype is given as a list of alleles, and a set of weights for each diplotype\n",
    "#     We provide \"inputs#_credences\" and \"inputs#_binarized\" as possible inputs\n",
    "#     which correspond to the credences we derived (see Section 3.2) and credences from Liu et al. respectively.\n",
    "# 2. The number of peptides in the vaccine design\n",
    "# 3. The marginal improvement of the utility as an array. The ith entry of the array should\n",
    "#   contain the marginal improvement for going from i to i+1\n",
    "# 4. The Levenshtein distance threshold. Peptides that are within this threshold of the peptides\n",
    "#   that have already been selected will not be considered for inclusion\n",
    "# 5. The batch size. This is required because of memory limitations on the GPU. The larger this value the better.\n",
    "#   It seems values between 10-50 work relatively well\n",
    "# 6. The list of devices to use. list(range(torch.cuda.device_count()))) should enumerate all available devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate designs using both credences that we derived and using 0-1 binarized credences that\n",
    "# match with those used in Liu et al.\n",
    "# Designs will be dumped in ./outputs/\n",
    "\n",
    "for threshold in range(3, 21):\n",
    "    # MHC Class 1, using credences\n",
    "    x = greedySelectionMulticore(*inputs1_credences,\n",
    "               151,\n",
    "               getMarginalImprovement( getThresholdUtility(threshold) ),\n",
    "               3,\n",
    "               20,\n",
    "               list(range(torch.cuda.device_count())))\n",
    "    with open(\"./outputs/mhc1_credences_threshold_{}.pkl\".format(threshold), 'wb') as fout:\n",
    "        pickle.dump(x, fout)\n",
    "    \n",
    "    # MHC Class 2, using credences\n",
    "    x = greedySelectionMulticore(*inputs2_credences,\n",
    "               151,\n",
    "               getMarginalImprovement( getThresholdUtility(threshold) ),\n",
    "               5,\n",
    "               20,\n",
    "               list(range(torch.cuda.device_count())))\n",
    "    with open(\"./outputs/mhc2_credences_threshold_{}.pkl\".format(threshold), 'wb') as fout:\n",
    "        pickle.dump(x, fout)\n",
    "        \n",
    "    # MHC Class 1, using binarized values\n",
    "    x = greedySelectionMulticore(*inputs1_binarized,\n",
    "               151,\n",
    "               getMarginalImprovement( getThresholdUtility(threshold) ),\n",
    "               3,\n",
    "               20,\n",
    "               list(range(torch.cuda.device_count())))\n",
    "    with open(\"./outputs/mhc1_binarized_threshold_{}.pkl\".format(threshold), 'wb') as fout:\n",
    "        pickle.dump(x, fout)\n",
    "    \n",
    "    # MHC Class 2, using binarized values\n",
    "    x = greedySelectionMulticore(*inputs2_binarized,\n",
    "               151,\n",
    "               getMarginalImprovement( getThresholdUtility(threshold) ),\n",
    "               5,\n",
    "               20,\n",
    "               list(range(torch.cuda.device_count())))\n",
    "    with open(\"./outputs/mhc2_binarized_threshold_{}.pkl\".format(threshold), 'wb') as fout:\n",
    "        pickle.dump(x, fout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}