(F) Calculate Number of Isolates with 100x Coverage for Specific Loci.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# This notebook was created to calculate the number of isolates that have 100x coverage in $\\ge 99\\%$ of sites given a certain locus."
   ]
  },
  {
   "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": 2,
   "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",
    "\n",
    "import time\n",
    "import sys\n",
    "import pickle\n",
    "from slurmpy import Slurm\n",
    "import shutil\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",
    "from Bio import Phylo\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",
    "import scipy\n",
    "from collections import Counter\n",
    "\n",
    "import allel # import scikit-allel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "################################################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [1] Write script 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": "markdown",
   "metadata": {},
   "source": [
    "################################################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tag = 'SAMN08732238'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import allel # import scikit-allel\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# genomic data directory\n",
    "rolling_DB_dir = '/n/data1/hms/dbmi/farhat/rollingDB/genomic_data/'\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))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "################################################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [2] Submit jobs to subset full *VCF* files to just the *Position* and *Valid Depth* and store result as Series"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "################################################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load INDEL genotype matrix *Isolate* Annotation Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load isolate annotation file (columns of Genotype Matrix)\n",
    "isolate_annotation_DF = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/Genotypes_Filtered_2/genotypes_isolate_annotation.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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>isolate_ID</th>\n",
       "      <th>lineage_1</th>\n",
       "      <th>lineage_2</th>\n",
       "      <th>lineage_3</th>\n",
       "      <th>lineage_4</th>\n",
       "      <th>lineage_5</th>\n",
       "      <th>lineage_6</th>\n",
       "      <th>lineage_7</th>\n",
       "      <th>lineage_8</th>\n",
       "      <th>lineage_9</th>\n",
       "      <th>lineage_10</th>\n",
       "      <th>lineage_11</th>\n",
       "      <th>lineage_call</th>\n",
       "      <th>group</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SAMN13051687</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>i3</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>2.2.1.1.1.i3</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SAMN09100245</td>\n",
       "      <td>4</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>i3</td>\n",
       "      <td>2</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>4.2.1.2.1.1.i3.2</td>\n",
       "      <td>4B</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SAMN08732238</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>2.2.1.1.1</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SAMN07658260</td>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>3.1.1</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SAMN03648003</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>2.2.1.1.1</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     isolate_ID lineage_1 lineage_2 lineage_3 lineage_4 lineage_5 lineage_6  \\\n",
       "0  SAMN13051687         2         2         1         1         1        i3   \n",
       "1  SAMN09100245         4         2         1         2         1         1   \n",
       "2  SAMN08732238         2         2         1         1         1       NaN   \n",
       "3  SAMN07658260         3         1         1       NaN       NaN       NaN   \n",
       "4  SAMN03648003         2         2         1         1         1       NaN   \n",
       "\n",
       "  lineage_7 lineage_8 lineage_9 lineage_10 lineage_11      lineage_call group  \n",
       "0       NaN       NaN       NaN        NaN        NaN      2.2.1.1.1.i3     2  \n",
       "1        i3         2       NaN        NaN        NaN  4.2.1.2.1.1.i3.2    4B  \n",
       "2       NaN       NaN       NaN        NaN        NaN         2.2.1.1.1     2  \n",
       "3       NaN       NaN       NaN        NaN        NaN             3.1.1     3  \n",
       "4       NaN       NaN       NaN        NaN        NaN         2.2.1.1.1     2  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "isolate_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Filter out isolates that have already been processed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the file names of the depth series (isolates that have been processed)\n",
    "isolate_depth_series_files = os.listdir('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/depth_series')\n",
    "\n",
    "# get the isolate tags for these depth series\n",
    "isolate_IDs_processed = set([isolate_depth_series_file.split('_')[0] for isolate_depth_series_file in isolate_depth_series_files])\n",
    "\n",
    "# get a boolean for the isolates that have NOT been processed\n",
    "isolate_IDs_not_processed_filter = [(isolate_ID not in isolate_IDs_processed) for isolate_ID in isolate_annotation_DF.isolate_ID]\n",
    "\n",
    "# subset isolate annotation for isolates that need to be processed and re-index\n",
    "isolate_annotation_DF = isolate_annotation_DF[isolate_IDs_not_processed_filter]\n",
    "isolate_annotation_DF.reset_index(drop = True, inplace = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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>isolate_ID</th>\n",
       "      <th>lineage_1</th>\n",
       "      <th>lineage_2</th>\n",
       "      <th>lineage_3</th>\n",
       "      <th>lineage_4</th>\n",
       "      <th>lineage_5</th>\n",
       "      <th>lineage_6</th>\n",
       "      <th>lineage_7</th>\n",
       "      <th>lineage_8</th>\n",
       "      <th>lineage_9</th>\n",
       "      <th>lineage_10</th>\n",
       "      <th>lineage_11</th>\n",
       "      <th>lineage_call</th>\n",
       "      <th>group</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SAMN08214998</td>\n",
       "      <td>4</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>4.2.1.2.2.1.1</td>\n",
       "      <td>4B</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     isolate_ID lineage_1 lineage_2 lineage_3 lineage_4 lineage_5 lineage_6  \\\n",
       "0  SAMN08214998         4         2         1         2         2         1   \n",
       "\n",
       "  lineage_7 lineage_8 lineage_9 lineage_10 lineage_11   lineage_call group  \n",
       "0         1       NaN       NaN        NaN        NaN  4.2.1.2.2.1.1    4B  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "isolate_annotation_DF.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 14)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(isolate_annotation_DF)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Subset VCF to POSITION & VALID DEPTH fields for isolates that had Full VCF files "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create BASH scripts to extract VCFs of **N** isolates per job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify N isolates per job\n",
    "N = 600"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "njobs = int( np.ceil( float( len( isolate_annotation_DF.isolate_ID ) ) / float(N) ) ) #number of jobs required if we split for every 600 isolates\n",
    "isolate_ID_series = pd.Series(isolate_annotation_DF.isolate_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "njobs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#genomic data directory\n",
    "rolling_DB_dir = '/n/data1/hms/dbmi/farhat/rollingDB/genomic_data/'\n",
    "\n",
    "bash_scripts = []\n",
    "\n",
    "#split jobs up into tasks of N\n",
    "for first_isolate_i_for_job in np.arange(0 , (njobs)*N , N):\n",
    "\n",
    "    last_isolate_i_for_job = first_isolate_i_for_job + N\n",
    "    \n",
    "    commands_list = []\n",
    "    \n",
    "    for isolate_i in range(first_isolate_i_for_job , last_isolate_i_for_job): \n",
    "\n",
    "        try:\n",
    "            tag = isolate_ID_series[isolate_i]\n",
    "            commands_list.append( 'python /home/rv76/Farhat_Lab/Python_Scripts/homoplasy_project/mmpR_BDQ_mutant_project/subset_fullVCF_to_POS_and_DEPTH.py {0}'.format(tag) )\n",
    "            \n",
    "        except KeyError: #applies only to the last job if it has < N isolates\n",
    "            continue\n",
    "        \n",
    "    bash_scripts.append(commands_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Submit each job to O2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "VCF_1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "submitted: Submitted batch job 25866808\n"
     ]
    }
   ],
   "source": [
    "job_num = 1\n",
    "for job_i_commands_list in bash_scripts:\n",
    "    \n",
    "    #append all commands in a single string to be submitted as a job\n",
    "    VCF_extract_job = ''\n",
    "    for command_i in job_i_commands_list:\n",
    "        VCF_extract_job  = VCF_extract_job  + '\\n' + command_i\n",
    "    \n",
    "    #directory where you want output + error files\n",
    "    os.chdir('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/depth_series/depth_series_jobs/')\n",
    "\n",
    "    job_name = 'VCF_' + str(job_num)\n",
    "\n",
    "    s = Slurm(job_name , {'partition':'short' , 'n':'1' , 't':'0-3:00:00' , 'mem':'4G' , 'mail-type':'FAIL' , 'mail-user':'roger_vargas@g.harvard.edu'})\n",
    "\n",
    "    #submits the job\n",
    "    job_id = s.run(VCF_extract_job)\n",
    "\n",
    "    print job_name\n",
    "    job_num += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "################################################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [3] Write script that takes in two coordinates corresponding to a region of H37Rv and calculates the # of isolates that have *X* coverage in $\\ge 99\\%$ of sites in that region"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "################################################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "H37Rv_start_coord = 778990\n",
    "H37Rv_stop_coord = 779487\n",
    "min_coverage_threshold = 100\n",
    "directory_name = 'mmpR_mod'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os\n",
    "import sys\n",
    "import shutil"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create directory for text files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "directory_for_txt_files = '/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/depth_series/prop_isolates_with_min_coverage_for_locus/{0}/'.format(directory_name)\n",
    "\n",
    "# create directory to store progress txt file & txt file with the number of isolates that have adequate coverate for locus\n",
    "if os.path.exists(directory_for_txt_files):\n",
    "    shutil.rmtree(directory_for_txt_files)\n",
    "    os.makedirs(directory_for_txt_files)\n",
    "elif not os.path.exists(directory_for_txt_files):\n",
    "    os.makedirs(directory_for_txt_files)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### *Function* that calculates proportion of a locus that has adequate coverage for an isolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_proportion_of_interval_covered(H37Rv_start, H37Rv_stop, coverage_threshold, isolate_ID):\n",
    "    '''\n",
    "    This function takes as INPUT start/stop coordinates for a region of H37Rv, a specific coverage threshold and an isolate tag.\n",
    "    The OUTPUT is the proportion of the interval for the specific isolate that had coverage >= coverage threshold.\n",
    "    '''\n",
    "    # load in the series for the depth across the genome\n",
    "    isolate_ID_depth_series = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/depth_series/{0}_POS_VD.pkl'.format(isolate_ID))\n",
    "\n",
    "    # get the valid depth values for a specific region of H37Rv\n",
    "    isolate_ID_depth_interval = isolate_ID_depth_series[(isolate_ID_depth_series.index >= H37Rv_start) & (isolate_ID_depth_series.index <= H37Rv_stop)]\n",
    "\n",
    "    # get the proportion of the interval that has a valid depth >= coverage threshold\n",
    "    prop_interval_covered = ( float(sum(isolate_ID_depth_interval >= coverage_threshold)) / float(len(isolate_ID_depth_interval)) ) * 100.0\n",
    "    \n",
    "    return prop_interval_covered"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load *mixed* INDEL Isolate Annotation File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load isolate annotation file (columns of Genotype Matrix)\n",
    "isolate_annotation_DF = pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/Genotypes_Mixed_only/genotypes_isolate_annotation.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate the # of isolates that have *X* coverage in $\\ge 99\\%$ of sites in the specific region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set threshold for adequate coverage for a single isolate (99% of the interval has adequate coverage)\n",
    "adequate_interval_covered_thresh = 99.0 \n",
    "\n",
    "# get proportion of isolates with adequate coverage in specified interval\n",
    "num_isolates_with_adeq_cov = 0\n",
    "\n",
    "# save isolate IDs for checking if mixed INDELs are present in high coverage isolates\n",
    "isolates_with_adeq_cov_series = pd.Series()\n",
    "\n",
    "for isolate_i, isolate_ID in zip(isolate_annotation_DF.index, isolate_annotation_DF.isolate_ID):\n",
    "    \n",
    "    if get_proportion_of_interval_covered(H37Rv_start_coord, H37Rv_stop_coord, min_coverage_threshold, isolate_ID) >= adequate_interval_covered_thresh:\n",
    "        num_isolates_with_adeq_cov += 1 #increment count by 1\n",
    "        isolates_with_adeq_cov_series.loc[num_isolates_with_adeq_cov] = isolate_ID #store isolate ID in series\n",
    "        \n",
    "    if isolate_i % 500 == 0:\n",
    "        with open('{0}progress_check.txt'.format(directory_for_txt_files), 'a') as myfile:\n",
    "            myfile.write( str( float(isolate_i)/float(len(isolate_annotation_DF.index)) ) + '\\n' )\n",
    "            \n",
    "# once we iterate through all 31,428 isolates, store number of isolates with adequate coverage in txt file\n",
    "file = open('{0}num_isolates_high_coverage.txt'.format(directory_for_txt_files), 'w')\n",
    "file.write(str(num_isolates_with_adeq_cov))\n",
    "file.close()\n",
    "\n",
    "# save the list of Isolates IDs as a pickled series & CSV file\n",
    "isolates_with_adeq_cov_series.to_pickle('{0}isolates_high_coverage.pkl'.format(directory_for_txt_files))\n",
    "isolates_with_adeq_cov_series.to_csv('{0}isolates_high_coverage.csv'.format(directory_for_txt_files))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "################################################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [4] Submit script above as a job to calculate the # of isolates that have 100x coverage in $\\ge 99\\%$ of sites in the *mmpL5-mmpS5-mmpR* operon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "################################################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **mmpL5** H37Rv start: 775586, H37Rv end: 778480\n",
    "- **mmpS5** H37Rv start: 778477, H37Rv end: 778905\n",
    "- **mmpR** H37Rv start: 778990, H37Rv end: 779487\n",
    "- **eis** H37Rv start: 2714124, H37Rv end: \t2715332\n",
    "- **whiB7** H37Rv start: 3568401, H37Rv end: 3568679\n",
    "- **ahpC** H37Rv start: 2726193, H37Rv end: 2726780"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "H37Rv_start_coord = '778990'\n",
    "H37Rv_stop_coord = '779487'\n",
    "min_coverage_threshold = '100'\n",
    "directory_name = 'mmpR'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mmpR\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "submitted: Submitted batch job 29568014\n"
     ]
    }
   ],
   "source": [
    "#create the job\n",
    "num_isolates_adequate_depth_cov_job = 'python /home/rv76/Farhat_Lab/Python_Scripts/homoplasy_project/mmpR_BDQ_mutant_project/prop_isolates_with_adequate_depth_for_locus.py {0} {1} {2} {3}'.format(H37Rv_start_coord, H37Rv_stop_coord, min_coverage_threshold, directory_name)\n",
    "\n",
    "#directory where you want output + error files\n",
    "os.chdir('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/depth_series/prop_isolates_with_min_coverage_for_locus/num_isolates_adequate_coverage_jobs')\n",
    "\n",
    "job_name = directory_name\n",
    "\n",
    "s = Slurm(job_name , {'partition':'medium' , 'n':'1' , 't':'1-12:00:00' , 'mem':'96G' , 'mail-type':'ALL' , 'mail-user':'roger_vargas@g.harvard.edu'})\n",
    "\n",
    "#submits the job\n",
    "job_id = s.run(num_isolates_adequate_depth_cov_job)\n",
    "\n",
    "print job_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "################################################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [5] Number of isolates with 100x coverage in that locus in *mmpR*, *mmpS5* and *mmpL5*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "################################################################################################################################################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "mmpS5_isolates = list(pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/depth_series/prop_isolates_with_min_coverage_for_locus/mmpS5/isolates_high_coverage.pkl'))\n",
    "mmpL5_isolates = list(pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/depth_series/prop_isolates_with_min_coverage_for_locus/mmpL5/isolates_high_coverage.pkl'))\n",
    "mmpR_isolates = list(pd.read_pickle('/n/data1/hms/dbmi/farhat/Roger/mmpR_BDQ_mutant_project/rolling_DB_scrape_indels/depth_series/prop_isolates_with_min_coverage_for_locus/mmpR/isolates_high_coverage.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**mmpR**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7435"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(mmpR_isolates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**mmpS5**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8949"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(mmpS5_isolates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**mmpL5**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6217"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(mmpL5_isolates)"
   ]
  }
 ],
 "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
}