bilayer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'writeandexecute' magic loaded.\n"
     ]
    }
   ],
   "source": [
    "%load_ext ipyext.writeandexecute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%%writeandexecute -i bilayer bilayer.py\n",
    "from random import seed, shuffle\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import mbuild as mb\n",
    "from mbuild import clone\n",
    "\n",
    "class Bilayer(mb.Compound):\n",
    "    \"\"\"Create a lipid bilayer and add solvent above and below.\n",
    "\n",
    "     This example is still pretty immature as it only represents some brief\n",
    "     scratch work. Feel free to flesh it out!\n",
    "\n",
    "    Attributes\n",
    "    ----------\n",
    "    lipids : list\n",
    "        List of tuples in format (lipid, frac) where frac is the fraction of\n",
    "        that lipid in the bilayer (lipid is a Compound)\n",
    "    ref_atoms : int\n",
    "        Indices of the atom in lipids to form the interface, one for each lipid\n",
    "        in lipids (i.e., this atom is shifted to the 'interface' level)\n",
    "    n_lipids_x : int\n",
    "        Number of lipids in the x-direction per layer.\n",
    "    n_lipids_y : int\n",
    "        Number of lipids in the x-direction per layer.\n",
    "    area_per_lipid : float\n",
    "        Area per lipid.\n",
    "    solvent : Compound\n",
    "        Compound to solvate the bilayer with. Typically, a pre-equilibrated box\n",
    "        of solvent.\n",
    "    lipid_box : Box, optional\n",
    "        A Box containing the lipids where no solvent will be added.\n",
    "    spacing_z : float, optional\n",
    "        Amount of space to add between opposing monolayers.\n",
    "    solvent_per_lipid : int, optional, default=\n",
    "        Number of solvent molecules per lipid\n",
    "    n_solvent : int, optional, default=None\n",
    "        *Total* number of solvent molecules\n",
    "    random_seed : int, optional, default=12345\n",
    "        Seed for random number generator for filling in lipids.\n",
    "    mirror : bool, optional, default=True\n",
    "        Make top and bottom layers mirrors of each other.\n",
    "\n",
    "    \"\"\"\n",
    "    def __init__(self, lipids, ref_atoms, n_lipids_x=10, n_lipids_y=10, \n",
    "                 area_per_lipid=1.0, solvent=None, lipid_box=None, \n",
    "                 spacing_z=0.5, solvent_per_lipid=None, n_solvent=None,\n",
    "                 random_seed=12345, mirror=True):\n",
    "        super(Bilayer, self).__init__()\n",
    "\n",
    "        # Santitize inputs.\n",
    "        if sum([lipid[1] for lipid in lipids]) != 1.0:\n",
    "            raise ValueError('Lipid fractions do not add up to 1.')\n",
    "        assert len(ref_atoms) == len(lipids)\n",
    "\n",
    "        self.lipids = lipids\n",
    "        self.ref_atoms = ref_atoms\n",
    "        self._lipid_box = lipid_box\n",
    "\n",
    "        # 2D Lipid locations.\n",
    "        self.n_lipids_x = n_lipids_x\n",
    "        self.n_lipids_y = n_lipids_y\n",
    "        self.apl = area_per_lipid\n",
    "        self.n_lipids_per_layer = self.n_lipids_x * self.n_lipids_y\n",
    "        self.pattern = mb.Grid2DPattern(n_lipids_x, n_lipids_y)\n",
    "        self.pattern.scale(np.sqrt(self.apl * self.n_lipids_per_layer))\n",
    "\n",
    "        # Solvent info.\n",
    "        self.solvent = solvent\n",
    "        self.n_solvent = n_solvent\n",
    "        self.solvent_per_lipid = solvent_per_lipid\n",
    "\n",
    "        # Other inputs.\n",
    "        self.spacing = np.array([0, 0, spacing_z])\n",
    "        self.random_seed = random_seed\n",
    "        self.mirror = mirror\n",
    "\n",
    "        self._number_of_each_lipid_per_layer = []\n",
    "        self._solvent_per_layer = None\n",
    "\n",
    "        # Containers for lipids and solvent.\n",
    "        self.lipid_components = mb.Compound()\n",
    "        self.solvent_components = mb.Compound()\n",
    "\n",
    "        # Assemble the lipid layers\n",
    "        seed(self.random_seed)\n",
    "        top_layer, top_lipid_labels = self.create_layer()\n",
    "        self.lipid_components.add(top_layer)\n",
    "        if self.mirror == True:\n",
    "            bottom_layer, bottom_lipid_labels = self.create_layer(lipid_indices=top_lipid_labels, flip_orientation=True)\n",
    "        else:\n",
    "            bottom_layer, bottom_lipid_labels = self.create_layer(flip_orientation=True)\n",
    "        self.lipid_components.add(bottom_layer)\n",
    "\n",
    "        # solvate the lipids\n",
    "        #self.solvate_bilayer()  # TODO: needs fixing\n",
    "\n",
    "        # add everything to the big list\n",
    "        self.add(self.lipid_components)\n",
    "        self.add(self.solvent_components)\n",
    "        print(self.number_of_each_lipid_per_layer)\n",
    "        # TODO(tim): shift everything so that the lipids are centered in the box?\n",
    "\n",
    "    def create_layer(self, lipid_indices=None, flip_orientation=False):\n",
    "        \"\"\"Create a monolayer of lipids.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        lipid_indices : list, optional, default=None\n",
    "            A list of indices associated with each lipid in the layer.\n",
    "        flip_orientation : bool, optional, default=False\n",
    "            Flip the orientation of the layer with respect to the z-dimension.\n",
    "\n",
    "        \"\"\"\n",
    "        layer = mb.Compound()\n",
    "        if not lipid_indices:\n",
    "            lipid_indices = list(range(self.n_lipids_per_layer))\n",
    "            shuffle(lipid_indices)\n",
    "\n",
    "        for n_type, n_of_lipid_type in enumerate(self.number_of_each_lipid_per_layer):\n",
    "            current_type = self.lipids[n_type][0]\n",
    "            for n_this_type, n_this_lipid_type in enumerate(range(n_of_lipid_type)):\n",
    "                lipids_placed = n_type + n_this_type\n",
    "                new_lipid = clone(current_type)\n",
    "                random_index = lipid_indices[lipids_placed]\n",
    "                position = self.pattern[random_index]\n",
    "\n",
    "                # Zero and space in z-direction\n",
    "                particles = list(new_lipid.particles())\n",
    "                ref_atom = self.ref_atoms[n_type]\n",
    "                mb.translate(new_lipid, -particles[ref_atom].pos + self.spacing)\n",
    "\n",
    "                # Move to point on pattern\n",
    "                if flip_orientation == True:\n",
    "                    # TODO: Function for this?\n",
    "                    # E.g., rotate_around_x_keep_com(compound, bool(3))\n",
    "                    center = new_lipid.center\n",
    "                    center[2] = 0.0\n",
    "                    mb.translate(new_lipid, -center)\n",
    "                    mb.rotate_around_x(new_lipid, np.pi)\n",
    "                    mb.translate(new_lipid, center)\n",
    "                mb.translate(new_lipid, position)\n",
    "                layer.add(new_lipid)\n",
    "        return layer, lipid_indices\n",
    "\n",
    "    def solvate_bilayer(self):\n",
    "        \"\"\"Solvate the constructed bilayer. \"\"\"\n",
    "        solvent_number_density = self.solvent.n_particles / np.prod(self.solvent.periodicity)\n",
    "\n",
    "        lengths = self.lipid_box.lengths\n",
    "        water_box_z = self.solvent_per_layer / (lengths[0] * lengths[1] * solvent_number_density)\n",
    "\n",
    "        mins = self.lipid_box.mins\n",
    "        maxs = self.lipid_box.maxs\n",
    "        bilayer_solvent_box = mb.Box(mins=[mins[0], mins[1], maxs[2]],\n",
    "                                     maxs=[maxs[0], maxs[1], maxs[2] + 2 * water_box_z])\n",
    "\n",
    "        self.solvent_components.add(mb.fill_box(self.solvent, bilayer_solvent_box))\n",
    "\n",
    "    @property\n",
    "    def solvent_per_layer(self):\n",
    "        \"\"\"Determine the number of solvent molecules per single layer.  \"\"\"\n",
    "        if self._solvent_per_layer:\n",
    "            return self._solvent_per_layer\n",
    "\n",
    "        assert not (self.solvent_per_lipid is None and self.n_solvent is None)\n",
    "        if self.solvent_per_lipid is not None:\n",
    "            assert self.n_solvent is None\n",
    "            self._solvent_per_layer = self.n_lipids_per_layer * self.solvent_per_lipid\n",
    "        elif self.n_solvent is not None:\n",
    "            assert self.solvent_per_lipid is None\n",
    "            self._solvent_per_layer = self.n_solvent / 2\n",
    "        return self._solvent_per_layer\n",
    "\n",
    "    @property\n",
    "    def number_of_each_lipid_per_layer(self):\n",
    "        \"\"\"The number of each lipid per layer. \"\"\"\n",
    "        if self._number_of_each_lipid_per_layer:\n",
    "            return self._number_of_each_lipid_per_layer\n",
    "\n",
    "        for lipid in self.lipids[:-1]:\n",
    "            self._number_of_each_lipid_per_layer.append(int(round(lipid[1] * self.n_lipids_per_layer)))\n",
    "\n",
    "        # TODO: give warning if frac * n different than actual\n",
    "        # Rounding errors may make this off by 1, so just do total - whats_been_added.\n",
    "        self._number_of_each_lipid_per_layer.append(self.n_lipids_per_layer - sum(self._number_of_each_lipid_per_layer))\n",
    "        assert len(self._number_of_each_lipid_per_layer) == len(self.lipids)\n",
    "        return self._number_of_each_lipid_per_layer\n",
    "\n",
    "    @property\n",
    "    def lipid_box(self):\n",
    "        \"\"\"The box containing all of the lipids. \"\"\"\n",
    "        if self._lipid_box:\n",
    "            return self._lipid_box\n",
    "        else:\n",
    "            self._lipid_box = self.lipid_components.boundingbox\n",
    "            # Add buffer around lipid box.\n",
    "            self._lipid_box.mins -= np.array([0.5*np.sqrt(self.apl),\n",
    "                                              0.5*np.sqrt(self.apl),\n",
    "                                              0.5*np.sqrt(self.apl)])\n",
    "            self._lipid_box.maxs += np.array([0.5*np.sqrt(self.apl),\n",
    "                                              0.5*np.sqrt(self.apl),\n",
    "                                              0.5*np.sqrt(self.apl)])\n",
    "            return self._lipid_box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[112, 113]\n",
      "<Bilayer 4947 particles, non-periodic, 5174 bonds, id: 4509429152>\n"
     ]
    }
   ],
   "source": [
    "from mbuild.utils.io import get_fn\n",
    "from mbuild.lib.moieties import H2O\n",
    "\n",
    "water = H2O()\n",
    "ecerns = mb.load(get_fn('ecer2.pdb'))\n",
    "\n",
    "chol = mb.load(get_fn('cg-chol.pdb'))\n",
    "# Orient along the z-direction.\n",
    "mb.rotate_around_x(chol, -135.0*np.pi/180)\n",
    "mb.rotate_around_y(chol, -45.0*np.pi/180)\n",
    "\n",
    "lipids = [(ecerns, 0.5), (chol, 0.5)]\n",
    "\n",
    "bilayer = Bilayer(lipids, n_lipids_x=15, n_lipids_y=15, area_per_lipid=1.4,\n",
    "                  solvent=water, ref_atoms=[1, 6],  spacing_z=0.7,\n",
    "                  solvent_per_lipid=20, mirror=False)\n",
    "\n",
    "print(bilayer)\n",
    "bilayer.visualize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.4.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}