Notebooks/guessbonds_benchmark.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# benchmarks for guess bonds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from MDAnalysis.topology.guessers import guess_bonds\n",
    "import numpy as np\n",
    "import MDAnalysis as mda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "u1,u2 = mda.Universe('small.gro'), mda.Universe('big.gro')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "bonds = guess_bonds(u1.atoms, u1.atoms.positions, box=u1.dimensions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8284"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(bonds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from MDAnalysis.topology import tables\n",
    "from MDAnalysis.lib import distances\n",
    "def old_guess_bonds(atoms, coords, box=None, **kwargs):\n",
    "    r\"\"\"Guess if bonds exist between two atoms based on their distance.\n",
    "\n",
    "    Bond between two atoms is created, if the two atoms are within\n",
    "\n",
    "    .. math::\n",
    "\n",
    "          d < f \\cdot (R_1 + R_2)\n",
    "\n",
    "    of each other, where :math:`R_1` and :math:`R_2` are the VdW radii\n",
    "    of the atoms and :math:`f` is an ad-hoc *fudge_factor*. This is\n",
    "    the `same algorithm that VMD uses`_.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    atoms : AtomGroup\n",
    "         atoms for which bonds should be guessed\n",
    "    coords : array\n",
    "         coordinates of the atoms (i.e., `AtomGroup.positions)`)\n",
    "    fudge_factor : float, optional\n",
    "        The factor by which atoms must overlap eachother to be considered a\n",
    "        bond.  Larger values will increase the number of bonds found. [0.72]\n",
    "    vdwradii : dict, optional\n",
    "        To supply custom vdwradii for atoms in the algorithm. Must be a dict\n",
    "        of format {type:radii}. The default table of van der Waals radii is\n",
    "        hard-coded as :data:`MDAnalysis.topology.tables.vdwradii`.  Any user\n",
    "        defined vdwradii passed as an argument will supercede the table\n",
    "        values. [``None``]\n",
    "    lower_bound : float, optional\n",
    "        The minimum bond length. All bonds found shorter than this length will\n",
    "        be ignored. This is useful for parsing PDB with altloc records where\n",
    "        atoms with altloc A and B maybe very close together and there should be\n",
    "        no chemical bond between them. [0.1]\n",
    "    box : array_like, optional\n",
    "        Bonds are found using a distance search, if unit cell information is\n",
    "        given, periodic boundary conditions will be considered in the distance\n",
    "        search. [``None``]\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    list\n",
    "        List of tuples suitable for use in Universe topology building.\n",
    "\n",
    "    Warnings\n",
    "    --------\n",
    "    No check is done after the bonds are guessed to see if Lewis\n",
    "    structure is correct. This is wrong and will burn somebody.\n",
    "\n",
    "    Raises\n",
    "    ------\n",
    "    :exc:`ValueError` if inputs are malformed or `vdwradii` data is missing.\n",
    "\n",
    "\n",
    "    .. _`same algorithm that VMD uses`:\n",
    "       http://www.ks.uiuc.edu/Research/vmd/vmd-1.9.1/ug/node26.html\n",
    "\n",
    "    .. versionadded:: 0.7.7\n",
    "    .. versionchanged:: 0.9.0\n",
    "       Updated method internally to use more :mod:`numpy`, should work\n",
    "       faster.  Should also use less memory, previously scaled as\n",
    "       :math:`O(n^2)`.  *vdwradii* argument now augments table list\n",
    "       rather than replacing entirely.\n",
    "    \"\"\"\n",
    "    # why not just use atom.positions?\n",
    "    if len(atoms) != len(coords):\n",
    "        raise ValueError(\"'atoms' and 'coord' must be the same length\")\n",
    "\n",
    "    fudge_factor = kwargs.get('fudge_factor', 0.72)\n",
    "\n",
    "    vdwradii = tables.vdwradii.copy()  # so I don't permanently change it\n",
    "    user_vdwradii = kwargs.get('vdwradii', None)\n",
    "    if user_vdwradii:  # this should make algo use their values over defaults\n",
    "        vdwradii.update(user_vdwradii)\n",
    "\n",
    "    # Try using types, then elements\n",
    "    atomtypes = atoms.types\n",
    "\n",
    "    # check that all types have a defined vdw\n",
    "    if not all(val in vdwradii for val in set(atomtypes)):\n",
    "        raise ValueError((\"vdw radii for types: \" +\n",
    "                          \", \".join([t for t in set(atomtypes) if\n",
    "                                     not t in vdwradii]) +\n",
    "                          \". These can be defined manually using the\" +\n",
    "                          \" keyword 'vdwradii'\"))\n",
    "\n",
    "    lower_bound = kwargs.get('lower_bound', 0.1)\n",
    "\n",
    "    if box is not None:\n",
    "        box = np.asarray(box)\n",
    "\n",
    "    # to speed up checking, calculate what the largest possible bond\n",
    "    # atom that would warrant attention.\n",
    "    # then use this to quickly mask distance results later\n",
    "    max_vdw = max([vdwradii[t] for t in atomtypes])\n",
    "\n",
    "    bonds = []\n",
    "\n",
    "    for i, atom in enumerate(atoms[:-1]):\n",
    "        vdw_i = vdwradii[atomtypes[i]]\n",
    "        max_d = (vdw_i + max_vdw) * fudge_factor\n",
    "\n",
    "        # using self_distance_array scales O(n^2)\n",
    "        # 20,000 atoms = 1.6 Gb memory\n",
    "        dist = distances.distance_array(coords[i][None, :], coords[i + 1:],\n",
    "                                        box=box)[0]\n",
    "        idx = np.where((dist > lower_bound) & (dist <= max_d))[0]\n",
    "\n",
    "        for a in idx:\n",
    "            j = i + 1 + a\n",
    "            atom_j = atoms[j]\n",
    "\n",
    "            if dist[a] < (vdw_i + vdwradii[atomtypes[j]]) * fudge_factor:\n",
    "                # because of method used, same bond won't be seen twice,\n",
    "                # so don't need to worry about duplicates\n",
    "                bonds.append((atom.index, atom_j.index))\n",
    "\n",
    "    return tuple(bonds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "bonds_old = old_guess_bonds(u1.atoms, u1.atoms.positions, box=u1.dimensions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8284"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(bonds_old)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.19 s ± 108 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit old_guess_bonds(u1.atoms, u1.atoms.positions, box=u1.dimensions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "394 ms ± 26 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit guess_bonds(u1.atoms, u1.atoms.positions, box=u1.dimensions)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}