{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openff.toolkit import ForceField\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": [
    "import mdtraj\n",
    "import nglview\n",
    "import openmm\n",
    "import openmm.unit\n",
    "from openff.toolkit import Molecule\n",
    "\n",
    "\n",
    "def viz(\n",
    "    smiles: str, force_field: ForceField = sage_with_example_virtual_sites\n",
    ") -> nglview.NGLWidget:\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, force_field: ForceField = sage_with_example_virtual_sites\n",
    ") -> nglview.NGLWidget:\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",
    "    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(\n",
    "        getPositions=True\n",
    "    ).getPositions()\n",
    "\n",
    "    return nglview.show_mdtraj(\n",
    "        mdtraj.load(\n",
    "            \"trajectory.dcd\",\n",
    "            top=mdtraj.Topology.from_openmm(\n",
    "                interchange.to_openmm_topology(),\n",
    "            ),\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"CC=O\")"
   ]
  },
  {
   "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\"})[\n",
    "    0\n",
    "].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "viz(\"CCCCCl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"CCCCCl\")"
   ]
  },
  {
   "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(\n",
    "    {\"name\": \"planar_carbonyl\"}\n",
    ")[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\")"
   ]
  },
  {
   "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\"})[\n",
    "    0\n",
    "].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"O\")"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sage_with_example_virtual_sites[\"VirtualSites\"].get_parameter(\n",
    "    {\"name\": \"trivalent_nitrogen\"}\n",
    ")[0].to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run(\"N\")"
   ]
  },
  {
   "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!"
   ]
  },
  {
   "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": [
    "# This crashes\n",
    "# run(\"C1=CC(=CC2=C1O[C@H](CN[S](=O)(=O)N)CO2)Cl\")"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}