uci-pharmsci/lectures/SMIRNOFF_simulations/mixture_simulations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Automated setup of mixtures\n",
    "\n",
    "We've been working on streamlining setup of simulations of arbitrary mixtures in AMBER/GROMACS/OpenMM and others for some of our own research. I thought I'd demo this really quick so you can get a feel for it and see if you're interested in contributing. It also allows quick setup and analysis of nontrivial liquid simulations, which can be a good opportunity to try out MDTraj and other analysis tools.\n",
    "\n",
    "*Before running the below*, you will need to have followed the [getting started instructions](https://github.com/MobleyLab/drug-computing/blob/master/uci-pharmsci/getting-started.md) for this course."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Warning: importing 'simtk.openmm' is deprecated.  Import 'openmm' instead.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finished build 0\n",
      "finished build 1\n",
      "finished build 2\n",
      "finished build 3\n",
      "finished build 4\n",
      "finished build 5\n",
      "Finished building\n"
     ]
    }
   ],
   "source": [
    "from openff.evaluator.protocols.coordinates import BuildCoordinatesPackmol\n",
    "from openff.evaluator.substances import Substance, Component, MoleFraction, ExactAmount\n",
    "from openeye import oeiupac\n",
    "from openeye import oechem\n",
    "import shutil, os\n",
    "\n",
    "#Number of solute/solvent molecules\n",
    "Nsolu = 3\n",
    "Nsolv = 100\n",
    "\n",
    "#solute names\n",
    "solutes = ['phenol', 'toluene', 'benzene', 'methane', 'ethanol', 'naphthalene']\n",
    "#Solvent names\n",
    "solvents = ['cyclohexane', 'cyclohexane', 'cyclohexane', 'octanol', 'octanol', 'octanol']\n",
    "\n",
    "# Generate SMILES for solutes and solvents and store for use in building mixtures\n",
    "solute_smiles = []\n",
    "solvent_smiles = []\n",
    "for name_pair in zip(solutes, solvents):\n",
    "    #print(name_pair)\n",
    "    solu_mol = oechem.OEMol()\n",
    "    oeiupac.OEParseIUPACName(solu_mol, name_pair[0])\n",
    "    solute_smiles.append(oechem.OECreateIsoSmiString(solu_mol))\n",
    "    solv_mol = oechem.OEMol()\n",
    "    oeiupac.OEParseIUPACName(solv_mol, name_pair[1])\n",
    "    solvent_smiles.append(oechem.OECreateIsoSmiString(solv_mol))\n",
    "    \n",
    "\n",
    "# Storage for the mixtures we've built\n",
    "mixtures = []\n",
    "outdir = 'coordinate_files'\n",
    "if not os.path.isdir(outdir): os.mkdir(outdir)\n",
    "    \n",
    "# Loop and build mixtures\n",
    "for idx in range( len( solutes) ):\n",
    "    # Define new mixture\n",
    "    mixture_build = BuildCoordinatesPackmol(\"\")\n",
    "    substance = Substance()\n",
    "    # Add solute and solvent\n",
    "    substance.add_component(Component(solute_smiles[idx], role=Component.Role.Solute), ExactAmount(Nsolu))\n",
    "    substance.add_component(Component(solvent_smiles[idx], role=Component.Role.Solvent), ExactAmount(Nsolv))\n",
    "    #substance.add_component(Component(\"Cc1ccccc1\", role=Component.Role.Solvent), MoleFraction(0.1))\n",
    "    #substance.add_component(Component(\"C1CCCCC1\", role=Component.Role.Solvent), MoleFraction(0.9))\n",
    "    # Note you can optionally specify mole fraction instead, or a mix of numbers/mole fractions, etc.\n",
    "    \n",
    "    #substance.max_molecules = 150\n",
    "    mixture_build.substance = substance\n",
    "\n",
    "    #build\n",
    "    mixture_build.execute()\n",
    "    \n",
    "    # Do file bookkeeping so the output files don't overwrite one another\n",
    "    outfile = os.path.join(outdir, f'coordinate_file{idx}.pdb')\n",
    "    shutil.copy( mixture_build.coordinate_file_path, outfile )\n",
    "    mixture_build.coordinate_file_path = outfile\n",
    "    \n",
    "    #Store details\n",
    "    mixtures.append(mixture_build)\n",
    "\n",
    "    print(f\"finished build {idx}\")\n",
    "    \n",
    "    \n",
    "print(\"Finished building\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's try and see if we can do a quick visualization of one of the systems via mdtraj just to make sure it looks right"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7ccdbe76c21f44079fe6e5bd148429e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e108f00aa524401bd7db74645456fb4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Import MDTraj\n",
    "import mdtraj as md\n",
    "#Load \"trajectory\" (structures)\n",
    "#You can load from either format\n",
    "#traj = md.load( 'data/amber/phenol_cyclohexane_3_100.inpcrd', top = 'data/amber/phenol_cyclohexane_3_100.prmtop' )\n",
    "traj = md.load( os.path.join(outdir, 'coordinate_file0.pdb'))#'data/gromacs/phenol_cyclohexane_3_100.gro')\n",
    "\n",
    "#Input viewer\n",
    "import nglview\n",
    "\n",
    "#Set up view of structure\n",
    "view = nglview.show_mdtraj(traj)\n",
    "\n",
    "#Try some of the following to modify representations\n",
    "view.clear_representations()\n",
    "view.add_licorice('all')\n",
    "# Select the first three \"residues\" of chain A for special display\n",
    "# (Evaluator seems to make the first substance to be chain A, and then individual molecules are residues within that chain)\n",
    "view.add_licorice('1:A or 2:A or 3:A', color = \"blue\") #NGLview has a whole selection lanuage\n",
    "view.add_surface('1:A or 2:A or 3:A', opacity=0.3)\n",
    "\n",
    "#Show the view. Note that this needs to be the last command used to manipulate the view, i.e. if you modify the\n",
    "#representation after this, your view will be empty.\n",
    "view\n",
    "\n",
    "#VIEWER USAGE:\n",
    "# - Use your typical zoom command/gesture (i.e. pinch) to zoom in and out \n",
    "# - Click and drag to reorient\n",
    "# - Click on specific atoms/residues to find out details of what they are (and how they could be selected)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's use a SMIRNOFF forcefield to parameterize the system, minimize, and run dynamics\n",
    "\n",
    "(This requires `openforcefield`, which you will have conda-installed if you've followed the getting started info.)\n",
    "\n",
    "First we handle imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the SMIRNOFF forcefield engine and some useful tools\n",
    "from openff.toolkit.topology import Molecule, Topology\n",
    "from openff.toolkit.typing.engines.smirnoff import ForceField\n",
    "import openeye.oechem as oechem #Here we'll use OpenEye tookits, but RDKIt use is also possible\n",
    "\n",
    "# We use PDBFile to get OpenMM topologies from PDB files\n",
    "try:\n",
    "    from openmm.app import PDBFile\n",
    "except:\n",
    "    from simtk.openmm.app import PDBFile\n",
    "\n",
    "# MDTraj for working with trajectories; time for timing\n",
    "import time\n",
    "import mdtraj"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now we handle assignment of force field parameters and generation of an OpenMM System"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify names of molecules that are components of the system\n",
    "molnames = ['phenol', 'cyclohexane']\n",
    "\n",
    "# Get molecules for components - required as input for OpenFF\n",
    "oemols = []\n",
    "for name in molnames:\n",
    "    mol = oechem.OEGraphMol()\n",
    "    oeiupac.OEParseIUPACName(mol, name)\n",
    "    oemols.append(mol)\n",
    "\n",
    "# Build set of OpenFF mols from OE molecules\n",
    "OFFmols = []\n",
    "for mol in oemols:\n",
    "    OFFmols.append( Molecule.from_openeye(mol))\n",
    "    \n",
    "# Load OpenFF 2.0 force field\n",
    "ff = ForceField('openff-2.0.0.offxml') \n",
    "\n",
    "# Get OpenMM topology for mixture of phenol and cyclohexane from where they were created\n",
    "# it on disk (the first mixture we built)\n",
    "pdbfile = PDBFile(mixtures[0].coordinate_file_path)\n",
    "\n",
    "# Create OpenFF Topology\n",
    "off_topology = Topology.from_openmm(openmm_topology = pdbfile.topology, unique_molecules = OFFmols)\n",
    "\n",
    "# Assign SMIRNOFF parameters and create system\n",
    "system = ff.create_openmm_system( off_topology)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finally we energy minimize and run dynamics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting simulation\n",
      "Elapsed time 14.19 seconds\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "from datetime import datetime\n",
    "try:\n",
    "    import openmm\n",
    "    from openmm import app, unit\n",
    "    from openmm.app import PDBFile\n",
    "except ImportError:\n",
    "    from simtk import openmm, app, unit\n",
    "    from simtk.openmm.app import PDBFile\n",
    "\n",
    "# Set how many steps we'll run and other run parameters\n",
    "num_steps=10000\n",
    "trj_freq = 100 #Trajectory output frequency\n",
    "data_freq = 100 #Energy/data output frequency\n",
    "temperature = 300*unit.kelvin #Temperature\n",
    "time_step = 2.*unit.femtoseconds\n",
    "friction = 1./unit.picosecond #Langevin friction constant\n",
    "\n",
    "# Bookkeeping -- if you run this more than once and perhaps encountered an exception, we need to make sure the reporter is closed\n",
    "try: \n",
    "    reporter.close()\n",
    "except: pass\n",
    "    \n",
    "# Set up integrator, platform for running simulation    \n",
    "integrator = openmm.LangevinIntegrator(temperature, friction, time_step)\n",
    "platform = openmm.Platform.getPlatformByName('Reference')\n",
    "simulation = app.Simulation(pdbfile.topology, system, integrator)\n",
    "# Set positions, velocities\n",
    "simulation.context.setPositions(pdbfile.positions)\n",
    "simulation.context.setVelocitiesToTemperature(temperature)\n",
    "\n",
    "# Before doing dynamics, energy minimize (initial geometry will be strained)\n",
    "simulation.minimizeEnergy()\n",
    "\n",
    "# Set up reporter for output\n",
    "reporter = mdtraj.reporters.HDF5Reporter('mixture.h5', trj_freq)\n",
    "simulation.reporters=[]\n",
    "simulation.reporters.append(reporter)\n",
    "simulation.reporters.append(app.StateDataReporter('data.csv', data_freq, step=True, potentialEnergy=True, temperature=True, density=True))\n",
    "\n",
    "# Run the dynamics\n",
    "print(\"Starting simulation\")\n",
    "start = datetime.now()\n",
    "simulation.step(num_steps)\n",
    "end = datetime.now()\n",
    "\n",
    "print(\"Elapsed time %.2f seconds\" % (end-start).total_seconds())\n",
    "#netcdf_reporter.close()\n",
    "reporter.close()\n",
    "print(\"Done!\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's make a movie of our simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4dd45991187948169ce4df99f26acd42",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget(max_frame=99)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import nglview\n",
    "traj=mdtraj.load('mixture.h5')\n",
    "view = nglview.show_mdtraj(traj)\n",
    "\n",
    "#Try some of the following to modify representations\n",
    "view.clear_representations()\n",
    "view.add_licorice('all')\n",
    "view.add_licorice('1:A or 2:A or 3:A', color = \"blue\") \n",
    "view.add_surface('1:A', opacity=0.3)\n",
    "view.add_surface('2:A or 3:A', color = 'red', opacity=0.3)\n",
    "\n",
    "view #Note that if you view a movie and keep it playing, your notebook will run a hair slow..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other possibly interesting things to try:\n",
    "* Find the average distance from phenol to phenol (would need longer simulation)\n",
    "* Calculate the density or volume of the system (would need to add barostat)\n",
    "* etc.\n",
    "\n",
    "Note: I would need to check that I've configured PBC properly above, and if I have, update the wrapping so the trajectory visualization correctly handles it.\n",
    "\n",
    "(Drawing on MDTraj - see docs online)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use this box to try additional things"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "drugcomp",
   "language": "python",
   "name": "drugcomp"
  },
  "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.12"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": false
  },
  "widgets": {
   "state": {
    "b486def970134f80a2bd2b8fa0f0b158": {
     "views": [
      {
       "cell_index": 3
      }
     ]
    },
    "fccd5a48296e490c88d379241a9e3309": {
     "views": [
      {
       "cell_index": 13
      }
     ]
    }
   },
   "version": "1.2.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}