In [1]:
from __future__ import print_function

import math

%load_ext Cython

from Bio import PDB

In [2]:
repository = PDB.PDBList()
parser = PDB.PDBParser()
repository.retrieve_pdb_file('1TUP', pdir='.')
p53_1tup = parser.get_structure('P 53', 'pdb1tup.ent')


Structure exists: './pdb1tup.ent' 
/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain A is discontinuous at line 6146.
  PDBConstructionWarning)
/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain B is discontinuous at line 6147.
  PDBConstructionWarning)
/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain C is discontinuous at line 6148.
  PDBConstructionWarning)
/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain E is discontinuous at line 6149.
  PDBConstructionWarning)
/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain F is discontinuous at line 6171.
  PDBConstructionWarning)
/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain A is discontinuous at line 6185.
  PDBConstructionWarning)
/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain B is discontinuous at line 6383.
  PDBConstructionWarning)
/home/tra/anaconda3/envs/book2/lib/python2.7/site-packages/Bio/PDB/StructureBuilder.py:86: PDBConstructionWarning: WARNING: Chain C is discontinuous at line 6453.
  PDBConstructionWarning)

In [3]:
def get_distance(atoms):
    atoms = list(atoms)  # not great
    natoms = len(atoms)
    for i in range(natoms - 1):
        xi, yi, zi = atoms[i].coord
        for j in range(i + 1, natoms):
            xj, yj, zj = atoms[j].coord
            my_dist = math.sqrt((xi - xj)**2 + (yi - yj)**2 + (zi - zj)**2)

In [4]:
%timeit get_distance(p53_1tup.get_atoms())


1 loops, best of 3: 2min 9s per loop

In [5]:
%%cython
import math
def get_distance_cython_0(atoms):
    atoms = list(atoms)
    natoms = len(atoms)
    for i in range(natoms - 1):
        xi, yi, zi = atoms[i].coord
        for j in range(i + 1, natoms):
            xj, yj, zj = atoms[j].coord
            my_dist = math.sqrt((xi - xj)**2 + (yi - yj)**2 + (zi - zj)**2)

In [6]:
%timeit get_distance_cython_0(p53_1tup.get_atoms())


1 loops, best of 3: 2min 7s per loop

In [7]:
%%cython
cimport cython
from libc.math cimport sqrt, pow

cdef double get_dist_cython(double xi, double yi, double zi,
                     double xj, double yj, double zj):
    return sqrt(pow(xi - xj, 2) + pow(yi - yj, 2) + pow(zi - zj, 2))

def get_distance_cython_1(object atoms):
    natoms = len(atoms)
    cdef double x1, xj, yi, yj, zi, zj
    for i in range(natoms - 1):
        xi, yi, zi = atoms[i]
        for j in range(i + 1, natoms):
            xj, yj, zj = atoms[j]
            my_dist = get_dist_cython(xi, yi, zi, xj, yj, zj)

In [8]:
%timeit get_distance_cython_1([atom.coord for atom in p53_1tup.get_atoms()])


1 loops, best of 3: 14.4 s per loop

Numba


In [9]:
from numba import float_
from numba.decorators import jit

In [10]:
get_distance_numba_0 = jit(get_distance)

In [11]:
%timeit get_distance_numba_0(p53_1tup.get_atoms())


1 loops, best of 3: 2min 4s per loop

In [12]:
@jit
def get_dist_numba(xi, yi, zi, xj, yj, zj):
    return math.sqrt((xi - xj)**2 + (yi - yj)**2 + (zi - zj)**2) 

def get_distance_numba_1(atoms):
    natoms = len(atoms)
    for i in range(natoms - 1):
        xi, yi, zi = atoms[i]
        for j in range(i + 1, natoms):
            xj, yj, zj = atoms[j]
            my_dist = get_dist_numba(xi, yi, zi, xj, yj, zj)

In [13]:
%timeit get_distance_numba_1([atom.coord for atom in p53_1tup.get_atoms()])


1 loops, best of 3: 25.2 s per loop

In [ ]: