4. Data preprocessing - ORF and transcript characteristics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Data preprocessing - ORF and transcript characteristics\n",
    "\n",
    "The initial data preprocessing steps calculate relevant data and organizes them in a format that is easy for later analysis: tab- and comma-separated files compatible with the DataFrame format (\"\\_main.df\" and \"\\_profile.df\"). These may be easily read by either R or Python (using the Pandas package).\n",
    "\n",
    "This data preprocessing step collates ORF characteristics (sequence features, RNA-seq and ribosome profiling data) for downstream analyses.\n",
    "\n",
    "ORF parameters are fully described later, but include calculations for initiation context Weighted Relative ENTropy (WRENT) score, seconday structure, length, number of reads and position within transcript.\n",
    "\n",
    "Before ORF parameters can be calculated, the position specific scoring matrix for determining WRENT scores must first be constructed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both the WRENT scoring matrix computation and the collation of ORF characteristics take substantial computational time (a few hours) on a desktop computer / laptop. Copies of this notebook may be concurrently opened to run these analyses on multiple datasets (by commenting / uncommenting the cell below), especially for computers with multi-core/multi-thread capabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "species = \"mm\"\n",
    "stage = \"mES\"\n",
    "ASSEMBLY = \"GRCm38_ens\"\n",
    "\n",
    "# species = \"dr\"\n",
    "# stage = \"Shield\"\n",
    "# ASSEMBLY = \"Zv9_ens\"\n",
    "\n",
    "# species = \"hs\"\n",
    "# stage = \"HeLa\"\n",
    "# ASSEMBLY = \"GRCh37_ens\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports, File Locations and Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# IMPORTS\n",
    "import RNA\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import corebio\n",
    "import weblogolib\n",
    "\n",
    "\n",
    "from Bio import SeqIO\n",
    "from numpy import log1p, log2\n",
    "from ast import literal_eval\n",
    "from pandas import Series, DataFrame\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "'\\_canonical.trpedf' means the .trpedf file was generated based on a transcriptome that used the longest transcript with longest coding sequence (historically the \"canonical\" transcript in transcript annotations). \n",
    "\n",
    "Because the ORF_starts, ORF_ends, RPF_csvProfile and CDS columns in the .trpedf files are comma-separated values, **literal_eval** is used to parse them.\n",
    "\n",
    "Corresponding fasta file is indexed using Biopython's SeqIO module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# FILENAME PARAMETERS\n",
    "DATA_DIR = \"./data/\" + species + \"/\"\n",
    "ANNOTATIONS_DIR = \"./annotations/\"\n",
    "TRPE_FILE = DATA_DIR + stage + \"_canonical.trpedf\"\n",
    "CONVERTERS = {i:literal_eval for i in (\"ORF_starts\", \"ORF_ends\", \"RPF_csvProfile\", \"CDS\")}\n",
    "\n",
    "FASTA_FILE = ANNOTATIONS_DIR + ASSEMBLY + \"_genes_canonical.fasta\"\n",
    "SEQS = SeqIO.index(FASTA_FILE, \"fasta\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More parameters defined. Transcripts were required to have a minimum 5' UTR and RNA-seq expression value (in FPKM) to be used in determining the weighted relative entropy (WRENT) scoring matrix.\n",
    "\n",
    "The size/position of the initiation context used for WRENT scoring and secondary structure prediction are defined around the start codon and start position respectively.\n",
    "\n",
    "When reads over ORFs are considered, a number of positions upstream of the stop are discarded, as determined from biases detected in metagene profiles.\n",
    "\n",
    "For plotting ribosome profiling reads around the start and stop codon of a metagene (for quality control purposes), minimum CDS and UTR lengths are required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# OTHER PARAMETERS\n",
    "UTR5_LENGTH_MIN = 25\n",
    "FPKM_MIN = 1\n",
    "\n",
    "WRENT_CONTEXT = (10, 10) # BEFORE START CODON, AFTER START CODON\n",
    "SS_CONTEXT_L = (25, 10)     # BEFORE ORF START, AFTER ORF START\n",
    "SS_CONTEXT_R = (-1, 36)\n",
    "\n",
    "ORF_END_TRIM = 10 # NUMBER OF POSITIONS BEFORE uORF END TO\n",
    "                  # DISCARD READS DUE TO ARTIFACTUAL EXPERIMENTAL\n",
    "                  # ACCUMULATION OF RPF READS AT uORF END.\n",
    "\n",
    "uORF_LENGTH_MIN = 21\n",
    "uORF_FROM_TSS_MIN = 25\n",
    "\n",
    "NT_COLOR = ('#00d700', '#df1f00', '#0226cc', '#ffb700')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some helper functions to retrieve sequences of initiation contexts for determining weighted Kozak scores and RNA secondary structure, as well as calculating translational efficiency (TE) of an ORF and scoring the weighted Kozak score of a sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def wrent_context_seq(seq, ORF, WRENT_CONTEXT):\n",
    "    '''Takes transcript sequence and ORF (in transcript coordinates),\n",
    "    along with weighted Kozak context size as inputs,\n",
    "    returns sequence'''\n",
    "    left_flank, right_flank = WRENT_CONTEXT\n",
    "    seq_to_return = str(seq[ORF[0] - left_flank:ORF[0]] + seq[ORF[0] + 3: ORF[0] + right_flank + 3])\n",
    "    if len(seq_to_return) == left_flank + right_flank:\n",
    "        return seq_to_return\n",
    "    else: return \"A\"\n",
    "\n",
    "def ss_context_seq(seq, ORF, SS_CONTEXT):\n",
    "    '''Takes transcript sequence and ORF (in transcript coordinates),\n",
    "    along with secondary structure context size as inputs,\n",
    "    returns sequence'''\n",
    "    left_flank, right_flank = SS_CONTEXT\n",
    "    seq_to_return = str(seq[ORF[0] - left_flank : ORF[0] + right_flank])\n",
    "    if len(seq_to_return) == left_flank + right_flank:\n",
    "        return seq_to_return\n",
    "    else: return \"A\"\n",
    "\n",
    "def TE(RPF_csvProfile, ORF, expression, ORF_END_TRIM):\n",
    "    '''Normalizes the density of ribosome profiling reads over ORF\n",
    "    by the expression of the transcript'''\n",
    "    trimmed_ORF_RPF_reads = sum(RPF_csvProfile[ORF[0]:ORF[1] - ORF_END_TRIM])\n",
    "    trimmed_ORF_length = max(((ORF[1] - ORF[0]) - ORF_END_TRIM), 1)\n",
    "    return float(trimmed_ORF_RPF_reads) / trimmed_ORF_length / expression\n",
    "\n",
    "def wrent_score_seq(en_score, seq_to_score):\n",
    "    '''Scores a sequence based on weighted entropy score matrix'''\n",
    "    return sum([en_score[pos][nt] for pos, nt in enumerate(seq_to_score)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def UTR5_uORF_segments(CDS, uORFs):\n",
    "    is_uORF = [False for i in xrange(int(CDS[0]))]\n",
    "    for uORF in uORFs:\n",
    "        for pos in xrange(uORF[0], min(uORF[1], CDS[0])):\n",
    "            is_uORF[pos] = True\n",
    "    return is_uORF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating weighted relative entropy (WRENT) scoring matrix\n",
    "\n",
    "Using translational efficiencies as weights, we determine the sequence motif for efficient translation initiation based on information content, with the nucleotide frequencies at the initiation context as background.\n",
    "\n",
    "To do this, we read in the frequencies (weighted and unweighted) of nucleotides at each position independently. Position-specific scoring matrices (PSSMs) are initiated as a pandas DataFrames.\n",
    "\n",
    "A file iterator is created to read the .trpedf file line-by-line, and a counter is initiated to count the number of transcripts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# INITIALIZE PSSMS FOR RELATIVE ENTROPY SCORING MATRICES\n",
    "en_pssm = DataFrame(index=[nt for nt in 'ATCGN'], columns=range(sum(WRENT_CONTEXT)))\n",
    "en_pssm.fillna(0., inplace=True)\n",
    "en_pssm_unweighted = DataFrame(index=[nt for nt in 'ATCGN'], columns=range(sum(WRENT_CONTEXT)))\n",
    "en_pssm_unweighted.fillna(0., inplace=True)\n",
    "\n",
    "uORF_en_pssm = DataFrame(index=[nt for nt in 'ATCGN'], columns=range(sum(WRENT_CONTEXT)))\n",
    "uORF_en_pssm.fillna(0., inplace=True)\n",
    "uORF_en_pssm_unweighted = DataFrame(index=[nt for nt in 'ATCGN'], columns=range(sum(WRENT_CONTEXT)))\n",
    "uORF_en_pssm_unweighted.fillna(0., inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# INITIALIZE FILE ITERATOR, LIST OF TRANSCRIPTS\n",
    "trpedf_file_iterator = pd.read_table(TRPE_FILE, converters=CONVERTERS, chunksize=1)\n",
    "transcript_list = []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The .trpedf file is read line-by-line. ORFs are merged, and uORFs are defined as beginning before the CDS.\n",
    "\n",
    "If the transcript passes the filter parameters defined earlier (minimum 5' UTR length and RNA-seq expression), the initiation context sequence is added to the PSSMs, either unweighted (i.e. weight of 1) weighted by the transcript CDS's translational efficiency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# ITERATES OVER TRANSCRIPTS IN .trpedf FILE\n",
    "\n",
    "CDS_motif_count = 0\n",
    "uORF_motif_count = 0\n",
    "\n",
    "for transcript in trpedf_file_iterator:\n",
    "    transcript_list.append(transcript[\"Transcript\"][0])\n",
    "    \n",
    "    seq = SEQS[transcript[\"Transcript\"][0]].seq        # get transcript sequence\n",
    "    \n",
    "    expression = transcript[\"Gene_Expression_FPKM\"][0]\n",
    "    RPF_csvProfile = transcript[\"RPF_csvProfile\"][0]\n",
    "    CDS = transcript[\"CDS\"][0]\n",
    "\n",
    "    ORF_starts = transcript[\"ORF_starts\"][0]           \n",
    "    ORF_ends = transcript[\"ORF_ends\"][0]\n",
    "    if type(ORF_starts) is np.int64:                   # corrects for single-entry\n",
    "        ORF_starts = (ORF_starts,)\n",
    "        ORF_ends = (ORF_ends,)\n",
    "    ORFs = zip(ORF_starts, ORF_ends)                   # zips starts and stops into ORF\n",
    "    uORFs = [ORF for ORF in ORFs if (ORF[0] < CDS[0])] # uORFs defined as beginning before CDS\n",
    "\n",
    "    if CDS[0] >= UTR5_LENGTH_MIN and expression >= FPKM_MIN:\n",
    "        if len(uORFs) == 0:                            # filter for transcripts with\n",
    "                                                       # minimum 5' UTR length and expression\n",
    "            weight = log1p(TE(RPF_csvProfile, CDS, expression, ORF_END_TRIM))\n",
    "            context_seq = wrent_context_seq(seq, CDS, WRENT_CONTEXT)\n",
    "            for pos1, nt1 in enumerate(context_seq):\n",
    "                en_pssm[pos1][nt1] += weight\n",
    "                en_pssm_unweighted[pos1][nt1] += 1\n",
    "            CDS_motif_count += 1\n",
    "        \n",
    "        elif (len(uORFs) == 1)\\\n",
    "        and (uORFs[0][1] < CDS[0])\\\n",
    "        and (uORFs[0][1] - uORFs[0][0] > uORF_LENGTH_MIN)\\\n",
    "        and (uORFs[0][0] > uORF_FROM_TSS_MIN):             # 1 uORF, non-overlapping, min uORF length, min distance from TSS\n",
    "            weight = log1p(TE(RPF_csvProfile, uORFs[0], expression, ORF_END_TRIM))\n",
    "            context_seq = wrent_context_seq(seq, uORFs[0], WRENT_CONTEXT)\n",
    "            for pos1, nt1 in enumerate(context_seq):\n",
    "                uORF_en_pssm[pos1][nt1] += weight\n",
    "                uORF_en_pssm_unweighted[pos1][nt1] += 1\n",
    "            uORF_motif_count += 1\n",
    "        \n",
    "transcript_count = len(transcript_list)\n",
    "print \"%g transcripts used for CDS WRENT motif\" % CDS_motif_count\n",
    "print \"%g transcripts used for uORF WRENT motif\" % uORF_motif_count"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first determine relative entropy: unweighted, then weighted for TE.\n",
    "\n",
    "'N's are ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# IGNORE N\n",
    "en_pssm_unweighted.drop('N', inplace=True)\n",
    "en_pssm.drop('N', inplace=True)\n",
    "uORF_en_pssm_unweighted.drop('N', inplace=True)\n",
    "uORF_en_pssm.drop('N', inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PSSMs are stored for later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "en_pssm_unweighted.to_csv(DATA_DIR + stage + \"_pssm_unweighted.df\", sep=\"\\t\")\n",
    "en_pssm.to_csv(DATA_DIR + stage + \"_pssm.df\", sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "uORF_en_pssm_unweighted.to_csv(DATA_DIR + stage + \"_pssm_uORF_unweighted.df\", sep=\"\\t\")\n",
    "uORF_en_pssm.to_csv(DATA_DIR + stage + \"_pssm_uORF.df\", sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For calculating both unweighted and unweighted relative entropies, the background nucleotide distribution is from the total nucleotide frequency from the unweighted PSSM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "en_pssm = DataFrame.from_csv(DATA_DIR + stage + \"_pssm.df\", sep=\"\\t\")\n",
    "en_pssm_unweighted = DataFrame.from_csv(DATA_DIR + stage + \"_pssm_unweighted.df\", sep=\"\\t\")\n",
    "uORF_en_pssm = DataFrame.from_csv(DATA_DIR + stage + \"_pssm_uORF.df\", sep=\"\\t\")\n",
    "uORF_en_pssm_unweighted = DataFrame.from_csv(DATA_DIR + stage + \"_pssm_uORF_unweighted.df\", sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "en_unweighted_nt_freq = en_pssm_unweighted.mean(axis='columns')  # background model is the unweighted PSSM's total nucleotide freq\n",
    "en_unweighted_nt_prob = en_unweighted_nt_freq / en_unweighted_nt_freq.sum(axis='rows')\n",
    "\n",
    "en_pssm_unweighted_prob = en_pssm_unweighted / en_pssm_unweighted.sum(axis=\"rows\")\n",
    "en_unweighted_score = en_pssm_unweighted_prob.divide(en_unweighted_nt_prob, axis='rows').apply(log2)\n",
    "\n",
    "en_pssm_prob = en_pssm / en_pssm.sum(axis='rows')             # normalized nucleotide probabilities\n",
    "en_score = en_pssm_prob.divide(en_unweighted_nt_prob, axis='rows').apply(log2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Stores the scoring matrices (weighted and unweighted) for later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "en_score.to_csv(DATA_DIR + stage + \"_en_score.df\", sep=\"\\t\")\n",
    "en_unweighted_score.to_csv(DATA_DIR + stage + \"_en_unweighted_score.df\", sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "en_score = DataFrame.from_csv(DATA_DIR + stage + \"_en_score.df\", sep=\"\\t\")\n",
    "en_unweighted_score = DataFrame.from_csv(DATA_DIR + stage + \"_en_unweighted_score.df\", sep=\"\\t\")\n",
    "\n",
    "en_score.columns = [int(i) for i in en_score.columns]\n",
    "en_unweighted_score.columns = [int(i) for i in en_unweighted_score.columns]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**en_score** is subsequently used to score all ORFs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Precalculating weighted relative entropy (WRENT) scores, secondary structure ensemble free energies\n",
    "For faster downstream analysis, weighted relative entropy (WRENT) scores and the computationally-predicted ensemble free energies of secondary structure around initiation contexts of all ORFs are precalculated and stored.\n",
    "\n",
    "Ribosome profiling reads contained within ORFs, as well as other ORF and transcript parameters are also determined and stored.\n",
    "\n",
    "Dataframe columns elaborated in following table:\n",
    "<table>\n",
    "<tr><td>**Column**</td><td>**Description**</td></tr>\n",
    "<tr><td>Transcript</td><td>Transcript ID</td></tr>\n",
    "<tr><td>Gene</td><td>Gene ID</td></tr>\n",
    "<tr><td>Gene_Name</td><td>Gene Name</td></tr>\n",
    "<tr><td>Gene_Expression_FPKM</td><td>Expression at gene level (from corresponding RNA-seq data; Tophat + Cufflinks)</td></tr>\n",
    "\n",
    "<tr><td>UTR5_length</td><td>Length of 5' UTR</td></tr>\n",
    "<tr><td>UTR3_length</td><td>Length of 3' UTR</td></tr>\n",
    "<tr><td>UTR5_reads</td><td>Ribosome profiling reads in 5' UTR</td></tr>\n",
    "<tr><td>UTR5_reads_trunc</td><td>Ribosome profiling reads in 5' UTR, truncated</td></tr>\n",
    "<tr><td>UTR5_GC</td><td>Total GC content of 5' leader</td></tr>\n",
    "<tr><td>num_uORFs</td><td>Number of uORFs (ORFs beginning before CDS)</td></tr>\n",
    "<tr><td>uORFs_read</td><td>Number of reads in each uORF, comma-separated values (CSV)</td></tr>\n",
    "<tr><td>UTR5_uORF_reads</td><td>Number of reads in 5' leader within uORFs</td></tr>\n",
    "<tr><td>UTR5_uORF_tlength</td><td>Total length of uORFs within 5' leader</td></tr>\n",
    "<tr><td>uORFs_length</td><td>Length of uORFs, CSV</td></tr>\n",
    "<tr><td>uORFs_wrent_score</td><td>Calculated WRENT score of uORFs, CSV</td></tr>\n",
    "<tr><td>uORFs_urent_score</td><td>Calculated URENT (unweighted) score of uORFs, CSV</td></tr>\n",
    "<tr><td>uORFs_wrent_seq</td><td>Sequences used to calculate WRENT scores, CSV</td></tr>\n",
    "<tr><td>uORFs_sec_struct_EFE_L</td><td>Computationally predicted ensemble free energy of secondary structure around uORF start left (-25), CSV</td></tr>\n",
    "<tr><td>uORFs_sec_struct_EFE_R</td><td>Computationally predicted ensemble free energy of secondary structure around uORF start right (+1), CSV</td></tr>\n",
    "<tr><td>uORFs_start_pos_wrt_tss</td><td>Start position of uORFs with respect to 5' end of transcript, CSV</td></tr>\n",
    "<tr><td>uORFs_end_pos_wrt_CDS</td><td>End position of uORFs with respect to start of CDS, CSV</td></tr>\n",
    "<tr><td>CDS_read</td><td>Number of reads in CDS</td></tr>\n",
    "<tr><td>CDS_GC</td><td>Total GC content of CDS</td></tr>\n",
    "<tr><td>CDS_length</td><td>Length of CDS</td></tr>\n",
    "<tr><td>CDS_wrent_score</td><td>Calculated WRENT score of CDS</td></tr>\n",
    "<tr><td>CDS_urent_score</td><td>Calculated URENT score of CDS</td></tr>\n",
    "<tr><td>CDS_wrent_seq</td><td>Sequence used to calculate WRENT score, CSV</td></tr>\n",
    "<tr><td>CDS_sec_struct_EFE_L</td><td>Computationally predicted ensemble free energy of secondary structure around CDS start, left (-25)</td></tr>\n",
    "<tr><td>CDS_sec_struct_EFE_R</td><td>Computationally predicted ensemble free energy of secondary structure around CDS start, right (+1)</td></tr>\n",
    "<tr><td>ORF_5CI3</td><td>Annotation of ORF within transcript, if begins in **5**' UTR, is **C**DS, begins with**I**n CDS, or in **3**' UTR, string of letters</td></tr>\n",
    "<tr><td>ORFs_wrent_score</td><td>Calculated WRENT score of ORFs, CSV</td></tr>\n",
    "<tr><td>ORFs_urent_score</td><td>Calculated URENT score of ORFs, CSV</td></tr>\n",
    "<tr><td>ORFs_wrent_seq</td><td>Sequences used to calculate WRENT score, CSV</td></tr>\n",
    "<tr><td>ORFs_sec_struct_EFE_L</td><td>Computationally predicted ensemble free energy of secondary structure around ORF start, left (-25), CSV</td></tr>\n",
    "<tr><td>ORFs_sec_struct_EFE_R</td><td>Computationally predicted ensemble free energy of secondary structure around ORF start, right (+1), CSV</td></tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "INDEX = \"Transcript\"\n",
    "COLUMNS = [\"Gene\", \"Gene_Name\", \"Gene_Expression_FPKM\", \"UTR5_length\", \"UTR3_length\", \"UTR5_reads\", \"UTR5_reads_trunc\",\n",
    "           \"UTR5_GC\", \"num_uORFs\", \"uORFs_reads\", \"UTR5_uORF_reads\", \"UTR5_uORF_tlength\", \"uORFs_length\",\n",
    "           \"uORFs_wrent_score\", \"uORFs_urent_score\", \"uORFs_wrent_seq\", \"uORFs_sec_struct_EFE_L\", \"uORFs_sec_struct_EFE_R\",\n",
    "           \"uORFs_start_pos_wrt_tss\", \"uORFs_end_pos_wrt_CDS\",\n",
    "           \"CDS_reads\", \"CDS_length\", \"CDS_GC\",\n",
    "           \"CDS_wrent_score\", \"CDS_urent_score\", \"CDS_wrent_seq\", \"CDS_sec_struct_EFE_L\", \"CDS_sec_struct_EFE_R\",\n",
    "           \"ORFs_5CI3\", \"ORFs_wrent_score\", \"ORFs_urent_score\", \"ORFs_wrent_seq\",\n",
    "           \"ORFs_sec_struct_EFE_L\", \"ORFs_sec_struct_EFE_R\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculated values are first stored in a dictionary, and a DataFrame is constructed afterward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# INITIALIZE DICTIONARIES FOR STORAGE\n",
    "df_main = DataFrame(columns=COLUMNS)\n",
    "df_main.index.name = INDEX\n",
    "\n",
    "# FILE ITERATOR FOR .trpedf\n",
    "trpedf_file_iterator = pd.read_table(TRPE_FILE, converters=CONVERTERS, chunksize=1)\n",
    "transcript_count2 = 0      # counter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A counter counts every 500 transcripts and prints to output.\n",
    "\n",
    "Initiation contexts for weighted Kozak scoring and secondary structure prediction are first determined using the earlier helper functions. uORFs, CDSes, then all ORFs are scored, with ORFs annotated for whether they begin in the 5' UTR, are the CDS, begin within the CDS, or in the 3' UTR.\n",
    "\n",
    "An additional dictionary stores the profile of ribosome profiling reads around the starts and ends of transcript CDSes, with equal weight to each transcript, to generate a metaprofile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# ITERATE OVER .trpedf\n",
    "for transcript in trpedf_file_iterator:\n",
    "    \n",
    "    # Counter for iPython Notebook display, every 100 transcripts\n",
    "    transcript_count2 += 1\n",
    "    if transcript_count2 % 100 == 0:\n",
    "        print \"%d transcripts read\" % transcript_count2\n",
    "    \n",
    "    # Reads in data from each transcript\n",
    "    RPF_csvProfile = transcript[\"RPF_csvProfile\"][0]\n",
    "    seq = SEQS[transcript[\"Transcript\"][0]].seq     # get transcript sequence\n",
    "    if \"N\" in seq: continue\n",
    "    ORF_starts = transcript[\"ORF_starts\"][0]\n",
    "    ORF_ends = transcript[\"ORF_ends\"][0]\n",
    "    if type(ORF_starts) is np.int64:      # corrects for single-entry\n",
    "        ORF_starts = (ORF_starts,)\n",
    "        ORF_ends = (ORF_ends,)\n",
    "    CDS = transcript[\"CDS\"][0]\n",
    "    ORFs = zip(ORF_starts, ORF_ends)\n",
    "    uORFs = [ORF for ORF in ORFs if ORF[0] < CDS[0]]  # uORFs defined as beginning before CDS\n",
    "\n",
    "    # Determines initiation contexts for weighted Kozak scoring of:\n",
    "    # uORFs\n",
    "    uORF_wrent_seqs = [wrent_context_seq(seq, uORF, WRENT_CONTEXT) for uORF in uORFs]\n",
    "        \n",
    "    # CDS\n",
    "    CDS_wrent_seq = wrent_context_seq(seq, CDS, WRENT_CONTEXT)\n",
    "    \n",
    "    # All ORFs\n",
    "    ORF_wrent_seqs = [wrent_context_seq(seq, ORF, WRENT_CONTEXT) for ORF in ORFs]\n",
    "        \n",
    "    # Determines initiation contexts for secondary structure prediction of:\n",
    "    # uORFs\n",
    "    uORF_ss_seqs_l = [ss_context_seq(seq, uORF, SS_CONTEXT_L) for uORF in uORFs]\n",
    "    uORF_ss_seqs_r = [ss_context_seq(seq, uORF, SS_CONTEXT_R) for uORF in uORFs]\n",
    "    \n",
    "    # CDS\n",
    "    CDS_ss_seq_l = ss_context_seq(seq, CDS, SS_CONTEXT_L)\n",
    "    CDS_ss_seq_r = ss_context_seq(seq, CDS, SS_CONTEXT_R)\n",
    "    \n",
    "    # All ORFs\n",
    "    ORF_ss_seqs_l = [ss_context_seq(seq, ORF, SS_CONTEXT_L) for ORF in ORFs]\n",
    "    ORF_ss_seqs_r = [ss_context_seq(seq, ORF, SS_CONTEXT_R) for ORF in ORFs]\n",
    "    \n",
    "    # Calculate and store values in main dictionary\n",
    "    entry = {}\n",
    "    for j in (\"Gene\", \"Gene_Name\", \"Gene_Expression_FPKM\"):\n",
    "        entry[j] = transcript[j][0]\n",
    "    entry[\"UTR5_length\"] = CDS[0]\n",
    "    entry[\"UTR3_length\"] = len(RPF_csvProfile) - CDS[1]\n",
    "    entry[\"UTR5_reads\"] = sum(RPF_csvProfile[:CDS[0]])\n",
    "    entry[\"UTR5_reads_trunc\"] = sum(RPF_csvProfile[:CDS[0]-ORF_END_TRIM])\n",
    "    entry[\"UTR5_GC\"] = 0.0 if CDS[0] == 0 else float(seq[:CDS[0]].count(\"G\") + seq[:CDS[0]].count(\"C\")) / CDS[0]\n",
    "    entry[\"num_uORFs\"] = len(uORFs)\n",
    "\n",
    "    entry[\"uORFs_reads\"] = [sum(RPF_csvProfile[uORF[0]:uORF[1]-ORF_END_TRIM]) for uORF in uORFs]\n",
    "    \n",
    "    transcript_UTR5_uORF_segments = UTR5_uORF_segments(CDS, uORFs)\n",
    "    entry[\"UTR5_uORF_reads\"] = sum([i for i, j in zip(RPF_csvProfile, transcript_UTR5_uORF_segments) if j])\n",
    "    entry[\"UTR5_uORF_tlength\"] = sum(transcript_UTR5_uORF_segments)\n",
    "    \n",
    "    entry[\"uORFs_length\"] = [uORF[1] - uORF[0] for uORF in uORFs]\n",
    "    entry[\"uORFs_wrent_score\"] = [wrent_score_seq(en_score, init_seq) \\\n",
    "                                            for init_seq in uORF_wrent_seqs]\n",
    "    entry[\"uORFs_urent_score\"] = [wrent_score_seq(en_unweighted_score, init_seq) \\\n",
    "                                            for init_seq in uORF_wrent_seqs]\n",
    "    entry[\"uORFs_wrent_seq\"] = uORF_wrent_seqs\n",
    "    entry[\"uORFs_sec_struct_EFE_L\"] = [RNA.pf_fold(init_seq)[1] for init_seq in uORF_ss_seqs_l]\n",
    "    entry[\"uORFs_sec_struct_EFE_R\"] = [RNA.pf_fold(init_seq)[1] for init_seq in uORF_ss_seqs_r]\n",
    "    entry[\"uORFs_start_pos_wrt_tss\"] = [uORF[0] for uORF in uORFs]\n",
    "    entry[\"uORFs_end_pos_wrt_CDS\"] = [uORF[1] - CDS[0] for uORF in uORFs]\n",
    "\n",
    "    entry[\"CDS_reads\"] = sum(RPF_csvProfile[CDS[0]:CDS[1]-ORF_END_TRIM])\n",
    "    entry[\"CDS_GC\"] = 0.0 if CDS[0] == CDS[1] else float(seq[CDS[0]:CDS[1]].count(\"G\") \\\n",
    "                                                         + seq[CDS[0]:CDS[1]].count(\"C\")) \\\n",
    "                                                         / (CDS[1] - CDS[0])\n",
    "    entry[\"CDS_length\"] = CDS[1] - CDS[0]\n",
    "    entry[\"CDS_wrent_score\"] = wrent_score_seq(en_score, CDS_wrent_seq)\n",
    "    entry[\"CDS_urent_score\"] = wrent_score_seq(en_unweighted_score, CDS_wrent_seq)\n",
    "    entry[\"CDS_wrent_seq\"] = CDS_wrent_seq\n",
    "    entry[\"CDS_sec_struct_EFE_L\"] = RNA.pf_fold(CDS_ss_seq_l)[1]\n",
    "    entry[\"CDS_sec_struct_EFE_R\"] = RNA.pf_fold(CDS_ss_seq_r)[1]\n",
    "\n",
    "    # 5CIS refers to annotation of ORF as beginning in 5' UTR ('5'),\n",
    "    # being the CDS ('C'),\n",
    "    # beginning in the CDS ('I'), or\n",
    "    # beginning in the 3' UTR ('3')\n",
    "    entry[\"ORFs_5CI3\"] = \"\".join(['C' if ORF_start == CDS[0] else \\\n",
    "                                  '5' if ORF_start < CDS[0] else \\\n",
    "                                  'I' if CDS[0] < ORF_start < CDS[1] else \\\n",
    "                                  '3' for ORF_start in ORF_starts])\n",
    "    entry[\"ORFs_wrent_score\"] = [wrent_score_seq(en_score, init_seq) \\\n",
    "                                           for init_seq in ORF_wrent_seqs]\n",
    "    entry[\"ORFs_urent_score\"] = [wrent_score_seq(en_unweighted_score, init_seq) \\\n",
    "                                           for init_seq in ORF_wrent_seqs]\n",
    "    entry[\"ORFs_wrent_seq\"] = ORF_wrent_seqs\n",
    "    entry[\"ORFs_sec_struct_EFE_L\"] = [RNA.pf_fold(init_seq)[1] for init_seq in ORF_ss_seqs_l]\n",
    "    entry[\"ORFs_sec_struct_EFE_R\"] = [RNA.pf_fold(init_seq)[1] for init_seq in ORF_ss_seqs_r]\n",
    "    \n",
    "    df_main.loc[transcript[\"Transcript\"][0]] = Series(entry)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "df_main.to_csv(DATA_DIR + stage + \"_main.df\", sep=\"\\t\")"
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}