{"cells":[{"metadata":{},"cell_type":"markdown","source":"Distance matrix between atoms of a single molecule are used in many public kernels to compute features, for instance Coulomb interaction, Van de Walls interaction, and Yukawa interactions, \n\nThis kernel shows how to speed up the distance matrix computation tremendously compared to the already fast version used in [coulomb_interaction - speed up!](https://www.kaggle.com/rio114/coulomb-interaction-speed-up).  The code from that kernel takes 2 minutes for all molecules.  \n\nHere we provide a code that runs in 3 seconds, i.e. 40 times faster.\n\nThis speedup is nice, but the code optimization technique used in this kernel is rather generic and can be reused in other context.\n\nV4 update.  @jmtest has suggested a nice improvement using einssum.  I added his version.  This brings down time to about 1.2 second.  We can do even better using numba, which brings down time to about 0.4 second. This is more than 250 faster than original code."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load in \n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nfrom tqdm import tqdm_notebook\n\n# Input data files are available in the \"../input/\" directory.\n# For example, running this (by clicking run or pressing Shift+Enter) will list the files in the input directory\n\nimport os\nprint(os.listdir(\"../input\"))\n\n# Any results you write to the current directory are saved as output.","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"structures = pd.read_csv('../input/structures.csv')\n\nstructures","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> An efficient way of computing distance matrix is provided in the [coulomb_interaction - speed up!](https://www.kaggle.com/rio114/coulomb-interaction-speed-up) kernel from Ryoji Nomura.  Please drop me a comment if the code is from another author and I'll correct the attribution."},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"structures_idx = structures.set_index('molecule_name')\n\ndef get_dist_matrix(df_structures_idx, molecule):\n    df_temp = df_structures_idx.loc[molecule]\n    locs = df_temp[['x','y','z']].values\n    num_atoms = len(locs)\n    loc_tile = np.tile(locs.T, (num_atoms,1,1))\n    dist_mat = np.sqrt((loc_tile - loc_tile.T)**2).sum(axis=1)\n    return dist_mat","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's see how much time this takes for all molecules"},{"metadata":{"trusted":true},"cell_type":"code","source":"for molecule in tqdm_notebook(structures.molecule_name.unique()):\n    get_dist_matrix(structures_idx, molecule)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This takes 2 minutes.\n\nCan we do better?\n\nA way to accelerate code is to replace pandas operations by numpy operations.  The code above already does this in a very clever way by using the numpy tile function and vectorized distance computation.\n\nStill, it creates a pandas dataframe for each molecule, then converts it into numpy array.  This is slow.  Also, filtering by the molecule name to extract the structure information for a given molecule is expensive.  We can speed up both by using numpy arrays up front.  Let's create the array we need."},{"metadata":{"trusted":true},"cell_type":"code","source":"xyz = structures[['x','y','z']].values\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This will remove the need to first create a data frame then convert it into a numpy array.  This is nice, but we need a way to extract the part relevant to a given molecule from this single numpy array.  \n\nThe original structures data frame is sorted by molecule, meaning that information for a given molecule is in a consecutive set of rows. We can then precompute the indices of each molecule segment.  \n\nWe first compute the number of rows for each molecule, then using cumsum to get indices for molecule changes.  We then store this as a numpy array where we add 0 as the first value.  The molecule names can be retrieved as the index of the series."},{"metadata":{"trusted":true},"cell_type":"code","source":"ss = structures.groupby('molecule_name').size()\nss = ss.cumsum()\nss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"ssx = np.zeros(len(ss) + 1, 'int')\nssx[1:] = ss\nssx","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Then, to retrieve information for the i-th molecule, we slice the xyz array using the indices we just compuited.  For instance, to get the information for the 20th molecule we do:"},{"metadata":{"trusted":true},"cell_type":"code","source":"molecule_id = 20\nprint(ss.index[molecule_id])\nstart_molecule = ssx[molecule_id]\nend_molecule = ssx[molecule_id+1]\nxyz[start_molecule:end_molecule]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can compare with the information we get from the original pandas dataframe"},{"metadata":{"trusted":true},"cell_type":"code","source":"structures_idx.loc['dsgdb9nsd_000022'][['x', 'y', 'z']].values","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Looks good.\n\nWe can now rewrite our function using our arrays.\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_fast_dist_matrix(xyz, ssx, molecule_id):\n    start_molecule, end_molecule = ssx[molecule_id], ssx[molecule_id+1]\n    locs = xyz[start_molecule:end_molecule]    \n    num_atoms = end_molecule - start_molecule\n    loc_tile = np.tile(locs.T, (num_atoms,1,1))\n    dist_mat = np.sqrt((loc_tile - loc_tile.T)**2).sum(axis=1)\n    return dist_mat","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's check we get the same result with both techniques.  We use the smallest molecule for the sake of simplicity. It is water."},{"metadata":{"trusted":true},"cell_type":"code","source":"molecule_id = 2\nmolecule = ss.index[molecule_id]\nprint(molecule)\nget_fast_dist_matrix(xyz, ssx, molecule_id)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"get_dist_matrix(structures_idx, molecule)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Looks good.\n\nWe can now benchmark our version."},{"metadata":{"trusted":true},"cell_type":"code","source":"for molecule_id in tqdm_notebook(range(structures.molecule_name.nunique())):\n    get_fast_dist_matrix(xyz, ssx, molecule_id)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We are down to 3 seconds from 2 minutes!\n\nA nice improvement was proposed by @jmtest in the comments.  In order to benchmark our code and his, we need to use a more precise way than tqdm.  Let's use the `%time` magic. It requires us to wrap the code in a function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def ultra_fast_dist_matrices(xyz, ssx):\n    for molecule_id in range(structures.molecule_name.nunique()):\n        get_fast_dist_matrix(xyz, ssx, molecule_id)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%time ultra_fast_dist_matrices(xyz, ssx)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's do the same for @jmtest code"},{"metadata":{"trusted":true},"cell_type":"code","source":"def sofast_dist(xyz, ssx, molecule_id):\n    start_molecule, end_molecule = ssx[molecule_id], ssx[molecule_id+1]\n    locs = xyz[start_molecule:end_molecule]     \n    d=locs[:,None,:]-locs\n    return np.sqrt(np.einsum('ijk,ijk->ij',d,d))\n\ndef sofast_dist_matrices(xyz, ssx):\n    for molecule_id in range(structures.molecule_name.nunique()):\n        sofast_dist(xyz, ssx, molecule_id)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%time sofast_dist_matrices(xyz, ssx)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"almost a 3x speedup.  \n\nCan we do better?  The ultimate weapon to accelerate vectorized operaitons in Python is to use a compiler.  Numba is my favorite because we stay in Python.  Alternatives are Cython compiler, or calling C/C++ code.  In my experience Numba provides almost the same speed as compiled C code if there are no indirect array access, which is the case here.  \n\nNumba works by annotating code."},{"metadata":{"trusted":true},"cell_type":"code","source":"from numba import jit\nfrom math import sqrt\n\n@jit\ndef numba_dist_matrix(xyz, ssx, molecule_id):\n    start_molecule, end_molecule = ssx[molecule_id], ssx[molecule_id+1]\n    locs = xyz[start_molecule:end_molecule]     \n   # return locs\n    num_atoms = end_molecule - start_molecule\n    dmat = np.zeros((num_atoms, num_atoms))\n    for i in range(num_atoms):\n        for j in range(i+1, num_atoms):\n            d = sqrt((locs[i,0] - locs[j,0])**2 + (locs[i,1] - locs[j,1])**2 + (locs[i,2] - locs[j,2])**2)\n            dmat[i,j] = d\n            dmat[j,i] = d\n    return dmat\n\ndef numba_dist_matrices(xyz, ssx):\n    for molecule_id in range(structures.molecule_name.nunique()):\n        numba_dist_matrix(xyz, ssx, molecule_id)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's make sure our code is correct by comparing with the previous one."},{"metadata":{"trusted":true},"cell_type":"code","source":"molecule_id = 2\nmolecule = ss.index[molecule_id]\nprint(molecule)\nnumba_dist_matrix(xyz, ssx, molecule_id)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sofast_dist(xyz, ssx, molecule_id)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Looks good"},{"metadata":{"trusted":true},"cell_type":"code","source":"%time numba_dist_matrices(xyz, ssx)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"A further 3x speedup!  We are now about 250 times faster than the original code...\n"},{"metadata":{},"cell_type":"markdown","source":""},{"metadata":{},"cell_type":"markdown","source":"It was suggested in comments that scipy distance would be faster.  Let's have a look.  scipy requires two functions to output a square matrix."},{"metadata":{"trusted":true},"cell_type":"code","source":"from scipy.spatial.distance import pdist, squareform\n\ndef scipy_dist_matrix(xyz, ssx, molecule_id):\n    start_molecule, end_molecule = ssx[molecule_id], ssx[molecule_id+1]\n    locs = xyz[start_molecule:end_molecule]     \n    cmat = pdist(locs)\n    dmat = squareform(cmat, force='tomatrix')\n    return dmat","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":""},{"metadata":{},"cell_type":"markdown","source":"Let's check that we get the same result."},{"metadata":{"trusted":true},"cell_type":"code","source":"scipy_dist_matrix(xyz, ssx, molecule_id)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can now benchmark."},{"metadata":{"trusted":true},"cell_type":"code","source":"def scipy_dist_matrices(xyz, ssx):\n    for molecule_id in range(structures.molecule_name.nunique()):\n        scipy_dist_matrix(xyz, ssx, molecule_id)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%time scipy_dist_matrices(xyz, ssx)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Results are disappointing.\n\nIf this kernel gets enough votes then I'll share another kernel with ultra fast way to compute angles and dighedral angles with numpy ;)"},{"metadata":{},"cell_type":"markdown","source":"@ilyivanchenko points out that there is an issue in the get_dist_matrix thagt I borrowed from [coulomb_interaction - speed up!](https://www.kaggle.com/rio114/coulomb-interaction-speed-up)."},{"metadata":{"trusted":true},"cell_type":"code","source":"epsilon = 1e-5\n\ndef get_dist_matrix_assert(df_structures_idx, molecule):\n    df_temp = df_structures_idx.loc[molecule]\n    locs = df_temp[['x','y','z']].values\n    num_atoms = len(locs)\n    loc_tile = np.tile(locs.T, (num_atoms,1,1))\n    dist_mat = np.sqrt((loc_tile - loc_tile.T)**2).sum(axis=1)\n    assert np.abs(dist_mat[0,1] - np.linalg.norm(locs[0] - locs[1])) < epsilon\n    return dist_mat\n\nfor molecule in tqdm_notebook(structures.molecule_name.unique()[660:]):\n    try:\n        get_dist_matrix_assert(structures_idx, molecule)\n    except: \n        print('assertion error on', molecule)\n        break\n        \n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's see if we have the issue in our code."},{"metadata":{"trusted":true},"cell_type":"code","source":"@jit\ndef numba_dist_matrix_ssert(xyz, ssx, molecule_id):\n    start_molecule, end_molecule = ssx[molecule_id], ssx[molecule_id+1]\n    locs = xyz[start_molecule:end_molecule]     \n   # return locs\n    num_atoms = end_molecule - start_molecule\n    dmat = np.zeros((num_atoms, num_atoms))\n    for i in range(num_atoms):\n        for j in range(i+1, num_atoms):\n            d = sqrt((locs[i,0] - locs[j,0])**2 + (locs[i,1] - locs[j,1])**2 + (locs[i,2] - locs[j,2])**2)\n            dmat[i,j] = d\n            dmat[j,i] = d\n    assert np.abs(dmat[0,1] - np.linalg.norm(locs[0] - locs[1])) < epsilon\n    return dmat\n\nfor molecule_id in range(structures.molecule_name.nunique()):\n    numba_dist_matrix_ssert(xyz, ssx, molecule_id)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The code runs fine."},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}