(A) Filter and Process 12 isolates with eis C-14T mutation and AG MICs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This notebook was created to include and process (i.e. genotype) an additional 12 *eis* C-14T mutants with AG MICs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>.container { width:100% !important; }<|/style>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:100% !important; }<|/style>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import vcf\n",
    "\n",
    "%matplotlib inline\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import matplotlib.ticker as ticker\n",
    "from itertools import compress\n",
    "from pylab import MaxNLocator\n",
    "import seaborn as sns; sns.set()\n",
    "from matplotlib.colors import LogNorm\n",
    "from matplotlib import gridspec\n",
    "import ast\n",
    "import itertools\n",
    "import seaborn as sns\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "import fastcluster\n",
    "from sklearn import cluster, datasets\n",
    "import scipy.cluster.hierarchy as hier\n",
    "from sklearn.cluster import KMeans\n",
    "import time\n",
    "import sys\n",
    "\n",
    "import Bio\n",
    "from Bio.Alphabet import IUPAC\n",
    "from Bio.Blast.Applications import NcbiblastnCommandline\n",
    "from Bio.Blast import NCBIXML\n",
    "from Bio.Seq import Seq\n",
    "from Bio.SeqRecord import SeqRecord\n",
    "from Bio.SeqFeature import SeqFeature, FeatureLocation\n",
    "from Bio import pairwise2\n",
    "from Bio import SeqIO\n",
    "from Bio.Graphics import GenomeDiagram\n",
    "from Bio.SeqUtils import GC\n",
    "\n",
    "from Bio.Align.Applications import MuscleCommandline\n",
    "from StringIO import StringIO\n",
    "from Bio import AlignIO\n",
    "from Bio.Align import AlignInfo\n",
    "from Bio.Seq import MutableSeq\n",
    "import itertools\n",
    "import gzip\n",
    "\n",
    "import networkx as nx\n",
    "import scipy\n",
    "import pickle\n",
    "\n",
    "from collections import Counter\n",
    "import allel # import scikit-allel\n",
    "\n",
    "#for exporting to Adobe Illustrator\n",
    "mpl.rcParams['pdf.fonttype'] = 42\n",
    "mpl.rcParams['ps.fonttype'] = 42"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [1] Get Coverage for Isolates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Isolate IDs listed below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "strain_tags = ['168-19','622-19','655-19','IT1070','IT123','IT184','IT230','IT233','IT524','IT634','IT77','IT947','IT952']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function that takes in the path to a full zipped VCF files, extracts just the *POS* and *VD* fields and stores this data as a Series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_depth_series(tag):\n",
    "    \n",
    "    # genomic data directory\n",
    "    rolling_DB_dir = '/n/data1/hms/dbmi/farhat/lfreschi/repos/megapipe/megapipe_snakemake/results_mmpR_eis/'\n",
    "\n",
    "    # path to full zipped VCF\n",
    "    zipped_VCF = rolling_DB_dir + tag + '/pilon/' + tag + '_full.vcf.gz'\n",
    "\n",
    "    # load in the Position & Valid Depth fields from the VCF file\n",
    "    POS_DP_fields = allel.read_vcf(zipped_VCF, fields=['variants/POS', 'variants/DP'])\n",
    "\n",
    "    # create a series with H37Rv Reference Positions as the Index & Depth as the values\n",
    "    depth_series = pd.Series(POS_DP_fields['variants/DP'], index = POS_DP_fields['variants/POS'])\n",
    "\n",
    "    # pickle series for downstream analyses\n",
    "    valid_depth_VCF_dir = '/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/depth_series/'\n",
    "    #depth_series.to_pickle('{0}{1}_POS_VD.pkl'.format(valid_depth_VCF_dir, tag))\n",
    "    \n",
    "    return depth_series"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Iterate through all tag names and get depth series for all strains"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "depth_series_per_isolate_dict = {}\n",
    "for tag in strain_tags:\n",
    "    depth_series_per_isolate_dict[tag] = get_depth_series(tag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_depth_series_per_isolate = [np.mean(depth_series_per_isolate_dict[depth_series_isolate_i]) for depth_series_isolate_i in strain_tags]\n",
    "avg_depth_series_per_isolate = pd.Series(avg_depth_series_per_isolate, index = strain_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "168-19    92.890217\n",
       "622-19    51.134228\n",
       "655-19    91.686259\n",
       "IT1070    58.677926\n",
       "IT123     61.067880\n",
       "IT184     62.892872\n",
       "IT230     18.216280\n",
       "IT233     76.325863\n",
       "IT524     40.771388\n",
       "IT634     28.441705\n",
       "IT77      43.649817\n",
       "IT947     67.702494\n",
       "IT952     58.931473\n",
       "dtype: float64"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "avg_depth_series_per_isolate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- isolate **IT230** from Milan has an average coverage of 18.2x across the genome. One of our filters to call an allele at a SNP site (from the VCF) requires at least a depth of 20x, so we won't be able to call the majority of alleles present at the 790k SNP sites we're analyzing in total, for this isolate, using our allele calling filters.\n",
    "\n",
    "- isolate **63** from Ukraine doesn't pass our contamination filter (used Kraken for this), the number of reads that map to the MTBC is ~80%. That level of contamination could lead to spurious heterozygous calls which would muddy the observations if we're looking at mixed indel calls for this isolate."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [2] Process SNP Genotypes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Side-Note: Some isolate have ~20x coverage and we require 20 reads at a given position to make a SNP call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "strain_tags = ['168-19','622-19','655-19','IT1070','IT123','IT184','IT233','IT524','IT634','IT77','IT947','IT952']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "12"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(strain_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [2.1] This section of the notebook prepares the script that iterates over all of the VCF files in rollingDB and outputs a genotype vector."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import vcf\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import sys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### load SNP annotation file (columns of Genotype Matrix for 31.5k strains) for indexing genotypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "SNV_annotation_DF = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/homoplasy_project/rolling_DB_scrape/Genotypes_Filtered_2/genotypes_SNP_annotation.pkl')\n",
    "\n",
    "#subset to columns that are needed\n",
    "SNV_annotation_DF = SNV_annotation_DF.loc[:, ['pos','ref','alt']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>48</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>64</td>\n",
       "      <td>G</td>\n",
       "      <td>[C]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>67</td>\n",
       "      <td>G</td>\n",
       "      <td>[A, T]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>69</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>71</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   pos ref     alt\n",
       "0   48   C     [T]\n",
       "1   64   G     [C]\n",
       "2   67   G  [A, T]\n",
       "3   69   C     [T]\n",
       "4   71   C     [T]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SNV_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Convert SNV annotation into dictionaries with keys = ROW indices of genotype array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#keys: H37Rv Reference Position , values: Genotype Row Index\n",
    "ref_pos_dict = dict( zip( list(SNV_annotation_DF.pos) , list(SNV_annotation_DF.index) ) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create set of H37Rv Reference Positions that are SNP sites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "H37Rv_ref_positions_of_interest = set(SNV_annotation_DF.pos)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Construct 'skeleton' Numpy Array that will store the supported base for all Reference Positions of interest for a single isolate in rollingDB\n",
    "\n",
    "- **0** for _A_\n",
    "- **1** for _C_\n",
    "- **2** for _G_\n",
    "- **3** for _T_\n",
    "- **9** for Bad Quality SNV calls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#number of Reference Positions we're interested in (where High Quality SNV calls were made)\n",
    "number_of_SNVs = len(ref_pos_dict.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "782565"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "number_of_SNVs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### *Function* to get SNP genotypes array from VCF file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_genotypes_array_from_full_VCF(isolate_ID):\n",
    "    \n",
    "    # INPUT - isolate ID\n",
    "    \n",
    "    ##################################################################\n",
    "    # In constructing the Genotype Vector, we're going to iterate through all Reference Positions of interest and find the supported base from the VCF file for each position\n",
    "    # create an empty numpy array to store genotypes\n",
    "    genotypes_array = np.full(number_of_SNVs , 9 , dtype = int)\n",
    "    \n",
    "    ##################################################################\n",
    "    # Iterate through the VCF file corresponding to the isolate in rollingDB and fill in genotypes vector\n",
    "\n",
    "    #genomic data directory\n",
    "    VCF_genotypes_dir = '/n/data1/hms/dbmi/farhat/lfreschi/repos/megapipe/megapipe_snakemake/results_mmpR_eis/'\n",
    "\n",
    "    #directory that stores files for each sequenced isolate\n",
    "    directory_for_sequenced_isolate = VCF_genotypes_dir + isolate_ID + '/pilon/'\n",
    "\n",
    "    #position-base dictionaries (order of output in Pilon)\n",
    "    base_order = {0:'A' , 1:'C' , 2:'G', 3:'T'}\n",
    "    base_order_r = {'A':0 , 'C':1 , 'G':2 , 'T':3}\n",
    "\n",
    "    #load in VCF file if one was found\n",
    "    vcf_reader = vcf.Reader( open(directory_for_sequenced_isolate  + isolate_ID + '_full.vcf.gz' , 'r') )\n",
    "\n",
    "    #iterate through each Variant Call \n",
    "    for record in vcf_reader:\n",
    "\n",
    "        #check to see if H37Rv Reference Position is one of the SNP sites of interest\n",
    "        if record.POS in H37Rv_ref_positions_of_interest:\n",
    "\n",
    "            #check to see if variant is SNP (length of Reference Allele is 1 base and there's only 1 alternate allele and length of alternate allele is 1 base is alternate allele exists)\n",
    "            if (len(record.REF) == 1) and (len(record.ALT) == 1) and ( (record.ALT[0] == None) or (len(str(record.ALT[0])) == 1) ):\n",
    "\n",
    "                #check to see if the call is a PASS by Pilon\n",
    "                if record.FILTER == []:\n",
    "\n",
    "                    ##### Retrieve Relevant information for filtering quality of Base Call #####\n",
    "                    # Mean Base Quality @ locus\n",
    "                    BQ = record.INFO['BQ']\n",
    "                    # Mean Mapping Quality @ locus\n",
    "                    MQ = record.INFO['MQ']\n",
    "                    # Number of Reads w/ Deletion \n",
    "                    DC = record.INFO['DC']\n",
    "                    # Number of Reads w/ Insertion\n",
    "                    IC = record.INFO['IC']\n",
    "                    # Depth of Valid Reads in Pileup\n",
    "                    VD = record.INFO['DP']\n",
    "\n",
    "                    ### Filtering Criteria\n",
    "                    #---> Mean Base Quality > 20\n",
    "                    #---> Mean Mapping Quality > 30\n",
    "                    #---> No Reads Supporting Insertions\n",
    "                    #---> No Reads Supporting Deletions\n",
    "                    #---> Number of High Quality Reads >= 20\n",
    "                    if (BQ > 20) and (MQ > 30) and (DC == 0) and (IC == 0) and (VD >= 20): \n",
    "\n",
    "                        # Quality-Percentage - Percentage of As, Cs, Gs, Ts weighted by Q & MQ at locus\n",
    "                        QP = record.INFO['QP']\n",
    "\n",
    "                        if (sum(np.array(QP) >= 75) == 1): #1 Allele is supported by at least 75% of the reads, store information for this Base Call\n",
    "\n",
    "                            #create a Boolean list for the Bases that have at least 75% of support from reads\n",
    "                            supported_bases_boolean = list( np.array( QP ) >= 75 )\n",
    "\n",
    "                            #position of Base [A,C,G,T] with at least 75% of support from reads\n",
    "                            supported_bases_i = [i for i,x in enumerate(supported_bases_boolean) if x == True] #returns list of length \n",
    "                            supported_base = base_order[supported_bases_i[0]]\n",
    "\n",
    "                            #write an exception in case a SNV (with a reference position that we're not interested in) 'sneaks' past the subsetting and filtering\n",
    "                            try:\n",
    "                                #store a '0', '1', '2', or '3' in Genotype Vector (Reference Position Index) referencing the Base with at least 90% support from the reads\n",
    "                                genotypes_array[ref_pos_dict[record.POS]] = base_order_r[supported_base]\n",
    "                            except KeyError:\n",
    "                                pass\n",
    "\n",
    "    ##################################################################\n",
    "    # Create directory for the isolate to store output for genotyping\n",
    "\n",
    "    isolate_genotype_dir = '/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes/Genotypes_per_isolate/' + isolate_ID\n",
    "    if os.path.exists(isolate_genotype_dir):\n",
    "        shutil.rmtree(isolate_genotype_dir)\n",
    "        os.makedirs(isolate_genotype_dir)\n",
    "    elif not os.path.exists(isolate_genotype_dir):\n",
    "        os.makedirs(isolate_genotype_dir)\n",
    "\n",
    "    ##################################################################\n",
    "    # Pickle Genotypes Array for downstream analysis\n",
    "    np.save('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes/Genotypes_per_isolate/' + isolate_ID + '/' + isolate_ID + '_genotypes' , genotypes_array , allow_pickle = True)\n",
    "\n",
    "    print(isolate_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "168-19\n",
      "622-19\n",
      "655-19\n",
      "IT1070\n",
      "IT123\n",
      "IT184\n",
      "IT233\n",
      "IT524\n",
      "IT634\n",
      "IT77\n",
      "IT947\n"
     ]
    }
   ],
   "source": [
    "for isolate_ID in strain_tags:\n",
    "    get_genotypes_array_from_full_VCF(isolate_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IT952\n"
     ]
    }
   ],
   "source": [
    "get_genotypes_array_from_full_VCF('IT952')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [2.2] Construct Genotype Matrix (from genotype vectors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import matplotlib.ticker as ticker\n",
    "from itertools import compress\n",
    "from matplotlib.colors import LogNorm\n",
    "from matplotlib import gridspec\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "strain_tags = ['168-19','622-19','655-19','IT1070','IT123','IT184','IT233','IT524','IT634','IT77','IT947','IT952']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "12"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(strain_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make sure all isolates have a corresponding _genotypes_ array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "isolate_id_series = pd.Series(strain_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    168-19\n",
       "1    622-19\n",
       "2    655-19\n",
       "3    IT1070\n",
       "4     IT123\n",
       "dtype: object"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "isolate_id_series.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(12,)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(isolate_id_series)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Concatenate genotypes arrays into a genotypes matrix for all isolates in rollingDB"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Store the __genotypes array__ for each isolate into a list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "genotypes_directory = '/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes/Genotypes_per_isolate/'\n",
    "genotype_arrays_per_isolate_list = []\n",
    "\n",
    "#iterate through the genotypes arrays for each isolate\n",
    "isolate_i = 0\n",
    "for isolate_ID in isolate_id_series:\n",
    "    \n",
    "    #load genotypes array\n",
    "    genotypes_array = np.load(genotypes_directory + isolate_ID + '/' + isolate_ID + '_genotypes.npy')\n",
    "    \n",
    "    #convert genotypes array from dtype 'int64' to 'int8' to save lots of memory\n",
    "    genotypes_array = genotypes_array.astype('int8')\n",
    "    \n",
    "    #store genotypes array for this isolate in a list with all other genotypes arrays\n",
    "    genotype_arrays_per_isolate_list.append(genotypes_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Stack genotype arrays for all isolates to create one __genotypes matrix__ for all isolates "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "genotypes_matrix = np.stack(genotype_arrays_per_isolate_list , axis = -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(782565, 12)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(genotypes_matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 1, 1, ..., 1, 1, 1],\n",
       "       [2, 2, 2, ..., 2, 2, 2],\n",
       "       [2, 2, 2, ..., 2, 2, 2],\n",
       "       ...,\n",
       "       [1, 1, 1, ..., 1, 1, 1],\n",
       "       [1, 1, 1, ..., 1, 1, 1],\n",
       "       [1, 1, 1, ..., 1, 1, 1]], dtype=int8)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genotypes_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9390780 bytes\n"
     ]
    }
   ],
   "source": [
    "print(\"%d bytes\" % (genotypes_matrix.size * genotypes_matrix.itemsize))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Genotypes Matrix is ~9.4 MB large"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save __Genotypes Matrix__ in case kernel crashes when filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "#save Genotypes Matrix\n",
    "np.save('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes/Genotypes_Filtered/genotypes_matrix' , genotypes_matrix , allow_pickle = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load __SNP annotation file__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "SNP_annotation_DF = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/homoplasy_project/rolling_DB_scrape/Genotypes_Filtered_2/genotypes_SNP_annotation.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "      <th>gene_id</th>\n",
       "      <th>gene_pos</th>\n",
       "      <th>gene_category</th>\n",
       "      <th>gene_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>48</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>47.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>64</td>\n",
       "      <td>G</td>\n",
       "      <td>[C]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>63.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>67</td>\n",
       "      <td>G</td>\n",
       "      <td>[A, T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>66.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>69</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>68.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>71</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>70.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   pos ref     alt gene_id  gene_pos gene_category gene_name\n",
       "0   48   C     [T]  Rv0001      47.0     Essential      dnaA\n",
       "1   64   G     [C]  Rv0001      63.0     Essential      dnaA\n",
       "2   67   G  [A, T]  Rv0001      66.0     Essential      dnaA\n",
       "3   69   C     [T]  Rv0001      68.0     Essential      dnaA\n",
       "4   71   C     [T]  Rv0001      70.0     Essential      dnaA"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SNP_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(782565, 7)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(SNP_annotation_DF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [3]  Lineage Typing using SNPs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [3.1] Load/Construct SNP genotype matrix and Annotation Files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Columns of Genotype Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "strain_tags = ['168-19','622-19','655-19','IT1070','IT123','IT184','IT233','IT524','IT634','IT77','IT947','IT952']\n",
    "isolate_annotation_DF = pd.Series(strain_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    168-19\n",
       "1    622-19\n",
       "2    655-19\n",
       "3    IT1070\n",
       "4     IT123\n",
       "dtype: object"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "isolate_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(12,)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(isolate_annotation_DF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rows of Genotype Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load SNP annotation file (rows of Genotype Matrix) with gene annotation information\n",
    "SNP_annotation_DF = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/homoplasy_project/rolling_DB_scrape/Genotypes_Filtered_2/genotypes_SNP_annotation.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "      <th>gene_id</th>\n",
       "      <th>gene_pos</th>\n",
       "      <th>gene_category</th>\n",
       "      <th>gene_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>48</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>47.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>64</td>\n",
       "      <td>G</td>\n",
       "      <td>[C]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>63.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>67</td>\n",
       "      <td>G</td>\n",
       "      <td>[A, T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>66.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>69</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>68.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>71</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>70.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   pos ref     alt gene_id  gene_pos gene_category gene_name\n",
       "0   48   C     [T]  Rv0001      47.0     Essential      dnaA\n",
       "1   64   G     [C]  Rv0001      63.0     Essential      dnaA\n",
       "2   67   G  [A, T]  Rv0001      66.0     Essential      dnaA\n",
       "3   69   C     [T]  Rv0001      68.0     Essential      dnaA\n",
       "4   71   C     [T]  Rv0001      70.0     Essential      dnaA"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SNP_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(782565, 7)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(SNP_annotation_DF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Genotype Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load Genotypes Matrix\n",
    "genotypes_array =  np.load('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes/Genotypes_Filtered/genotypes_matrix.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 1, 1, ..., 1, 1, 1],\n",
       "       [2, 2, 2, ..., 2, 2, 2],\n",
       "       [2, 2, 2, ..., 2, 2, 2],\n",
       "       ...,\n",
       "       [1, 1, 1, ..., 1, 1, 1],\n",
       "       [1, 1, 1, ..., 1, 1, 1],\n",
       "       [1, 1, 1, ..., 1, 1, 1]], dtype=int8)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genotypes_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(782565, 12)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(genotypes_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [3.2] Import  Luca's SNP barcode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Import 96-SNP barcode from Luca's work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_barcode_diagnostic_SNPs = pd.read_csv('/home/lf61/lf61/mic_assemblies/40-full-analysis/lin-sp-var-10k/results/validation_resistant_isolates/snp_scheme_freschi_experimental.tsv' , sep = '\\t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>#lineage</th>\n",
       "      <th>position</th>\n",
       "      <th>gene_coord</th>\n",
       "      <th>allele_change</th>\n",
       "      <th>codon_number</th>\n",
       "      <th>codon_change</th>\n",
       "      <th>aa_change</th>\n",
       "      <th>locus_id</th>\n",
       "      <th>gene_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3.1.1</td>\n",
       "      <td>635139</td>\n",
       "      <td>NaN</td>\n",
       "      <td>G/C</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3</td>\n",
       "      <td>342873</td>\n",
       "      <td>NaN</td>\n",
       "      <td>C/T</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3.2.2</td>\n",
       "      <td>16727</td>\n",
       "      <td>NaN</td>\n",
       "      <td>G/A</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3.2</td>\n",
       "      <td>63850</td>\n",
       "      <td>NaN</td>\n",
       "      <td>T/C</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>3.1</td>\n",
       "      <td>4044872</td>\n",
       "      <td>NaN</td>\n",
       "      <td>G/A</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  #lineage  position  gene_coord allele_change  codon_number codon_change  \\\n",
       "0    3.1.1    635139         NaN           G/C           NaN          NaN   \n",
       "1        3    342873         NaN           C/T           NaN          NaN   \n",
       "2    3.2.2     16727         NaN           G/A           NaN          NaN   \n",
       "3      3.2     63850         NaN           T/C           NaN          NaN   \n",
       "4      3.1   4044872         NaN           G/A           NaN          NaN   \n",
       "\n",
       "  aa_change locus_id gene_name  \n",
       "0       NaN      NaN       NaN  \n",
       "1       NaN      NaN       NaN  \n",
       "2       NaN      NaN       NaN  \n",
       "3       NaN      NaN       NaN  \n",
       "4       NaN      NaN       NaN  "
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_barcode_diagnostic_SNPs.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Drop redundant SNP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>#lineage</th>\n",
       "      <th>position</th>\n",
       "      <th>gene_coord</th>\n",
       "      <th>allele_change</th>\n",
       "      <th>codon_number</th>\n",
       "      <th>codon_change</th>\n",
       "      <th>aa_change</th>\n",
       "      <th>locus_id</th>\n",
       "      <th>gene_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>4.2.1.2.1</td>\n",
       "      <td>63771</td>\n",
       "      <td>NaN</td>\n",
       "      <td>C/T</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     #lineage  position  gene_coord allele_change  codon_number codon_change  \\\n",
       "20  4.2.1.2.1     63771         NaN           C/T           NaN          NaN   \n",
       "\n",
       "   aa_change locus_id gene_name  \n",
       "20       NaN      NaN       NaN  "
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_barcode_diagnostic_SNPs[all_barcode_diagnostic_SNPs.loc[: , '#lineage'] == '4.2.1.2.1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>#lineage</th>\n",
       "      <th>position</th>\n",
       "      <th>gene_coord</th>\n",
       "      <th>allele_change</th>\n",
       "      <th>codon_number</th>\n",
       "      <th>codon_change</th>\n",
       "      <th>aa_change</th>\n",
       "      <th>locus_id</th>\n",
       "      <th>gene_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>4.2.1.2</td>\n",
       "      <td>63771</td>\n",
       "      <td>NaN</td>\n",
       "      <td>C/T</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   #lineage  position  gene_coord allele_change  codon_number codon_change  \\\n",
       "27  4.2.1.2     63771         NaN           C/T           NaN          NaN   \n",
       "\n",
       "   aa_change locus_id gene_name  \n",
       "27       NaN      NaN       NaN  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_barcode_diagnostic_SNPs[all_barcode_diagnostic_SNPs.loc[: , '#lineage'] == '4.2.1.2']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_barcode_diagnostic_SNPs = all_barcode_diagnostic_SNPs[all_barcode_diagnostic_SNPs.loc[: , '#lineage'] != '4.2.1.2']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(96, 9)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(all_barcode_diagnostic_SNPs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Subset DataFrame to SNPs that we have present in our Genotypes Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "barcode_SNPs_in_Genotypes_Matrix = []\n",
    "barcode_SNPs_not_in_Genotypes_Matrix = []\n",
    "\n",
    "#iterate through each of 97 SNPs and find which ones are present in our dataset\n",
    "for SNP_i in all_barcode_diagnostic_SNPs.index:\n",
    "    \n",
    "    if int( all_barcode_diagnostic_SNPs.loc[SNP_i , 'position'] ) in list(SNP_annotation_DF.pos):\n",
    "        \n",
    "        barcode_SNPs_in_Genotypes_Matrix.append(SNP_i)\n",
    "        \n",
    "    else:\n",
    "        \n",
    "        barcode_SNPs_not_in_Genotypes_Matrix.append(SNP_i)\n",
    "\n",
    "#subset to SNPs that we have in Genotypes Matrix\n",
    "barcode_diagnostic_SNPs = all_barcode_diagnostic_SNPs.loc[barcode_SNPs_in_Genotypes_Matrix , :]\n",
    "barcode_diagnostic_SNPs_missing = all_barcode_diagnostic_SNPs.loc[barcode_SNPs_not_in_Genotypes_Matrix , :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>#lineage</th>\n",
       "      <th>position</th>\n",
       "      <th>gene_coord</th>\n",
       "      <th>allele_change</th>\n",
       "      <th>codon_number</th>\n",
       "      <th>codon_change</th>\n",
       "      <th>aa_change</th>\n",
       "      <th>locus_id</th>\n",
       "      <th>gene_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3.1.1</td>\n",
       "      <td>635139</td>\n",
       "      <td>NaN</td>\n",
       "      <td>G/C</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3</td>\n",
       "      <td>342873</td>\n",
       "      <td>NaN</td>\n",
       "      <td>C/T</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3.2.2</td>\n",
       "      <td>16727</td>\n",
       "      <td>NaN</td>\n",
       "      <td>G/A</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3.2</td>\n",
       "      <td>63850</td>\n",
       "      <td>NaN</td>\n",
       "      <td>T/C</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>3.1</td>\n",
       "      <td>4044872</td>\n",
       "      <td>NaN</td>\n",
       "      <td>G/A</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  #lineage  position  gene_coord allele_change  codon_number codon_change  \\\n",
       "0    3.1.1    635139         NaN           G/C           NaN          NaN   \n",
       "1        3    342873         NaN           C/T           NaN          NaN   \n",
       "2    3.2.2     16727         NaN           G/A           NaN          NaN   \n",
       "3      3.2     63850         NaN           T/C           NaN          NaN   \n",
       "4      3.1   4044872         NaN           G/A           NaN          NaN   \n",
       "\n",
       "  aa_change locus_id gene_name  \n",
       "0       NaN      NaN       NaN  \n",
       "1       NaN      NaN       NaN  \n",
       "2       NaN      NaN       NaN  \n",
       "3       NaN      NaN       NaN  \n",
       "4       NaN      NaN       NaN  "
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "barcode_diagnostic_SNPs.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(94, 9)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(barcode_diagnostic_SNPs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>#lineage</th>\n",
       "      <th>position</th>\n",
       "      <th>gene_coord</th>\n",
       "      <th>allele_change</th>\n",
       "      <th>codon_number</th>\n",
       "      <th>codon_change</th>\n",
       "      <th>aa_change</th>\n",
       "      <th>locus_id</th>\n",
       "      <th>gene_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>3.1.2</td>\n",
       "      <td>17842</td>\n",
       "      <td>NaN</td>\n",
       "      <td>G/C</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>4.2.1.1.1.2.1.1.1.1.2</td>\n",
       "      <td>126249</td>\n",
       "      <td>NaN</td>\n",
       "      <td>C/A</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                 #lineage  position  gene_coord allele_change  codon_number  \\\n",
       "8                   3.1.2     17842         NaN           G/C           NaN   \n",
       "17  4.2.1.1.1.2.1.1.1.1.2    126249         NaN           C/A           NaN   \n",
       "\n",
       "   codon_change aa_change locus_id gene_name  \n",
       "8           NaN       NaN      NaN       NaN  \n",
       "17          NaN       NaN      NaN       NaN  "
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "barcode_diagnostic_SNPs_missing.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 9)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(barcode_diagnostic_SNPs_missing)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The original 96 SNP barcode outlined in Luca's analysis corresponds to 96 branches, we have 94 of 96 SNPs in our Genotypes Matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Format DataFrame with 95 lineage-defining SNPs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "barcode_SNP_annot = barcode_diagnostic_SNPs.loc[: , ['#lineage' , 'position' , 'allele_change'] ]\n",
    "\n",
    "#create Reference & Alternate allele columns from Allele Change column\n",
    "barcode_SNP_annot['ref'] = [ref_alt_alleles.split('/')[0] for ref_alt_alleles in list( barcode_SNP_annot.loc[: , 'allele_change'] ) ]\n",
    "barcode_SNP_annot['alt'] = [ref_alt_alleles.split('/')[1] for ref_alt_alleles in list( barcode_SNP_annot.loc[: , 'allele_change'] ) ]\n",
    "\n",
    "#drop unnecessary Allele Change column\n",
    "barcode_SNP_annot = barcode_SNP_annot.loc[: , ['#lineage' , 'position' , 'ref' , 'alt'] ]\n",
    "\n",
    "#rename reference position column\n",
    "barcode_SNP_annot = barcode_SNP_annot.rename(index=str, columns={\"position\": \"pos\"})\n",
    "\n",
    "#reset index\n",
    "barcode_SNP_annot.reset_index(inplace = True , drop = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(94, 4)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(barcode_SNP_annot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>#lineage</th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3.1.1</td>\n",
       "      <td>635139</td>\n",
       "      <td>G</td>\n",
       "      <td>C</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3</td>\n",
       "      <td>342873</td>\n",
       "      <td>C</td>\n",
       "      <td>T</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3.2.2</td>\n",
       "      <td>16727</td>\n",
       "      <td>G</td>\n",
       "      <td>A</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3.2</td>\n",
       "      <td>63850</td>\n",
       "      <td>T</td>\n",
       "      <td>C</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>3.1</td>\n",
       "      <td>4044872</td>\n",
       "      <td>G</td>\n",
       "      <td>A</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  #lineage      pos ref alt\n",
       "0    3.1.1   635139   G   C\n",
       "1        3   342873   C   T\n",
       "2    3.2.2    16727   G   A\n",
       "3      3.2    63850   T   C\n",
       "4      3.1  4044872   G   A"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "barcode_SNP_annot.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a column that contains the diagnostic allele and another that contains the base code for the diagnostic allele"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Note__: The diagnostic allele is the Reference Allele (not the Alternate Allele) for these two SNPs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>#lineage</th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>90</th>\n",
       "      <td>4**</td>\n",
       "      <td>931123</td>\n",
       "      <td>T</td>\n",
       "      <td>C</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   #lineage     pos ref alt\n",
       "90      4**  931123   T   C"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "barcode_SNP_annot[barcode_SNP_annot.pos == 931123]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "barcode_SNP_annot.loc[90 , '#lineage'] = '4'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>#lineage</th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>89</th>\n",
       "      <td>4.2.1.1.1.1.1.1.i2**</td>\n",
       "      <td>1759252</td>\n",
       "      <td>G</td>\n",
       "      <td>T</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                #lineage      pos ref alt\n",
       "89  4.2.1.1.1.1.1.1.i2**  1759252   G   T"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "barcode_SNP_annot[barcode_SNP_annot.pos == 1759252]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "barcode_SNP_annot.loc[89 , '#lineage'] = '4.2.1.1.1.1.1.1.i2'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "#create dictionary for each base - code map\n",
    "base_to_code_map = {'A':0 , 'C':1 , 'G':2 , 'T':3}\n",
    "diag_allele = []\n",
    "diag_allele_code = []\n",
    "\n",
    "for diagnostic_SNP_i in barcode_SNP_annot.index:\n",
    "    \n",
    "    #default that the alternate allele is the diagnostic allele, except for two SNPs in which the diagnostic allele is the reference allele\n",
    "    if barcode_SNP_annot.loc[diagnostic_SNP_i,:].pos in [931123 , 1759252]:\n",
    "        \n",
    "        diagnostic_SNP_i_allele = barcode_SNP_annot.loc[diagnostic_SNP_i,:].ref\n",
    "        \n",
    "    else:\n",
    "        \n",
    "        diagnostic_SNP_i_allele = barcode_SNP_annot.loc[diagnostic_SNP_i,:].alt\n",
    "        \n",
    "        \n",
    "    diag_allele.append(diagnostic_SNP_i_allele)\n",
    "    diag_allele_code.append(base_to_code_map[diagnostic_SNP_i_allele])\n",
    "    \n",
    "#append columns to barcode SNP DF\n",
    "barcode_SNP_annot['diag_allele'] = diag_allele\n",
    "barcode_SNP_annot['diag_allele_code'] = diag_allele_code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>#lineage</th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "      <th>diag_allele</th>\n",
       "      <th>diag_allele_code</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>89</th>\n",
       "      <td>4.2.1.1.1.1.1.1.i2</td>\n",
       "      <td>1759252</td>\n",
       "      <td>G</td>\n",
       "      <td>T</td>\n",
       "      <td>G</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>90</th>\n",
       "      <td>4</td>\n",
       "      <td>931123</td>\n",
       "      <td>T</td>\n",
       "      <td>C</td>\n",
       "      <td>T</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>91</th>\n",
       "      <td>5</td>\n",
       "      <td>1799921</td>\n",
       "      <td>C</td>\n",
       "      <td>A</td>\n",
       "      <td>A</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>92</th>\n",
       "      <td>6</td>\n",
       "      <td>1816587</td>\n",
       "      <td>C</td>\n",
       "      <td>G</td>\n",
       "      <td>G</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>93</th>\n",
       "      <td>7</td>\n",
       "      <td>1137518</td>\n",
       "      <td>G</td>\n",
       "      <td>A</td>\n",
       "      <td>A</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              #lineage      pos ref alt diag_allele  diag_allele_code\n",
       "89  4.2.1.1.1.1.1.1.i2  1759252   G   T           G                 2\n",
       "90                   4   931123   T   C           T                 3\n",
       "91                   5  1799921   C   A           A                 0\n",
       "92                   6  1816587   C   G           G                 2\n",
       "93                   7  1137518   G   A           A                 0"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "barcode_SNP_annot.tail()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [3.3] Assign Isolates to Lineage and Sub-lineages based off of SNP barcode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Subset Genotypes matrix to SNP barcode SNPs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "#subset SNP annotation to the annotation for just the SNP barcode\n",
    "barcode_SNPs_filter = [SNP_i in list(barcode_SNP_annot.pos) for SNP_i in SNP_annotation_DF.pos]\n",
    "barcode_SNPs_in_genotypes_matrix = SNP_annotation_DF[barcode_SNPs_filter]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "      <th>gene_id</th>\n",
       "      <th>gene_pos</th>\n",
       "      <th>gene_category</th>\n",
       "      <th>gene_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>71</th>\n",
       "      <td>324</td>\n",
       "      <td>C</td>\n",
       "      <td>[G]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>323.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4524</th>\n",
       "      <td>15177</td>\n",
       "      <td>C</td>\n",
       "      <td>[T, G]</td>\n",
       "      <td>Rv0013</td>\n",
       "      <td>263.0</td>\n",
       "      <td>Non-Essential</td>\n",
       "      <td>trpG</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4720</th>\n",
       "      <td>15890</td>\n",
       "      <td>G</td>\n",
       "      <td>[A]</td>\n",
       "      <td>Rv0014c</td>\n",
       "      <td>300.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>pknB</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4904</th>\n",
       "      <td>16526</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0014c</td>\n",
       "      <td>936.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>pknB</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4958</th>\n",
       "      <td>16727</td>\n",
       "      <td>G</td>\n",
       "      <td>[A]</td>\n",
       "      <td>Rv0014c</td>\n",
       "      <td>1137.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>pknB</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        pos ref     alt  gene_id  gene_pos  gene_category gene_name\n",
       "71      324   C     [G]   Rv0001     323.0      Essential      dnaA\n",
       "4524  15177   C  [T, G]   Rv0013     263.0  Non-Essential      trpG\n",
       "4720  15890   G     [A]  Rv0014c     300.0      Essential      pknB\n",
       "4904  16526   C     [T]  Rv0014c     936.0      Essential      pknB\n",
       "4958  16727   G     [A]  Rv0014c    1137.0      Essential      pknB"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "barcode_SNPs_in_genotypes_matrix.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(94, 7)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape( barcode_SNPs_in_genotypes_matrix )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Merge lineage/sub-lineage, diagnostic allele and diagnostic allele codes (from the barcode SNP annotation above) information with barcode SNPs annotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "barcode_SNPs_in_genotypes_matrix = barcode_SNPs_in_genotypes_matrix.merge(barcode_SNP_annot.loc[: , ['diag_allele' , 'diag_allele_code' , '#lineage' , 'pos']] , how = 'inner' , on = 'pos')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "      <th>gene_id</th>\n",
       "      <th>gene_pos</th>\n",
       "      <th>gene_category</th>\n",
       "      <th>gene_name</th>\n",
       "      <th>diag_allele</th>\n",
       "      <th>diag_allele_code</th>\n",
       "      <th>#lineage</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>324</td>\n",
       "      <td>C</td>\n",
       "      <td>[G]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>323.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "      <td>G</td>\n",
       "      <td>2</td>\n",
       "      <td>4.2.1.1.1.2.1.1.1.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>15177</td>\n",
       "      <td>C</td>\n",
       "      <td>[T, G]</td>\n",
       "      <td>Rv0013</td>\n",
       "      <td>263.0</td>\n",
       "      <td>Non-Essential</td>\n",
       "      <td>trpG</td>\n",
       "      <td>G</td>\n",
       "      <td>2</td>\n",
       "      <td>1.1.1.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>15890</td>\n",
       "      <td>G</td>\n",
       "      <td>[A]</td>\n",
       "      <td>Rv0014c</td>\n",
       "      <td>300.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>pknB</td>\n",
       "      <td>A</td>\n",
       "      <td>0</td>\n",
       "      <td>2.2.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>16526</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0014c</td>\n",
       "      <td>936.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>pknB</td>\n",
       "      <td>T</td>\n",
       "      <td>3</td>\n",
       "      <td>1.2.1.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>16727</td>\n",
       "      <td>G</td>\n",
       "      <td>[A]</td>\n",
       "      <td>Rv0014c</td>\n",
       "      <td>1137.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>pknB</td>\n",
       "      <td>A</td>\n",
       "      <td>0</td>\n",
       "      <td>3.2.2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     pos ref     alt  gene_id  gene_pos  gene_category gene_name diag_allele  \\\n",
       "0    324   C     [G]   Rv0001     323.0      Essential      dnaA           G   \n",
       "1  15177   C  [T, G]   Rv0013     263.0  Non-Essential      trpG           G   \n",
       "2  15890   G     [A]  Rv0014c     300.0      Essential      pknB           A   \n",
       "3  16526   C     [T]  Rv0014c     936.0      Essential      pknB           T   \n",
       "4  16727   G     [A]  Rv0014c    1137.0      Essential      pknB           A   \n",
       "\n",
       "   diag_allele_code             #lineage  \n",
       "0                 2  4.2.1.1.1.2.1.1.1.2  \n",
       "1                 2              1.1.1.2  \n",
       "2                 0                2.2.2  \n",
       "3                 3              1.2.1.1  \n",
       "4                 0                3.2.2  "
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "barcode_SNPs_in_genotypes_matrix.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Subset Genotypes Matrix to barcode SNPs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "genotypes_array_barcode_SNPs = genotypes_array[ np.array(barcode_SNPs_filter) , :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(94, 12)"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(genotypes_array_barcode_SNPs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 1, 1, ..., 1, 1, 1],\n",
       "       [1, 1, 1, ..., 1, 1, 1],\n",
       "       [2, 2, 2, ..., 2, 2, 2],\n",
       "       ...,\n",
       "       [1, 1, 1, ..., 1, 1, 1],\n",
       "       [2, 2, 2, ..., 2, 2, 2],\n",
       "       [2, 2, 2, ..., 2, 2, 2]], dtype=int8)"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genotypes_array_barcode_SNPs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Iterate through each isolate and assign a lineage/sub-lineage based off of the most specific lineage call made for that isolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "#create DataFrame to store the lineage call for each isolate\n",
    "lineage_calls_for_isolates_DF = pd.DataFrame(index = isolate_annotation_DF.index , columns = ['lineage_call'])\n",
    "\n",
    "for isolate_i in list(isolate_annotation_DF.index):\n",
    "\n",
    "    genotypes_for_isolate = genotypes_array_barcode_SNPs[: , isolate_i]\n",
    "\n",
    "    #find Lineage(s) types by the isolate (match the genotypes for the isolate to the diagnostic alleles from the SNP barcode and get lineages for matching alleles)\n",
    "    lineage_calls = list( barcode_SNPs_in_genotypes_matrix[genotypes_for_isolate == barcode_SNPs_in_genotypes_matrix.diag_allele_code.values].loc[: , '#lineage'] )\n",
    "    \n",
    "    #sort by length of element...global lineage first, then sublineage, then sub-sub-lineage\n",
    "    lineage_calls = sorted(lineage_calls , key=len)\n",
    "    \n",
    "    #keep the 'most specific' lineage call if one exists\n",
    "    if len(lineage_calls) > 0:\n",
    "        \n",
    "        most_specific_lineage_call = lineage_calls[-1]\n",
    "    \n",
    "        #iterate through the lineage & sublineages for most specific lineage call and fill in DF\n",
    "\n",
    "        #check Bovis\n",
    "        if most_specific_lineage_call == 'BOV': \n",
    "            lineage_calls_for_isolates_DF.iloc[isolate_i , 0] = 'BOV'\n",
    "\n",
    "        #check Bovis-Africa\n",
    "        elif most_specific_lineage_call == 'BOV_AFRI': \n",
    "            lineage_calls_for_isolates_DF.iloc[isolate_i , 0] = 'BOV_AFRI'\n",
    "\n",
    "            #check to see if global lineage was also called\n",
    "            try: \n",
    "                lineage_calls_for_isolates_DF.iloc[isolate_i , 0] = lineage_calls[-2]\n",
    "                \n",
    "            except IndexError: #BOV_AFRI sublineage was called but Global Lineage was not\n",
    "                continue\n",
    "    \n",
    "        #any other lineage call can be split into the global lineage & sub-lineages\n",
    "        else:\n",
    "            #iterate through lineage & sub-lineages and fill in DF\n",
    "            lineage_calls_for_isolates_DF.iloc[isolate_i , 0] = most_specific_lineage_call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>lineage_call</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2.2.1.1.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>4.2.1.2.1.1.i4.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2.2.1.1.1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       lineage_call\n",
       "0         2.2.1.1.1\n",
       "1      2.2.1.1.1.i3\n",
       "2  4.2.1.2.1.1.i4.1\n",
       "3      2.2.1.1.1.i3\n",
       "4         2.2.1.1.1"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lineage_calls_for_isolates_DF.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Append Isolate Annotation Series to Lineage/Sub-lineage assignments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "lineage_calls_for_isolates_DF['isolate_ID'] = list( isolate_annotation_DF )\n",
    "isolate_annotation_DF = lineage_calls_for_isolates_DF.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>lineage_call</th>\n",
       "      <th>isolate_ID</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2.2.1.1.1</td>\n",
       "      <td>168-19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "      <td>622-19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>4.2.1.2.1.1.i4.1</td>\n",
       "      <td>655-19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "      <td>IT1070</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2.2.1.1.1</td>\n",
       "      <td>IT123</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       lineage_call isolate_ID\n",
       "0         2.2.1.1.1     168-19\n",
       "1      2.2.1.1.1.i3     622-19\n",
       "2  4.2.1.2.1.1.i4.1     655-19\n",
       "3      2.2.1.1.1.i3     IT1070\n",
       "4         2.2.1.1.1      IT123"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "isolate_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(12, 2)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(isolate_annotation_DF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save Isolate Annotation file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "isolate_annotation_DF.to_pickle('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes/Genotypes_Filtered/genotypes_isolate_annotation.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check to see that all isolates were genotyped for the *diagnostic allele* that types sub-lineage *2.2.1.1.1.i3*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "      <th>gene_id</th>\n",
       "      <th>gene_pos</th>\n",
       "      <th>gene_category</th>\n",
       "      <th>gene_name</th>\n",
       "      <th>diag_allele</th>\n",
       "      <th>diag_allele_code</th>\n",
       "      <th>#lineage</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>90</th>\n",
       "      <td>4238675</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv3792</td>\n",
       "      <td>743.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>aftA</td>\n",
       "      <td>T</td>\n",
       "      <td>3</td>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        pos ref  alt gene_id  gene_pos gene_category gene_name diag_allele  \\\n",
       "90  4238675   C  [T]  Rv3792     743.0     Essential      aftA           T   \n",
       "\n",
       "    diag_allele_code      #lineage  \n",
       "90                 3  2.2.1.1.1.i3  "
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "barcode_SNPs_in_genotypes_matrix[barcode_SNPs_in_genotypes_matrix.loc[:, '#lineage'] == '2.2.1.1.1.i3']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 3, 1, 3, 1, 1, 3, 1, 3, 3, 3, 1], dtype=int8)"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genotypes_array_barcode_SNPs[90, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [4] Process INDEL Genotypes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "strain_tags = ['168-19','622-19','655-19','IT1070','IT123','IT184','IT233','IT524','IT634','IT77','IT947','IT952']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "12"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(strain_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [4.1] This section of the notebook prepares the script that iterates over all of the VCF files in rollingDB and outputs a genotype vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import vcf\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import sys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### load INDEL-index annotation file (columns of Genotype Matrix for 31.5k strains) for indexing genotypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load INDEL annotation file (rows of Genotype Matrix) with gene annotation information\n",
    "INDEL_annotation_DF = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/homoplasy_project/rolling_DB_scrape_indels/Genotypes_Filtered_2/genotypes_INDEL_functional_annotation.pkl')\n",
    "INDEL_annotation_DF.reset_index(inplace = True , drop = False)\n",
    "\n",
    "#subset to columns that are needed\n",
    "INDEL_annotation_DF = INDEL_annotation_DF.loc[:, ['key','pos','ref','alt']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>key</th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ACCGACGAAG_313_A</td>\n",
       "      <td>313</td>\n",
       "      <td>ACCGACGAAG</td>\n",
       "      <td>A</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>TC_1549_T</td>\n",
       "      <td>1549</td>\n",
       "      <td>TC</td>\n",
       "      <td>T</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>T_1552_TAA</td>\n",
       "      <td>1552</td>\n",
       "      <td>T</td>\n",
       "      <td>TAA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>TAA_1552_T</td>\n",
       "      <td>1552</td>\n",
       "      <td>TAA</td>\n",
       "      <td>T</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>T_1552_TA</td>\n",
       "      <td>1552</td>\n",
       "      <td>T</td>\n",
       "      <td>TA</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                key   pos         ref  alt\n",
       "0  ACCGACGAAG_313_A   313  ACCGACGAAG    A\n",
       "1         TC_1549_T  1549          TC    T\n",
       "2        T_1552_TAA  1552           T  TAA\n",
       "3        TAA_1552_T  1552         TAA    T\n",
       "4         T_1552_TA  1552           T   TA"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "INDEL_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Convert INDEL annotation into dictionaries with keys = ROW indices of genotype array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "#keys: INDEL identifier , values: Genotype Row Index\n",
    "INDEL_dict = dict( zip( list(INDEL_annotation_DF.key) , list(INDEL_annotation_DF.index) ) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create set of H37Rv Reference Positions that are INDEL sites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "H37Rv_ref_positions_of_interest = set(INDEL_annotation_DF.pos)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Construct 'skeleton' Numpy Array that will store the supported base for all INDELs of interest for a single isolate in rollingDB\n",
    "\n",
    "- **0** INDEL not called in isolate (but good quality call otherwise)\n",
    "- **1** INDEL called in isolate (Reference & Alternate Alleles match) \n",
    "- **9** for Bad Quality calls\n",
    "\n",
    "#### In constructing the Genotype Vector, we're going to iterate through all Reference Positions of interest and find the supported variant call from the VCF file for each position, assume bad quality Base Call/Missing Data unless evidence from VCF file supports a high-quality variant call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "#number of INDELs we're interested in (where High Quality INDEL calls were made)\n",
    "number_of_INDELs = len(INDEL_dict.keys()) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "50260"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "number_of_INDELs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### *Function* to get INDEL genotypes array from VCF file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_genotypes_array_from_full_VCF(isolate_ID):\n",
    "    \n",
    "    # INPUT - isolate ID\n",
    "    \n",
    "    ##################################################################\n",
    "    # In constructing the Genotype Vector, we're going to iterate through all Reference Positions of interest and find the supported base from the VCF file for each position\n",
    "    #create an empty numpy array to store genotypes\n",
    "    genotypes_array = np.full(number_of_INDELs , 9 , dtype = int)\n",
    "    \n",
    "    ##################################################################\n",
    "    # Iterate through the VCF file corresponding to the isolate in rollingDB and fill in genotypes vector\n",
    "\n",
    "    #genomic data directory\n",
    "    VCF_genotypes_dir = '/n/data1/hms/dbmi/farhat/lfreschi/repos/megapipe/megapipe_snakemake/results_mmpR_eis/'\n",
    "\n",
    "    #directory that stores files for each sequenced isolate\n",
    "    directory_for_sequenced_isolate = VCF_genotypes_dir + isolate_ID + '/pilon/'\n",
    "\n",
    "    #load in VCF file if one was found\n",
    "    vcf_reader = vcf.Reader( open(directory_for_sequenced_isolate  + isolate_ID + '_full.vcf.gz' , 'r') )\n",
    "    \n",
    "    #iterate through each Variant Call \n",
    "    for record in vcf_reader:\n",
    "\n",
    "        #check to see if H37Rv Reference Position is one of the INDEL sites of interest\n",
    "        if record.POS in H37Rv_ref_positions_of_interest:\n",
    "\n",
    "            ref_pos = int( record.POS )\n",
    "            ref_allele = str( record.REF )  #allele on H37Rv\n",
    "            alt_allele = str( record.ALT[0] ) #alternate allele supported by reads\n",
    "            quality_metrics = record.INFO.keys()\n",
    "            variant_i_key = ref_allele + '_' + str(ref_pos) + '_' + alt_allele #variant key, unique identifier for a specific variant\n",
    "\n",
    "            #get the all relevant INDEL keys for possible INDELs @ this reference position\n",
    "            INDEL_keys_at_ref_pos = list( INDEL_annotation_DF[INDEL_annotation_DF.pos == ref_pos].key )\n",
    "\n",
    "             #BOOLEAN for whether INDEL was found\n",
    "            INDEL_FOUND = False\n",
    "\n",
    "            ###################################################################################################\n",
    "            #The variant call at this reference position was high quality and was an INDEL of interest\n",
    "            ###################################################################################################\n",
    "            #check to see if the call is a PASS by Pilon (PASS by Pilon ; there's exactly 1 alternate allele ; call supports the alternate allele)\n",
    "            if (record.FILTER == []) and (len(record.ALT) == 1) and (record.ALT != [None]):\n",
    "\n",
    "                #check to see if variant is INDEL (either Refeference or Alternate Alleles are > 1 base but not both ; 1 allele is 1bp, the other allele is > 1bp)\n",
    "                if ( (len(ref_allele) > 1) or (len(alt_allele) > 1) ) and ( (len(ref_allele) == 1) or (len(alt_allele) == 1) ):\n",
    "\n",
    "                    #keep only INDELs <= 10bp long (1 allele is 1bp ; the other allele is <= 11bp)\n",
    "                    if (len(ref_allele) <= 11) and (len(alt_allele) <= 11):\n",
    "\n",
    "                        #analyze only INDELs with Depth metrics (larger structural variants don't have this)\n",
    "                        if ('MQ' in quality_metrics) and ('DC' in quality_metrics) and ('IC' in quality_metrics) and ('TD' in quality_metrics) and ('DP' in quality_metrics):\n",
    "\n",
    "                            ##### Retrieve Relevant information for filtering quality of Base Call #####\n",
    "                            # Mean Mapping Quality @ locus\n",
    "                            MQ = record.INFO['MQ']\n",
    "                            # Number of Reads w/ Deletion \n",
    "                            DC = record.INFO['DC']\n",
    "                            # Number of Reads w/ Insertion\n",
    "                            IC = record.INFO['IC']\n",
    "                            # Depth of Valid Reads in Pileup\n",
    "                            VD = record.INFO['DP']\n",
    "                            # Total Depth in Pileup\n",
    "                            TD = record.INFO['TD']\n",
    "\n",
    "                            ### Filtering Criteria\n",
    "\n",
    "                            #---> Mean Mapping Quality > 30\n",
    "                            #---> Number of High Quality Reads >= 20\n",
    "                            if (MQ > 30) and (VD >= 20): #INDEL passed filtering criteria!\n",
    "\n",
    "                                # Calculate Alternate Allele Frequency\n",
    "\n",
    "                                #calculate INDEL (Alternate Allele) Frequency from Deletion/Insertion Count & Total Depth (all metrics include ALL reads)\n",
    "                                alt_allele_frequency = float(max(DC , IC)) / float(TD)\n",
    "\n",
    "                                # use only consensus INDELs (where alternate allele frequency >= 75%)\n",
    "                                if alt_allele_frequency >= 0.75:\n",
    "\n",
    "                                    #INDEL passed filters, find relevant INDEL index and store '1' and store '0' in for all other possible INDELs at this reference position\n",
    "                                    if variant_i_key in INDEL_dict.keys():\n",
    "\n",
    "                                        #BOOLEAN for whether INDEL was found\n",
    "                                        INDEL_FOUND = True\n",
    "\n",
    "                                        #store a '1' in Genotype Vector (INDEL index) referencing that the INDEL is supported in this isolate\n",
    "                                        genotypes_array[INDEL_dict[variant_i_key]] = 1\n",
    "\n",
    "                                        #store a '0' in Genotype Vector (INDEL index) for all other INDELs @ this reference position (if there were multiple indels here)\n",
    "                                        INDELs_not_detected = INDEL_keys_at_ref_pos.remove(variant_i_key)\n",
    "\n",
    "                                        #if list of not empty, then there are some INDEL @ this ref pos that were not detected in this isolate\n",
    "                                        if INDELs_not_detected != None: \n",
    "\n",
    "                                            for INDEL_i_key in INDELs_not_detected:\n",
    "\n",
    "                                                genotypes_array[INDEL_dict[INDEL_i_key]] = 0\n",
    "\n",
    "            ###################################################################################################\n",
    "            #The variant call at this reference position was high quality but was not an INDEL of interest\n",
    "            ###################################################################################################\n",
    "            if (INDEL_FOUND == False) and (record.FILTER == []) and ('MQ' in quality_metrics) and ('DP' in quality_metrics):\n",
    "\n",
    "                ##### Retrieve Relevant information for filtering quality of Base Call #####\n",
    "                # Mean Mapping Quality @ locus\n",
    "                MQ = record.INFO['MQ']\n",
    "                # Depth of Valid Reads in Pileup\n",
    "                VD = record.INFO['DP']\n",
    "\n",
    "                ### Filtering Criteria\n",
    "\n",
    "                #---> Mean Mapping Quality > 30\n",
    "                #---> Number of High Quality Reads >= 20\n",
    "                if (MQ > 30) and (VD >= 20): #variant passed filtering criteria!\n",
    "\n",
    "                    #store a '0' in Genotype Vector (INDEL index) for all INDELs @ this reference position\n",
    "                    for INDEL_i_key in INDEL_keys_at_ref_pos:\n",
    "\n",
    "                        genotypes_array[INDEL_dict[INDEL_i_key]] = 0                           \n",
    "                            \n",
    "    ##################################################################\n",
    "    # Create directory for the isolate to store output for genotyping\n",
    "    isolate_genotype_dir = '/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes_indels/Genotypes_per_isolate/' + isolate_ID\n",
    "    if os.path.exists(isolate_genotype_dir):\n",
    "        shutil.rmtree(isolate_genotype_dir)\n",
    "        os.makedirs(isolate_genotype_dir)\n",
    "    elif not os.path.exists(isolate_genotype_dir):\n",
    "        os.makedirs(isolate_genotype_dir)\n",
    "\n",
    "    ##################################################################\n",
    "    # Pickle Genotypes Array for downstream analysis\n",
    "    np.save('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes_indels/Genotypes_per_isolate/' + isolate_ID + '/' + isolate_ID + '_genotypes' , genotypes_array , allow_pickle = True)\n",
    "\n",
    "    print(isolate_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "168-19\n",
      "622-19\n",
      "655-19\n",
      "IT1070\n",
      "IT123\n",
      "IT184\n",
      "IT233\n",
      "IT524\n",
      "IT634\n",
      "IT77\n",
      "IT947\n",
      "IT952\n"
     ]
    }
   ],
   "source": [
    "for isolate_ID in strain_tags:\n",
    "    get_genotypes_array_from_full_VCF(isolate_ID)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [4.2] Construct Genotype Matrix (from genotype vectors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import matplotlib.ticker as ticker\n",
    "from itertools import compress\n",
    "from matplotlib.colors import LogNorm\n",
    "from matplotlib import gridspec\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "strain_tags = ['168-19','622-19','655-19','IT1070','IT123','IT184','IT233','IT524','IT634','IT77','IT947','IT952']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "12"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(strain_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make sure all isolates have a corresponding _genotypes_ array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "isolate_id_series = pd.Series(strain_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    168-19\n",
       "1    622-19\n",
       "2    655-19\n",
       "3    IT1070\n",
       "4     IT123\n",
       "dtype: object"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "isolate_id_series.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(12,)"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(isolate_id_series)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Concatenate genotypes arrays into a genotypes matrix for all isolates in rollingDB"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Store the __genotypes array__ for each isolate into a list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "genotypes_directory = '/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes_indels/Genotypes_per_isolate/'\n",
    "genotype_arrays_per_isolate_list = []\n",
    "\n",
    "#iterate through the genotypes arrays for each isolate\n",
    "isolate_i = 0\n",
    "for isolate_ID in isolate_id_series:\n",
    "    \n",
    "    #load genotypes array\n",
    "    genotypes_array = np.load(genotypes_directory + isolate_ID + '/' + isolate_ID + '_genotypes.npy')\n",
    "    \n",
    "    #convert genotypes array from dtype 'int64' to 'int8' to save lots of memory\n",
    "    genotypes_array = genotypes_array.astype('int8')\n",
    "    \n",
    "    #store genotypes array for this isolate in a list with all other genotypes arrays\n",
    "    genotype_arrays_per_isolate_list.append(genotypes_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Stack genotype arrays for all isolates to create one __genotypes matrix__ for all isolates "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "genotypes_matrix = np.stack(genotype_arrays_per_isolate_list , axis = -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50260, 12)"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(genotypes_matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 0, 0, ..., 0, 0, 0],\n",
       "       [0, 0, 0, ..., 0, 0, 0],\n",
       "       [0, 0, 0, ..., 0, 0, 0],\n",
       "       ...,\n",
       "       [0, 0, 0, ..., 0, 0, 0],\n",
       "       [0, 0, 0, ..., 9, 0, 0],\n",
       "       [0, 0, 0, ..., 0, 0, 0]], dtype=int8)"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genotypes_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "603120 bytes\n"
     ]
    }
   ],
   "source": [
    "print(\"%d bytes\" % (genotypes_matrix.size * genotypes_matrix.itemsize))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Genotypes Matrix is ~0.6 MB large"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save __Genotypes Matrix__ in case kernel crashes when filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "#save Genotypes Matrix\n",
    "np.save('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes_indels/Genotypes_Filtered/genotypes_matrix' , genotypes_matrix , allow_pickle = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load __INDEL annotation file__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load INDEL annotation file (rows of Genotype Matrix) with gene annotation information\n",
    "INDEL_annotation_DF = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/homoplasy_project/rolling_DB_scrape_indels/Genotypes_Filtered_2/genotypes_INDEL_functional_annotation.pkl')\n",
    "INDEL_annotation_DF.reset_index(inplace = True , drop = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>key</th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "      <th>gene_category</th>\n",
       "      <th>gene_name</th>\n",
       "      <th>gene_id</th>\n",
       "      <th>gene_pos</th>\n",
       "      <th>ins_del</th>\n",
       "      <th>INDEL_type</th>\n",
       "      <th>codon_pos</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ACCGACGAAG_313_A</td>\n",
       "      <td>313</td>\n",
       "      <td>ACCGACGAAG</td>\n",
       "      <td>A</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>313</td>\n",
       "      <td>del</td>\n",
       "      <td>inframe</td>\n",
       "      <td>105</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>TC_1549_T</td>\n",
       "      <td>1549</td>\n",
       "      <td>TC</td>\n",
       "      <td>T</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>Rv0001_Rv0002</td>\n",
       "      <td>None</td>\n",
       "      <td>del</td>\n",
       "      <td>frameshift</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>T_1552_TAA</td>\n",
       "      <td>1552</td>\n",
       "      <td>T</td>\n",
       "      <td>TAA</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>Rv0001_Rv0002</td>\n",
       "      <td>None</td>\n",
       "      <td>ins</td>\n",
       "      <td>frameshift</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>TAA_1552_T</td>\n",
       "      <td>1552</td>\n",
       "      <td>TAA</td>\n",
       "      <td>T</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>Rv0001_Rv0002</td>\n",
       "      <td>None</td>\n",
       "      <td>del</td>\n",
       "      <td>frameshift</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>T_1552_TA</td>\n",
       "      <td>1552</td>\n",
       "      <td>T</td>\n",
       "      <td>TA</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>Rv0001_Rv0002</td>\n",
       "      <td>None</td>\n",
       "      <td>ins</td>\n",
       "      <td>frameshift</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                key   pos         ref  alt gene_category gene_name  \\\n",
       "0  ACCGACGAAG_313_A   313  ACCGACGAAG    A     Essential      dnaA   \n",
       "1         TC_1549_T  1549          TC    T          None      None   \n",
       "2        T_1552_TAA  1552           T  TAA          None      None   \n",
       "3        TAA_1552_T  1552         TAA    T          None      None   \n",
       "4         T_1552_TA  1552           T   TA          None      None   \n",
       "\n",
       "         gene_id gene_pos ins_del  INDEL_type codon_pos  \n",
       "0         Rv0001      313     del     inframe       105  \n",
       "1  Rv0001_Rv0002     None     del  frameshift      None  \n",
       "2  Rv0001_Rv0002     None     ins  frameshift      None  \n",
       "3  Rv0001_Rv0002     None     del  frameshift      None  \n",
       "4  Rv0001_Rv0002     None     ins  frameshift      None  "
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "INDEL_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50260, 11)"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(INDEL_annotation_DF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [5] How to load SNP & INDEL genotypes for extra strains"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Isolate Annotation DataFrame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Columns of Genotype Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [],
   "source": [
    "isolate_annotation_DF = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes/Genotypes_Filtered/genotypes_isolate_annotation.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>lineage_call</th>\n",
       "      <th>isolate_ID</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2.2.1.1.1</td>\n",
       "      <td>168-19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "      <td>622-19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>4.2.1.2.1.1.i4.1</td>\n",
       "      <td>655-19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "      <td>IT1070</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2.2.1.1.1</td>\n",
       "      <td>IT123</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>2.2.1.1.1</td>\n",
       "      <td>IT184</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "      <td>IT233</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>4.2.1.1.1.1.1.1.i2</td>\n",
       "      <td>IT524</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "      <td>IT634</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "      <td>IT77</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "      <td>IT947</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>4.2.1.1.1.1.1.1.i2</td>\n",
       "      <td>IT952</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          lineage_call isolate_ID\n",
       "0            2.2.1.1.1     168-19\n",
       "1         2.2.1.1.1.i3     622-19\n",
       "2     4.2.1.2.1.1.i4.1     655-19\n",
       "3         2.2.1.1.1.i3     IT1070\n",
       "4            2.2.1.1.1      IT123\n",
       "5            2.2.1.1.1      IT184\n",
       "6         2.2.1.1.1.i3      IT233\n",
       "7   4.2.1.1.1.1.1.1.i2      IT524\n",
       "8         2.2.1.1.1.i3      IT634\n",
       "9         2.2.1.1.1.i3       IT77\n",
       "10        2.2.1.1.1.i3      IT947\n",
       "11  4.2.1.1.1.1.1.1.i2      IT952"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "isolate_annotation_DF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(12, 2)"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(isolate_annotation_DF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SNP Genotypes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rows of Genotype Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load SNP annotation file (rows of Genotype Matrix) with gene annotation information\n",
    "SNP_annotation_DF = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/homoplasy_project/rolling_DB_scrape/Genotypes_Filtered_2/genotypes_SNP_annotation.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "      <th>gene_id</th>\n",
       "      <th>gene_pos</th>\n",
       "      <th>gene_category</th>\n",
       "      <th>gene_name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>48</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>47.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>64</td>\n",
       "      <td>G</td>\n",
       "      <td>[C]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>63.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>67</td>\n",
       "      <td>G</td>\n",
       "      <td>[A, T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>66.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>69</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>68.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>71</td>\n",
       "      <td>C</td>\n",
       "      <td>[T]</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>70.0</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   pos ref     alt gene_id  gene_pos gene_category gene_name\n",
       "0   48   C     [T]  Rv0001      47.0     Essential      dnaA\n",
       "1   64   G     [C]  Rv0001      63.0     Essential      dnaA\n",
       "2   67   G  [A, T]  Rv0001      66.0     Essential      dnaA\n",
       "3   69   C     [T]  Rv0001      68.0     Essential      dnaA\n",
       "4   71   C     [T]  Rv0001      70.0     Essential      dnaA"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SNP_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(782565, 7)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(SNP_annotation_DF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Genotype Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load Genotypes Matrix\n",
    "genotypes_array =  np.load('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes/Genotypes_Filtered/genotypes_matrix.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 1, 1, ..., 1, 1, 1],\n",
       "       [2, 2, 2, ..., 2, 2, 2],\n",
       "       [2, 2, 2, ..., 2, 2, 2],\n",
       "       ...,\n",
       "       [1, 1, 1, ..., 1, 1, 1],\n",
       "       [1, 1, 1, ..., 1, 1, 1],\n",
       "       [1, 1, 1, ..., 1, 1, 1]], dtype=int8)"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genotypes_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(782565, 12)"
      ]
     },
     "execution_count": 95,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(genotypes_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### INDEL Genotypes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rows of Genotype Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load INDEL annotation file (rows of Genotype Matrix) with gene annotation information\n",
    "INDEL_annotation_DF = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/homoplasy_project/rolling_DB_scrape_indels/Genotypes_Filtered_2/genotypes_INDEL_functional_annotation.pkl')\n",
    "INDEL_annotation_DF.reset_index(inplace = True , drop = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>key</th>\n",
       "      <th>pos</th>\n",
       "      <th>ref</th>\n",
       "      <th>alt</th>\n",
       "      <th>gene_category</th>\n",
       "      <th>gene_name</th>\n",
       "      <th>gene_id</th>\n",
       "      <th>gene_pos</th>\n",
       "      <th>ins_del</th>\n",
       "      <th>INDEL_type</th>\n",
       "      <th>codon_pos</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ACCGACGAAG_313_A</td>\n",
       "      <td>313</td>\n",
       "      <td>ACCGACGAAG</td>\n",
       "      <td>A</td>\n",
       "      <td>Essential</td>\n",
       "      <td>dnaA</td>\n",
       "      <td>Rv0001</td>\n",
       "      <td>313</td>\n",
       "      <td>del</td>\n",
       "      <td>inframe</td>\n",
       "      <td>105</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>TC_1549_T</td>\n",
       "      <td>1549</td>\n",
       "      <td>TC</td>\n",
       "      <td>T</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>Rv0001_Rv0002</td>\n",
       "      <td>None</td>\n",
       "      <td>del</td>\n",
       "      <td>frameshift</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>T_1552_TAA</td>\n",
       "      <td>1552</td>\n",
       "      <td>T</td>\n",
       "      <td>TAA</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>Rv0001_Rv0002</td>\n",
       "      <td>None</td>\n",
       "      <td>ins</td>\n",
       "      <td>frameshift</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>TAA_1552_T</td>\n",
       "      <td>1552</td>\n",
       "      <td>TAA</td>\n",
       "      <td>T</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>Rv0001_Rv0002</td>\n",
       "      <td>None</td>\n",
       "      <td>del</td>\n",
       "      <td>frameshift</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>T_1552_TA</td>\n",
       "      <td>1552</td>\n",
       "      <td>T</td>\n",
       "      <td>TA</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>Rv0001_Rv0002</td>\n",
       "      <td>None</td>\n",
       "      <td>ins</td>\n",
       "      <td>frameshift</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                key   pos         ref  alt gene_category gene_name  \\\n",
       "0  ACCGACGAAG_313_A   313  ACCGACGAAG    A     Essential      dnaA   \n",
       "1         TC_1549_T  1549          TC    T          None      None   \n",
       "2        T_1552_TAA  1552           T  TAA          None      None   \n",
       "3        TAA_1552_T  1552         TAA    T          None      None   \n",
       "4         T_1552_TA  1552           T   TA          None      None   \n",
       "\n",
       "         gene_id gene_pos ins_del  INDEL_type codon_pos  \n",
       "0         Rv0001      313     del     inframe       105  \n",
       "1  Rv0001_Rv0002     None     del  frameshift      None  \n",
       "2  Rv0001_Rv0002     None     ins  frameshift      None  \n",
       "3  Rv0001_Rv0002     None     del  frameshift      None  \n",
       "4  Rv0001_Rv0002     None     ins  frameshift      None  "
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "INDEL_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50260, 11)"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(INDEL_annotation_DF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Genotype Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load Genotypes Matrix\n",
    "genotypes_array =  np.load('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/eis_promoter_mutant_strains_from_Milan/Genotypes_indels/Genotypes_Filtered/genotypes_matrix.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 0, 0, ..., 0, 0, 0],\n",
       "       [0, 0, 0, ..., 0, 0, 0],\n",
       "       [0, 0, 0, ..., 0, 0, 0],\n",
       "       ...,\n",
       "       [0, 0, 0, ..., 0, 0, 0],\n",
       "       [0, 0, 0, ..., 9, 0, 0],\n",
       "       [0, 0, 0, ..., 0, 0, 0]], dtype=int8)"
      ]
     },
     "execution_count": 100,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genotypes_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50260, 12)"
      ]
     },
     "execution_count": 101,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(genotypes_array)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}