{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining custom attributes with the new Topology system\n",
    "\n",
    "This notebook aims to explain how to augment an existing MDAnalysis Universe with additional data.  This is designed to allow more advanced users to define extra properties for AtomGroups which can then be queried or selected according to.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import MDAnalysis as mda\n",
    "from MDAnalysis.core.topologyattrs import AtomAttr\n",
    "from MDAnalysis.core.selection import StringSelection\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "u = mda.Universe('./4E43.pdb')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define an Attribute class for the layer property.  Here we subclass `AtomAttr` which is a shortcut for per-atom attributes.\n",
    "\n",
    "Also define an associated selection for layers.  This uses the `StringSelection` subclass which is how other selections such as `name` and `type` work.  Defining the Selection automatically adds it to the selection language."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Layer(AtomAttr):\n",
    "    attrname = 'layers'  # how it is called in AtomGroup\n",
    "    singular = 'layer'  # how it is called in Atom\n",
    "    \n",
    "class LayerSelection(StringSelection):\n",
    "    token = 'layer'  # how to refer to it in a selection string\n",
    "    field = 'layers'  # where to look in AtomGroup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Arbitrarily split the system into 'lower' and 'upper', obviously this could be done better...."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "a = 1877 // 2\n",
    "b = 1877 - a\n",
    "# Attributes must be able to handle numpy fancy indexing (ie a numpy array)\n",
    "layers = Layer(np.array(['lower'] * a + ['upper'] * b, dtype=object))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add the attribute to the Universe.  This populates the namespace of all Atom, AtomGroup etc classes with the `layer` related data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "u.add_TopologyAttr(layers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create an arbitrary AtomGroup and access the new property or select by it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['lower' 'lower' 'lower' 'lower' 'lower']\n",
      "<AtomGroup with 9 atoms>\n"
     ]
    }
   ],
   "source": [
    "ag = u.atoms[::10]\n",
    "print ag.layers[:5]\n",
    "print ag.select_atoms('name N and layer upper')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}