process_bams.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "#%matplotlib inline \n",
    "\n",
    "import warnings\n",
    "# both loading plastid and one of the lines below provoke some warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "from plastid import *\n",
    "from twobitreader import TwoBitFile\n",
    "from collections import defaultdict\n",
    "import six\n",
    "\n",
    "#to make executable\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "##All files are defined by their path rather than by a prefix.\n",
    "os.system('pwd')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# what are the ensembl ids for these genes?\n",
    "targeted_genes = pd.read_csv('target_gene_list.txt', header=None, names=['gene_name'])\n",
    "feature_names = pd.read_csv('./references/features.tsv.gz',\n",
    "                            sep='\\t',\n",
    "                            header=None,\n",
    "                            names=['gene_id', 'gene_name', 'feature_type'])\n",
    "targeted_genes = feature_names[feature_names['gene_name'].isin(targeted_genes['gene_name'])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Missing:\n",
      "Empty DataFrame\n",
      "Columns: [gene_name]\n",
      "Index: []\n"
     ]
    }
   ],
   "source": [
    "#check for genes that were not assigned an ensembl id\n",
    "tmp = pd.read_csv('target_gene_list.txt', header=None, names=['gene_name'])\n",
    "print('Missing:')\n",
    "print(pd.DataFrame(tmp)[~pd.DataFrame(tmp)['gene_name'].isin(targeted_genes['gene_name'])])\n",
    "#The gene names need to be updated for these genes. Search ensembl website to find desired gene name.\n",
    "#Append these genes MANUALLY to the end of target_gene_list.txt and rerun this script.\n",
    "#These genes must be added to target_gene_list.txt because this list is used by pick_targets.ipynb script.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write these to a text document\n",
    "np.savetxt('targeted_gene_ids.csv', targeted_genes['gene_id'].values, fmt='%s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# # extract the relevant transcripts from the genome reference\n",
    "os.system('grep -f targeted_gene_ids.csv /nvme/indices/refdata-cellranger-GRCh38-1.2.0/genes/genes.gtf > ./references/targeted_genes.gtf')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# # convert genome to 2bit format\n",
    "os.system('faToTwoBit /nvme/indices/refdata-cellranger-GRCh38-1.2.0/fasta/genome.fa ./references/genome.2bit')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this line produces a lot of warnings about duplicate tags...\n",
    "annotated_transcripts = {transcript.attr['transcript_id']: transcript for transcript \n",
    "                             in GTF2_TranscriptAssembler('./references/targeted_genes.gtf',return_type=Transcript)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make a table of all the unique transcripts for each gene\n",
    "targeted_transcripts = pd.Series({transcript_id: transcript.get_gene() for transcript_id, transcript in annotated_transcripts.iteritems()}).reset_index()\n",
    "targeted_transcripts.columns = ['transcript_id', 'gene_id']\n",
    "\n",
    "name_mapper = dict(zip(targeted_genes['gene_id'], targeted_genes['gene_name']))\n",
    "id_mapper = dict(zip(targeted_genes['gene_name'], targeted_genes['gene_id']))\n",
    "\n",
    "targeted_transcripts['gene_name'] = targeted_transcripts['gene_id'].map(name_mapper)\n",
    "targeted_transcripts = targeted_transcripts.sort_values('gene_name').reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct dictionary that groups transcript models by gene\n",
    "transcripts_by_gene = defaultdict(dict)\n",
    "\n",
    "for name, x in targeted_transcripts.iterrows():\n",
    "    transcript_id = x['transcript_id']\n",
    "    transcripts_by_gene[x['gene_name']][transcript_id] = annotated_transcripts[transcript_id]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find spanning segments that cover all transcripts for a particular gene\n",
    "spanning_segments = dict()\n",
    "\n",
    "# we will add some extra space at the end just in case 3' UTRs aren't fully annotated\n",
    "utr3_offset = 200\n",
    "\n",
    "for gene, transcripts in transcripts_by_gene.iteritems():\n",
    "    # note the spanning segments include annotated 3' UTRs\n",
    "    start = np.min([x.spanning_segment.start for x in transcripts.values()])\n",
    "    end = np.max([x.spanning_segment.end for x in transcripts.values()])\n",
    "    \n",
    "    sense = transcripts.values()[0].strand\n",
    "    \n",
    "    if sense == '+':\n",
    "        end = end + utr3_offset\n",
    "    elif sense == '-':\n",
    "        start = start - utr3_offset\n",
    "    \n",
    "    spanning_segments[gene] = SegmentChain(GenomicSegment(transcripts.values()[0].chrom,\n",
    "                                                          start, end,\n",
    "                                                          transcripts.values()[0].strand),\n",
    "                                          gene_id=id_mapper[gene])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cPickle as pickle\n",
    "\n",
    "def save_obj(obj, name ):\n",
    "    with open(name + '.pkl', 'wb') as f:\n",
    "        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "def load_obj(name ):\n",
    "    with open(name, 'rb') as f:\n",
    "        return pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_obj(spanning_segments, './counts/spanning_segments')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "bam_list = {'CR2_l{0}'.format(i): '/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l{0}/outs/possorted_genome_bam.bam'.format(i) for i in range(1,7)}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l4/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l5/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l6/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l1/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l2/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l3/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n",
      "================================================\n",
      "\tTOMM7\n",
      "\tASNA1\n",
      "\tTUBB\n",
      "\tORC1\n",
      "\tFOXO3\n",
      "\tRPS5\n",
      "\tATF4\n",
      "\tCDH3\n",
      "\tSEC61A1\n"
     ]
    }
   ],
   "source": [
    "for pretty_name, bam_file in bam_list.iteritems():\n",
    "    alignments = BAMGenomeArray(bam_file)\n",
    "    print(bam_file)\n",
    "    print('================================================')\n",
    "    \n",
    "    # this will put weight 1/N on each nucleotide covered by a read, where N is the length of the read\n",
    "    alignments.set_mapping(CenterMapFactory())\n",
    "\n",
    "    # now accumulate read density\n",
    "    count_vectors = dict()\n",
    "    for gene, segment in spanning_segments.iteritems():\n",
    "        print('\\t{0}'.format(gene))\n",
    "        count_vectors[gene] = segment.get_counts(alignments)\n",
    "        \n",
    "    save_obj(count_vectors, './counts/count_vectors_{0}.count'.format(pretty_name))\n",
    "\n",
    "    # this will put weight 1 at the 5' end of the mapped read\n",
    "    alignments.set_mapping(FivePrimeMapFactory())\n",
    "\n",
    "    print('================================================')\n",
    "    # now accumulate read density\n",
    "    start_count_vectors = dict()\n",
    "    for gene, segment in spanning_segments.iteritems():\n",
    "        print('\\t{0}'.format(gene))\n",
    "        start_count_vectors[gene] = segment.get_counts(alignments)\n",
    "        \n",
    "    save_obj(start_count_vectors, './counts/start_count_vectors_{0}.count'.format(pretty_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l4/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l5/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l6/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l1/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l2/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l3/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n"
     ]
    }
   ],
   "source": [
    "# this is transcript-resolved mappings 5' end mappings, though with short reads many of these will end up randomly assigned\n",
    "for pretty_name, bam_file in bam_list.iteritems():\n",
    "    alignments = BAMGenomeArray(bam_file)\n",
    "    alignments.set_mapping(FivePrimeMapFactory())\n",
    "    \n",
    "    print(bam_file)\n",
    "    print('================================================')\n",
    "\n",
    "    transcript_count_vectors = defaultdict(dict)\n",
    "\n",
    "    for gene, transcripts in transcripts_by_gene.iteritems():\n",
    "        print('\\t{0}'.format(gene))\n",
    "        for transcript_id, transcript in transcripts.iteritems():\n",
    "            print('\\t\\t{0}'.format(transcript_id))\n",
    "            transcript_count_vectors[gene][transcript_id] = transcript.get_counts(alignments)\n",
    "            \n",
    "    save_obj(transcript_count_vectors, './counts/transcript_count_vectors_{0}.count'.format(pretty_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l4/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l5/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l6/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l1/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l2/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n",
      "/home/jmr/direct_capture/doubles/perturbseq/count/CR2_l3/outs/possorted_genome_bam.bam\n",
      "================================================\n",
      "\tTOMM7\n",
      "\t\tENST00000483581\n",
      "\t\tENST00000372879\n",
      "\t\tENST00000463284\n",
      "\t\tENST00000358435\n",
      "\t\tENST00000405021\n",
      "\t\tENST00000496129\n",
      "\t\tENST00000621567\n",
      "\tASNA1\n",
      "\t\tENST00000590633\n",
      "\t\tENST00000586561\n",
      "\t\tENST00000611688\n",
      "\t\tENST00000357332\n",
      "\t\tENST00000591090\n",
      "\tTUBB\n",
      "\t\tENST00000396389\n",
      "\t\tENST00000327892\n",
      "\t\tENST00000396384\n",
      "\t\tENST00000330914\n",
      "\tORC1\n",
      "\t\tENST00000371566\n",
      "\t\tENST00000371568\n",
      "\tFOXO3\n",
      "\t\tENST00000406360\n",
      "\t\tENST00000540898\n",
      "\t\tENST00000343882\n",
      "\tRPS5\n",
      "\t\tENST00000598807\n",
      "\t\tENST00000598495\n",
      "\t\tENST00000596046\n",
      "\t\tENST00000196551\n",
      "\t\tENST00000598098\n",
      "\t\tENST00000596314\n",
      "\t\tENST00000601521\n",
      "\t\tENST00000599909\n",
      "\t\tENST00000599232\n",
      "\tATF4\n",
      "\t\tENST00000396680\n",
      "\t\tENST00000404241\n",
      "\t\tENST00000337304\n",
      "\tCDH3\n",
      "\t\tENST00000569080\n",
      "\t\tENST00000569036\n",
      "\t\tENST00000567674\n",
      "\t\tENST00000542274\n",
      "\t\tENST00000569117\n",
      "\t\tENST00000429102\n",
      "\t\tENST00000264012\n",
      "\t\tENST00000568292\n",
      "\t\tENST00000566808\n",
      "\t\tENST00000565453\n",
      "\tSEC61A1\n",
      "\t\tENST00000498837\n",
      "\t\tENST00000491668\n",
      "\t\tENST00000243253\n",
      "\t\tENST00000481210\n",
      "\t\tENST00000424880\n",
      "\t\tENST00000464451\n",
      "\t\tENST00000483956\n"
     ]
    }
   ],
   "source": [
    "for pretty_name, bam_file in bam_list.iteritems():\n",
    "    alignments = BAMGenomeArray(bam_file)\n",
    "    alignments.set_mapping(CenterMapFactory())\n",
    "    \n",
    "    print(bam_file)\n",
    "    print('================================================')\n",
    "\n",
    "    transcript_count_vectors = defaultdict(dict)\n",
    "\n",
    "    for gene, transcripts in transcripts_by_gene.iteritems():\n",
    "        print('\\t{0}'.format(gene))\n",
    "        for transcript_id, transcript in transcripts.iteritems():\n",
    "            print('\\t\\t{0}'.format(transcript_id))\n",
    "            transcript_count_vectors[gene][transcript_id] = transcript.get_counts(alignments)\n",
    "            \n",
    "    save_obj(transcript_count_vectors, './counts/transcript_count_vectors_{0}.count'.format(pretty_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_files(pattern):\n",
    "    file_list = glob.glob(pattern)\n",
    "    all_count_vectors = dict()\n",
    "    \n",
    "    for file_name in file_list:\n",
    "        all_count_vectors[file_name] = load_obj(file_name)\n",
    "    return all_count_vectors\n",
    "\n",
    "def merge_count_files(pattern):\n",
    "    data = get_files(pattern)\n",
    "    \n",
    "    merged = dict()\n",
    "\n",
    "    # initialize dictionary with values from first file\n",
    "    first = data.pop(data.keys()[0])\n",
    "    for gene_name, values in first.iteritems():\n",
    "        merged[gene_name] = values.copy()\n",
    "\n",
    "    # add on the rest\n",
    "    for file, arrays in data.iteritems():\n",
    "        for gene_name, values in arrays.iteritems():\n",
    "            merged[gene_name] += values\n",
    "            \n",
    "    return merged\n",
    "\n",
    "def merge_transcript_count_files(pattern):\n",
    "    data = get_files(pattern)\n",
    "    \n",
    "    merged = defaultdict(dict)\n",
    "\n",
    "    # initialize dictionary with values from first file\n",
    "    first = data.pop(data.keys()[0])\n",
    "    for gene_name, transcript_dict in first.iteritems():\n",
    "        for transcript_id, values in transcript_dict.iteritems():\n",
    "            merged[gene_name][transcript_id] = values.copy()\n",
    "\n",
    "    # add on the rest\n",
    "    for file, arrays in data.iteritems():\n",
    "        for gene_name, transcript_dict in arrays.iteritems():\n",
    "            for transcript_id, values in transcript_dict.iteritems():\n",
    "                merged[gene_name][transcript_id] += values\n",
    "            \n",
    "    return merged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_obj(merge_count_files('./counts/count_vectors*'), './counts/merged_count_vectors.count')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_obj(merge_count_files('./counts/start_count_vectors*'), './counts/merged_start_count_vectors.count')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_obj(merge_transcript_count_files('./counts/transcript_count_vectors*'), './counts/merged_transcript_count_vectors.count')"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}