{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Prepare a GAFF ligand with `openmmforcefields` and Interchange\n",
    "\n",
    "[`openmmforcefields`](https://github.com/openmm/openmmforcefields) is a Python package that provides OpenMM implementations of some small molecule and biopolymer force fields via template generators. It provides classes that port SMIRNOFF, GAFF, and Espaloma force fields into a format that be used alongside the many force fields that ship with OpenMM by default."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Comparing to GAFF parametrization\n",
    "\n",
    "`openmmforcefields` provides a (partial) implementation of GAFF/GAFF2, using AM1-BCC in place of RESP to generate partial charges. The API for using `GAFFTemplateGenerator` is similar to `SMIRNOFFTemplateGenerator`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmm.app\n",
    "from openff.toolkit import Molecule, Quantity, unit\n",
    "from openmmforcefields.generators import GAFFTemplateGenerator\n",
    "from rich.pretty import pprint\n",
    "\n",
    "from openff.interchange import Interchange\n",
    "from openff.interchange.drivers.openmm import (\n",
    "    _get_openmm_energies,\n",
    "    _process,\n",
    "    get_openmm_energies,\n",
    ")\n",
    "\n",
    "molecule = Molecule.from_smiles(\"O=S(=O)(N)c1c(Cl)cc2c(c1)S(=O)(=O)NCN2\")\n",
    "molecule.generate_conformers(n_conformers=1)\n",
    "\n",
    "topology = molecule.to_topology()\n",
    "topology.box_vectors = Quantity([4, 4, 4], unit.nanometer)\n",
    "\n",
    "gaff = GAFFTemplateGenerator(molecules=molecule)\n",
    "\n",
    "# If using this alongside another force field, we'd load that in here. But\n",
    "# in this case, we're only using GAFF for this one molecule, so a \"blank\"\n",
    "# force field is a sufficient starting point\n",
    "forcefield_gaff = openmm.app.ForceField()\n",
    "forcefield_gaff.registerTemplateGenerator(gaff.generator)\n",
    "\n",
    "system_gaff = forcefield_gaff.createSystem(\n",
    "    topology=topology.to_openmm(),\n",
    "    nonbondedMethod=openmm.app.PME,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%env INTERCHANGE_EXPERIMENTAL=1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "From here, one could use `from_openmm` to import this `openmm.System` into an `Interchange` object, which then could be combined with other `Interchange` objects generated by different paths."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "imported = Interchange.from_openmm(\n",
    "    topology=topology.to_openmm(),\n",
    "    system=system_gaff,\n",
    "    positions=molecule.conformers[0].to_openmm(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "This `Interchange` contains an internal representation of approximately everything contained in the `openmm.System`. Topological information (atoms, their elements, bonds, the graph they compose, etc.), forces (including non-bonded settings), and box vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "imported.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "We can use Interchange's machinery to evaluate each of these objects - the `openmm.System` created by `openmmforcefields` and the `Interchange` object created from it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_openmmforcefields = _process(\n",
    "    _get_openmm_energies(\n",
    "        system=system_gaff,\n",
    "        box_vectors=topology.to_openmm().getPeriodicBoxVectors(),\n",
    "        positions=molecule.conformers[0].to_openmm(),\n",
    "        round_positions=None,\n",
    "        platform=\"Reference\",\n",
    "    ),\n",
    "    system=system_gaff,\n",
    "    combine_nonbonded_forces=True,\n",
    "    detailed=False,\n",
    ")\n",
    "\n",
    "pprint(energy_openmmforcefields)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_imported = get_openmm_energies(imported, detailed=False)\n",
    "\n",
    "pprint(energy_imported)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "Except for some small differences in non-bonded energies due to PME errors, the energies evaluated from each of these objects should match quite closely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_imported.compare(\n",
    "    energy_openmmforcefields,\n",
    "    {\"Nonbonded\": abs(energy_openmmforcefields[\"Nonbonded\"] * 1e-5)},\n",
    ")\n",
    "\n",
    "energy_imported.diff(energy_openmmforcefields)"
   ]
  }
 ],
 "metadata": {
  "category": "parametrization_evaluation",
  "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "86c9b142c8dc60dd36d17e2a57efabbd2ed015b9d3db80dd77f3e0894d5aea85"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}