tip3p-indexedset-lookup.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import mbuild as mb\n",
    "\n",
    "import gmso\n",
    "from gmso import ForceField\n",
    "from gmso.external import convert_mbuild\n",
    "from gmso.formats.lammpsdata import write_lammpsdata\n",
    "import warnings; warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile spce.mol2\n",
    "@<TRIPOS>MOLECULE\n",
    "RES\n",
    "3 0 1 0 1\n",
    "SMALL\n",
    "NO_CHARGES\n",
    "@<TRIPOS>CRYSIN\n",
    "    3.0130     3.0130     3.0130    90.0000    90.0000    90.0000  1  1\n",
    "@<TRIPOS>ATOM\n",
    "       1 O            0.0000     0.0000     0.0000 O             1 RES     \n",
    "       2 H           -0.6126    -0.7357     0.0000 H             1 RES     \n",
    "       3 H           -0.5469     0.7762     0.0000 H             1 RES     \n",
    "@<TRIPOS>BOND\n",
    "     1     1     2    1\n",
    "     2     1     3    1\n",
    "@<TRIPOS>SUBSTRUCTURE\n",
    "       1 RES             1 RESIDUE    0 **** ROOT      0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile spce.xml\n",
    "<?xml version='1.0' encoding='UTF-8'?>\n",
    "<ForceField name=\"spce\" version=\"0.0.1\">\n",
    "  <FFMetaData>\n",
    "    <Units energy=\"kcal/mol\" mass=\"amu\" charge=\"elementary_charge\" distance=\"nm\"/>\n",
    "  </FFMetaData>\n",
    "  <AtomTypes expression=\"4 * epsilon * ((sigma/r)**12 - (sigma/r)**6)\">\n",
    "    <ParametersUnitDef parameter=\"epsilon\" unit=\"kJ/mol\"/>\n",
    "    <ParametersUnitDef parameter=\"sigma\" unit=\"nm\"/>\n",
    "    <AtomType name=\"opls_116\" atomclass=\"OW\" element=\"O\" charge=\"-0.8476\" mass=\"15.999\" definition=\"O\">\n",
    "      <Parameters>\n",
    "        <Parameter name=\"epsilon\" value=\"0.650194\"/>\n",
    "        <Parameter name=\"sigma\" value=\"0.316557\"/>\n",
    "      </Parameters>\n",
    "    </AtomType>\n",
    "    <AtomType name=\"opls_117\" atomclass=\"HW\" element=\"H\" charge=\"0.4238\" mass=\"1.00784\" definition=\"H\">\n",
    "      <Parameters>\n",
    "        <Parameter name=\"epsilon\" value=\"0.0\"/>\n",
    "        <Parameter name=\"sigma\" value=\"0.1\"/>\n",
    "      </Parameters>\n",
    "    </AtomType>\n",
    "  </AtomTypes>\n",
    "  <BondTypes expression=\"0.5 * k * (r-r_eq)**2\">\n",
    "    <ParametersUnitDef parameter=\"k\" unit=\"kJ/mol/nm**2\"/>\n",
    "    <ParametersUnitDef parameter=\"r_eq\" unit=\"nm\"/>\n",
    "    <BondType name=\"BondType-Harmonic-1\" type1=\"opls_116\" type2=\"opls_117\">\n",
    "      <Parameters>\n",
    "        <Parameter name=\"k\" value=\"345000.0\"/>\n",
    "        <Parameter name=\"r_eq\" value=\"0.1\"/>\n",
    "      </Parameters>\n",
    "    </BondType>\n",
    "  </BondTypes>\n",
    "  <AngleTypes expression=\"0.5 * k * (theta - theta_eq)**2\">\n",
    "    <ParametersUnitDef parameter=\"k\" unit=\"kJ/mol/degree**2\"/>\n",
    "    <ParametersUnitDef parameter=\"theta_eq\" unit=\"degree\"/>\n",
    "    <AngleType name=\"AngleType-Harmonic-1\" type1=\"opls_116\" type2=\"opls_117\" type3=\"opls_117\">\n",
    "      <Parameters>\n",
    "        <Parameter name=\"k\" value=\"383.0\"/>\n",
    "        <Parameter name=\"theta_eq\" value=\"1.91061193\"/>\n",
    "      </Parameters>\n",
    "    </AngleType>\n",
    "  </AngleTypes>\n",
    "</ForceField>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load in single water (SPC/E) structure\n",
    "\n",
    "water = mb.load(\"spce.mol2\")\n",
    "water = water.children[0]\n",
    "water.name = \"water\"\n",
    "\n",
    "# element_map = which site name corresponds to which atom_type name. \n",
    "# In the future atomtyping will be done through foyer. \n",
    "element_map = {\"O\": \"opls_116\",\n",
    "               \"H\": \"opls_117\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load in topology forcefield\n",
    "\n",
    "forcefield = ForceField(\"spce.xml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to set up a water box with n waters\n",
    "def setup_sys(n=100):\n",
    "    water_box = mb.fill_box(\n",
    "        compound=water,     \n",
    "        n_compounds=n,\n",
    "        density=1000,\n",
    "    )\n",
    "    \n",
    "    top = convert_mbuild.from_mbuild(water_box)\n",
    "    \n",
    "    # Assign atom types\n",
    "    for atom in top.sites:\n",
    "        atom.atom_type = forcefield.atom_types[element_map[atom.name]]\n",
    "\n",
    "    # Assign bond types\n",
    "    for bond in top.bonds:\n",
    "        bond.bond_type = bond.connection_type = forcefield.bond_types[\"opls_116~opls_117\"]\n",
    "\n",
    "    # Create angles with correct atom type and add to top\n",
    "    for subtop in top.subtops:\n",
    "        angle = gmso.core.angle.Angle(\n",
    "            connection_members=[site for site in subtop.sites],\n",
    "            name=\"opls_116~opls_117~opls_117\",\n",
    "            connection_type=forcefield.angle_types[\"opls_116~opls_117~opls_117\"]\n",
    "        )\n",
    "        top.add_connection(angle, update_types=False)\n",
    "\n",
    "    top.update_topology()\n",
    "    \n",
    "    return top"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Benchmark IndexedSet lookup\n",
    "\n",
    "n_list = [1, 5, 10, 20, 50, 100]\n",
    "lookup_times = []\n",
    "n_trials = 10\n",
    "n_reps = 100\n",
    "\n",
    "for n in n_list:\n",
    "    print(n)\n",
    "    current_lookup_times = []\n",
    "    for trial in range(n_trials):\n",
    "        top = setup_sys(n)\n",
    "        site = top.sites[random.randint(0, top.n_sites-1)]\n",
    "        start = time.time()\n",
    "        for rep in range(n_reps):\n",
    "            top.sites.index(site)\n",
    "        end_lookup = time.time()\n",
    "        current_lookup_times.append(end_lookup-start)\n",
    "    lookup_times.append(np.mean(current_lookup_times))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(n_list, lookup_times)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}