cellgrid-example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cellgrid\n",
    "import MDAnalysis as mda\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# show  capped_self_distance_array  error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "box = np.ones(3).astype(np.float32) * 100\n",
    "# equally distributed points (a tough case!)\n",
    "a = np.random.uniform(low=0, high=box[0], size=(1000, 3))\n",
    "a = np.array(a, order='F', dtype=np.float32)\n",
    "maxdist = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "capped_dist = cellgrid.capped_self_distance_array(a, maxdist, box=box)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[269, 727],\n",
       "        [269, 610],\n",
       "        [269, 760],\n",
       "        ...,\n",
       "        [207, 539],\n",
       "        [207, 439],\n",
       "        [207, 758]]),\n",
       " array([14.953247, 16.304543, 10.556907, ..., 18.55288 , 20.035353,\n",
       "        16.36398 ], dtype=float32))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "capped_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "check = mda.lib.distances.distance_array(a, a, box)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "row, col = np.where(check < maxdist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(5126, (13427,))"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(row), capped_dist[1].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of rows found by a brute force distance search should match up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Give Performance example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import sparse\n",
    "from MDAnalysis.lib.distances import distance_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "def distances(a, maxdist, box, nbins_grid=5):\n",
    "    # estimate binsize of cellgrid for optimal performance\n",
    "    binsize = (1 / nbins_grid) * (a.max() - a.min())\n",
    "    cg = cellgrid.CellGrid(box, binsize, a)\n",
    "    N = len(a)\n",
    "    out = sparse.lil_matrix((N, N))\n",
    "\n",
    "    for cell in cg:\n",
    "        # check for particles in cell\n",
    "        if len(cell) > 1:\n",
    "            d = distance_array(cell.coordinates,\n",
    "                               cell.coordinates,\n",
    "                               box)\n",
    "            # overwrite self distance\n",
    "            i, j = np.where(d < maxdist)\n",
    "            out[cell.indices[i], cell.indices[j]] = d[(i, j)]\n",
    "        for addr in cell.all_neighbours:\n",
    "            other = cg[addr]\n",
    "            if not other:\n",
    "                continue\n",
    "            d = distance_array(cell.coordinates,\n",
    "                               other.coordinates,\n",
    "                               box)\n",
    "            i, j = np.where(d < maxdist)\n",
    "            out[cell.indices[i], other.indices[j]] = d[(i, j)]\n",
    "    return out.tocoo()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "box = np.ones(3).astype(np.float32) * 100\n",
    "# equally distributed points (a tough case!)\n",
    "a = np.random.uniform(low=0, high=box[0], size=(1000, 3))\n",
    "a = np.array(a, order='F', dtype=np.float32)\n",
    "maxdist = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_distance = distances(a, maxdist, box)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check for correctness \n",
    "check = distance_array(a, a, box)\n",
    "# the cellgrid version doesn't \n",
    "check += 2*maxdist*np.eye(len(a))\n",
    "row, col = np.where(check < maxdist)\n",
    "\n",
    "np.testing.assert_equal(row, max_distance.row)\n",
    "np.testing.assert_equal(col, max_distance.col)\n",
    "np.testing.assert_almost_equal(check[row, col], max_distance.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "200 ms ± 5.65 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit distances(a, maxdist, box)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 47s, sys: 137 ms, total: 1min 47s\n",
      "Wall time: 1min 47s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(array([[222, 148],\n",
       "        [590, 266],\n",
       "        [888, 429],\n",
       "        [815, 675],\n",
       "        [927, 454],\n",
       "        [378, 333],\n",
       "        [464, 166],\n",
       "        [434, 597],\n",
       "        [712, 844],\n",
       "        [929, 209],\n",
       "        [839, 345],\n",
       "        [742, 837],\n",
       "        [366, 667]]),\n",
       " array([1.3200516 , 1.6378901 , 1.7164932 , 2.3929284 , 1.2950763 ,\n",
       "        1.8017806 , 1.5394814 , 1.9913622 , 1.4522259 , 1.9939705 ,\n",
       "        0.77306837, 2.4229352 , 1.6182246 ], dtype=float32))"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%time cellgrid.capped_self_distance_array(a, maxdist, box)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:distances]",
   "language": "python",
   "name": "conda-env-distances-py"
  },
  "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.4"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}