pick_targets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "#for executable\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\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 docx import Document\n",
    "from docx.shared import Inches\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": [
       "256"
      ]
     },
     "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')\n",
    "##clear out figures and fasta for new parameters\n",
    "os.system('rm ./figures/*')\n",
    "os.system('rm ./fastas/*')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# keep adding FASTAs for additional transcripts (in descending order of abundance)\n",
    "# until this fraction of total reads is accounted for\n",
    "GENE_COVERAGE_THRESHOLD = 0.8 \n",
    "# fraction of reads to try to capture in each transcript\n",
    "TRANSCRIPT_COVERAGE_THRESHOLD = 0.8\n",
    "# kernel size for median filter used when smoothing read densities\n",
    "KERNEL_SIZE = 49\n",
    "# amount to pad the 5' end of each suggested capture region\n",
    "PAD5 = 25\n",
    "# amount to pad the 3' end of each suggested capture region\n",
    "PAD3 = 400\n",
    "# amount from end of transcripts to target for genes with inadequate sequencing\n",
    "FAILED_LENGTH = 800\n",
    "# maximum length, all sequences will be clipped to this with the clipping beginning from the 3' end\n",
    "MAXIMUM_LENGTH = 2000\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Algorithm outline**\n",
    "\n",
    "1. Using K562 ENCODE data, estimate the abundance of different transcript isoforms. Target a sufficient number of isoforms such that at least the fraction `GENE_COVERAGE_THRESHOLD` of reads (from the ENCODE data) would be captured.\n",
    "\n",
    "2. In each targeted isoform, we perform a peak finding procedure to estimate the target region. Take all reads that are compatible with that isoform and then smooth them using a median filter (with width `KERNEL_SIZE`). Then set a threshold such that at least the fraction `TRANSCRIPT_COVERAGE_THRESHOLD` of reads is accounted for. Then pad the selected region on the 5' and 3' sides by the constants `PAD5` and `PAD3` (this will not extend target sequences past the annotated transcript ends). If the resulting sequence is too long (this happens e.g. if there is an extraneous peak of density early in a transcript), clip the sequence on the side closer to the beginning of the transcript such that the total targeted length does not exceed `MAXIMUM_LENGTH`. For isoforms that have insufficient sequencing coverage, just target `FAILED_LENGTH` nt from the annotated 3' ends. Diagram below:\n",
    "\n",
    "    ```\n",
    "                     /----\\\n",
    "    5'              /      \\      _____         3'\n",
    "                   /        \\    /     \\\n",
    "    _______/-\\____/          ----       \\______\n",
    "    \n",
    "                         ||\n",
    "                         ||\n",
    "                         \\/\n",
    "                         \n",
    "                |    /----\\                    |\n",
    "                |   /      \\      _____        |\n",
    "                |  /        \\    /     \\       |\n",
    "    ______/-\\___|_/          ----       \\______|\n",
    "                    \n",
    "    ```\n",
    "3. For each gene, compare the regions chosen across different isoforms. If one of the smaller regions is a strict subset of one of the bigger ones, eliminate the smaller region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "genome = TwoBitFile('./references/genome.2bit') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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": 6,
   "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": 7,
   "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": 8,
   "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": "markdown",
   "metadata": {},
   "source": [
    "# Load sequencing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load aggregated sequencing data (from process_bams notebook)\n",
    "\n",
    "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": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "spanning_segments = load_obj('./counts/spanning_segments.pkl')\n",
    "count_vectors = load_obj('./counts/merged_count_vectors.count.pkl')\n",
    "start_count_vectors = load_obj('./counts/merged_start_count_vectors.count.pkl')\n",
    "transcript_count_vectors = load_obj('./counts/merged_transcript_count_vectors.count.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ASNA1     62355.0\n",
       "ATF4     386623.0\n",
       "CDH3       1363.0\n",
       "FOXO3    269350.0\n",
       "ORC1      85226.0\n",
       "dtype: float64"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "start_read_counts = pd.Series({name: np.sum(c) for name, c in start_count_vectors.iteritems()})\n",
    "start_read_counts.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Find isoforms to target for each gene"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ENCODE K562 transcript isoform estimates from RSEM\n",
    "transcript_abundances = pd.read_csv('./references/ENCFF717EVE.tsv', sep='\\t')\n",
    "transcript_abundances['base_transcript_id'] = transcript_abundances['transcript_id'].map(lambda x: x.split('.')[0])\n",
    "transcript_abundances.set_index('base_transcript_id', inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# order transcripts for each gene according to RSEM estimates of abundance\n",
    "transcript_sort_order = dict()\n",
    "\n",
    "for gene, transcripts in transcripts_by_gene.iteritems():\n",
    "    transcript_sort_order[gene] = transcript_abundances.loc[transcripts.keys()].sort_values('TPM', ascending=False)['TPM']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# select the top transcripts necessary to account for the majority of reads\n",
    "transcript_cumulative_fraction = dict()\n",
    "no_data = list()\n",
    "\n",
    "for gene, abundance in transcript_sort_order.iteritems():\n",
    "    data = abundance.copy()\n",
    "    data = data/data.sum()\n",
    "    data = data.cumsum()\n",
    "    if (abundance == 0).all():\n",
    "        no_data.extend([gene,])\n",
    "        transcript_cumulative_fraction[gene] = data\n",
    "        continue\n",
    "    num_needed = np.max([1, len(data[data < GENE_COVERAGE_THRESHOLD]) + 1])\n",
    "    transcript_cumulative_fraction[gene] = data.iloc[0:num_needed]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# genes that have no representation in ENCODE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CDH3\n"
     ]
    }
   ],
   "source": [
    "for gene in no_data:\n",
    "    print(gene)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finding target sequences for top isoforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal import medfilt\n",
    "\n",
    "def get_threshold_for_coverage(density, coverage_threshold=0.95):\n",
    "    levels = np.unique(density)\n",
    "    density_by_level = pd.Series({level: density[density >= level].sum() for level in levels})\n",
    "    thresh = density_by_level[density_by_level >= coverage_threshold].argmin()\n",
    "    return thresh\n",
    "\n",
    "def get_thresholded_density(density, coverage_threshold=0.95):\n",
    "    thresh = get_threshold_for_coverage(density, coverage_threshold=coverage_threshold)\n",
    "    thresholded_density = density.copy()\n",
    "    thresholded_density[thresholded_density < thresh] = 0\n",
    "    return thresholded_density\n",
    "\n",
    "def find_runs(a):\n",
    "    # Create an array that is 1 where a is `value`, and pad each end with an extra 0.\n",
    "    isvalue = np.concatenate(([0], (a != 0).view(np.int8), [0]))\n",
    "    absdiff = np.abs(np.diff(isvalue))\n",
    "    # Runs start and end where absdiff is 1.\n",
    "    ranges = np.where(absdiff == 1)[0].reshape(-1, 2)\n",
    "    return ranges\n",
    "\n",
    "def get_transcript_density(gene_name, transcript_id, coverage_threshold=TRANSCRIPT_COVERAGE_THRESHOLD, kernel_size=KERNEL_SIZE):\n",
    "    counts = transcript_count_vectors[gene_name][transcript_id]\n",
    "    filtered = medfilt(counts, kernel_size=kernel_size)\n",
    "    density = filtered/filtered.sum()\n",
    "    thresholded_density = get_thresholded_density(density, coverage_threshold=coverage_threshold)\n",
    "    return thresholded_density, density\n",
    "\n",
    "def get_transcript_cover(thresholded_density, pad3=PAD3, pad5=PAD5):\n",
    "    runs = find_runs(thresholded_density)\n",
    "    thresholded_start = runs[0][0]\n",
    "    thresholded_end = runs[-1][1]\n",
    "    transcript_end = len(thresholded_density) - 1\n",
    "    padded_start = np.max([thresholded_start - pad5, 0])\n",
    "    padded_end = np.min([thresholded_end + pad3, transcript_end])\n",
    "    \n",
    "    cover_length = padded_end - padded_start\n",
    "    if cover_length > MAXIMUM_LENGTH:\n",
    "        padded_start = padded_end - MAXIMUM_LENGTH\n",
    "    \n",
    "    return padded_start, padded_end\n",
    "\n",
    "def get_transcript_subsequence(gene_name, transcript_id, start, end):\n",
    "    transcript = transcripts_by_gene[gene_name][transcript_id]\n",
    "    return transcript.get_sequence(genome)[start:end]\n",
    "\n",
    "def get_transcript_cover_sequence(gene_name, transcript_id, coverage_threshold=TRANSCRIPT_COVERAGE_THRESHOLD, kernel_size=KERNEL_SIZE, pad3=PAD3, pad5=PAD5):\n",
    "    thresholded_density, density = get_transcript_density(gene_name,\n",
    "                                                          transcript_id,\n",
    "                                                          coverage_threshold=coverage_threshold,\n",
    "                                                          kernel_size=kernel_size)\n",
    "    start, end = get_transcript_cover(thresholded_density, pad3=pad3, pad5=pad5)\n",
    "    return get_transcript_subsequence(gene_name, transcript_id, start, end)\n",
    "\n",
    "def get_transcript_cover_position_list(gene_name, transcript_id, coverage_threshold=TRANSCRIPT_COVERAGE_THRESHOLD, kernel_size=KERNEL_SIZE, pad3=PAD3, pad5=PAD5):\n",
    "    thresholded_density, density = get_transcript_density(gene_name,\n",
    "                                                          transcript_id,\n",
    "                                                          coverage_threshold=coverage_threshold,\n",
    "                                                          kernel_size=kernel_size)\n",
    "    start, end = get_transcript_cover(thresholded_density, pad3=pad3, pad5=pad5)\n",
    "    transcript = transcripts_by_gene[gene_name][transcript_id]\n",
    "    if transcript.strand == '+':\n",
    "        pos_list = transcript.get_position_list()[start:end]\n",
    "    else:\n",
    "        pos_list = transcript.get_position_list()[::-1][start:end][::-1]\n",
    "    return pos_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TOMM7\n",
      "=================================\n",
      "\tENST00000358435...\n",
      "ASNA1\n",
      "=================================\n",
      "\tENST00000357332...\n",
      "TUBB\n",
      "=================================\n",
      "\tENST00000327892...\n",
      "ORC1\n",
      "=================================\n",
      "\tENST00000371568...\n",
      "FOXO3\n",
      "=================================\n",
      "\tENST00000343882...\n",
      "\tENST00000406360...\n",
      "RPS5\n",
      "=================================\n",
      "\tENST00000196551...\n",
      "ATF4\n",
      "=================================\n",
      "\tENST00000396680...\n",
      "CDH3\n",
      "=================================\n",
      "\tENST00000569080...\n",
      "\tFAILED\n",
      "\tENST00000569036...\n",
      "\tFAILED\n",
      "\tENST00000567674...\n",
      "\tFAILED\n",
      "\tENST00000542274...\n",
      "\tFAILED\n",
      "\tENST00000569117...\n",
      "\tENST00000429102...\n",
      "\tFAILED\n",
      "\tENST00000264012...\n",
      "\tFAILED\n",
      "\tENST00000568292...\n",
      "\tFAILED\n",
      "\tENST00000566808...\n",
      "\tFAILED\n",
      "\tENST00000565453...\n",
      "\tFAILED\n",
      "SEC61A1\n",
      "=================================\n",
      "\tENST00000243253...\n"
     ]
    }
   ],
   "source": [
    "cover_positions = defaultdict(dict)\n",
    "cover_sequences = defaultdict(dict)\n",
    "no_cover = list()\n",
    "\n",
    "for gene, transcripts in transcript_cumulative_fraction.iteritems():\n",
    "    print(gene)\n",
    "    print('=================================')\n",
    "    for transcript_id in transcripts.index:\n",
    "        print('\\t{0}...'.format(transcript_id))\n",
    "        try:\n",
    "            cover_positions[gene][transcript_id] = get_transcript_cover_position_list(gene, transcript_id)\n",
    "            cover_sequences[gene][transcript_id] = get_transcript_cover_sequence(gene, transcript_id)\n",
    "        except ValueError:\n",
    "            no_cover.extend([(gene, transcript_id),])\n",
    "            print('\\tFAILED')\n",
    "            transcript = transcripts_by_gene[gene][transcript_id]\n",
    "            cover_positions[gene][transcript_id] = transcript.get_position_list()[-FAILED_LENGTH:]\n",
    "            cover_sequences[gene][transcript_id] = transcript.get_sequence(genome)[-FAILED_LENGTH:]\n",
    "            continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('CDH3', 'ENST00000569080'),\n",
       " ('CDH3', 'ENST00000569036'),\n",
       " ('CDH3', 'ENST00000567674'),\n",
       " ('CDH3', 'ENST00000542274'),\n",
       " ('CDH3', 'ENST00000429102'),\n",
       " ('CDH3', 'ENST00000264012'),\n",
       " ('CDH3', 'ENST00000568292'),\n",
       " ('CDH3', 'ENST00000566808'),\n",
       " ('CDH3', 'ENST00000565453')]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# transcripts where too little sequencing so just took piece from annotated end\n",
    "no_cover"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Removing target sequences that overlap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "\n",
    "def find_redundant_sequences(seqs):\n",
    "    # order sequences by length as smaller must be in larger\n",
    "    seq_lens = pd.Series({transcript_id: len(seq) for transcript_id, seq in seqs.iteritems()}).sort_values()\n",
    "    # make all pairwise comparisons of smaller to larger ((1, 2), (1, 3), (1, 4), ..., (2, 3), (2, 4), ...)\n",
    "    ordered_pairs = list(itertools.combinations(np.arange(len(seq_lens)), 2))    \n",
    "    redundant_pairs = list()\n",
    "    \n",
    "    for pair in ordered_pairs:\n",
    "        name1 = seq_lens.index[pair[0]]\n",
    "        name2 = seq_lens.index[pair[1]]\n",
    "        len1 = seq_lens.loc[name1]\n",
    "        len2 = seq_lens.loc[name2]\n",
    "        seq1 = seqs[name1]\n",
    "        seq2 = seqs[name2]\n",
    "\n",
    "        if seq1 in seq2:\n",
    "            print('\\t{0} (length: {1}) is contained in {2} (length: {3})'.format(name1, len1, name2, len2))\n",
    "            redundant_pairs.extend([(name1, name2),])\n",
    "    return redundant_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TOMM7\n",
      "=================================\n",
      "ASNA1\n",
      "=================================\n",
      "TUBB\n",
      "=================================\n",
      "ORC1\n",
      "=================================\n",
      "FOXO3\n",
      "=================================\n",
      "RPS5\n",
      "=================================\n",
      "ATF4\n",
      "=================================\n",
      "CDH3\n",
      "=================================\n",
      "SEC61A1\n",
      "=================================\n"
     ]
    }
   ],
   "source": [
    "redundant_sequences = dict()\n",
    "for gene, seqs in cover_sequences.iteritems():\n",
    "    print(gene)\n",
    "    print('=================================')\n",
    "    \n",
    "    redundant_pairs = find_redundant_sequences(seqs)\n",
    "    if len(redundant_pairs) > 0:\n",
    "        redundant_sequences[gene] = redundant_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "redundant_covers = list()\n",
    "\n",
    "for pairs in redundant_sequences.itervalues():\n",
    "    for pair in pairs:\n",
    "        redundant_covers.append(pair[0])\n",
    "\n",
    "# list to keep for plots\n",
    "redundant_covers = np.unique(redundant_covers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# makes names summarizing all of the transcripts targeted by overlapping sequences\n",
    "dominance_relations = dict()\n",
    "redundant_names = dict()\n",
    "\n",
    "for gene, pairs in redundant_sequences.iteritems():\n",
    "    dominant_covers = np.unique([pair[1] for pair in pairs])\n",
    "    for d in dominant_covers:\n",
    "        dominance_relations[d] = reduce(np.union1d, [pair for pair in pairs if pair[1] == d])\n",
    "        redundant_names[d] = '_'.join(dominance_relations[d])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this has all sequences that are encapsulated in another removed\n",
    "# and the name has been changed to be all transcripts targeted mashed together\n",
    "filtered_cover_sequences = defaultdict(dict)\n",
    "\n",
    "for gene, covers in cover_sequences.iteritems():\n",
    "    for transcript_id, seq in covers.iteritems():\n",
    "        if transcript_id not in redundant_covers:\n",
    "            if transcript_id not in redundant_names.keys():\n",
    "                filtered_cover_sequences[gene][transcript_id] = seq\n",
    "            else:\n",
    "                filtered_cover_sequences[gene][redundant_names[transcript_id]] = seq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "vec_start = np.vectorize(lambda x: x.start)\n",
    "vec_end = np.vectorize(lambda x: x.end)\n",
    "\n",
    "def transcript_to_segments(transcript):\n",
    "    starts = vec_start(transcript.segments)\n",
    "    ends = vec_end(transcript.segments)\n",
    "    return starts, ends\n",
    "\n",
    "def plot_genome_coverage(gene_name, transcript_id, xlim=None, ylim=None):\n",
    "    transcript = transcripts_by_gene[gene_name][transcript_id]\n",
    "    strand = transcript.strand\n",
    "    if (strand == '+'):\n",
    "        positions = transcript.get_position_list()\n",
    "    elif (strand == '-'):\n",
    "        positions = transcript.get_position_list()[::-1]\n",
    "    \n",
    "    plt.step(positions, transcript_count_vectors[gene_name][transcript_id], linewidth=0.5, alpha=0.5)\n",
    "    \n",
    "    starts, ends = transcript_to_segments(transcript)\n",
    "\n",
    "    for i in xrange(len(starts)):\n",
    "        plt.plot([starts[i], ends[i]], [0, 0], linewidth=5, color='gray')\n",
    "\n",
    "    utr3 = transcript.get_utr3()\n",
    "    if len(utr3) > 0:\n",
    "        starts, ends = transcript_to_segments(utr3)\n",
    "        for i in xrange(len(starts)):\n",
    "            plt.plot([starts[i], ends[i]], [0, 0], linewidth=5, color='r', alpha=0.6)\n",
    "        \n",
    "    ax = plt.gca()\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['left'].set_visible(False)\n",
    "    \n",
    "    if xlim is not None:\n",
    "        plt.xlim(xlim)\n",
    "    if ylim is not None:\n",
    "        plt.ylim(ylim)\n",
    "        \n",
    "    if transcript_id in cover_positions[gene_name].keys():\n",
    "        cov_pos = np.array(cover_positions[gene_name][transcript_id])\n",
    "        x = np.arange(cov_pos[0], cov_pos[-1] + 1)\n",
    "        y = np.zeros(*x.shape)\n",
    "        ymin, ymax = plt.gca().get_ylim()\n",
    "        y[cov_pos - cov_pos[0]] = ymax\n",
    "        color = 'g' if transcript_id not in redundant_covers else 'lightgreen'\n",
    "        plt.fill_between(x, y, step='pre', alpha=0.3, facecolor=color)\n",
    "\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_transcript_coverage(gene_name, transcript_id, xlim=None, ylim=None):\n",
    "    transcript = transcripts_by_gene[gene_name][transcript_id]\n",
    "    strand = transcript.strand\n",
    "\n",
    "    genome_positions = transcript.get_position_list()\n",
    "    transcript_length = len(genome_positions)\n",
    "    positions = np.arange(transcript_length)\n",
    "\n",
    "    if (strand == '+'):\n",
    "        # line them all up at right side\n",
    "        positions = positions + genome_positions[-1] - transcript_length\n",
    "    elif (strand == '-'):\n",
    "        # line them up at left side\n",
    "        positions = positions[::-1] + genome_positions[0]\n",
    "\n",
    "    plt.step(positions, transcript_count_vectors[gene_name][transcript_id], linewidth=0.5, alpha=0.5, color='red')\n",
    "\n",
    "    transcript_to_genome = dict(enumerate(genome_positions))\n",
    "    pos_mapper = {v: positions[k] for k, v in transcript_to_genome.iteritems()}\n",
    "\n",
    "    starts, ends = transcript_to_segments(transcript)\n",
    "    \n",
    "    for i in xrange(len(starts)):\n",
    "        plt.plot([pos_mapper[starts[i]], pos_mapper[ends[i] - 1]], [0, 0], linewidth=5, color='gray')\n",
    "\n",
    "    ax = plt.gca()\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['left'].set_visible(False)\n",
    "        \n",
    "    if xlim is not None:\n",
    "        plt.xlim(xlim)\n",
    "    if ylim is not None:\n",
    "        plt.ylim(ylim)\n",
    "        \n",
    "    if transcript_id in cover_positions[gene_name].keys():\n",
    "        cov_pos = np.array(cover_positions[gene_name][transcript_id]) \n",
    "        x = np.arange(cov_pos[0], cov_pos[-1] + 1)\n",
    "        y = np.zeros(*x.shape)\n",
    "        ymin, ymax = plt.gca().get_ylim()\n",
    "        y[cov_pos - cov_pos[0]] = ymax\n",
    "        if strand == '+':\n",
    "            genome_pos_mapper = dict(zip(genome_positions, positions))\n",
    "        else:\n",
    "            genome_pos_mapper = dict(zip(genome_positions, positions[::-1])) \n",
    "        x = np.vectorize(genome_pos_mapper.get, otypes=[float])(x)\n",
    "        \n",
    "        color = 'g' if transcript_id not in redundant_covers else 'lightgreen'\n",
    "        plt.fill_between(x, y, step='pre', alpha=0.3, facecolor=color)\n",
    "        \n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_aggregated_genome_coverage(gene_name, xlim=None, ylim=None):\n",
    "    strand = spanning_segments[gene_name].strand\n",
    "    chrom = spanning_segments[gene_name].chrom\n",
    "\n",
    "    if (strand == '+'):\n",
    "        positions = spanning_segments[gene_name].get_position_list()\n",
    "    elif (strand == '-'):\n",
    "        positions = spanning_segments[gene_name].get_position_list()[::-1]\n",
    "      \n",
    "    plt.step(positions, start_count_vectors[gene_name], linewidth=0.5, alpha=0.5)\n",
    "    plt.plot(positions, count_vectors[gene_name], linewidth=1)\n",
    "\n",
    "    if xlim is not None:\n",
    "        plt.xlim(xlim)\n",
    "    if ylim is not None:\n",
    "        plt.ylim(ylim)\n",
    "    \n",
    "    ax = plt.gca()\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_all_transcript_coverage(gene_name, sort=True, filename=None):\n",
    "    transcripts = transcripts_by_gene[gene_name]\n",
    "    strand = spanning_segments[gene_name].strand\n",
    "    \n",
    "    num_transcripts = len(transcripts)\n",
    "    xlim = [spanning_segments[gene_name].get_position_list()[0], spanning_segments[gene_name].get_position_list()[-1]]\n",
    "    \n",
    "    max_counts = np.max([np.max(counts) for counts in transcript_count_vectors[gene_name].itervalues()])\n",
    "    ylim = [0, max_counts]\n",
    "    \n",
    "    plt.figure(figsize=[15, num_transcripts])\n",
    "    \n",
    "    if not sort:\n",
    "        keys = transcripts.iterkeys()\n",
    "    else:\n",
    "        keys = transcript_sort_order[gene_name].index.values\n",
    "    \n",
    "    for i, transcript_id in enumerate(keys):\n",
    "        plt.subplot(num_transcripts, 1, i+1)\n",
    "        ax = plot_transcript_coverage(gene_name, transcript_id, xlim=xlim, ylim=ylim)\n",
    "\n",
    "        if sort:\n",
    "            min_x, max_x = ax.get_xlim()\n",
    "            x_range = max_x - min_x\n",
    "            max_y = ax.get_ylim()[1]\n",
    "            ax.text(max_x - 0.05*x_range, 0.7*max_y, transcript_sort_order[gene_name].loc[transcript_id])\n",
    "        \n",
    "        if i == 0:\n",
    "            plt.title('{0} - {1} ({2})'.format(gene_name, id_mapper[gene_name], strand))\n",
    "            \n",
    "        if i < num_transcripts - 1:\n",
    "            plt.xticks([])\n",
    "    if filename is not None:\n",
    "        plt.savefig('./figures/{0}.png'.format(filename))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_all_genome_coverage(gene_name, sort=True, filename=None):\n",
    "    transcripts = transcripts_by_gene[gene_name]\n",
    "    strand = spanning_segments[gene_name].strand\n",
    "    \n",
    "    num_transcripts = len(transcripts)\n",
    "    xlim = [spanning_segments[gene_name].get_position_list()[0], spanning_segments[gene_name].get_position_list()[-1]]\n",
    "    \n",
    "    max_counts = np.max([np.max(counts) for counts in transcript_count_vectors[gene_name].itervalues()])\n",
    "    ylim = [0, max_counts]\n",
    "    \n",
    "    plt.figure(figsize=[15, 1*(num_transcripts + 1)])\n",
    "    plt.subplot(num_transcripts + 1, 1, 1)\n",
    "    plot_aggregated_genome_coverage(gene_name, xlim=xlim)\n",
    "    plt.title('{0} - {1} ({2})'.format(gene_name, id_mapper[gene_name], strand))\n",
    "    plt.xticks([])\n",
    "    \n",
    "    if not sort:\n",
    "        keys = transcripts.iterkeys()\n",
    "    else:\n",
    "        keys = transcript_sort_order[gene_name].index.values\n",
    "    \n",
    "    for i, transcript_id in enumerate(keys):\n",
    "        plt.subplot(num_transcripts + 1, 1, i+2)\n",
    "        ax = plot_genome_coverage(gene_name, transcript_id, xlim=xlim, ylim=ylim)\n",
    "        if sort:\n",
    "            min_x, max_x = ax.get_xlim()\n",
    "            x_range = max_x - min_x\n",
    "            max_y = ax.get_ylim()[1]\n",
    "            ax.text(max_x - 0.05*x_range, 0.7*max_y, transcript_sort_order[gene_name].loc[transcript_id])\n",
    "            \n",
    "        if i < num_transcripts - 1:\n",
    "            plt.xticks([])\n",
    "    if filename is not None:\n",
    "        plt.savefig('./figures/{0}.png'.format(filename))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "from docx import Document\n",
    "from docx.shared import Inches\n",
    "from docx.enum.section import WD_SECTION\n",
    "from docx.enum.section import WD_ORIENT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "document = Document()\n",
    "summary_data = defaultdict(dict)\n",
    "\n",
    "section = document.sections[-1]\n",
    "new_width, new_height = section.page_height, section.page_width\n",
    "section.orientation = WD_ORIENT.LANDSCAPE\n",
    "section.page_width = new_width\n",
    "section.page_height = new_height\n",
    "\n",
    "for name, row in targeted_genes.sort_values('gene_name').iterrows():\n",
    "    gene_name = row['gene_name']\n",
    "    gene_id = row['gene_id']\n",
    "    plot_all_genome_coverage(gene_name, filename=gene_name + '_genome')\n",
    "    plot_all_transcript_coverage(gene_name, filename=gene_name + '_transcriptome')\n",
    "    \n",
    "    spanning_start = spanning_segments[gene_name].segments[0].start\n",
    "    spanning_end = spanning_segments[gene_name].segments[0].end\n",
    "    spanning_chrom = spanning_segments[gene_name].segments[0].chrom\n",
    "    spanning_strand = spanning_segments[gene_name].segments[0].strand\n",
    "    \n",
    "    document.add_heading('{0} - {1}'.format(gene_name, gene_id), 0)\n",
    "    document.add_picture('./figures/{0}_genome.png'.format(gene_name), width=Inches(9))\n",
    "    document.add_picture('./figures/{0}_transcriptome.png'.format(gene_name), width=Inches(9))\n",
    "    \n",
    "    p = document.add_paragraph()\n",
    "    run = p.add_run('Position: ' + str(spanning_chrom) + ':' + \\\n",
    "                              str(spanning_start) + '-' + \\\n",
    "                              str(spanning_end) + ' (' + \\\n",
    "                              str(spanning_strand) + ')')\n",
    "    font = run.font\n",
    "    font.name = 'Courier New'\n",
    "    \n",
    "    for transcript_target, seq in filtered_cover_sequences[gene_name].iteritems():\n",
    "        document.add_heading(transcript_target, 1)\n",
    "        p = document.add_paragraph()\n",
    "        run = p.add_run(seq)\n",
    "        font = run.font\n",
    "        font.name = 'Courier New'\n",
    "        \n",
    "        try:\n",
    "            fasta_name = '{0}__{1}'.format(gene_id, transcript_target)\n",
    "            fasta_file = open('./fastas/{0}.fasta'.format(fasta_name), 'wb')\n",
    "            fasta_file.write('>{0}\\n'.format(fasta_name))\n",
    "            fasta_file.write(seq)\n",
    "            fasta_file.close()\n",
    "        except IOError as ioerr:\n",
    "            if ioerr.errno == 36:\n",
    "                fasta_name = '{0}__{1}'.format(gene_id, transcript_target[0:3])\n",
    "                fasta_file = open('./fastas/{0}.fasta'.format(fasta_name), 'wb')\n",
    "                fasta_file.write('>{0}\\n'.format(fasta_name))\n",
    "                fasta_file.write(seq)\n",
    "                fasta_file.close()\n",
    "                \n",
    "        \n",
    "        # summary data spreadsheet\n",
    "        summary_data[transcript_target]['gene_name'] =  gene_name\n",
    "        summary_data[transcript_target]['gene_id'] =  gene_id\n",
    "        summary_data[transcript_target]['strand'] =  spanning_strand\n",
    "        summary_data[transcript_target]['length'] =  len(seq)\n",
    "        summary_data[transcript_target]['seq'] =  seq\n",
    " \n",
    "    document.add_page_break()\n",
    "        \n",
    "document.save('target_regions.docx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame(summary_data).T.sort_values('gene_name').to_excel('./target_regions_summary.xlsx')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total length:\n",
      "15736\n"
     ]
    }
   ],
   "source": [
    "print('Total length:')\n",
    "print(pd.DataFrame(summary_data).T.sort_values('gene_name')['length'].sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>gene_id</th>\n",
       "      <th>gene_name</th>\n",
       "      <th>length</th>\n",
       "      <th>seq</th>\n",
       "      <th>strand</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>ENST00000396680</th>\n",
       "      <td>ENSG00000128272</td>\n",
       "      <td>ATF4</td>\n",
       "      <td>411</td>\n",
       "      <td>ATAGGAGCCTCCCATCTCCAGGTGTTCTCTGTGGGTCTGCCCGTCC...</td>\n",
       "      <td>+</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                         gene_id gene_name length  \\\n",
       "ENST00000396680  ENSG00000128272      ATF4    411   \n",
       "\n",
       "                                                               seq strand  \n",
       "ENST00000396680  ATAGGAGCCTCCCATCTCCAGGTGTTCTCTGTGGGTCTGCCCGTCC...      +  "
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(summary_data).T[pd.DataFrame(summary_data).T['gene_name']=='ATF4']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}