diff --git a/notebooks/tutorials/Use OpenMM to Verify Your MM Functional Form Implementation.ipynb b/notebooks/tutorials/Use OpenMM to Verify Your MM Functional Form Implementation.ipynb
new file mode 100644
index 0000000..079bb4c
--- /dev/null
+++ b/notebooks/tutorials/Use OpenMM to Verify Your MM Functional Form Implementation.ipynb
@@ -0,0 +1,570 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "62ee8371",
+ "metadata": {},
+ "source": [
+ "## Use OpenMM to Verify Your MM Functional Form Implementation"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "024e7580",
+ "metadata": {},
+ "source": [
+ "Recently, there have been efforts where MM functional forms are implemented within machine learning packages using tensor-accelerating framework to assist the differentiable parametrization of force fields, such as [Espaloma](https://github.com/choderalab/espaloma), [JAX-MD](https://github.com/jax-md/jax-md), and [TimeMachine](https://github.com/proteneer/timemachine/tree/master).\n",
+ "When you implement your own MM functional forms, you might want to make sure that they match what are used in, say, OpenMM, exactly.\n",
+ "Since the nonbonded part of the energy requires a lot more love to implement exactly, in this tutorial we first show you how to check your bonded energy implementation."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "7808e632",
+ "metadata": {},
+ "source": [
+ "For simplicity, we only work with small molecules for now, which means that we would need to import `openff.toolkit` for chemoinformatics modeling."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 238,
+ "id": "6347e55b",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import numpy as np\n",
+ "import openmm\n",
+ "from openff.toolkit.topology import Molecule"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "23dbd437",
+ "metadata": {},
+ "source": [
+ "### Get a Toy Water Molecule System"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "d9e3fa86",
+ "metadata": {},
+ "source": [
+ "We use a water molecule to focus first on bonds and angles."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 150,
+ "id": "079df807",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "molecule = Molecule.from_smiles(\"[H]O[H]\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 151,
+ "id": "960825b2",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "image/svg+xml": [
+ ""
+ ],
+ "text/plain": [
+ ""
+ ]
+ },
+ "execution_count": 151,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "molecule.visualize()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f915510d",
+ "metadata": {},
+ "source": [
+ "It's also worth noting the index-element correspondance in this molecule, where the 0th and 2nd atoms are hydrogens whereas the 1st is oxygen."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 152,
+ "id": "38f4b37d",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "[Atom(name=, atomic number=1),\n",
+ " Atom(name=, atomic number=8),\n",
+ " Atom(name=, atomic number=1)]"
+ ]
+ },
+ "execution_count": 152,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "molecule.atoms"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "a5f83c48",
+ "metadata": {},
+ "source": [
+ "To create OpenMM systems, we use an OpenFF force field."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 153,
+ "id": "6c5c757b",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from openff.toolkit.typing.engines.smirnoff import ForceField\n",
+ "forcefield = ForceField(\"openff_unconstrained-1.2.0.offxml\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 154,
+ "id": "e168c674",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/Users/wangy1/miniconda3/envs/malt/lib/python3.10/site-packages/openff/interchange/components/interchange.py:339: UserWarning: Automatically up-converting BondHandler from version 0.3 to 0.4. Consider manually upgrading this BondHandler (or section in an OFFXML file) to 0.4 or newer. For more details, see https://openforcefield.github.io/standards/standards/smirnoff/#bonds.\n",
+ " warnings.warn(\n"
+ ]
+ }
+ ],
+ "source": [
+ "molecule.assign_partial_charges(\"mmff94\")\n",
+ "system = forcefield.create_openmm_system(molecule.to_topology(), charge_from_molecules=[molecule])\n",
+ "topology = molecule.to_topology().to_openmm()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "d0deedd0",
+ "metadata": {},
+ "source": [
+ "### Get Energy Components with OpenMM"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "951dc8ab",
+ "metadata": {},
+ "source": [
+ "Querying the energy contributions from OpenMM is somewhat non-intuitive.\n",
+ "We would need to set different force groups first."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 155,
+ "id": "3fd0e951",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "for idx, force in enumerate(system.getForces()):\n",
+ " force.setForceGroup(idx)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "af83459f",
+ "metadata": {},
+ "source": [
+ "Next, we create a simulation so that we can inspect each energy component."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 156,
+ "id": "31d99d79",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "integrator = openmm.VerletIntegrator(0.0)\n",
+ "simulation = openmm.app.Simulation(topology,system, integrator)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "bab7fa2a",
+ "metadata": {},
+ "source": [
+ "Let's also set the initial positions to one of the conformations of the water molecule."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 157,
+ "id": "b320c078",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "molecule.generate_conformers() # generate molecule conformers\n",
+ "\n",
+ "# NOTE:\n",
+ "# OpenMM and OpenFF use two sets of unit systems,\n",
+ "# hence the awkward conversion\n",
+ "simulation.context.setPositions(\n",
+ " openmm.unit.nanometer * molecule._conformers[0].magnitude,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 162,
+ "id": "65898ea4",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Bond Energy 368275.4375 kJ/mol\n",
+ "Angle Energy 1.3650074005126953 kJ/mol\n"
+ ]
+ }
+ ],
+ "source": [
+ "for idx, force in enumerate(system.getForces()):\n",
+ " state = simulation.context.getState(getEnergy=True, groups={idx})\n",
+ " if \"HarmonicBondForce\" in force.getName():\n",
+ " bond_energy = state.getPotentialEnergy()\n",
+ " if \"HarmonicAngleForce\" in force.getName():\n",
+ " angle_energy = state.getPotentialEnergy()\n",
+ "\n",
+ "print(\"Bond Energy\", bond_energy)\n",
+ "print(\"Angle Energy\", angle_energy)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "a9555d75",
+ "metadata": {},
+ "source": [
+ "### Implement our own distance and angle functions\n",
+ "Here we show a very simple example.\n",
+ "This is usually done in an elaborate tensor-accelerating framework."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "6d5d0f5e",
+ "metadata": {},
+ "source": [
+ "The bond length can be computed as the root of the squared difference between two bonded atoms."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 166,
+ "id": "b466f554",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def get_distance_vector(x0, x1):\n",
+ " return x1 - x0"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 194,
+ "id": "94e76b8c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def get_norm(x01):\n",
+ " return np.linalg.norm(x01, ord=2, axis=-1)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 195,
+ "id": "abe40b29",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def get_distance(x0, x1):\n",
+ " return get_norm(get_distance_vector(x0, x1))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 208,
+ "id": "cc038688",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def get_angle(x0, x1, x2):\n",
+ " x10 = get_distance_vector(x1, x0)\n",
+ " x12 = get_distance_vector(x1, x2)\n",
+ " c012 = np.arctan2(\n",
+ " get_norm(np.cross(x10, x12)),\n",
+ " (x10 * x12).sum(-1)\n",
+ " )\n",
+ " return c012"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "638fa9f4",
+ "metadata": {},
+ "source": [
+ "### Implement our own MM forces"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f42fbbd0",
+ "metadata": {},
+ "source": [
+ "Let's implement harmonic bond and angle forces from scratch:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 209,
+ "id": "25c71478",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def harmonic(x, k, r0):\n",
+ " return 0.5 * k * (x - r0) ** 2"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "2f422194",
+ "metadata": {},
+ "source": [
+ "### Compute angle and bond energy"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "2137744d",
+ "metadata": {},
+ "source": [
+ "One of the nice things about OpenMM is that it works with units, whereas most tensor-accelerating frameworks don't.\n",
+ "Here our crappy implementation doesn't take care of units either.\n",
+ "So we work with numeric values only---everything has the default of OpenMM units: `nanometer` for distances, `radian` for angles, and `kJ/mol` for energies."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ad357717",
+ "metadata": {},
+ "source": [
+ "We first compute the bonds and angles using the functions defined above."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 210,
+ "id": "e23a808d",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "conformer = molecule._conformers[0].magnitude\n",
+ "x0, x1, x2 = conformer"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 227,
+ "id": "e6abf994",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "r01 = get_distance(x0, x1)\n",
+ "r12 = get_distance(x1, x2)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 214,
+ "id": "b3057adc",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "c012 = get_angle(x0, x1, x2)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "7c6d91c2",
+ "metadata": {},
+ "source": [
+ "Then we grab the parameters: bond equilibrium length `r0`, bond force constant `k_r`, equilibrium angle `theta_0`, angle force constant `k_theta`, from the OpenMM `System`.\n",
+ "Note that here we know that there is only one kind of bond and one kind of angle, saving us the need for bookkeeping.\n",
+ "In real life one would need to carefully manage bond and angle types."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 230,
+ "id": "44ac59b5",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "for force in system.getForces():\n",
+ " if \"HarmonicBond\" in force.getName():\n",
+ " _, __, r0, k_r = force.getBondParameters(0)\n",
+ " if \"HarmonicAngle\" in force.getName():\n",
+ " _, __, ___, theta0, k_theta = force.getAngleParameters(0)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "31117ca9",
+ "metadata": {},
+ "source": [
+ "We have to erase the units, unfortunately."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 231,
+ "id": "2769d77d",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "r0, k_r, theta0, k_theta = r0._value, k_r._value, theta0._value, k_theta._value"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 239,
+ "id": "aadb3640",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "368275.4577602855"
+ ]
+ },
+ "execution_count": 239,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "bond_energy_computed = harmonic(r01, k_r, r0) + harmonic(r12, k_r, r0)\n",
+ "bond_energy_computed"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 240,
+ "id": "c753c122",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "1.3650163207594144"
+ ]
+ },
+ "execution_count": 240,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "angle_energy_computed = harmonic(c012, k_theta, theta0)\n",
+ "angle_energy_computed"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1e4a33be",
+ "metadata": {},
+ "source": [
+ "### Check consistency"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 246,
+ "id": "e8080be0",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "assert np.allclose(bond_energy_computed, bond_energy._value)\n",
+ "assert np.allclose(angle_energy_computed, angle_energy._value)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "839ea7ab",
+ "metadata": {},
+ "source": [
+ "### Postlude\n",
+ "Here we have only checked the _bond_ and _angle_ energy consistency for _one_ molecule with _one_ type of bond and angle each.\n",
+ "But with the help of careful index matching and bookkeeping with `PyTree` containers, we are able to scale things up dramatically to check the consistency for other terms among considerable chemical and conformational spaces.\n",
+ "Check out unit tests [here](https://github.com/choderalab/espaloma/blob/master/espaloma/mm/tests/test_openmm_consistency.py) to see how this is done."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "08590304",
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "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.10.4"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}