#!/usr/bin/env python
# coding: utf-8

# In[1]:


from __future__ import print_function

import math

get_ipython().run_line_magic('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')


# 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]:


get_ipython().run_line_magic('timeit', 'get_distance(p53_1tup.get_atoms())')


# In[5]:


get_ipython().run_cell_magic('cython', '', 'import math\ndef 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) \n')


# In[6]:


get_ipython().run_line_magic('timeit', 'get_distance_cython_0(p53_1tup.get_atoms())')


# In[7]:


get_ipython().run_cell_magic('cython', '', 'cimport cython\nfrom libc.math cimport sqrt, pow\n\ncdef 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\ndef 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)\n')


# In[8]:


get_ipython().run_line_magic('timeit', 'get_distance_cython_1([atom.coord for atom in p53_1tup.get_atoms()])')


# ## Numba

# In[9]:


from numba import float_
from numba.decorators import jit


# In[10]:


get_distance_numba_0 = jit(get_distance)


# In[11]:


get_ipython().run_line_magic('timeit', 'get_distance_numba_0(p53_1tup.get_atoms())')


# 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]:


get_ipython().run_line_magic('timeit', 'get_distance_numba_1([atom.coord for atom in p53_1tup.get_atoms()])')


# In[ ]: