212/ssbio_01.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SSBIO: Collection of functions used throughout a structural reconstruction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# ESSENTIAL\n",
    "import os\n",
    "import ast\n",
    "import json\n",
    "import pickle\n",
    "import operator\n",
    "import warnings\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import collections\n",
    "from collections import defaultdict\n",
    "from dateutil.parser import parse as dateparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# CLEARING OUTPUT\n",
    "import sys \n",
    "import time\n",
    "from IPython.display import display, clear_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# SUPPRESSING ALL OUTPUT\n",
    "from contextlib import contextmanager\n",
    "import sys, os\n",
    " \n",
    "@contextmanager\n",
    "def suppress_stdout():\n",
    "    with open(os.devnull, \"w\") as devnull:\n",
    "        old_stdout = sys.stdout\n",
    "        sys.stdout = devnull\n",
    "        try:  \n",
    "            yield\n",
    "        finally:\n",
    "            sys.stdout = old_stdout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import sys, time\n",
    "import itertools\n",
    "try:\n",
    "    from IPython.display import clear_output\n",
    "    have_ipython = True\n",
    "except ImportError:\n",
    "    have_ipython = False\n",
    "\n",
    "class ProgressBar:\n",
    "    def __init__(self, iterations):\n",
    "        self.iterations = iterations\n",
    "        self.prog_bar = '[]'\n",
    "        self.fill_char = '*'\n",
    "        self.width = 40\n",
    "        self.__update_amount(0)\n",
    "        if have_ipython:\n",
    "            self.animate = self.animate_ipython\n",
    "        else:\n",
    "            self.animate = self.animate_noipython\n",
    "\n",
    "    def animate_ipython(self, iter):\n",
    "        print '\\r', self,\n",
    "        sys.stdout.flush()\n",
    "        self.update_iteration(iter + 1)\n",
    "\n",
    "    def update_iteration(self, elapsed_iter):\n",
    "        self.__update_amount((elapsed_iter / float(self.iterations)) * 100.0)\n",
    "        self.prog_bar += '  %d of %s complete' % (elapsed_iter, self.iterations)\n",
    "\n",
    "    def __update_amount(self, new_amount):\n",
    "        percent_done = int(round((new_amount / 100.0) * 100.0))\n",
    "        all_full = self.width - 2\n",
    "        num_hashes = int(round((percent_done / 100.0) * all_full))\n",
    "        self.prog_bar = '[' + self.fill_char * num_hashes + ' ' * (all_full - num_hashes) + ']'\n",
    "        pct_place = (len(self.prog_bar) // 2) - len(str(percent_done))\n",
    "        pct_string = '%d%%' % percent_done\n",
    "        self.prog_bar = self.prog_bar[0:pct_place] + \\\n",
    "            (pct_string + self.prog_bar[pct_place + len(pct_string):])\n",
    "\n",
    "    def __str__(self):\n",
    "        return str(self.prog_bar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "cobra/solvers/__init__.py:63 \u001b[1;31mUserWarning\u001b[0m: No LP solvers found\n",
      "@> ProDy is configured: verbosity='none'\n",
      "INFO:.prody:ProDy is configured: verbosity='none'\n"
     ]
    }
   ],
   "source": [
    "# COBRAPY\n",
    "# used to read GEMs into python\n",
    "from cobra import Model, Reaction, Metabolite\n",
    "from cobra.io.sbml import create_cobra_model_from_sbml_file\n",
    "from cobra.io.mat import load_matlab_model\n",
    "\n",
    "# LIBSBML\n",
    "# used to directly read SBML files\n",
    "from libsbml import *\n",
    "reader = SBMLReader()\n",
    "\n",
    "# REQUESTS\n",
    "# used for interfacing with REST\n",
    "import requests\n",
    "import urllib2\n",
    "import gzip\n",
    "\n",
    "# STRINGIO\n",
    "import StringIO as sio\n",
    "\n",
    "# BIOSERVICES\n",
    "# http://bioinformatics.oxfordjournals.org/content/29/24/3241\n",
    "from bioservices.uniprot import UniProt\n",
    "bsup = UniProt()\n",
    "\n",
    "# PRODY\n",
    "# http://bioinformatics.oxfordjournals.org/content/27/11/1575\n",
    "# used for protein structure information\n",
    "import prody as pr\n",
    "pr.confProDy(verbosity='none')\n",
    "\n",
    "# BIOPYTHON\n",
    "# used for sequence alignment\n",
    "from Bio import SeqIO\n",
    "from Bio import AlignIO\n",
    "from Bio.Seq import Seq\n",
    "from Bio.SeqRecord import SeqRecord\n",
    "from Bio.Alphabet import IUPAC\n",
    "from Bio.PDB.Polypeptide import *\n",
    "from Bio.PDB.PDBParser import PDBParser\n",
    "from Bio.PDB.MMCIFParser import MMCIFParser\n",
    "from Bio.PDB import PDBIO\n",
    "from Bio import PDB\n",
    "\n",
    "# BEAUTIFULSOUP\n",
    "from bs4 import BeautifulSoup\n",
    "\n",
    "# TEMPFILE - for storing files temporarily (cross-platform!)\n",
    "import tempfile\n",
    "\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# PANDAS DISPLAY OPTIONS\n",
    "pd.set_option('display.max_rows', 500)\n",
    "pd.set_option('display.max_columns', 500)\n",
    "# pd.set_option('display.width', 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# MATPLOTLIB FIGURES\n",
    "import seaborn as sns\n",
    "sns.set_style(\"dark\")\n",
    "sns.set_context(\"poster\")\n",
    "\n",
    "from matplotlib import rcParams\n",
    "rcParams['figure.dpi'] = 300\n",
    "rcParams['lines.linewidth'] = 2\n",
    "rcParams['axes.facecolor'] = 'white'\n",
    "rcParams['patch.edgecolor'] = 'white'\n",
    "\n",
    "from matplotlib import font_manager as fm\n",
    "proptease = fm.FontProperties()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: flatlist_dropdup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def flatlist_dropdup(list_of_lists):\n",
    "    return list(set([str(item) for sublist in list_of_lists for item in sublist]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: de_unicodeify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def de_unicodeify(data):\n",
    "    if isinstance(data, basestring):\n",
    "        return str(data)\n",
    "    elif isinstance(data, collections.Mapping):\n",
    "        return dict(map(de_unicodeify, data.iteritems()))\n",
    "    elif isinstance(data, collections.Iterable):\n",
    "        return type(data)(map(de_unicodeify, data))\n",
    "    else:\n",
    "        return data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "## MAPPING TOOLS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: kegg_mapper\n",
    "#### KEGG ORGANISM & MAP TO DATABASE ID -> DICTIONARY OF MAPPED GENE IDS\n",
    "This function simply creates a dictionary that maps KEGG IDs to either the **NCBI Entrez Gene ID** or the **UniProt ID**. It works using the KEGG REST API documented here: http://www.kegg.jp/kegg/docs/keggapi.html\n",
    "\n",
    "**INPUTS**:\n",
    "1. kegg_organism\n",
    "    * Pick the one the SBML model is built off of from here: http://www.genome.jp/kegg/catalog/org_list.html\n",
    "    * E. coli - 'eco'\n",
    "    * Homo sapiens - 'hsa'\n",
    "    * Thermotoga maritima (1999) - 'tma'\n",
    "2. map_db\n",
    "    * UniProt - 'uniprot'\n",
    "    * NCBI Entrez - 'ncbi-geneid'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def kegg_mapper(kegg_organism, map_db='uniprot'):\n",
    "    '''\n",
    "    Generates a dictionary that maps KEGG gene IDs to NCBI Entrez gene IDs ('ncbi-geneid') or UniProt ('uniprot').\n",
    "    Input:  kegg_organism - the KEGG organism name which you can determine from http://www.genome.jp/kegg/catalog/org_list.html\n",
    "            map_db - ncbi-geneid OR uniprot (default): the database you want to map to\n",
    "    Output: id_mapper - a dictionary of {KEGG_ID: mapped_ID}\n",
    "    '''\n",
    "\n",
    "    id_mapper = {}\n",
    "    r = requests.post('http://rest.kegg.jp/conv/%s/%s' % (map_db,kegg_organism))\n",
    "\n",
    "    for line in r.content.split('\\n'):\n",
    "        if not line: continue\n",
    "\n",
    "        idents = line.split('\\t')\n",
    "\n",
    "        orig = idents[0].split(':')\n",
    "        orig_database = orig[0]\n",
    "        orig_id = orig[1]\n",
    "\n",
    "        conv = idents[1].split(':')\n",
    "        conv_database = conv[0]\n",
    "        conv_id = conv[1]\n",
    "\n",
    "        id_mapper[orig_id] = conv_id\n",
    "\n",
    "    return id_mapper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: bioservices_uniprot_mapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": false,
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "def bioservices_uniprot_mapper(map_from, map_to, ident):\n",
    "    return de_unicodeify(dict(bsup.mapping(fr=map_from, to=map_to, query=ident)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: uniprot_reviewed_checker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# status can be \"unreviewed\" or \"reviewed\"\n",
    "\n",
    "def uniprot_reviewed_checker(uniprot_ids):\n",
    "    # splitting query up into managable sizes (200 IDs each)\n",
    "    Nmax = 200\n",
    "    N, rest = divmod(len(uniprot_ids), Nmax)\n",
    "\n",
    "    uni_rev_dict = {}\n",
    "\n",
    "    if rest>0:\n",
    "        N+=1\n",
    "    for i in range(0,N):\n",
    "        i1 = i*Nmax\n",
    "        i2 = (i+1)*Nmax\n",
    "        if i2>len(uniprot_ids):\n",
    "            i2 = len(uniprot_ids)\n",
    "\n",
    "        query = uniprot_ids[i1:i2]\n",
    "\n",
    "        query_string = ''\n",
    "        for x in query:\n",
    "            query_string += 'id:' + x + '+OR+'\n",
    "        query_string = query_string.strip('+OR+')\n",
    "    \n",
    "        uni_rev_raw = sio.StringIO(bsup.search(query_string, columns='id,reviewed', frmt='tab'))\n",
    "        uni_rev_df = pd.read_table(uni_rev_raw, sep='\\t', index_col=0)\n",
    "        uni_rev_dict_adder = uni_rev_df.to_dict()['Status']\n",
    "        uni_rev_dict.update(uni_rev_dict_adder)\n",
    "        \n",
    "        clear_output(wait=True)\n",
    "        sys.stdout.flush()\n",
    "    clear_output(wait=True)\n",
    "    sys.stdout.flush()\n",
    "    \n",
    "    return de_unicodeify(uni_rev_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: pdb_current_checker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "theoretical_pdbs_link = \"ftp://ftp.wwpdb.org/pub/pdb/data/structures/models/index/titles.idx\"\n",
    "\n",
    "response = urllib2.urlopen(theoretical_pdbs_link)\n",
    "theoretical_pdbs_raw = sio.StringIO(response.read())\n",
    "theoretical_pdbs_df = pd.read_table(theoretical_pdbs_raw, sep='\\t', header=None)\n",
    "\n",
    "PDB_THEORETICAL = theoretical_pdbs_df[0].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "PDB_OBSOLETE_MAPPING = {}\n",
    "\n",
    "req = urllib2.Request('ftp://ftp.wwpdb.org/pub/pdb/data/status/obsolete.dat')\n",
    "response = urllib2.urlopen(req)\n",
    "for line in response:\n",
    "    entry = line.split()\n",
    "    if entry[0] == 'LIST':\n",
    "        continue\n",
    "    if len(entry) == 3:\n",
    "        PDB_OBSOLETE_MAPPING[entry[2].upper()] = 'obsolete'\n",
    "    if len(entry) >= 4:\n",
    "        new = [y.upper() for y in entry[3:]]\n",
    "        PDB_OBSOLETE_MAPPING[entry[2].upper()] = new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# status can be \"theoretical\", \"current\", or a PDB ID giving the non-obsolete entry\n",
    "\n",
    "def pdb_current_checker(pdb_ids):\n",
    "    if isinstance(pdb_ids, str):\n",
    "        pdb_ids = [pdb_ids]\n",
    "    \n",
    "    pdb_status = {}\n",
    "    \n",
    "    theoretical = list(set(PDB_THEORETICAL).intersection(pdb_ids))\n",
    "    obsolete = list(set(PDB_OBSOLETE_MAPPING.keys()).intersection(pdb_ids))\n",
    "    current = list(set(pdb_ids).difference(theoretical,obsolete))\n",
    "    \n",
    "    for t in theoretical:\n",
    "        pdb_status[t] = 'theoretical'\n",
    "    for o in obsolete:\n",
    "        pdb_status[o] = PDB_OBSOLETE_MAPPING[o]\n",
    "    for c in current:\n",
    "        pdb_status[c] = 'current'\n",
    "        \n",
    "    return pdb_status"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Internal use: SIFTS file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import os.path, time\n",
    "now = time.time()\n",
    "twodays_ago = now - 60*60*24*2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "baseURL = \"ftp://ftp.ebi.ac.uk/pub/databases/msd/sifts/flatfiles/csv/\"\n",
    "filename = \"pdb_chain_uniprot.csv.gz\"\n",
    "\n",
    "final_filename = tempfile.gettempdir() + '/' + filename.split('.gz')[0]\n",
    "\n",
    "if os.path.exists(final_filename):\n",
    "    final_filename_creation_time = os.path.getmtime(final_filename)\n",
    "    \n",
    "    if final_filename_creation_time < twodays_ago:\n",
    "        response = urllib2.urlopen(baseURL + filename)\n",
    "        compressedFile = sio.StringIO(response.read())\n",
    "        decompressedFile = gzip.GzipFile(fileobj=compressedFile)\n",
    "\n",
    "        with open(final_filename, 'w') as outfile:\n",
    "            outfile.write(decompressedFile.read())\n",
    "\n",
    "        SIFTS = pd.read_csv(final_filename, skiprows=1, index_col=['PDB','CHAIN'])\n",
    "        \n",
    "    else:\n",
    "        SIFTS = pd.read_csv(final_filename, skiprows=1, index_col=['PDB','CHAIN'])\n",
    "else:\n",
    "    response = urllib2.urlopen(baseURL + filename)\n",
    "    compressedFile = sio.StringIO(response.read())\n",
    "    decompressedFile = gzip.GzipFile(fileobj=compressedFile)\n",
    "\n",
    "    with open(final_filename, 'w') as outfile:\n",
    "        outfile.write(decompressedFile.read())\n",
    "\n",
    "    SIFTS = pd.read_csv(final_filename, skiprows=1, index_col=['PDB','CHAIN'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "SIFTS['PDB_BEG_INT'] = SIFTS['PDB_BEG'].replace(to_replace=r'[^\\d-]+', value='', regex=True)\n",
    "SIFTS['SP_BEG_INT'] = SIFTS['SP_BEG'].replace(to_replace=r'[^\\d-]+', value='', regex=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "SIFTS['OFFSET'] = SIFTS['SP_BEG_INT'].astype(int) - SIFTS['PDB_BEG_INT'].astype(int)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: sifts_pdb_chain_to_uniprot\n",
    "#### PDB ID & CHAIN ID -> UNIPROT ID\n",
    "This function uses the SIFTS mapping service to map PDB chains to their corresponding UniProt entries. PDB files have their own 'dbref' entry, however sometimes these map to obsolete UniProt entries.\n",
    "\n",
    "**INPUTS**:\n",
    "1. pdb\n",
    "    * a string - PDB ID\n",
    "2. chain\n",
    "    * a string - chain ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def sifts_pdb_chain_to_uniprot(pdb,chain):\n",
    "    return SIFTS.ix[(pdb.lower(),chain.upper())]['SP_PRIMARY'].unique().tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div style=\"max-height:1000px;max-width:1500px;overflow:auto;\">\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>SP_PRIMARY</th>\n",
       "      <th>RES_BEG</th>\n",
       "      <th>RES_END</th>\n",
       "      <th>PDB_BEG</th>\n",
       "      <th>PDB_END</th>\n",
       "      <th>SP_BEG</th>\n",
       "      <th>SP_END</th>\n",
       "      <th>PDB_BEG_INT</th>\n",
       "      <th>SP_BEG_INT</th>\n",
       "      <th>OFFSET</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>CHAIN</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>A</th>\n",
       "      <td> P25714</td>\n",
       "      <td> 1</td>\n",
       "      <td> 548</td>\n",
       "      <td> 1</td>\n",
       "      <td> 548</td>\n",
       "      <td> 1</td>\n",
       "      <td> 548</td>\n",
       "      <td> 1</td>\n",
       "      <td> 1</td>\n",
       "      <td> 0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Z</th>\n",
       "      <td> P68699</td>\n",
       "      <td> 1</td>\n",
       "      <td>  79</td>\n",
       "      <td> 1</td>\n",
       "      <td>  79</td>\n",
       "      <td> 1</td>\n",
       "      <td>  79</td>\n",
       "      <td> 1</td>\n",
       "      <td> 1</td>\n",
       "      <td> 0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      SP_PRIMARY  RES_BEG  RES_END PDB_BEG PDB_END  SP_BEG  SP_END  \\\n",
       "CHAIN                                                                \n",
       "A         P25714        1      548       1     548       1     548   \n",
       "Z         P68699        1       79       1      79       1      79   \n",
       "\n",
       "      PDB_BEG_INT  SP_BEG_INT  OFFSET  \n",
       "CHAIN                                  \n",
       "A               1           1       0  \n",
       "Z               1           1       0  "
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "SIFTS.ix['4utq']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "##UNIPROT RELATED"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: uniprot_valid_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def uniprot_valid_id(test_id):\n",
    "    # regex from: http://www.uniprot.org/help/accession_numbers\n",
    "    valid_id = re.compile(\"[OPQ][0-9][A-Z0-9]{3}[0-9]|[A-NR-Z][0-9]([A-Z][A-Z0-9]{2}[0-9]){1,2}\")\n",
    "    if valid_id.match(test_id): return True\n",
    "    else: return False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: uniprot_metadata\n",
    "#### UNIPROT ID -> UNIPROT METADATA\n",
    "\n",
    "This function retrieves data for a UniProt entry and returns a dictionary of that information.\n",
    "\n",
    "**INPUT**:\n",
    "1. uniprot_id\n",
    "\n",
    "**OUTPUT**:\n",
    "1. uniprot_metadata_dict\n",
    "    * { 'UNIPROT': id,\n",
    "        'DESCRIPTION': protein function,\n",
    "        'KEGG': associated KEGG IDs,\n",
    "        'SEQ': uniprot canonical sequence,\n",
    "        'PDB': list of PDB IDs\n",
    "      }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def parse_uniprot_txt_file(cache_txt):\n",
    "    \"\"\"\n",
    "    Parses the text of metadata retrieved from uniprot.org.\n",
    "\n",
    "    Only a few fields have been parsed, but this provides a\n",
    "    template for the other fields.\n",
    "\n",
    "    A single description is generated from joining alternative\n",
    "    descriptions.\n",
    "\n",
    "    Returns a dictionary with the main UNIPROT ACC as keys.\n",
    "    \"\"\"\n",
    "\n",
    "    tag = None\n",
    "    uniprot_id = None\n",
    "    metadata_by_seqid = {}\n",
    "    for l in cache_txt.splitlines():\n",
    "        test_tag = l[:5].strip()\n",
    "        if test_tag and test_tag != tag:\n",
    "            tag = test_tag\n",
    "        line = l[5:].strip()\n",
    "        words = line.split()\n",
    "        if tag == \"ID\":\n",
    "            uniprot_id = words[0]\n",
    "            is_reviewed = words[1].startswith('Reviewed')\n",
    "            length = int(words[2])\n",
    "            metadata_by_seqid[uniprot_id] = {\n",
    "                'id': uniprot_id,\n",
    "                'is_reviewed': is_reviewed,\n",
    "                'length': length,\n",
    "                'sequence': '',\n",
    "                'accs': [],\n",
    "            }\n",
    "            entry = metadata_by_seqid[uniprot_id]\n",
    "        if tag == \"DT\":\n",
    "            #DT   01-OCT-1996, integrated into UniProtKB/Swiss-Prot.\n",
    "            #DT   17-OCT-2006, sequence version 3.\n",
    "            #DT   22-JUL-2015, entry version 166.\n",
    "            comma_split = line.split(',')\n",
    "            if 'sequence version' in comma_split[1]:\n",
    "#                 print 'sequence_version', comma_split[0], dateparse(comma_split[0]).date()\n",
    "                entry['sequence_version'] = str(dateparse(comma_split[0]).date())\n",
    "            elif 'entry version' in comma_split[1]:\n",
    "#                 print 'entry_version', comma_split[0], dateparse(comma_split[0]).date()\n",
    "                entry['entry_version'] = str(dateparse(comma_split[0]).date())\n",
    "        if tag == \"SQ\":\n",
    "            if words[0] != \"SEQUENCE\":\n",
    "                entry['sequence'] += ''.join(words)\n",
    "        if tag == \"AC\":\n",
    "            accs = [w.replace(\";\", \"\") for w in words]\n",
    "            entry['accs'].extend(accs)\n",
    "        if tag == \"DR\":\n",
    "            if 'PDB' in words[0]:\n",
    "                if 'pdb' not in entry:\n",
    "                    entry['pdb'] = words[1][:-1]\n",
    "                if 'pdbs' not in entry:\n",
    "                    entry['pdbs'] = []\n",
    "                entry['pdbs'].append(words[1][:-1])\n",
    "            if 'RefSeq' in words[0]:\n",
    "                if 'refseq' not in entry:\n",
    "                    entry['refseq'] = []\n",
    "                ids = [w[:-1] for w in words[1:]]\n",
    "                entry['refseq'].extend(ids)\n",
    "            if 'KEGG' in words[0]:\n",
    "                if 'kegg' not in entry:\n",
    "                    entry['kegg'] = []\n",
    "                ids = [w[:-1] for w in words[1:]]\n",
    "                ids = filter(lambda w: len(w) > 1, ids)\n",
    "                entry['kegg'].extend(ids)\n",
    "            if 'GO' in words[0]:\n",
    "                if 'go' not in entry:\n",
    "                    entry['go'] = []\n",
    "                entry['go'].append(' '.join(words[1:]))\n",
    "            if 'Pfam' in words[0]:\n",
    "                if 'pfam' not in entry:\n",
    "                    entry['pfam'] = []\n",
    "                entry['pfam'].append(words[1][:-1])\n",
    "        if tag == \"GN\":\n",
    "            if 'gene' not in entry and len(words) > 0:\n",
    "                pieces = words[0].split(\"=\")\n",
    "                if len(pieces) > 1 and 'name' in pieces[0].lower():\n",
    "                    entry['gene'] = pieces[1].replace(';', '').replace(',', '')\n",
    "        if tag == \"OS\":\n",
    "            # OS   Homo sapiens (Human).\n",
    "            if 'organism' not in entry:\n",
    "                entry['organism'] = \"\"\n",
    "            entry['organism'] += line\n",
    "        if tag == \"DE\":\n",
    "            if 'descriptions' not in entry:\n",
    "                entry['descriptions'] = []\n",
    "            entry['descriptions'].append(line)\n",
    "        if tag == \"CC\":\n",
    "            if 'comment' not in entry:\n",
    "                entry['comment'] = ''\n",
    "            entry['comment'] += line + '\\n'\n",
    "\n",
    "    for entry in metadata_by_seqid.values():\n",
    "        descriptions = entry['descriptions']\n",
    "        for i in reversed(range(len(descriptions))):\n",
    "            description = descriptions[i]\n",
    "            if 'Short' in description or 'Full' in description or 'EC=' in description:\n",
    "                if 'Short' in description or 'Full' in description:\n",
    "                    j = description.find('=')\n",
    "                    descriptions[i] = description[j+1:].replace(';', '')\n",
    "                    if 'description' not in entry:\n",
    "                        entry['description'] = []\n",
    "                    entry['description'].append(descriptions[i])\n",
    "                if 'EC=' in description:\n",
    "                    j = description.find('=')\n",
    "                    descriptions[i] = description[j+1:].replace(';', '')\n",
    "                    if '{' in descriptions[i]:\n",
    "                        descriptions[i] = descriptions[i].split(' {')[0]\n",
    "                    if 'ec' not in entry:\n",
    "                        entry['ec'] = []\n",
    "                    entry['ec'].append(descriptions[i])\n",
    "            else:\n",
    "                del descriptions[i]\n",
    "    \n",
    "    return metadata_by_seqid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def uniprot_metadata(uniprot_ids):\n",
    "    '''\n",
    "    Input: UniProt ID or IDs\n",
    "    Output: dictionary of metadata associated with the UniProt IDs\n",
    "    '''\n",
    "    counter = 1\n",
    "    \n",
    "    if isinstance(uniprot_ids, str):\n",
    "        uniprot_ids = [uniprot_ids]\n",
    "        single = True\n",
    "    else: single = False\n",
    "    \n",
    "    uniprot_metadata_raw = bsup.retrieve(uniprot_ids, frmt='txt')\n",
    "    uniprot_metadata_final = {}\n",
    "    \n",
    "    for uniprot_id in uniprot_ids:\n",
    "        \n",
    "        clear_output(wait=True)\n",
    "        print('***PROGRESS: %s - PARSED %d/%d UNIPROT IDS***\\n' % (uniprot_id, counter, len(uniprot_ids)))\n",
    "        counter += 1\n",
    "        sys.stdout.flush()\n",
    "        \n",
    "        uniprot_metadata_dict = {}\n",
    "        \n",
    "        if single:\n",
    "            metadata = parse_uniprot_txt_file(uniprot_metadata_raw)\n",
    "        else:\n",
    "            metadata = parse_uniprot_txt_file(uniprot_metadata_raw[uniprot_ids.index(uniprot_id)])\n",
    "            \n",
    "        metadata_key = metadata.keys()[0]\n",
    "    \n",
    "        uniprot_metadata_dict['u_uniprot_acc'] = uniprot_id\n",
    "        uniprot_metadata_dict['u_seq'] = metadata[metadata_key]['sequence']\n",
    "        uniprot_metadata_dict['u_seq_len'] = len(str(metadata[metadata_key]['sequence']))\n",
    "        uniprot_metadata_dict['u_reviewed'] = metadata[metadata_key]['is_reviewed']\n",
    "        uniprot_metadata_dict['u_seq_version'] = metadata[metadata_key]['sequence_version']\n",
    "        uniprot_metadata_dict['u_entry_version'] = metadata[metadata_key]['entry_version']\n",
    "        if 'gene' in metadata[metadata_key]:\n",
    "            uniprot_metadata_dict['u_gene_name'] = metadata[metadata_key]['gene']\n",
    "        if 'description' in metadata[metadata_key]:\n",
    "            uniprot_metadata_dict['u_description'] = metadata[metadata_key]['description']\n",
    "        if 'refseq' in metadata[metadata_key]:\n",
    "            uniprot_metadata_dict['u_refseq'] = metadata[metadata_key]['refseq']\n",
    "        if 'kegg' in metadata[metadata_key]:\n",
    "            uniprot_metadata_dict['u_kegg_id'] = metadata[metadata_key]['kegg']\n",
    "        if 'go' in metadata[metadata_key]:\n",
    "            uniprot_metadata_dict['u_go'] = metadata[metadata_key]['go']\n",
    "        if 'ec' in metadata[metadata_key]:\n",
    "            uniprot_metadata_dict['u_ec_number'] = metadata[metadata_key]['ec']\n",
    "        if 'pfam' in metadata[metadata_key]:\n",
    "            uniprot_metadata_dict['u_pfam'] = metadata[metadata_key]['pfam']\n",
    "        \n",
    "        uniprot_metadata_final[uniprot_id] = uniprot_metadata_dict\n",
    "    return de_unicodeify(uniprot_metadata_final)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: uniprot_reviewed_mapper_plus_metadata\n",
    "#### ANY ID & INPUT TYPE -> UNIPROT METADATA\n",
    "\n",
    "This function maps an ID to UniProt, retrieves the metadata for the mapped ID, and returns a dictionary of that information.\n",
    "\n",
    "**INPUT**:\n",
    "1. gene_id\n",
    "2. id_type\n",
    "\n",
    "**OUTPUT**:\n",
    "1. uniprot_metadata_dict\n",
    "    * { 'UNIPROT': id,\n",
    "        'DESCRIPTION': protein function,\n",
    "        'KEGG': associated KEGG IDs,\n",
    "        'SEQ': uniprot canonical sequence,\n",
    "        'PDB': list of PDB IDs\n",
    "      }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def uniprot_reviewed_mapper_plus_metadata(gene_id, id_type):\n",
    "    '''\n",
    "    Input: Any ID type specified from: http://www.uniprot.org/faq/28#id_mapping_examples\n",
    "    reviewed=True for only reviewed UniProt entries\n",
    "    '''\n",
    "\n",
    "    uniprot_metadata_dict = {}\n",
    "\n",
    "    reviewed_uniprots = uniprot_mapper(gene_id, id_type, reviewed=True)\n",
    "\n",
    "    if len(reviewed_uniprots) > 1:\n",
    "        warnings.warn('%s: MORE THAN 1 REVIEWED UNIPROT ENTRY' % gene_id)\n",
    "\n",
    "    uniprot_id = reviewed_uniprots[0]\n",
    "\n",
    "    return uniprot_metadata(uniprot_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: uniprot_ec\n",
    "#### UNIPROT ID -> EC NUMBER\n",
    "\n",
    "This function returns the EC number associated with a UniProt ID\n",
    "\n",
    "**INPUT**:\n",
    "1. uniprot_id\n",
    "\n",
    "**OUTPUT**:\n",
    "1. EC number (as a string)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def uniprot_ec(uniprot_id):\n",
    "    r = requests.post('http://www.uniprot.org/uniprot/?query=%s&columns=ec&format=tab' % uniprot_id)\n",
    "\n",
    "    ec = r.content.splitlines()[1]\n",
    "\n",
    "    if len(ec) == 0:\n",
    "        ec = None\n",
    "\n",
    "    return ec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: uniprot_sites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def uniprot_sites(uniprot_id):\n",
    "    r = requests.post('http://www.uniprot.org/uniprot/%s.gff' % uniprot_id)\n",
    "    gff = sio.StringIO(r.content)\n",
    "    \n",
    "    gff_df = pd.read_table(gff, sep='\\t',skiprows=2,header=None)\n",
    "    gff_df.drop([0,1,5,6,7,9], axis=1, inplace=True)\n",
    "    gff_df.columns = ['type','seq_start','seq_end','notes']\n",
    "    \n",
    "    return gff_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "##PDB RELATED"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from Bio.PDB.PDBParser import PDBParser\n",
    "from Bio.PDB import PDBIO\n",
    "keep_alt=0\n",
    "# create class to select non disordered atoms and those with the correct alternate location ID\n",
    "class NotDisordered(PDB.Select):\n",
    "  def __init__(self,alt):\n",
    "    self.alt=alt\n",
    "  def accept_atom(self,atom):\n",
    "    if not atom.is_disordered():\n",
    "      return True\n",
    "    elif atom.get_altloc() == self.alt:\n",
    "      if not keep_alt:\n",
    "        atom.set_altloc(' ')\n",
    "      return True\n",
    "    else:\n",
    "      return False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: clean_pdb\n",
    "\n",
    "Adds chains and occupancies to incomplete PDB files. Used for automatically generated homology models that don't have these, since STRIDE, DSSP, biopython, and others will complain about these missing components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def clean_pdb(input_pdb):\n",
    "    \n",
    "    # default is to keep atoms with alternate location ID 'A'\n",
    "    alt = 'A'\n",
    "    # default is to not keep \"alt\" identifier\n",
    "    keep_alt = 0\n",
    "    \n",
    "    parser = PDB.PDBParser()\n",
    "    \n",
    "    filename = input_pdb.split('.pdb')[0]\n",
    "    occ_beta_added = filename + '_fix.pdb'\n",
    "\n",
    "    struct = parser.get_structure(filename, input_pdb)\n",
    "\n",
    "    # adding chain X if the chain column is empty (http://comments.gmane.org/gmane.comp.python.bio.devel/10639)\n",
    "    for x in struct.get_chains():\n",
    "        if not x.id.strip(): # chain could be an empty string ' ' so strip it!\n",
    "            x.id = 'X'\n",
    "\n",
    "    # adding occupancies if there are none (http://comments.gmane.org/gmane.comp.python.bio.general/6289)\n",
    "    for atom in struct.get_atoms():\n",
    "        atom.set_occupancy(1)\n",
    "\n",
    "    # create PDBIO class for writing\n",
    "    io=PDB.PDBIO()\n",
    "    io.set_structure(struct)\n",
    "    # save it\n",
    "    io.save(occ_beta_added, select=NotDisordered(alt))\n",
    "    \n",
    "    return occ_beta_added"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: get_pdb_seq\n",
    "#### 2015-02-10 in progress\n",
    "\n",
    "trying to test biopython \n",
    "\n",
    "resources:\n",
    "\n",
    "* http://mmcif.wwpdb.org/pdbx-mmcif-home-page.html\n",
    "* http://biopython.org/wiki/The_Biopython_Structural_Bioinformatics_FAQ#I.27d_like_to_have_some_more_low_level_access_to_an_mmCIF_file..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_pdb_seq(pdb_id, pdb_file):\n",
    "    '''\n",
    "    this uses biopython to get the sequence of the residues that are resolved in the structure. \n",
    "    IT ADDS AN X IN POSITIONS THAT HAVE NO RESOLVED RESIDUE!\n",
    "    IT EVEN WORKS WITH CIF FILES!\n",
    "    '''\n",
    "    \n",
    "    if pdb_file.endswith('.pdb'):\n",
    "        parser = PDBParser()\n",
    "        structure = parser.get_structure(pdb_id, pdb_file)\n",
    "    elif pdb_file.endswith('.cif'):\n",
    "        parser = MMCIFParser()\n",
    "        try:\n",
    "            structure = parser.get_structure(pdb_id, pdb_file)\n",
    "        except:\n",
    "            return None\n",
    "        \n",
    "    structure_seqs = []\n",
    "    for chain in structure[0]:\n",
    "        chain_seq = SeqRecord(Seq('', IUPAC.protein), name=pdb_id, id=pdb_id+'.'+chain.get_id(), description='PDB structure residues')\n",
    "        first = True\n",
    "        tracker = 0\n",
    "        for res in chain.get_residues():\n",
    "            if is_aa(res, standard=True):\n",
    "                full_id = res.get_full_id()\n",
    "                end_tracker = full_id[3][1]\n",
    "                i_code = full_id[3][2]\n",
    "                aa = three_to_one(res.get_resname())\n",
    "                \n",
    "                # a bit odd -- shouldnt i check if end_tracker = tracker first?\n",
    "                # then if the i_code is ' ', do the same thing\n",
    "                # and if it is not, it is because the last residue had an i_code so we just need to add this residue\n",
    "                # right??\n",
    "                # as it currently is it adds -1 X's (none) if there is a previous insertion code, which works\n",
    "                if end_tracker != (tracker + 1) and first == False:\n",
    "                    if i_code != ' ':\n",
    "                        chain_seq += aa\n",
    "                        tracker = end_tracker + 1\n",
    "                        continue\n",
    "\n",
    "                    else:\n",
    "                        chain_seq += 'X'*(end_tracker - tracker - 1)\n",
    "                        \n",
    "                chain_seq += aa\n",
    "                first = False\n",
    "                tracker = end_tracker\n",
    "\n",
    "        structure_seqs.append(chain_seq)\n",
    "\n",
    "    return structure_seqs\n",
    "\n",
    "\n",
    "## ALTERNATE BIOPYTHON STRUCTURE SEQUENCE\n",
    "# using the polypeptide builder to get the residues\n",
    "# works BUT if you cannot build a polypeptide (ie there is a \"floating\" residue somewhere), it ignores that residue\n",
    "# example is with mmCIF but can use a pdb too\n",
    "\n",
    "# from Bio.PDB import *\n",
    "# parser = MMCIFParser()\n",
    "# structure = parser.get_structure('1seh', '1SEH.cif')\n",
    "\n",
    "# ppb=PPBuilder()\n",
    "# for pp in ppb.build_peptides(structure): \n",
    "#     print pp\n",
    "#     seq = pp.get_sequence()\n",
    "#     print seq\n",
    "\n",
    "# for model in structure:\n",
    "#     for chain in model:\n",
    "#         chain_seq = ''\n",
    "\n",
    "#         first = True\n",
    "#         end_tracker = 0\n",
    "        \n",
    "#         for pp in ppb.build_peptides(chain):\n",
    "#             start = pp[0].get_id()[1]\n",
    "#             end = pp[-1].get_id()[1]\n",
    "#             print pp, start, end\n",
    "#             seq = pp.get_sequence()\n",
    "#             print seq\n",
    "            \n",
    "#             if start != end_tracker - 1 and first == False:\n",
    "#                 print start - end_tracker - 1\n",
    "#                 chain_seq += 'X'*(start - end_tracker - 1)\n",
    "            \n",
    "#             chain_seq += seq\n",
    "                \n",
    "#             first=False\n",
    "#             end_tracker = end\n",
    "            \n",
    "#         print chain_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_pdb_seq_prody(pdbid):\n",
    "    # TODO: ABOUT TO DEPRECATE THIS ONE\n",
    "    '''\n",
    "    This function downloads a PDB ID into the working directory and also\n",
    "    returns the sequence that is represented in the structure (what residues are visible).\n",
    "    The sequence is a SeqRecord Biopython object.\n",
    "\n",
    "    Input: pdbid - a 4 letter string\n",
    "    Output: structure_seqs - a Seq object that represents what residues can be seen in the structure\n",
    "            resolution - the crystal structure resolution\n",
    "    '''\n",
    "    \n",
    "    try:\n",
    "        pdb = pr.parsePDB(pdbid, header=False, biomol=False, secondary=False)\n",
    "        \n",
    "        # if multiple biomols choose the first one (this could be improved)\n",
    "        if type(pdb) == list:\n",
    "            pdb = pdb[0]\n",
    "            \n",
    "    # if PDB isn't able to be downloaded, try downloading the mmCIF\n",
    "    # TODO: how to parse structure information from mmCIF? returning nothing for now, just downloading it\n",
    "    except IOError:\n",
    "        try:\n",
    "            pr.fetchPDBviaFTP(pdbid,format='cif')\n",
    "            return None\n",
    "        except IOError:\n",
    "            return None\n",
    "\n",
    "    structure_seqs = []\n",
    "    for chain in pdb.iterChains():\n",
    "        ch = chain.getChid()\n",
    "        seq = SeqRecord(Seq(chain.getSequence(), IUPAC.protein), id=pdbid+'.'+ch, description='PDB structure residues')\n",
    "        structure_seqs.append(seq)\n",
    "\n",
    "    return structure_seqs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_pdb_res_starts(pdb_id, pdb_file):\n",
    "    '''\n",
    "    this uses biopython to get the first residue number in a pdb file. \n",
    "    returns a list of tuples of pdb, chain, resnum. \n",
    "\n",
    "    '''\n",
    "    if pdb_file.endswith('.pdb'):\n",
    "        parser = PDBParser()\n",
    "        structure = parser.get_structure(pdb_id, pdb_file)\n",
    "    if pdb_file.endswith('.cif'):\n",
    "        parser = MMCIFParser()\n",
    "        structure = parser.get_structure(pdb_id, pdb_file)\n",
    "        \n",
    "    start_residues = []\n",
    "    for chain in structure[0]:\n",
    "        residues = chain.get_residues()\n",
    "        start_residues.append((pdb_id, chain.get_id(), residues.next().get_id()[1]))\n",
    "    \n",
    "    return start_residues"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def pdb_metadata_and_download(pdbids):\n",
    "    final_dict = {}\n",
    "    counter = 1\n",
    "    \n",
    "    if isinstance(pdbids, str):\n",
    "        pdbids = [pdbids]\n",
    "    \n",
    "    for pdbid in pdbids:\n",
    "        \n",
    "        clear_output(wait=True)\n",
    "        print('***PROGRESS: PARSED & DOWNLOADED %d/%d PDB IDS***\\n' % (counter, len(pdbids)))\n",
    "        counter += 1\n",
    "        sys.stdout.flush()\n",
    "        \n",
    "        try:\n",
    "            pr.fetchPDB(pdbid, compressed=False)\n",
    "            header = pr.parsePDBHeader(pdbid)\n",
    "\n",
    "        # if PDB isn't able to be downloaded\n",
    "        except IOError:\n",
    "            try:\n",
    "                header = pr.parsePDBHeader(pr.fetchPDBviaFTP(pdbid,format='cif', compressed=False))\n",
    "            except IOError:\n",
    "                continue\n",
    "        \n",
    "        appender = defaultdict(list)\n",
    "        for prop,val in header.iteritems():\n",
    "            if isinstance(val, pr.proteins.header.Chemical):\n",
    "                appender['p_chemicals'].append(val.resname)\n",
    "            elif isinstance(val, pr.proteins.header.Polymer):\n",
    "                appender['p_chains'].append(val.chid)\n",
    "                if val.ec:\n",
    "                    appender['p_ec_numbers'].append((val.chid,val.ec))\n",
    "            elif prop == 'reference':\n",
    "                if 'doi' in val:\n",
    "                    appender['p_doi'] = val['doi']\n",
    "                if 'pmid' in val:\n",
    "                    appender['p_pmid'] = val['pmid']\n",
    "            elif prop in ['resolution','space_group','experiment','deposition_date']:\n",
    "                appender['p_' + prop] = val\n",
    "        \n",
    "        tmp = {}\n",
    "        for chain in appender['p_chains']:\n",
    "            try:\n",
    "                uniprot_mapping = sifts_pdb_chain_to_uniprot(pdbid, chain)\n",
    "            except KeyError:\n",
    "                uniprot_mapping = ['PDB-'+chain]\n",
    "            tmp[chain] = uniprot_mapping\n",
    "        appender['p_chain_uniprot_map'] = tmp\n",
    "\n",
    "        final_dict[pdbid] = dict(appender)\n",
    "            \n",
    "    return de_unicodeify(final_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "AAdict={'MSE': 'M', 'ALA': 'A', 'LEU': 'L', 'MET': 'M', 'ARG': 'R', 'LYS': 'K', 'GLN': 'Q', 'GLU': 'E', 'ILE': 'I', 'SER': 'S', 'TRP': 'W', 'TYR': 'Y', 'PHE': 'F', 'VAL': 'V', 'THR': 'T', 'HIS': 'H', 'CYS': 'C', 'ASN': 'N', 'ASP': 'D', 'GLY': 'G', 'PRO': 'P'}\n",
    "\n",
    "AAdict2 = {'C':'CYS',\n",
    " 'I':'ILE',\n",
    " 'G':'GLY', \n",
    " 'S':'SER', \n",
    " 'Q':'GLN', \n",
    " 'K':'LYS',\n",
    " 'N':'ASN',\n",
    " 'P':'PRO', \n",
    " 'D':'ASP', \n",
    " 'T':'THR', \n",
    " 'F':'PHE', \n",
    " 'A':'ALA', \n",
    " 'M':'MET', \n",
    " 'H':'HIS', \n",
    " 'L':'LEU', \n",
    " 'R':'ARG', \n",
    " 'W':'TRP', \n",
    " 'V':'VAL', \n",
    " 'E':'GLU', \n",
    " 'Y':'TYR'}\n",
    "\n",
    "keep_atom_list = ['N', 'C', 'O', 'CA']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def setup_pdb_for_amber(df, name, file_name):\n",
    "    \n",
    "    ''' takes in a PDB or homology file and (1) performs site-directed mutagenesis given a user-defined sequence change,\n",
    "    (2) removes all hetero atoms (e.g. water and metals/cofactors) and (3) removes disordered atoms\n",
    "    '''\n",
    "    \n",
    "    parser = PDBParser()\n",
    "    structure = parser.get_structure(name, file_name)\n",
    "    model = structure[0]\n",
    "    chains = [i.id for i in model.child_list]\n",
    "\n",
    "    for chain_id in chains:\n",
    "        chain = model[chain_id]\n",
    "        residue_list = chain.child_list\n",
    "\n",
    "        for residue in list(residue_list):\n",
    "            # mutate residues according to sequence alignment data\n",
    "#             print residue.id[1]\n",
    "            if len(df)>0:\n",
    "                if residue.id[1] in df[df.id_b.str.endswith(chain_id)].id_b_start.tolist():\n",
    "                \n",
    "    #                 print name, residue.id[1], chain_id\n",
    "    #                 print df\n",
    "                    res_w = df[df.id_b.str.endswith(chain_id)][df.id_b_start == residue.id[1]].id_b_aa.values[0]\n",
    "                    res_mut = df[df.id_b.str.endswith(chain_id)][df.id_b_start == residue.id[1]].id_a_aa.values[0]\n",
    "\n",
    "                    if res_mut not in AAdict2.keys():\n",
    "                        warnings.warn('***UNKNOWN AMINO ACID IN UNIPROT SEQUENCE. SKIPPING***') ### TO CHANGE!!!!!!!!!!!!!!!\n",
    "                        continue\n",
    "\n",
    "                    print df[df.id_b.str.endswith(chain_id)][df.id_b_start == residue.id[1]].values \n",
    "                    print '\\n'\n",
    "                    print residue.id[1], residue.resname.upper(),\" | \", res_w, \"mutate to:\", res_mut\n",
    "\n",
    "                    # Remove all atoms except protein backbone atoms:\n",
    "                    for atom in residue.child_dict.keys():\n",
    "                        if atom in keep_atom_list:\n",
    "                            pass\n",
    "                        else:\n",
    "                            residue.detach_child(atom)\n",
    "\n",
    "                    residue.resname = AAdict2[res_mut]\n",
    "                    # if residue is non-standard, make it a non-heteroatom\n",
    "                    if residue.id[0] != ' ':\n",
    "                        resnum = residue.id[1]\n",
    "                        residue.id = (' ',resnum, ' ')\n",
    "                    # check that its the residue you expect and rename the residue:  \n",
    "    #                 if AAdict[residue.resname] == res_w:#AAdict[res_w]:\n",
    "    #                     print 'matches \\n'\n",
    "    #                     residue.resname = AAdict2[res_mut]\n",
    "    #                     print residue.resname\n",
    "    #                 else:\n",
    "    #                     raise ValueError(\"Unrecognised residue %r\" % residue.resname)\n",
    "            \n",
    "        # Remove all hydrogens and heteroatom residues (e.g. WAT or metal) from structure:  \n",
    "        for residue in list(chain):\n",
    "            id = residue.id\n",
    "            if id[0] != ' ':\n",
    "                chain.detach_child(id)\n",
    "            if len(chain) == 0:\n",
    "                model.detach_child(chain.id)\n",
    "            for atom in residue.get_list():\n",
    "            #print residue.resname, residue.id[1], atom.element\n",
    "                if atom.element == 'H':\n",
    "                    residue.detach_child(atom.id)\n",
    "                \n",
    "\n",
    "    return structure\n",
    "\n",
    "def write_to_pdb(struct, new_file_name):\n",
    "    w = PDBIO()\n",
    "    w.set_structure(struct)\n",
    "    w.save(new_file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def final_mutation_function(corrected_df, pdb_id, pdb_file):\n",
    "    # default is to keep atoms with alternate location ID 'A'\n",
    "    alt = 'A'\n",
    "    # default is to not keep \"alt\" identifier\n",
    "    keep_alt = 0\n",
    "    \n",
    "    output = pdb_id + '_modified.pdb'\n",
    "#     output = pdb_id + '.pdb'\n",
    "    \n",
    "    #read in structure to parser\n",
    "    parser = PDB.PDBParser()\n",
    "    struct = parser.get_structure(pdb_id, pdb_file)\n",
    "    \n",
    "    # create PDBIO class for writing\n",
    "    io = PDB.PDBIO()\n",
    "    io.set_structure(struct)\n",
    "    # save it\n",
    "    io.save('/tmp/%s_no_disorder.pdb' % pdb_id, select=NotDisordered(alt))\n",
    "    \n",
    "    # choose only parts of the protein to mutate\n",
    "    df = corrected_df[corrected_df.type == 'mutation']\n",
    "    # load structure (no disorder) into biopython & perform modifications\n",
    "    structure_new = setup_pdb_for_amber(df, pdb_id, '/tmp/%s_no_disorder.pdb' % pdb_id)\n",
    "    #  write out new structure\n",
    "    write_to_pdb(structure_new, output)\n",
    "    \n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def setup_pdb_for_amber_monomer(df, name, file_name):\n",
    "    \n",
    "    ''' takes in a PDB or homology file and (1) performs site-directed mutagenesis given a user-defined sequence change,\n",
    "    (2) removes all hetero atoms (e.g. water and metals/cofactors) and (3) removes disordered atoms\n",
    "    MOD: keep only 1 chain (df is already the \"corrected\" df with only the gene of interest anyway)\n",
    "    '''\n",
    "    \n",
    "    parser = PDBParser()\n",
    "    structure = parser.get_structure(name, file_name)\n",
    "    model = structure[0]\n",
    "#     chains = [i.id for i in model.child_list]\n",
    "\n",
    "    # get mutations, if any\n",
    "    mutation_df = df[df.type == 'mutation']\n",
    "    \n",
    "    # choosing a representative chain\n",
    "    sorted_by_match_and_chain = df[df['type']=='match'].sort(['count','chain'], ascending=[False, True]).set_index(['chain','count']).index.tolist()\n",
    "    chain_id = str(sorted_by_match_and_chain[0][0])\n",
    "    \n",
    "    # operate only on this chain\n",
    "    chain = model[chain_id]\n",
    "    residue_list = chain.child_list\n",
    "\n",
    "    for residue in list(residue_list):\n",
    "        # mutate residues according to sequence alignment data\n",
    "#             print residue.id[1]\n",
    "        if len(mutation_df)>0:\n",
    "            if residue.id[1] in df[df.id_b.str.endswith(chain_id)].id_b_start.tolist():\n",
    "\n",
    "#                 print name, residue.id[1], chain_id\n",
    "#                 print df\n",
    "                res_w = df[df.id_b.str.endswith(chain_id)][df.id_b_start == residue.id[1]].id_b_aa.values[0]\n",
    "                res_mut = df[df.id_b.str.endswith(chain_id)][df.id_b_start == residue.id[1]].id_a_aa.values[0]\n",
    "\n",
    "                if res_mut not in AAdict2.keys():\n",
    "                    warnings.warn('***UNKNOWN AMINO ACID IN UNIPROT SEQUENCE. SKIPPING***') ### TO CHANGE!!!!!!!!!!!!!!!\n",
    "                    continue\n",
    "\n",
    "                print df[df.id_b.str.endswith(chain_id)][df.id_b_start == residue.id[1]].values \n",
    "                print '\\n'\n",
    "                print residue.id[1], residue.resname.upper(),\" | \", res_w, \"mutate to:\", res_mut\n",
    "\n",
    "                # Remove all atoms except protein backbone atoms:\n",
    "                for atom in residue.child_dict.keys():\n",
    "                    if atom in keep_atom_list:\n",
    "                        pass\n",
    "                    else:\n",
    "                        residue.detach_child(atom)\n",
    "\n",
    "                residue.resname = AAdict2[res_mut]\n",
    "                # if residue is non-standard, make it a non-heteroatom\n",
    "                if residue.id[0] != ' ':\n",
    "                    resnum = residue.id[1]\n",
    "                    residue.id = (' ',resnum, ' ')\n",
    "                # check that its the residue you expect and rename the residue:  \n",
    "#                 if AAdict[residue.resname] == res_w:#AAdict[res_w]:\n",
    "#                     print 'matches \\n'\n",
    "#                     residue.resname = AAdict2[res_mut]\n",
    "#                     print residue.resname\n",
    "#                 else:\n",
    "#                     raise ValueError(\"Unrecognised residue %r\" % residue.resname)\n",
    "\n",
    "    # Remove all hydrogens and heteroatom residues (e.g. WAT or metal) from structure:  \n",
    "    for residue in list(chain):\n",
    "        id = residue.id\n",
    "        if id[0] != ' ':\n",
    "            chain.detach_child(id)\n",
    "        if len(chain) == 0:\n",
    "            model.detach_child(chain.id)\n",
    "        for atom in residue.get_list():\n",
    "        #print residue.resname, residue.id[1], atom.element\n",
    "            if atom.element == 'H':\n",
    "                residue.detach_child(atom.id)\n",
    "\n",
    "\n",
    "    return chain\n",
    "\n",
    "def write_to_pdb(struct, new_file_name):\n",
    "    w = PDBIO()\n",
    "    w.set_structure(struct)\n",
    "    w.save(new_file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def monomer_or_mutation_function(corrected_df, pdb_id, pdb_file):\n",
    "    # default is to keep atoms with alternate location ID 'A'\n",
    "    alt = 'A'\n",
    "    # default is to not keep \"alt\" identifier\n",
    "    keep_alt = 0\n",
    "    \n",
    "    output = pdb_id + '_modified.pdb'\n",
    "#     output = pdb_id + '.pdb'\n",
    "    \n",
    "    #read in structure to parser\n",
    "    parser = PDB.PDBParser()\n",
    "    struct = parser.get_structure(pdb_id, pdb_file)\n",
    "    \n",
    "    # create PDBIO class for writing\n",
    "    io = PDB.PDBIO()\n",
    "    io.set_structure(struct)\n",
    "    # save it\n",
    "    io.save('/tmp/%s_no_disorder.pdb' % pdb_id, select=NotDisordered(alt))\n",
    "    \n",
    "    # load structure (no disorder) into biopython & perform modifications\n",
    "    structure_new = setup_pdb_for_amber_monomer(corrected_df, pdb_id, '/tmp/%s_no_disorder.pdb' % pdb_id)\n",
    "    #  write out new structure\n",
    "    write_to_pdb(structure_new, output)\n",
    "    \n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: pdb_chain_stoichiometry_biomolone\n",
    "#### PDB ID -> DOWNLOAD/PARSE PDB & CHAIN STOICHIOMETRY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def pdb_chain_stoichiometry_biomolone(pdbid):\n",
    "    '''\n",
    "    Takes in a PDB ID and returns the stoichiometry of the chains in biological assembly 1 as a dictionary.\n",
    "    Steps taken are:\n",
    "    1) download PDB and parse header, make biomolecule if provided\n",
    "    2) count how many times each chain appears in biomolecule #1\n",
    "    3) convert chain id to uniprot id\n",
    "    4) return final dictionary\n",
    "    Input: pdbid - a 4 character PDB ID\n",
    "    Output: a dictionary of {(ChainID,UniProtID): # occurences}\n",
    "    '''\n",
    "    \n",
    "    try:\n",
    "        # change to biomol=True if you want biological assemblies\n",
    "        pdb, header = pr.parsePDB(pdbid, header=True, biomol=True, secondary=False)\n",
    "\n",
    "        # if multiple biomols choose the first one, also if NMR choose the first coord set\n",
    "        if type(pdb) == list:\n",
    "            pdb = pdb[0]\n",
    "\n",
    "    except IOError:\n",
    "        return None, None\n",
    "    # if there are no biological assemblies\n",
    "    except ValueError:\n",
    "        pdb, header = pr.parsePDB(pdbid, header=True, biomol=False, secondary=False)\n",
    "    \n",
    "    # count the occurences of each chain\n",
    "    chain_stoich = defaultdict(int)\n",
    "    hier = pdb.getHierView()\n",
    "    for chain in hier:\n",
    "        chain_stoich[chain.getChid()] += 1\n",
    "    \n",
    "    # DBREF entry in PDB file sometimes contains obsolete UniProt entries\n",
    "    # chain_to_uniprot = {}\n",
    "    # for chain in header['polymers']:\n",
    "        # for dbref in chain.dbrefs:\n",
    "            # if dbref.database.lower() == 'uniprot':\n",
    "                # chain_to_uniprot[chain.chid] = dbref.accession\n",
    "    \n",
    "    # convert chain IDs to uniprot IDs\n",
    "    chain_to_uniprot = {}\n",
    "    for chain in header['polymers']:\n",
    "        try:\n",
    "            chain_to_uniprot[chain.chid] = sifts_pdb_chain_to_uniprot(pdbid.lower(), chain.chid)\n",
    "        except KeyError:\n",
    "            chain_to_uniprot[chain.chid] = ['PDB-'+chain.chid]\n",
    "                \n",
    "    # keep both chain ID and uniprot ID (this is the final dictionary)\n",
    "    combined = {}\n",
    "    for k,v in chain_to_uniprot.iteritems():\n",
    "        for uni in v:\n",
    "            combined[(k,uni)] = chain_stoich[k]\n",
    "    \n",
    "    return combined"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: pisa_complex_information\n",
    "#### PDB ID -> PREDICTED PISA COMPLEXES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def pisa_complex_information(pdb_id):\n",
    "    \n",
    "    pdb_id = pdb_id.lower()\n",
    "    pisa = {}\n",
    "    \n",
    "    # request the xml file for a PDB\n",
    "    r = requests.post('http://www.ebi.ac.uk/pdbe/pisa/cgi-bin/multimers.pisa?%s' % pdb_id)\n",
    "    soup = BeautifulSoup(r.content)\n",
    "    \n",
    "    # if PISA can't calculate an assembly...\n",
    "    if 'not found' in str(soup.pdb_entry.status.string) or 'No symmetry' in str(soup.pdb_entry.status.string) or 'Overlapping structures' in str(soup.pdb_entry.status.string):\n",
    "        pisa[pdb_id] = 'ERROR'\n",
    "        return pisa\n",
    "    \n",
    "    # if it is a monomer...\n",
    "    num_complexes = int(soup.total_asm.string)\n",
    "    if num_complexes == 0:\n",
    "        pisa[pdb_id] = 'MONOMER'\n",
    "        return pisa\n",
    "        \n",
    "    # otherwise, get the complex information!\n",
    "    elif num_complexes > 0:\n",
    "        \n",
    "        # all \"assembly sets\" (see PISA sets for more info)\n",
    "        sets = soup.findAll('asm_set')\n",
    "        \n",
    "        for s in sets:\n",
    "            \n",
    "            set_id = int(s.ser_no.string)\n",
    "            \n",
    "            # all assemblies\n",
    "            complexes = s.findAll('assembly')\n",
    "            \n",
    "            for cplx in complexes:\n",
    "                \n",
    "                ############################################################################################\n",
    "                # this part tells you the actual composition of the predicted complex (chains and ligands) #\n",
    "                parts = cplx.findAll('molecule')\n",
    "\n",
    "                chains = defaultdict(int)\n",
    "\n",
    "                for part in parts:\n",
    "                    part_id = part.chain_id.string\n",
    "                    if part_id.startswith('['):\n",
    "                        part_id = 'LIG_' + part_id.split(']')[0].strip('[')\n",
    "                    chains[str(part_id)] += 1\n",
    "\n",
    "                ligands = {}\n",
    "\n",
    "                for key in chains.keys():\n",
    "                    if key.startswith('LIG_'):\n",
    "                        ligands[str(key.split('_')[1])] = chains.pop(key)\n",
    "                        \n",
    "                chains_final = {}\n",
    "                for k,v in chains.iteritems():\n",
    "                    try:\n",
    "                        chain_to_uniprot = sifts_pdb_chain_to_uniprot(pdb_id.lower(), k)\n",
    "                    except KeyError:\n",
    "                        chain_to_uniprot = ['PDB-' + k]\n",
    "                    for m in chain_to_uniprot:\n",
    "                        chains_final[(str(k), str(m))] = v\n",
    "                ############################################################################################\n",
    "                \n",
    "                # this part give you something to add to a dataframe\n",
    "                adder = {}\n",
    "\n",
    "                cplx_id = int(cplx.id.string)\n",
    "                cplx_composition = str(cplx.composition.string)\n",
    "\n",
    "                d_g_diss = float(cplx.diss_energy.string)\n",
    "                d_g_int = float(cplx.int_energy.string)\n",
    "\n",
    "                pdb_biomol = int(cplx.r350.string)\n",
    "                \n",
    "                if d_g_diss >= 0:\n",
    "                    stable = True\n",
    "                else:\n",
    "                    stable = False\n",
    "\n",
    "                adder['cplx_composition'] = cplx_composition.strip()\n",
    "                adder['cplx_chains'] = chains_final\n",
    "                adder['cplx_ligands'] = ligands\n",
    "                adder['stable'] = stable\n",
    "                adder['d_g_diss'] = d_g_diss\n",
    "                adder['d_g_int'] = d_g_int\n",
    "                adder['pdb_biomol'] = pdb_biomol\n",
    "                \n",
    "                pisa[(set_id,cplx_id)] = adder\n",
    "    \n",
    "        return pisa"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "## SEQUENCE TOOLS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: write_fasta_file\n",
    "#### This writes a fasta file for a SeqRecord object. It also checks if the file exists already and returns the filename.\n",
    "\n",
    "**INPUTS**:\n",
    "1. sequence\n",
    "    * A Biopython SeqRecord object (that could contain multiple sequences in one)\n",
    "2. identification\n",
    "    * A string representing the ID (usually UniProt or PDB)\n",
    "    \n",
    "**OUTPUTS**\n",
    "1. outfile\n",
    "    * The filename of the fasta (.faa) file generated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def write_fasta_file(sequence, identification):\n",
    "    '''\n",
    "    This writes a fasta file for a SeqRecord object. It also checks if the file exists already and returns the filename.\n",
    "    \n",
    "    Input: sequence - Biopython SeqRecord object, identification - ID of the sequence.\n",
    "    Output: Filename of fasta file\n",
    "    '''\n",
    "    \n",
    "    from Bio import SeqIO\n",
    "    import os.path\n",
    "    \n",
    "    outfile = \"%s.faa\" % identification\n",
    "    if os.path.isfile(outfile):\n",
    "        print 'FASTA file already exists %s' % outfile\n",
    "        return outfile\n",
    "    else:\n",
    "        SeqIO.write(sequence, outfile, \"fasta\")\n",
    "        return outfile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: run_alignment\n",
    "#### This runs the EMBOSS needle application (Needleman–Wunsch algorithm) to align a uniprot sequence to a PDB fasta sequence. It also checks if the file exists already and returns the filename.\n",
    "\n",
    "**INPUTS**:\n",
    "1. uniprot_faa\n",
    "    * A Biopython SeqRecord object (usually just one sequence)\n",
    "2. pdb_faa\n",
    "    * A Biopython SeqRecord object (can have multiple because of multiple chains in a PDB structure)\n",
    "    \n",
    "**OUTPUTS**\n",
    "1. alignment_file\n",
    "    * The filename of the alignment file generated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def run_alignment(uni_id, uniprot_faa, pdb_id, pdb_faa):\n",
    "    '''\n",
    "    Runs the needle alignment program and writes the result to a file. Returns the filename.\n",
    "    \n",
    "    Input:  uniprot_faa - uniprot fasta file name\n",
    "            pdb_faa - pdb fasta file name\n",
    "    Output: alignment_file - file name of alignment\n",
    "    '''\n",
    "    \n",
    "    from Bio.Emboss.Applications import NeedleCommandline\n",
    "    import os.path\n",
    "\n",
    "    alignment_file = \"%s_%s_align.txt\" % (uni_id, pdb_id)\n",
    "    \n",
    "    if os.path.isfile(alignment_file):\n",
    "        print 'Alignment %s file already exists' % alignment_file\n",
    "        return alignment_file\n",
    "\n",
    "    else:\n",
    "        # print '**RUNNING ALIGNMENT FOR %s AND %s**' % (uni_id, pdb_id)\n",
    "        needle_cline = NeedleCommandline(asequence=uniprot_faa, bsequence=pdb_faa, gapopen=10, gapextend=0.5, outfile=alignment_file)\n",
    "        stdout, stderr = needle_cline()\n",
    "        return alignment_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# SEE: https://www.biostars.org/p/91124/\n",
    "\n",
    "def run_alignment2(a_id, a_seq, b_id, b_seq):\n",
    "    '''\n",
    "    Runs the needle alignment program and returns a raw text dump of the alignment\n",
    "    \n",
    "    Input:  a_id - sequence ID #1 (string)\n",
    "            a_seq - sequence #1 (string)\n",
    "            b_id - sequence ID #2 (string)\n",
    "            b_seq - sequence #2 (string)\n",
    "    Output: alignment_file - file name of alignment\n",
    "    \n",
    "    DEPENDENCIES:\n",
    "    get_alignment_allpos_df\n",
    "    '''\n",
    "    \n",
    "    from Bio.Emboss.Applications import NeedleCommandline\n",
    "    from Bio import AlignIO\n",
    "    import os.path\n",
    "\n",
    "    alignment_file = \"/tmp/%s_%s_align.txt\" % (a_id, b_id)\n",
    "\n",
    "    needle_cline = NeedleCommandline(asequence=\"asis::\"+a_seq, bsequence=\"asis::\"+b_seq, gapopen=10, gapextend=0.5, outfile=alignment_file)\n",
    "    stdout, stderr = needle_cline()\n",
    "    \n",
    "    return get_alignment_allpos_df(alignment_file, a_id, b_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: run_alignment_needleall\n",
    "#### TODO: runs needleall!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def run_alignment_needleall(a_id, a_faa, b_id, b_faa):\n",
    "    \n",
    "    from Bio.Emboss.Applications import NeedleallCommandline\n",
    "    import os.path\n",
    "\n",
    "    alignment_file = \"%s-%s_align.txt\" % (a_id, b_id)\n",
    "    \n",
    "    if os.path.isfile(alignment_file):\n",
    "        print 'Alignment %s file already exists' % alignment_file\n",
    "        return alignment_file\n",
    "\n",
    "    else:\n",
    "        print '**RUNNING ALIGNMENT FOR %s AND %s**' % (a_id, b_id)\n",
    "        needle_cline = NeedleallCommandline(asequence=a_faa, bsequence=b_faa, gapopen=10, gapextend=0.5, outfile=alignment_file)\n",
    "        stdout, stderr = needle_cline()\n",
    "        return alignment_file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: needle_reader\n",
    "#### Reads in a needle alignment file and spits out statistics of the alignment.\n",
    "\n",
    "**INPUTS**:\n",
    "1. fl\n",
    "    * Filename of the alignment file\n",
    "\n",
    "**OUTPUTS**\n",
    "1. alignment_properties\n",
    "    * Dictionary of dictionaries representing statistics of the alignment { 'PDB_A': {'gaps': ##, 'identity': ##, 'similarity': ##, 'score': ####}, 'PDB_B': ...}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def needle_reader(fl):\n",
    "    '''\n",
    "    Reads in a needle alignment file and spits out statistics of the alignment.\n",
    "    \n",
    "    Input: fl - alignment file name\n",
    "    Output: alignment_properties - a dictionary telling you the number of gaps, identity, etc.\n",
    "    '''\n",
    "    \n",
    "    alignments = list(AlignIO.parse(fl, \"emboss\"))\n",
    "    \n",
    "    f=open(fl)\n",
    "    alignment_properties = defaultdict(dict)\n",
    "    line = f.readline()\n",
    "\n",
    "    for i in range(len(alignments)):\n",
    "        #print alignments[i]\n",
    "        while line.rstrip() != \"#=======================================\":\n",
    "            line = f.readline()\n",
    "            if not line:\n",
    "                raise StopIteration\n",
    "\n",
    "        while line[0] == \"#\":\n",
    "            #Read in the rest of this alignment header,\n",
    "            #try and discover the number of records expected\n",
    "            #and their length\n",
    "            parts = line[1:].split(\":\", 1)\n",
    "            key = parts[0].lower().strip()\n",
    "            if key == '2':\n",
    "                pdb_id = parts[1].strip()\n",
    "            if key == 'identity':\n",
    "                ident_parse = parts[1].strip().replace('(','').replace(')','').replace('%','').split()\n",
    "                ident_num = int(ident_parse[0].split('/')[0])\n",
    "                ident_percent = float(ident_parse[1])\n",
    "                alignment_properties[pdb_id]['identity'] = ident_num\n",
    "                alignment_properties[pdb_id]['identity_percent'] = ident_percent\n",
    "            if key == 'similarity':\n",
    "                sim_parse = parts[1].strip().replace('(','').replace(')','').replace('%','').split()\n",
    "                sim_num = int(sim_parse[0].split('/')[0])\n",
    "                sim_percent = float(sim_parse[1])\n",
    "                alignment_properties[pdb_id]['similarity'] = sim_num\n",
    "                alignment_properties[pdb_id]['similarity_percent'] = sim_percent\n",
    "            if key == 'gaps':\n",
    "                gap_parse = parts[1].strip().replace('(','').replace(')','').replace('%','').split()\n",
    "                gap_num = int(gap_parse[0].split('/')[0])\n",
    "                gap_percent = float(gap_parse[1])\n",
    "                alignment_properties[pdb_id]['gaps'] = gap_num\n",
    "                alignment_properties[pdb_id]['gaps_percent'] = gap_percent\n",
    "            if key == 'score':\n",
    "                score = float(parts[1].strip())\n",
    "                alignment_properties[pdb_id]['score'] = score\n",
    "            \n",
    "            #And read in another line...\n",
    "            line = f.readline()\n",
    "\n",
    "    return alignment_properties"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: get_alignment_df\n",
    "#### Reads in a needle alignment file and creates a pandas dataframe telling you what the mutations are.\n",
    "\n",
    "**INPUTS**:\n",
    "1. TODO\n",
    "\n",
    "\n",
    "**OUTPUTS**\n",
    "1. TODO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_alignment_df(alignment_file):\n",
    "    alignments = list(AlignIO.parse(alignment_file, \"emboss\"))\n",
    "    alignment_df = pd.DataFrame(columns = ['id_a', 'id_b', 'type', 'id_a_start', 'id_a_stop', 'count', 'id_a_aa', 'id_b_aa'])\n",
    "\n",
    "    for alignment in alignments:\n",
    "    #         if not switch:\n",
    "        a_seq_id = list(alignment)[0].id\n",
    "        a_seq = str(list(alignment)[0].seq)\n",
    "        b_seq_id = list(alignment)[1].id\n",
    "        b_seq = str(list(alignment)[1].seq)\n",
    "    #         else:\n",
    "    #             a_seq_id = list(alignment)[1].id\n",
    "    #             a_seq = str(list(alignment)[1].seq)\n",
    "    #             b_seq_id = list(alignment)[0].id\n",
    "    #             b_seq = str(list(alignment)[0].seq)\n",
    "\n",
    "        idx_start = 1\n",
    "\n",
    "        uniprot_shift = 0\n",
    "        insertion_counter = 0\n",
    "        for i, (a,b) in enumerate(zip(a_seq,b_seq)):\n",
    "            if uniprot_shift != 0:\n",
    "                new_i = i-uniprot_shift\n",
    "            else:\n",
    "                new_i = i\n",
    "\n",
    "            if a == b and a != '-' and b != '-':\n",
    "                aa_flag = 'match'\n",
    "            if a != b and a == '-' and b != '-':\n",
    "                aa_flag = 'insertion'\n",
    "            if a != b and a != '-' and b == '-':\n",
    "                aa_flag = 'deletion'\n",
    "            if a != b and a != '-' and b == 'X':\n",
    "                aa_flag = 'unresolved'\n",
    "            elif a != b and a != '-' and b != '-':\n",
    "                aa_flag = 'mutation'\n",
    "            if i == 0:\n",
    "                aa_flag_tracker = aa_flag\n",
    "\n",
    "            if aa_flag == 'insertion':\n",
    "                uniprot_shift += 1\n",
    "                insertion_counter += 1\n",
    "\n",
    "\n",
    "            if aa_flag != aa_flag_tracker:\n",
    "                idx_end = new_i\n",
    "                if aa_flag_tracker == 'match' or aa_flag_tracker == 'deletion' or aa_flag_tracker == 'unresolved':\n",
    "                    appender = {}\n",
    "                    appender['id_a'] = a_seq_id\n",
    "                    appender['id_b'] = b_seq_id\n",
    "                    appender['type'] = aa_flag_tracker\n",
    "                    appender['id_a_start'] = idx_start\n",
    "                    appender['id_a_stop'] = idx_end\n",
    "                    appender['count'] = idx_end - idx_start + 1\n",
    "#                     print '%s from %d to %d with %d members' % (appender['type'], appender['start'], appender['stop'], appender['count'])\n",
    "                    alignment_df = alignment_df.append(appender, ignore_index = True)\n",
    "\n",
    "                if aa_flag_tracker == 'insertion':\n",
    "                    appender = {}\n",
    "                    appender['id_a'] = a_seq_id\n",
    "                    appender['id_b'] = b_seq_id\n",
    "                    appender['type'] = aa_flag_tracker\n",
    "                    appender['id_a_start'] = new_i\n",
    "                    appender['id_a_stop'] = new_i\n",
    "                    appender['count'] = insertion_counter\n",
    "#                     print '%s of length %d' % (aa_flag_tracker, insertion_counter)\n",
    "                    alignment_df = alignment_df.append(appender, ignore_index = True)\n",
    "                    insertion_counter = 0\n",
    "                idx_start = new_i + 1\n",
    "\n",
    "            # HEY. THIS NEEDS TO BE OUTSIDE THE FOR LOOP\n",
    "            if aa_flag == 'mutation':\n",
    "                appender = {}\n",
    "                appender['id_a'] = a_seq_id\n",
    "                appender['id_b'] = b_seq_id\n",
    "                appender['type'] = aa_flag\n",
    "                appender['id_a_start'] = new_i + 1\n",
    "                appender['id_a_stop'] = new_i + 1\n",
    "                appender['count'] = 1\n",
    "                appender['id_a_aa'] = a\n",
    "                appender['id_b_aa'] = b\n",
    "#                     print '%s at %d' % (aa_flag, new_i + 1)\n",
    "                alignment_df = alignment_df.append(appender, ignore_index = True)\n",
    "                idx_start = new_i + 1\n",
    "                \n",
    "\n",
    "\n",
    "            elif (i+1) == len(a_seq) and aa_flag_tracker == aa_flag:\n",
    "                idx_end = new_i + 1\n",
    "                if aa_flag_tracker != 'mutation' and aa_flag_tracker != 'insertion':\n",
    "                    appender = {}\n",
    "                    appender['id_a'] = a_seq_id\n",
    "                    appender['id_b'] = b_seq_id\n",
    "                    appender['type'] = aa_flag_tracker\n",
    "                    appender['id_a_start'] = idx_start\n",
    "                    appender['id_a_stop'] = idx_end\n",
    "                    appender['count'] = idx_end - idx_start + 1\n",
    "#                     print '%s from %d to %d with %d members' % (appender['type'], appender['start'], appender['stop'], appender['count'])\n",
    "                    alignment_df = alignment_df.append(appender, ignore_index = True)\n",
    "\n",
    "            elif (i+1) == len(a_seq) and aa_flag_tracker != aa_flag:\n",
    "                idx_end = new_i + 1\n",
    "                idx_start = new_i + 1\n",
    "                if aa_flag_tracker != 'mutation' and aa_flag_tracker != 'insertion':\n",
    "                    appender = {}\n",
    "                    appender['id_a'] = a_seq_id\n",
    "                    appender['id_b'] = b_seq_id\n",
    "                    appender['type'] = aa_flag\n",
    "                    appender['id_a_start'] = idx_start\n",
    "                    appender['id_a_stop'] = idx_end\n",
    "                    appender['count'] = idx_end - idx_start + 1\n",
    "#                     print '%s from %d to %d with %d members' % (appender['type'], appender['start'], appender['stop'], appender['count'])\n",
    "                    alignment_df = alignment_df.append(appender, ignore_index = True)\n",
    "#             print new_i + 1, ':', a, b, 'last=%s, now=%s' % (aa_flag_tracker,aa_flag) \n",
    "\n",
    "            aa_flag_tracker = aa_flag\n",
    "            \n",
    "    return alignment_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_alignment_allpos_df(alignment_file, a_seq_id=None, b_seq_id=None):\n",
    "    alignments = list(AlignIO.parse(alignment_file, \"emboss\"))\n",
    "\n",
    "    appender = defaultdict(dict)\n",
    "    idx = 0\n",
    "    for alignment in alignments:\n",
    "    #         if not switch:\n",
    "        if not a_seq_id:\n",
    "            a_seq_id = list(alignment)[0].id\n",
    "        a_seq = str(list(alignment)[0].seq)\n",
    "        if not b_seq_id:\n",
    "            b_seq_id = list(alignment)[1].id\n",
    "        b_seq = str(list(alignment)[1].seq)\n",
    "\n",
    "        a_idx = 1\n",
    "        b_idx = 1\n",
    "\n",
    "        for i, (a,b) in enumerate(zip(a_seq,b_seq)):\n",
    "            if a == b and a != '-' and b != '-':\n",
    "                aa_flag = 'match'\n",
    "            if a != b and a == '-' and b != '-':\n",
    "                aa_flag = 'insertion'\n",
    "            if a != b and a != '-' and b == '-':\n",
    "                aa_flag = 'deletion'\n",
    "            if a != b and a != '-' and b == 'X':\n",
    "                aa_flag = 'unresolved'\n",
    "            if a != b and b != '-' and a == 'X':\n",
    "                aa_flag = 'unresolved'\n",
    "            elif a != b and a != '-' and b != '-':\n",
    "                aa_flag = 'mutation'\n",
    "                \n",
    "            appender[idx]['id_a'] = a_seq_id\n",
    "            appender[idx]['id_b'] = b_seq_id\n",
    "            appender[idx]['type'] = aa_flag\n",
    "            \n",
    "            if aa_flag == 'match' or aa_flag == 'unresolved' or aa_flag == 'mutation':\n",
    "                appender[idx]['id_a_aa'] = a\n",
    "                appender[idx]['id_a_pos'] = a_idx\n",
    "                appender[idx]['id_b_aa'] = b\n",
    "                appender[idx]['id_b_pos'] = b_idx\n",
    "                a_idx += 1\n",
    "                b_idx += 1\n",
    "\n",
    "            if aa_flag == 'deletion':\n",
    "                appender[idx]['id_a_aa'] = a\n",
    "                appender[idx]['id_a_pos'] = a_idx\n",
    "                a_idx += 1\n",
    "\n",
    "            if aa_flag == 'insertion':\n",
    "                appender[idx]['id_b_aa'] = b\n",
    "                appender[idx]['id_b_pos'] = b_idx\n",
    "                b_idx += 1\n",
    "            \n",
    "            idx += 1\n",
    "\n",
    "    alignment_df = pd.DataFrame.from_dict(appender, orient='index')\n",
    "    alignment_df = alignment_df[['id_a', 'id_b', 'type', 'id_a_aa', 'id_a_pos', 'id_b_aa', 'id_b_pos']].fillna(value=np.nan)\n",
    "    \n",
    "#     alignment_df = alignment_df.join(alignment_df['id_a'].apply(lambda x: pd.Series(x.split('.')[1])))\n",
    "#     alignment_df = alignment_df.rename(columns={0:'id_a_chain'})\n",
    "    \n",
    "#     alignment_df = alignment_df.join(alignment_df['id_b'].apply(lambda x: pd.Series(x.split('.')[1])))\n",
    "#     alignment_df = alignment_df.rename(columns={0:'id_b_chain'})\n",
    "    \n",
    "    return alignment_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "## COBRA & SBML RELATED"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def is_spontaneous(gene):\n",
    "    spont = re.compile(\"[Ss](_|)0001\")\n",
    "    if spont.match(gene):\n",
    "        return True\n",
    "    else:\n",
    "        return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def true_num_genes(model):\n",
    "    true_num = 0\n",
    "    for gene in model.genes:\n",
    "        if not is_spontaneous(gene.id):\n",
    "            true_num += 1\n",
    "    return true_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def true_num_reactions(model):\n",
    "    true_num = 0\n",
    "    for rxn in model.reactions:\n",
    "        genes = [x.id for x in rxn.genes]\n",
    "        if len(genes) == 0:\n",
    "            continue\n",
    "        if len(genes) == 1 and is_spontaneous(genes[0]):\n",
    "            continue\n",
    "        else:\n",
    "            true_num += 1\n",
    "    return true_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def adj_num_reactions(model, missing_genes):\n",
    "    adj_num = 0\n",
    "    for rxn in model.reactions:\n",
    "        genes = [x.id for x in rxn.genes]\n",
    "        if len(genes) == 0:\n",
    "            continue\n",
    "        if len(genes) == 1 and (is_spontaneous(genes[0]) or genes[0] in missing_genes):\n",
    "            continue\n",
    "        else:\n",
    "            adj_num += 1\n",
    "    return adj_num  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def gene_name_to_id(full_model_xml):\n",
    "    '''\n",
    "    Input: the .xml filename of the 'full' model (in my case Recon2) which has the gene names and associated gene IDs listed as species\n",
    "    Output: a dictionary with the following structure - {GeneName: GeneID, ...}\n",
    "    '''\n",
    "\n",
    "    full_model_sbml = reader.readSBML(full_model_xml)\n",
    "    m = full_model_sbml.getModel()\n",
    "\n",
    "    gene_dict = defaultdict(dict)\n",
    "\n",
    "    # 'species' includes the genes\n",
    "    for i in m.getListOfSpecies():\n",
    "        entrez = i.getId()\n",
    "\n",
    "        # some start with _NM\n",
    "        if entrez.startswith('_NM'):\n",
    "            name = i.getName()\n",
    "            annotation = i.getAnnotation()\n",
    "            gene_dict[name] = entrez\n",
    "\n",
    "        # all other genes start with _\n",
    "        elif entrez.startswith('_'):\n",
    "            name = i.getName()\n",
    "            annotation = i.getAnnotation()\n",
    "            idz = entrez.split('_')\n",
    "            idz.pop(0)\n",
    "\n",
    "            # accounting for genes that are connected\n",
    "            newidz = idz[0::3]\n",
    "            namez = name.split(':')\n",
    "            for x in range(len(newidz)):\n",
    "                if namez[x] == 'null':\n",
    "                    continue\n",
    "                gene_dict[namez[x]] = newidz[x]\n",
    "\n",
    "    return gene_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def full_to_reduced(full_dict, reduced_list_of_genes):\n",
    "    '''\n",
    "    Input:\n",
    "        full_dict: the dictionary returned by function nameToId\n",
    "        reduced_list_of_genes: the list of genes of your reduced model (in my case iAB-RBC)\n",
    "    Output:\n",
    "        gene_mapping: the gene mapping for your reduced model (a dictionary of dictionaries with form {GeneName: {'ENTREZ': GeneID, 'UNIPROT': UniprotAcc#}, ...})\n",
    "        missing_in_full: genes that were not in the full model but in the reduced - stuff that needs to be looked at manually\n",
    "    '''\n",
    "    missing_in_full = []\n",
    "    gene_mapping = {}\n",
    "\n",
    "    for gene in reduced_list_of_genes:\n",
    "        genename = gene.name.strip('|&').split('.')[0].upper()\n",
    "        if genename in full_dict:\n",
    "            gene_mapping[genename] = full_dict[genename]\n",
    "        else:\n",
    "            missing_in_full.append(genename)\n",
    "\n",
    "    return gene_mapping, list(set(filter(None, missing_in_full)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "## MISCELLANEOUS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def write_pickle(filePath,d):\n",
    "    f=open(filePath,'w')\n",
    "    newData = pickle.dumps(d, 1)\n",
    "    f.write(newData)\n",
    "    f.close()\n",
    "\n",
    "def read_pickle(filePath):\n",
    "    f=open(filePath,'r')\n",
    "    data = pickle.load(f)\n",
    "    f.close()\n",
    "    return data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "***\n",
    "## DEPRECATED"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# UNIPROT\n",
    "# DEPRECATED: replaced by bioservices\n",
    "# used for uniprot information: https://github.com/boscoh/uniprot\n",
    "import uniprot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: general_mapper\n",
    "#### ANY ID & INPUT TYPE & OUTPUT TYPE -> MAPPED ID\n",
    "This function uses the UniProt mapping tool to map any ID to another ID. You need to specify the ID types though."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def general_mapper(map_from, map_to, gene_id):\n",
    "    '''\n",
    "    Test out a mapping here: http://www.uniprot.org/mapping/\n",
    "    Specify ID types(!) from here: http://www.uniprot.org/faq/28#id_mapping_examples\n",
    "    reviewed=True for only reviewed UniProt entries\n",
    "    '''\n",
    "    \n",
    "    pairs = uniprot.batch_uniprot_id_mapping_pairs(map_from, map_to, [gene_id])\n",
    "\n",
    "    maps = []\n",
    "    for idx in range(len(pairs)):\n",
    "        maps.append(str(pairs[idx][1]))\n",
    "\n",
    "    return maps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: uniprot_mapper\n",
    "#### ANY ID & INPUT TYPE -> UNIPROT ID\n",
    "\n",
    "This function will use the UniProt mapping tool and return a list of UniProt IDs associated with another ID.\n",
    "\n",
    "You need to know where your original ID comes from, and enter the code from http://www.uniprot.org/faq/28#id_mapping_examples\n",
    "\n",
    "By default it will return all UniProt IDs, including unreviewed ones. Put reviewed=True if you want only reviewed ones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def uniprot_mapper(gene_id, id_type, reviewed=False):\n",
    "    '''\n",
    "    Test out a mapping here: http://www.uniprot.org/mapping/\n",
    "    Specify ID type from here: http://www.uniprot.org/faq/28#id_mapping_examples\n",
    "    reviewed=True for only reviewed UniProt entries\n",
    "    '''\n",
    "    \n",
    "    pairs = uniprot.batch_uniprot_id_mapping_pairs(id_type, 'ACC', [gene_id])\n",
    "\n",
    "    unis = []\n",
    "    for idx in range(len(pairs)):\n",
    "        unis.append(str(pairs[idx][1]))\n",
    "\n",
    "    if reviewed == True:\n",
    "        meta = uniprot.batch_uniprot_metadata(unis)\n",
    "        for uni_id in meta.keys():\n",
    "            if meta[uni_id]['is_reviewed'] == False:\n",
    "                del meta[uni_id]\n",
    "        return meta.keys()\n",
    "\n",
    "    else:\n",
    "        return unis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: get_pdb_seq_and_rez\n",
    "#### PDB ID -> DOWNLOAD/PARSE PDB & STRUCTURE SEQUENCE & RESOLUTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_pdb_seq_and_rez(pdbid):\n",
    "    '''\n",
    "    This function downloads a PDB ID (biological assembly) into the working directory and also\n",
    "    returns the sequence that is represented in the structure (what residues are visible).\n",
    "    The sequence is a SeqRecord Biopython object.\n",
    "\n",
    "    Input: pdbid - a 4 letter string\n",
    "    Output: structure_seqs - a Seq object that represents what residues can be seen in the structure\n",
    "            resolution - the crystal structure resolution\n",
    "    '''\n",
    "\n",
    "    try:\n",
    "        # change to biomol=True if you want biological assemblies\n",
    "        pdb, header = pr.parsePDB(pdbid, header=True, biomol=False, secondary=False)\n",
    "\n",
    "        if 'NMR' in header['experiment'] or 'nmr' in header['experiment']:\n",
    "            warnings.warn('******WARNING: NMR entry %s******' % pdbid)\n",
    "\n",
    "        # if multiple biomols choose the first one (this could be improved)\n",
    "        if type(pdb) == list:\n",
    "            pdb = pdb[0]\n",
    "\n",
    "    except IOError:\n",
    "        return None, None\n",
    "    # if there are no biological assemblies\n",
    "    except ValueError:\n",
    "        pdb, header = pr.parsePDB(pdbid, header=True, biomol=False, secondary=False)\n",
    "\n",
    "    if 'resolution' not in header:\n",
    "        resolution = None\n",
    "    else:\n",
    "        resolution = float(header['resolution'])\n",
    "\n",
    "    ## if doing bio assembly NEED TO CONSIDER \"SEGMENTS\"!!\n",
    "    if pdb.numSegments() > 0:\n",
    "        structure_seqs = []\n",
    "        for seg in pdb.iterSegments():\n",
    "            sg = seg.getSegname()\n",
    "            for chain in seg.iterChains():\n",
    "                ch = chain.getChid()\n",
    "                seq = SeqRecord(Seq(chain.getSequence(), IUPAC.protein), id=pdbid+'.'+sg+'.'+ch, description='PDB structure residues, biomol 1')\n",
    "                structure_seqs.append(seq)\n",
    "    else:\n",
    "        structure_seqs = []\n",
    "        for chain in pdb.iterChains():\n",
    "            ch = chain.getChid()\n",
    "            seq = SeqRecord(Seq(chain.getSequence(), IUPAC.protein), id=pdbid+'.'+ch, description='PDB structure residues')\n",
    "            structure_seqs.append(seq)\n",
    "\n",
    "    return structure_seqs, resolution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "****"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NEW FUNCTIONS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: get_org_from_pdb\n",
    "#### PDB ID -> DOWNLOAD/PARSE PDB & FIND WHICH ORGANISM THE TEMPLATE IS DERIVED FROM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from bs4 import *\n",
    "import urllib2\n",
    "from urllib2 import Request\n",
    "\n",
    "def get_pdb_info(pdbid):\n",
    "    tmp_url='http://www.rcsb.org/pdb/rest/describeMol?structureId=%s'%pdbid\n",
    "    req = Request(tmp_url)\n",
    "    res = urllib2.urlopen(req)\n",
    "    soup = BeautifulSoup(urllib2.urlopen(req))\n",
    "    tmp_org = str(soup.contents[0].taxonomy.get(\"id\"))\n",
    "    return tmp_org\n",
    "    \n",
    "def get_org_from_pdb(pdb_file):\n",
    "    pdb = get_pdb_info(pdb_file)\n",
    "    uniprot_url = 'http://www.uniprot.org/taxonomy/%s'%pdb\n",
    "    req = Request(uniprot_url)\n",
    "    res = urllib2.urlopen(req)\n",
    "    soup = BeautifulSoup(urllib2.urlopen(req))\n",
    "    if 'unidentified' in str(soup.title.get_text()):\n",
    "        tmp ='unidentified'\n",
    "    else:\n",
    "        tmp = soup.title.get_text().split()[0]+\" \"+soup.title.get_text().split()[1]\n",
    "    return tmp.encode('utf8')\n",
    "\n",
    "# Usage: organism = get_org_from_pdb(pdbid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function: blast_pdb\n",
    "#### seq -> dictionary of blast results from blasting entire pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import xmltodict as xd\n",
    "\n",
    "def blast_pdb(seq, evalue = 0.001):\n",
    "    '''\n",
    "    Returns the first BLASTp hit with a default evalue search param of 0.001. \n",
    "    returns a dict of..  %TODO: lots of stuff!\n",
    "    up to u if u think its a good hit.\n",
    "    '''\n",
    "    print 'http://www.rcsb.org/pdb/rest/getBlastPDB1?sequence=%s&eCutOff=%s&matrix=BLOSUM62&outputFormat=XML' % (seq, evalue)\n",
    "    req = urllib2.Request('http://www.rcsb.org/pdb/rest/getBlastPDB1?sequence=%s&eCutOff=%s&matrix=BLOSUM62&outputFormat=XML' % (seq, evalue))\n",
    "    response = urllib2.urlopen(req)\n",
    "    \n",
    "    if response.getcode() == 200:\n",
    "        info_dict = {}\n",
    "        len_orig = len(seq)\n",
    "\n",
    "        raw = sio.StringIO(response.read())\n",
    "        doc = de_unicodeify(xd.parse(raw))\n",
    "        \n",
    "        if 'Iteration_hits' not in doc['BlastOutput']['BlastOutput_iterations']['Iteration'].keys():\n",
    "            return None\n",
    "        else:\n",
    "            if isinstance(doc['BlastOutput']['BlastOutput_iterations']['Iteration']['Iteration_hits']['Hit'], list):\n",
    "                best_hit = doc['BlastOutput']['BlastOutput_iterations']['Iteration']['Iteration_hits']['Hit'][0]\n",
    "            else:\n",
    "                best_hit = doc['BlastOutput']['BlastOutput_iterations']['Iteration']['Iteration_hits']['Hit']\n",
    "            hit_pdb = best_hit['Hit_def'].split('|')[0].split(':')[0]\n",
    "            hit_pdb_chain = best_hit['Hit_def'].split('|')[0].split(':')[2]\n",
    "\n",
    "            hit_dict = de_unicodeify(best_hit['Hit_hsps']['Hsp'])\n",
    "\n",
    "            if isinstance(hit_dict, list):\n",
    "                hit_dict = hit_dict[0]\n",
    "\n",
    "            info_dict['input_seq'] = seq\n",
    "            info_dict['input_seq_len'] = len_orig\n",
    "            info_dict['hit_pdb'] = hit_pdb\n",
    "            info_dict['hit_pdb_chain'] = hit_pdb_chain\n",
    "            info_dict['hit_seq'] = hit_dict['Hsp_hseq']\n",
    "            info_dict['hit_num_ident'] = hit_dict['Hsp_identity']\n",
    "            info_dict['hit_percent_ident'] = float(hit_dict['Hsp_identity'])/float(len_orig)\n",
    "            info_dict['hit_evalue'] = hit_dict['Hsp_evalue']\n",
    "            info_dict['hit_score'] = hit_dict['Hsp_score']\n",
    "    \n",
    "            return info_dict\n",
    "        \n",
    "    else:\n",
    "        return None"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}