Custom TopologyAttrs in MDAnalysis

custom_attrs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Attributes in MDAnalysis\n",
    "\n",
    "This notebooks shows how to add custom topology attributes to a Universe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import MDAnalysis as mda\n",
    "from MDAnalysisTests.datafiles import GRO\n",
    "\n",
    "u = mda.Universe(GRO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simple Topology attributes\n",
    "\n",
    "A simple attribute can be added by defining a class that inherits from `AtomAttr`, `ResidueAttr` or `SegmentAttr`.  This class should define the `attrname` and `singular` attributes, these are used as the resulting attributes on an `AtomGroup`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from MDAnalysis.core.topologyattrs import AtomAttr, ResidueAttr\n",
    "import numpy as np\n",
    "\n",
    "class XRayAttr(AtomAttr):\n",
    "    attrname = 'scatters'\n",
    "    singular = 'scatter'\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use the `add_TopologyAttr` method of a `Universe` to add this attribute to an existing Universe.\n",
    "\n",
    "The resulting data is available via the `.scatters` attribute of any `AtomGroup` of this `Universe`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0., 0., 0., ..., 0., 0., 0.])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u.add_TopologyAttr(\"scatters\", values=np.zeros(len(u.atoms)))\n",
    "\n",
    "u.atoms.scatters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Group methods from an Attribute\n",
    "\n",
    "Topology attributes can also give new methods to Group objects.  Here we define a method which adds the `absolute_scattering` method to all `ResidueGroup` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from MDAnalysis.core.groups import ResidueGroup\n",
    "\n",
    "class ResScatter(ResidueAttr):\n",
    "    attrname = 'res_scatters'\n",
    "    singular = 'res_scatter'\n",
    "    \n",
    "    # define a method\n",
    "    # here the 'self' argument will refer to the ResidueGroup\n",
    "    def abs_scatter(res_group):\n",
    "        return np.abs(res_group.res_scatters)\n",
    "    \n",
    "    # add to the transplants dict, Universe inspects this and adds\n",
    "    # methods to the AtomGroup functionality\n",
    "    transplants = defaultdict(list)\n",
    "    transplants[ResidueGroup].append(('absolute_scattering', abs_scatter))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.16754531, 0.17478824, 0.44387846, ..., 0.76977245, 0.98876158,\n",
       "       0.86706427])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u.add_TopologyAttr('res_scatters', values=np.random.random(len(u.residues)))\n",
    "\n",
    "u.residues.absolute_scattering()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom aggregation of attributes\n",
    "\n",
    "Finally we can control how topology attributes are returned at different levels of the system resolution.\n",
    "\n",
    "For example here we override the default `get_residues` method to make it return the mean value of all atoms.  (By default this would return an array of all of the values per Atom in the Residue)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from MDAnalysis.core.groups import Atom, AtomGroup, Residue, ResidueGroup\n",
    "import numbers\n",
    "\n",
    "class ComplexScatter(AtomAttr):\n",
    "    attrname = 'c_scatters'\n",
    "    singular = 'c_scatter'\n",
    "    \n",
    "    # define which classes this Attribute gets stuck onto\n",
    "    target_classes = [Atom, AtomGroup, Residue, ResidueGroup]\n",
    "    \n",
    "    # redefine the get_residues functionality\n",
    "    def get_residues(self, res_group):\n",
    "        # grab the atoms per Residue\n",
    "        resatoms = self.top.tt.residues2atoms_2d(res_group.ix)\n",
    "\n",
    "        if isinstance(res_group.ix, numbers.Integral):\n",
    "            # for a single residue\n",
    "            vals = self.values[resatoms].mean()\n",
    "        else:\n",
    "            # for a residuegroup\n",
    "            vals = np.empty(len(res_group))\n",
    "            for i, row in enumerate(resatoms):\n",
    "                vals[i] = self.values[row].mean()\n",
    "\n",
    "        return vals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "u.add_TopologyAttr('c_scatters', values=np.random.random(len(u.atoms)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.69798291, 0.37982618, 0.34488251, 0.94202816, 0.79299702,\n",
       "       0.50775793, 0.63592951, 0.66233932, 0.73865624, 0.98336715,\n",
       "       0.27424914, 0.3482525 , 0.50392121, 0.7671037 , 0.83558858,\n",
       "       0.43518562, 0.49788706, 0.13559358, 0.22358973])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u.residues[0].atoms.c_scatters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.563533581188531"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u.residues[0].atoms.c_scatters.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.563533581188531"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u.residues[0].c_scatter"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}