{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "import math\n",
    "\n",
    "%load_ext Cython\n",
    "\n",
    "from Bio import PDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Structure exists: './pdb1tup.ent' \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain A is discontinuous at line 6146.\n",
      "  PDBConstructionWarning)\n",
      "/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain B is discontinuous at line 6147.\n",
      "  PDBConstructionWarning)\n",
      "/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain C is discontinuous at line 6148.\n",
      "  PDBConstructionWarning)\n",
      "/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain E is discontinuous at line 6149.\n",
      "  PDBConstructionWarning)\n",
      "/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain F is discontinuous at line 6171.\n",
      "  PDBConstructionWarning)\n",
      "/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain A is discontinuous at line 6185.\n",
      "  PDBConstructionWarning)\n",
      "/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain B is discontinuous at line 6383.\n",
      "  PDBConstructionWarning)\n",
      "/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain C is discontinuous at line 6453.\n",
      "  PDBConstructionWarning)\n"
     ]
    }
   ],
   "source": [
    "repository = PDB.PDBList()\n",
    "parser = PDB.PDBParser()\n",
    "repository.retrieve_pdb_file('1TUP', pdir='.')\n",
    "p53_1tup = parser.get_structure('P 53', 'pdb1tup.ent')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def get_distance(atoms):\n",
    "    atoms = list(atoms)  # not great\n",
    "    natoms = len(atoms)\n",
    "    for i in range(natoms - 1):\n",
    "        xi, yi, zi = atoms[i].coord\n",
    "        for j in range(i + 1, natoms):\n",
    "            xj, yj, zj = atoms[j].coord\n",
    "            my_dist = math.sqrt((xi - xj)**2 + (yi - yj)**2 + (zi - zj)**2) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 loops, best of 3: 2min 9s per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit get_distance(p53_1tup.get_atoms())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%cython\n",
    "import math\n",
    "def get_distance_cython_0(atoms):\n",
    "    atoms = list(atoms)\n",
    "    natoms = len(atoms)\n",
    "    for i in range(natoms - 1):\n",
    "        xi, yi, zi = atoms[i].coord\n",
    "        for j in range(i + 1, natoms):\n",
    "            xj, yj, zj = atoms[j].coord\n",
    "            my_dist = math.sqrt((xi - xj)**2 + (yi - yj)**2 + (zi - zj)**2) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 loops, best of 3: 2min 7s per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit get_distance_cython_0(p53_1tup.get_atoms())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%cython\n",
    "cimport cython\n",
    "from libc.math cimport sqrt, pow\n",
    "\n",
    "cdef double get_dist_cython(double xi, double yi, double zi,\n",
    "                     double xj, double yj, double zj):\n",
    "    return sqrt(pow(xi - xj, 2) + pow(yi - yj, 2) + pow(zi - zj, 2))\n",
    "\n",
    "def get_distance_cython_1(object atoms):\n",
    "    natoms = len(atoms)\n",
    "    cdef double x1, xj, yi, yj, zi, zj\n",
    "    for i in range(natoms - 1):\n",
    "        xi, yi, zi = atoms[i]\n",
    "        for j in range(i + 1, natoms):\n",
    "            xj, yj, zj = atoms[j]\n",
    "            my_dist = get_dist_cython(xi, yi, zi, xj, yj, zj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 loops, best of 3: 14.4 s per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit get_distance_cython_1([atom.coord for atom in p53_1tup.get_atoms()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from numba import float_\n",
    "from numba.decorators import jit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "get_distance_numba_0 = jit(get_distance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 loops, best of 3: 2min 4s per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit get_distance_numba_0(p53_1tup.get_atoms())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "@jit\n",
    "def get_dist_numba(xi, yi, zi, xj, yj, zj):\n",
    "    return math.sqrt((xi - xj)**2 + (yi - yj)**2 + (zi - zj)**2) \n",
    "\n",
    "def get_distance_numba_1(atoms):\n",
    "    natoms = len(atoms)\n",
    "    for i in range(natoms - 1):\n",
    "        xi, yi, zi = atoms[i]\n",
    "        for j in range(i + 1, natoms):\n",
    "            xj, yj, zj = atoms[j]\n",
    "            my_dist = get_dist_numba(xi, yi, zi, xj, yj, zj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 loops, best of 3: 25.2 s per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit get_distance_numba_1([atom.coord for atom in p53_1tup.get_atoms()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}