<commit_before>{
 "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",<commit_msg>Bilayer example notebook (#406)<commit_after>{
 "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",
    "                new_lipid.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",
    "                    new_lipid.translate(new_lipid, -center)\n",
    "                    mb.rotate_around_x(new_lipid, np.pi)\n",
    "                    new_lipid.translate(new_lipid, center)\n",
    "                new_lipid.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",
{'fraction_non_alphanumeric': 0.14852294154619736, 'fraction_numerical': 0.00458830923947203, 'mean_word_length': 1.8534791965566715, 'pattern_counts': {'":': 32, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7375353', 'n_tokens_mistral': 5697, 'n_tokens_neox': 5336, 'n_words': 1090}