barycenters.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load_ext cython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import cython\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from MDAnalysisTests.datafiles import PSF, DCD\n",
    "import MDAnalysis as mda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "u = mda.Universe(PSF, DCD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%cython\n",
    "\n",
    "cimport cython\n",
    "cimport numpy as np\n",
    "import numpy as np\n",
    "\n",
    "@cython.boundscheck(False)\n",
    "@cython.wraparound(False)\n",
    "def cy_barycenters(np.ndarray[dtype=np.float32_t, ndim=2] positions,\n",
    "                int ngroups,\n",
    "                np.ndarray[dtype=np.int64_t, ndim=1] groupsizes):\n",
    "    cdef int i, j, k, size, offset\n",
    "    cdef np.ndarray[dtype=np.float32_t, ndim=2] output = np.empty((ngroups, 3), dtype=np.float32)\n",
    "    \n",
    "    offset = 0\n",
    "    \n",
    "    for i in range(ngroups):\n",
    "        size = groupsizes[i]\n",
    "        \n",
    "        for j in range(size):\n",
    "            for k in range(3):\n",
    "                output[i, k] += positions[offset + j, k]\n",
    "        \n",
    "        for k in range(3):\n",
    "            output[i, k] /= size\n",
    "        \n",
    "        offset += size\n",
    "\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def barycenters(rg):\n",
    "    # center of geometry for each residues\n",
    "    ngroups = len(rg)\n",
    "    sizes = np.array([len(r.atoms) for r in rg])\n",
    "    positions = rg.atoms.positions\n",
    "    \n",
    "    return cy_barycenters(positions, ngroups, sizes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def manual(rg):\n",
    "    return np.array([r.atoms.center_of_geometry() for r in rg])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100 loops, best of 3: 3.59 ms per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit a = barycenters(u.residues)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100 loops, best of 3: 10.1 ms per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit a = manual(u.residues)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100 loops, best of 3: 3.06 ms per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit [r.atoms for r in u.residues]"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda env:mdadev]",
   "language": "python",
   "name": "conda-env-mdadev-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}