{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a variety of different virtual sites\n",
    "\n",
    "In this example, several different types of SMIRNOFF virtual sites are used to model different chemistries, including\n",
    "* 4-site water\n",
    "* 5-site water\n",
    "* sigma holes\n",
    "* lone pairs of oxygens planar carbonyls\n",
    "* lone pairs of nitrogens in ammonia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mdtraj\n",
    "import nglview\n",
    "import openmm\n",
    "import openmm.unit\n",
    "from openff.toolkit import ForceField, Molecule\n",
    "\n",
    "sage_with_example_virtual_sites = ForceField(\n",
    "    \"openff-2.1.0.offxml\",\n",
    "    \"example-vsites-parameters-forcefield.offxml\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def viz(\n",
    "    smiles: str,\n",
    "    force_field: ForceField = sage_with_example_virtual_sites,\n",
    ") -> nglview.NGLWidget:\n",
    "    \"\"\"Visualize a molecule with virtual sites.\"\"\"\n",
    "    molecule = Molecule.from_smiles(smiles, allow_undefined_stereo=True)\n",
    "    molecule.generate_conformers(n_conformers=1)\n",
    "\n",
    "    interchange = force_field.create_interchange(molecule.to_topology())\n",
    "\n",
    "    return interchange._visualize_nglview(include_virtual_sites=True)\n",
    "\n",
    "\n",
    "def run(\n",
    "    smiles: str,\n",
    "    name: str,\n",
    "    force_field: ForceField = sage_with_example_virtual_sites,\n",
    ") -> nglview.NGLWidget:\n",
    "    \"\"\"Run and visualize a short simulation of a molecule with virtual sites.\"\"\"\n",
    "    molecule = Molecule.from_smiles(smiles, allow_undefined_stereo=True)\n",
    "    molecule.generate_conformers(n_conformers=1)\n",
    "\n",
    "    interchange = force_field.create_interchange(molecule.to_topology())\n",
    "\n",
    "    # interchange.minimize()\n",
    "\n",
    "    integrator = openmm.LangevinIntegrator(\n",
    "        300 * openmm.unit.kelvin,\n",
    "        1 / openmm.unit.picosecond,\n",
    "        1 * openmm.unit.femtoseconds,\n",
    "    )\n",
    "\n",
    "    simulation = interchange.to_openmm_simulation(integrator=integrator)\n",
    "\n",
    "    dcd_reporter = openmm.app.DCDReporter(\"trajectory.dcd\", 10)\n",
    "    simulation.reporters.append(dcd_reporter)\n",
    "\n",
    "    simulation.context.computeVirtualSites()\n",
    "    simulation.minimizeEnergy()\n",
    "    simulation.context.setVelocitiesToTemperature(openmm.unit.kelvin * 300)\n",
    "    simulation.step(1000)\n",
    "\n",
    "    interchange.positions = simulation.context.getState(getPositions=True).getPositions()\n",
    "\n",
    "    open(f\"{name}.xml\", \"w\").write(openmm.XmlSerializer.serialize(interchange.to_openmm()))\n",
    "\n",
    "    return nglview.show_mdtraj(\n",
    "        mdtraj.load(\n",
    "            \"trajectory.dcd\",\n",
    "            top=mdtraj.Topology.from_openmm(\n",
    "                interchange.to_openmm_topology(collate=True),\n",
    "            ),\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first parameter is of type `BondCharge`, which places a virtual site along the axis of a bond. This was originally intended for use with halogens to better model sigma holes. In this case, a  virtual site is added $1.45 \\AA$ \"outside\" a carbon-chlorine bond."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sage_with_example_virtual_sites[\"VirtualSites\"].get_parameter({\"name\": \"sigma_hole\"})[0].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz(\"CCCCCl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"CCCCCl\", name=\"sigma_hole\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next is a parameter using `MonovalentLonePair`. In this case, a virtual site is added near the oxygen in a carbonyl. It is placed at an angle (`inPlaneAngle`) formed by the oxygen and carbon atoms of the carbonyl and in the plane defined by those atoms and the alpha carbon. If the parameter `outOfPlaneAngle` were non-zero, it would be placed at an angle out of that plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sage_with_example_virtual_sites[\"VirtualSites\"].get_parameter({\"name\": \"planar_carbonyl\"})[0].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz(\"COC1=C(C=CC(=C1)C=O)O\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"COC1=C(C=CC(=C1)C=O)O\", name=\"planar_carbonyl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next parameter completes a so-called four-site water model like TIP4P or OPC. A single virtual site is placed near the oxygen in the plane of the three atoms. This is implemented with the type `DivalentLonePair`, though it is possible to also implement it with `MonovalentLonePair`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sage_with_example_virtual_sites[\"VirtualSites\"].get_parameter({\"name\": \"4_site_water\"})[0].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"O\", name=\"4_site_water\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz(\"O\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are also 5-site water models, such as TIP5P. These are more complex as they require placing virtual sites _outside_ the plane formed by the atoms in the molecule. To avoid squeezing two water models into a single force field, this portion uses a force field from a different file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tip5p = ForceField(\"tip5p.offxml\")\n",
    "\n",
    "tip5p[\"VirtualSites\"].get_parameter({\"name\": \"5_site_water\"})[0].to_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next parameter uses `TrivalentLonePair` to model the lone pair of a trivalent nitrogen. It is written to match only ammonia - the other capturing atoms are all hydrogens - but a SMIRKS pattern could be written to match trivalent nitrogens in other chemistries. A virtual site is placed $5 \\AA$ away from the nitrogen atom, opposite a plane defined by the three hydrogen atoms. (You may need to rotate the molecule, using your cursor, to see the virtual site)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz(\"O\", force_field=tip5p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"O\", force_field=tip5p, name=\"5_site_water\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sage_with_example_virtual_sites[\"VirtualSites\"].get_parameter({\"name\": \"trivalent_nitrogen\"})[0].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"N\", name=\"trivalent_nitrogen\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each of the molecules shown so far was selected to only have a single virtual site added, but larger molecules may have multiple copies of different types of virtual sites added. For example, this force field will give hexachlorobenzene 6 virtual sites, a dialdehyde two, or both types of virtual sites to something with a mixture of these chemistries. (Note that these parameters are made-up for instructive purposes, not the result of a force field fit, and likely to produce crashes or bad results for realistic ligands!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz(\"c1(Cl)c(Cl)c(Cl)c(Cl)c(Cl)c1Cl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz(\"O=CCCCCC=O\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz(\"ClCCCC(Cl)(CC=O)CC=O\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"O=C1c2ccc(Cl)cc2C(=O)N1[C@H]3CCC(=O)NC3=O\", name=\"ligand\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}