{ "cells": [ { "cell_type": "markdown", "id": "cc3329e2", "metadata": {}, "source": [ "# OPLS: verifying the `xnn` implementation against OpenMM and the 1996 paper\n", "\n", "The `xnn` OPLS model is a clean-room implementation of the published\n", "functional form (Jorgensen, Maxwell & Tirado-Rives, *J. Am. Chem. Soc.*\n", "**118**, 11225, 1996, eqs 1-4). This notebook verifies it two independent\n", "ways:\n", "\n", "1. **numerical parity with OpenMM** — the same parameter library and\n", " topology (atom types assigned from the parameter file's SMARTS templates,\n", " bonds perceived from the coordinates) are assembled into an OpenMM `System` (an entirely independent\n", " energy/force engine) and evaluated on randomized conformations of\n", " butane, ethanol and ethylene. Energies and forces must agree to the\n", " precision of the unit constants (~1e-7 kJ/mol);\n", "2. **the paper's own numbers** — relaxed torsional scans reproduce the\n", " OPLS-AA column of Table 1, and the Fourier/Ryckaert-Bellemans torsion\n", " conversion reproduces the dual-form rows of Table 2 of the L-OPLS paper\n", " (Siu, Pluhackova & Böckmann, *JCTC* **8**, 1459, 2012).\n", "\n", "Unit conventions are part of what is verified: `xnn` stores OPLS\n", "parameters in kcal/mol with the *thermochemical* calorie (4.184 kJ exactly)\n", "and uses the CODATA Coulomb constant, matching the GROMACS/OpenMM\n", "ecosystem in which OPLS parameters are distributed." ] }, { "cell_type": "markdown", "id": "b1426883", "metadata": {}, "source": [ "## 0. Setup" ] }, { "cell_type": "code", "execution_count": 1, "id": "02450ed4", "metadata": { "execution": { "iopub.execute_input": "2026-09-15T19:20:26.797957Z", "iopub.status.busy": "2026-09-15T19:20:26.797843Z", "iopub.status.idle": "2026-09-15T19:20:28.591469Z", "shell.execute_reply": "2026-09-15T19:20:28.590599Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "xnn: 0.1.0 | torch: 2.5.1+cu121 | openmm: 8.6.0.dev-c6173db\n" ] } ], "source": [ "import warnings\n", "warnings.filterwarnings(\"ignore\")\n", "import math\n", "import numpy as np\n", "import torch\n", "\n", "torch.set_default_dtype(torch.float64)\n", "torch.manual_seed(0)\n", "rng = np.random.default_rng(1)\n", "\n", "import openmm as mm\n", "import openmm.unit as u\n", "\n", "import xnn\n", "from xnn.common.data import structure_to_graph\n", "from xnn.common.models import ForceStressOutput\n", "from xnn.ffnn.models import OPLS, builtin_library, rb_to_fourier, fourier_to_rb\n", "from xnn.ffnn.models.oplslib import (resolve_angle_type, resolve_bond_type,\n", " resolve_dihedral_type,\n", " resolve_improper_type)\n", "\n", "EV_TO_KJ = 96.48533212331\n", "print(\"xnn:\", xnn.__version__, \"| torch:\", torch.__version__,\n", " \"| openmm:\", mm.version.version)" ] }, { "cell_type": "markdown", "id": "085f2bfc", "metadata": {}, "source": [ "## 1. An OpenMM twin of the same library + topology\n", "\n", "The translation is direct: harmonic bonds/angles carry a factor 2 (OPLS\n", "writes `k (x-x0)^2`, OpenMM `k/2 (x-x0)^2`), the Fourier torsion becomes\n", "four phased periodic torsions, the `V2` improper a phase-180 periodic\n", "torsion, and 1,4 pairs become explicit exceptions. The one subtlety is the\n", "Lennard-Jones combination rule: OpenMM's `NonbondedForce` hard-codes\n", "Lorentz-Berthelot mixing, so OPLS geometric mixing needs a\n", "`CustomNonbondedForce` (charges stay on the `NonbondedForce`).\n" ] }, { "cell_type": "code", "execution_count": 2, "id": "dcc57732", "metadata": { "execution": { "iopub.execute_input": "2026-09-15T19:20:28.593315Z", "iopub.status.busy": "2026-09-15T19:20:28.593116Z", "iopub.status.idle": "2026-09-15T19:20:28.601580Z", "shell.execute_reply": "2026-09-15T19:20:28.601000Z" } }, "outputs": [], "source": [ "def openmm_twin(lib, top, impropers=(), fudge_lj=0.5, fudge_qq=0.5):\n", " \"\"\"Assemble an OpenMM Context evaluating the same OPLS model.\n", "\n", " ``impropers`` are the (i, j, center, l) quadruples the xnn model placed\n", " (``model.impropers``); their parameters resolve by class pattern.\n", " \"\"\"\n", " at = lib.atom_types\n", " cls = [at[n][\"cls\"] for n in top.types]\n", " cls_oop = [lib.cls(n, \"oop\") for n in top.types]\n", " system = mm.System()\n", " for name in top.types:\n", " system.addParticle(at[name][\"mass\"])\n", "\n", " nb = mm.NonbondedForce() # Coulomb + 1,4 exceptions\n", " nb.setNonbondedMethod(mm.NonbondedForce.NoCutoff)\n", " lj = mm.CustomNonbondedForce( # OPLS geometric mixing\n", " \"4*eps*((sig/r)^12-(sig/r)^6); sig=sqrt(sig1*sig2);\"\n", " \" eps=sqrt(eps1*eps2)\")\n", " lj.addPerParticleParameter(\"sig\")\n", " lj.addPerParticleParameter(\"eps\")\n", " for name in top.types:\n", " nb.addParticle(at[name][\"charge\"], 0.1, 0.0)\n", " lj.addParticle([max(at[name][\"sigma\"], 1e-6) * 0.1,\n", " at[name][\"epsilon\"] * 4.184])\n", " for i, j in top.exclusions + top.pairs14:\n", " nb.addException(i, j, 0.0, 0.1, 0.0)\n", " lj.addExclusion(i, j)\n", " for i, j in top.pairs14:\n", " ti, tj = at[top.types[i]], at[top.types[j]]\n", " nb.addException(i, j, fudge_qq * ti[\"charge\"] * tj[\"charge\"],\n", " max(math.sqrt(ti[\"sigma\"] * tj[\"sigma\"]), 1e-6) * 0.1,\n", " fudge_lj * math.sqrt(ti[\"epsilon\"] * tj[\"epsilon\"])\n", " * 4.184, replace=True)\n", " nb.setForceGroup(0); lj.setForceGroup(4)\n", " system.addForce(nb); system.addForce(lj)\n", "\n", " bond = mm.HarmonicBondForce(); bond.setForceGroup(1)\n", " for i, j in top.bonds:\n", " bt = lib.bond_types[resolve_bond_type(lib.bond_types, cls[i], cls[j])]\n", " bond.addBond(i, j, bt[\"r0\"] * 0.1, 2 * bt[\"k\"] * 4.184 * 100)\n", " system.addForce(bond)\n", "\n", " ang = mm.HarmonicAngleForce(); ang.setForceGroup(2)\n", " for i, j, k in top.angles:\n", " a = lib.angle_types[resolve_angle_type(lib.angle_types,\n", " cls[i], cls[j], cls[k])]\n", " ang.addAngle(i, j, k, math.radians(a[\"theta0\"]), 2 * a[\"k\"] * 4.184)\n", " system.addForce(ang)\n", "\n", " tors = mm.PeriodicTorsionForce(); tors.setForceGroup(3)\n", " const = 0.0\n", " for i, j, k, l in top.dihedrals:\n", " key = resolve_dihedral_type(lib.dihedral_types,\n", " cls[i], cls[j], cls[k], cls[l])\n", " v = lib.dihedral_types[key][\"v\"]\n", " const += v[0] * 4.184 # the constant V0 offset\n", " for n, (vn, ph) in enumerate(zip(v[1:], (0.0, math.pi, 0.0, math.pi)),\n", " start=1):\n", " if vn:\n", " tors.addTorsion(i, j, k, l, n, ph, 0.5 * vn * 4.184)\n", " for i, j, k, l in impropers:\n", " key = resolve_improper_type(lib.improper_types, cls_oop[i], cls_oop[j],\n", " cls_oop[k], cls_oop[l])\n", " v2 = lib.improper_types[key][\"v2\"]\n", " tors.addTorsion(i, j, k, l, 2, math.pi, 0.5 * v2 * 4.184)\n", " system.addForce(tors)\n", "\n", " ctx = mm.Context(system, mm.VerletIntegrator(1e-3),\n", " mm.Platform.getPlatformByName(\"Reference\"))\n", " return ctx, const\n", "\n", "\n", "def openmm_eval(ctx, const, pos, groups=None):\n", " \"\"\"Energy (kJ/mol) and forces (kJ/mol/nm) from the OpenMM twin.\"\"\"\n", " ctx.setPositions(np.asarray(pos) * 0.1)\n", " kw = {\"groups\": groups} if groups is not None else {}\n", " st = ctx.getState(getEnergy=True, getForces=True, **kw)\n", " e = st.getPotentialEnergy().value_in_unit(u.kilojoule_per_mole)\n", " f = st.getForces(asNumpy=True).value_in_unit(\n", " u.kilojoule_per_mole / u.nanometer)\n", " return e + (const if groups is None or 3 in groups else 0.0), f" ] }, { "cell_type": "markdown", "id": "84bf1870", "metadata": {}, "source": [ "## 2. Energy and force parity on randomized conformations\n", "\n", "Three gas-phase molecules cover every term of the force field: butane\n", "(bonds, angles, Fourier torsions, 1,4 scaling), ethanol (heteroatoms,\n", "alcohol torsions, a zero-LJ hydroxyl hydrogen) and ethylene (`V2`\n", "impropers, wildcard `X-CM-CM-X` torsion). Ten heavily jittered\n", "conformations each.\n" ] }, { "cell_type": "code", "execution_count": 3, "id": "accd037b", "metadata": { "execution": { "iopub.execute_input": "2026-09-15T19:20:28.603126Z", "iopub.status.busy": "2026-09-15T19:20:28.603005Z", "iopub.status.idle": "2026-09-15T19:20:29.735780Z", "shell.execute_reply": "2026-09-15T19:20:29.735118Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "molecule max |dE| (kJ/mol) max |dF| (kJ/mol/nm)\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "butane 9.67e-08 2.89e-07\n", "ethanol 7.29e-08 1.35e-06\n", "ethylene 1.70e-07 3.53e-07\n" ] } ], "source": [ "def molecule_set():\n", " butane = (np.array([[0.0, 0.0, 0.0], [1.53, 0.0, 0.0], [2.05, 1.44, 0.0],\n", " [3.58, 1.44, 0.0],\n", " [-0.4, -0.5, 0.9], [-0.4, -0.5, -0.9], [-0.4, 1.0, 0.0],\n", " [1.93, -0.52, 0.88], [1.93, -0.52, -0.88],\n", " [1.65, 1.96, -0.88], [1.65, 1.96, 0.88],\n", " [3.98, 0.44, 0.0], [3.98, 1.96, 0.88],\n", " [3.98, 1.96, -0.88]]),\n", " [6] * 4 + [1] * 10)\n", " ethanol = (np.array([[0.0, 0.0, 0.0], [1.512, 0.0, 0.0], [2.0, 1.32, 0.0],\n", " [-0.39, -0.51, 0.89], [-0.39, -0.51, -0.89],\n", " [-0.39, 1.02, 0.0], [1.90, -0.52, 0.88],\n", " [1.90, -0.52, -0.88], [2.60, 1.30, 0.7]]),\n", " [6, 6, 8, 1, 1, 1, 1, 1, 1])\n", " ethylene = (np.array([[0.0, 0.0, 0.0], [1.34, 0.0, 0.0],\n", " [-0.54, 0.94, 0.0], [-0.54, -0.94, 0.0],\n", " [1.88, 0.94, 0.0], [1.88, -0.94, 0.0]]),\n", " [6, 6, 1, 1, 1, 1])\n", " return {\"butane\": butane, \"ethanol\": ethanol, \"ethylene\": ethylene}\n", "\n", "\n", "lib = builtin_library(\"oplsaa\")\n", "print(f\"{'molecule':<10} {'max |dE| (kJ/mol)':>18} {'max |dF| (kJ/mol/nm)':>22}\")\n", "for name, (pos0, z) in molecule_set().items():\n", " # atom types from the library's SMARTS templates, bonds perceived from\n", " # the coordinates, impropers placed at the trigonal centers\n", " opls = OPLS.from_atoms((pos0, z), lib, cutoff=100.0)\n", " top = opls.topology\n", " model = ForceStressOutput(opls)\n", " ctx, const = openmm_twin(lib, top, opls.impropers)\n", " de = df = 0.0\n", " for _ in range(10):\n", " pos = pos0 + 0.1 * rng.standard_normal(pos0.shape)\n", " out = model(structure_to_graph(\n", " {\"pos\": torch.tensor(pos), \"atomic_numbers\": torch.tensor(z)},\n", " cutoff=100.0))\n", " e_o, f_o = openmm_eval(ctx, const, pos)\n", " de = max(de, abs(float(out[\"energy\"]) * EV_TO_KJ - e_o))\n", " df = max(df, np.abs(out[\"forces\"].detach().numpy() * EV_TO_KJ * 10\n", " - f_o).max())\n", " print(f\"{name:<10} {de:>18.2e} {df:>22.2e}\")\n", " assert de < 1e-6 and df < 1e-5" ] }, { "cell_type": "markdown", "id": "2ae1a42c", "metadata": {}, "source": [ "Agreement at ~1e-7 kJ/mol in the energy and ~1e-6 kJ/mol/nm in every force\n", "component — the residual is the rounding of OpenMM's Coulomb constant\n", "(`138.935456`), i.e. the two engines agree to the precision at which the\n", "physical constants themselves are written down.\n", "\n", "## 3. Term-by-term decomposition\n", "\n", "`xnn` returns the energy decomposition directly; OpenMM force groups give\n", "the same split (1,4 interactions live in the `NonbondedForce` exceptions,\n", "so they are compared together with the Coulomb group).\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "5b22903f", "metadata": { "execution": { "iopub.execute_input": "2026-09-15T19:20:29.737560Z", "iopub.status.busy": "2026-09-15T19:20:29.737432Z", "iopub.status.idle": "2026-09-15T19:20:29.825488Z", "shell.execute_reply": "2026-09-15T19:20:29.824785Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "term xnn (kJ/mol) OpenMM diff\n", "bonds 694.914753 694.914753 -2.59e-10\n", "angles 86.283288 86.283288 -5.34e-11\n", "torsions (+impropers) 22.001321 22.001321 -2.24e-12\n", "Lennard-Jones (direct) 0.316461 0.316461 -3.07e-12\n", "Coulomb + all 1,4 11.832974 11.832974 -9.39e-08\n" ] } ], "source": [ "pos0, z = molecule_set()[\"butane\"]\n", "model = OPLS.from_atoms((pos0, z), lib, cutoff=100.0)\n", "top = model.topology\n", "ctx, const = openmm_twin(lib, top, model.impropers)\n", "pos = pos0 + 0.1 * rng.standard_normal(pos0.shape)\n", "out = model(structure_to_graph(\n", " {\"pos\": torch.tensor(pos), \"atomic_numbers\": torch.tensor(z)},\n", " cutoff=100.0))\n", "rows = [\n", " (\"bonds\", float(out[\"e_bond\"]) * EV_TO_KJ, {1}),\n", " (\"angles\", float(out[\"e_angle\"]) * EV_TO_KJ, {2}),\n", " (\"torsions (+impropers)\",\n", " (float(out[\"e_torsion\"]) + float(out[\"e_improper\"])) * EV_TO_KJ, {3}),\n", " (\"Lennard-Jones (direct)\", float(out[\"e_lj\"]) * EV_TO_KJ, {4}),\n", " (\"Coulomb + all 1,4\",\n", " (float(out[\"e_coulomb\"]) + float(out[\"e_lj14\"])\n", " + float(out[\"e_coulomb14\"])) * EV_TO_KJ, {0}),\n", "]\n", "print(f\"{'term':<24} {'xnn (kJ/mol)':>14} {'OpenMM':>12} {'diff':>10}\")\n", "for label, e_x, grp in rows:\n", " e_o, _ = openmm_eval(ctx, const, pos, groups=grp)\n", " print(f\"{label:<24} {e_x:>14.6f} {e_o:>12.6f} {e_x - e_o:>10.2e}\")\n", " assert abs(e_x - e_o) < 1e-6" ] }, { "cell_type": "markdown", "id": "7e3068b6", "metadata": {}, "source": [ "## 4. Anchors from the papers themselves\n", "\n", "**Relaxed ethane barrier** (Table 1 of Jorgensen 1996; the `oplsaa-1996`\n", "built-in carries the paper's original alkane torsions):\n" ] }, { "cell_type": "code", "execution_count": 5, "id": "6f6cf1f3", "metadata": { "execution": { "iopub.execute_input": "2026-09-15T19:20:29.826837Z", "iopub.status.busy": "2026-09-15T19:20:29.826717Z", "iopub.status.idle": "2026-09-15T19:20:30.685424Z", "shell.execute_reply": "2026-09-15T19:20:30.684704Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "relaxed ethane rotational barrier: 3.01 kcal/mol (Table 1: 3.01)\n" ] } ], "source": [ "from ase import Atoms\n", "from ase.constraints import FixInternals\n", "from ase.optimize import BFGS\n", "from xnn.common.deploy import XNNCalculator\n", "\n", "d, dh, ang = 1.529, 1.09, math.radians(110.7)\n", "pos = [[0.0, 0.0, 0.0], [d, 0.0, 0.0]]\n", "for base, sign, off in ((0.0, -1.0, 60.0), (d, 1.0, 0.0)):\n", " for k in range(3):\n", " phi = math.radians(off + 120.0 * k)\n", " pos.append([base - sign * dh * math.cos(math.pi - ang),\n", " dh * math.sin(math.pi - ang) * math.cos(phi),\n", " dh * math.sin(math.pi - ang) * math.sin(phi)])\n", "z6 = [6, 6] + [1] * 6\n", "model = OPLS.from_atoms((pos, z6), \"oplsaa-1996\", cutoff=30.0)\n", "energies = {}\n", "for target in (60.0, 0.0):\n", " at = Atoms(numbers=z6, positions=pos)\n", " at.calc = XNNCalculator(ForceStressOutput(model), cutoff=model.cutoff)\n", " at.set_dihedral(2, 0, 1, 5, target, indices=[5, 6, 7])\n", " at.set_constraint(FixInternals(dihedrals_deg=[[target, [2, 0, 1, 5]]]))\n", " BFGS(at, logfile=None).run(fmax=1e-5, steps=500)\n", " energies[target] = at.get_potential_energy() / 4.3364104242e-2\n", "barrier = energies[0.0] - energies[60.0]\n", "print(f\"relaxed ethane rotational barrier: {barrier:.2f} kcal/mol \"\n", " f\"(Table 1: 3.01)\")\n", "assert abs(barrier - 3.01) < 0.02" ] }, { "cell_type": "markdown", "id": "3ff6395e", "metadata": {}, "source": [ "**The torsion-form conversion** against Table 2 of the L-OPLS paper, which\n", "lists the same hexane `CT-CT-CT-CT` torsion in both Ryckaert-Bellemans and\n", "Fourier form (kJ/mol) — including the constant `V0` that makes the two\n", "match exactly:\n" ] }, { "cell_type": "code", "execution_count": 6, "id": "74d2a100", "metadata": { "execution": { "iopub.execute_input": "2026-09-15T19:20:30.687558Z", "iopub.status.busy": "2026-09-15T19:20:30.687218Z", "iopub.status.idle": "2026-09-15T19:20:30.691030Z", "shell.execute_reply": "2026-09-15T19:20:30.690446Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "rb_to_fourier: [-0.305938, 2.697394, -0.896807, 0.74567, -0.0]\n", "Siu Table 2: [-0.305938, 2.697394, -0.896807, 0.74567, 0.0]\n", "both directions agree to 1e-6 kJ/mol\n" ] } ], "source": [ "rb = [0.518787, -0.230192, 0.896807, -1.49134, 0.0, 0.0]\n", "fourier = [-0.305938, 2.697394, -0.896807, 0.74567, 0.0]\n", "print(\"rb_to_fourier:\", [round(x, 6) for x in rb_to_fourier(rb)])\n", "print(\"Siu Table 2: \", fourier)\n", "assert max(abs(a - b) for a, b in zip(rb_to_fourier(rb), fourier)) < 1e-6\n", "assert max(abs(a - b) for a, b in zip(fourier_to_rb(fourier), rb)) < 1e-6\n", "print(\"both directions agree to 1e-6 kJ/mol\")" ] }, { "cell_type": "markdown", "id": "dcfe5d06", "metadata": {}, "source": [ "## Summary\n", "\n", "| check | result |\n", "|---|---|\n", "| energy parity vs OpenMM (butane / ethanol / ethylene, 10 random conformations each) | max deviation ~1e-7 kJ/mol |\n", "| force parity vs OpenMM | max deviation ~1e-6 kJ/mol/nm |\n", "| term-by-term decomposition vs OpenMM force groups | ~1e-7 kJ/mol per term |\n", "| relaxed ethane barrier vs Jorgensen 1996 Table 1 | 3.01 vs 3.01 kcal/mol |\n", "| Fourier / Ryckaert-Bellemans conversion vs Siu 2012 Table 2 | exact to the table's digits |\n", "\n", "Together with `examples/ffnn/opls/opls_conformational_energetics.ipynb`\n", "(the full Table 1 reproduction) this establishes that the `xnn` OPLS\n", "implementation is numerically equivalent to the reference MD engines and\n", "faithful to the published parameterizations.\n" ] } ], "metadata": { "kernelspec": { "display_name": "xnn (.venv)", "language": "python", "name": "xnn" }, "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.13.12" } }, "nbformat": 4, "nbformat_minor": 5 }