SimonBoothroyd · SimonBoothroyd · Sep 15, 2023 · Sep 15, 2023 · Sep 15, 2023
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,19 +1,38 @@
-files: \.py$
 repos:
+  - repo: local
+    hooks:
+      - id: isort
+        name: "[Package] Import formatting"
+        language: system
+        entry: isort
+        files: \.py$
 
-- repo: local
-  hooks:
-    - id: isort
-      name: Import formatting
-      language: system
-      entry: isort
+      - id: black
+        name: "[Package] Code formatting"
+        language: system
+        entry: black
+        files: \.py$
 
-    - id: black
-      name: Code formatting
-      language: system
-      entry: black
+      - id: flake8
+        name: "[Package] Linting"
+        language: system
+        entry: flake8
+        files: \.py$
 
-    - id: flake8
-      name: Linting
-      language: system
-      entry: flake8
+      - id: isort-examples
+        name: "[Examples] Import formatting"
+        language: system
+        entry: nbqa isort
+        files: examples/.+\.ipynb$
+
+      - id: black-examples
+        name: "[Examples] Code formatting"
+        language: system
+        entry: nbqa black
+        files: examples/.+\.ipynb$
+
+      - id: flake8-examples
+        name: "[Examples] Linting"
+        language: system
+        entry: nbqa flake8 --ignore=E402
+        files: examples/.+\.ipynb$
diff --git a/Makefile b/Makefile
@@ -1,7 +1,7 @@
 PACKAGE_NAME  := smirnoffee
 CONDA_ENV_RUN := conda run --no-capture-output --name $(PACKAGE_NAME)
 
-.PHONY: env lint format test
+.PHONY: env lint format test test-examples
 
 env:
 	mamba create     --name $(PACKAGE_NAME)
@@ -10,14 +10,20 @@ env:
 	$(CONDA_ENV_RUN) pre-commit install || true
 
 lint:
-	$(CONDA_ENV_RUN) isort --check-only $(PACKAGE_NAME) examples
-	$(CONDA_ENV_RUN) black --check      $(PACKAGE_NAME) examples
-	$(CONDA_ENV_RUN) flake8             $(PACKAGE_NAME) examples
+	$(CONDA_ENV_RUN) isort --check-only $(PACKAGE_NAME)
+	$(CONDA_ENV_RUN) black --check      $(PACKAGE_NAME)
+	$(CONDA_ENV_RUN) flake8             $(PACKAGE_NAME)
+	$(CONDA_ENV_RUN) nbqa isort   --check-only  examples
+	$(CONDA_ENV_RUN) nbqa black   --check       examples
+	$(CONDA_ENV_RUN) nbqa flake8  --ignore=E402 examples
 
 format:
-	$(CONDA_ENV_RUN) isort  $(PACKAGE_NAME) examples
-	$(CONDA_ENV_RUN) black  $(PACKAGE_NAME) examples
-	$(CONDA_ENV_RUN) flake8 $(PACKAGE_NAME) examples
+	$(CONDA_ENV_RUN) isort  $(PACKAGE_NAME)
+	$(CONDA_ENV_RUN) black  $(PACKAGE_NAME)
+	$(CONDA_ENV_RUN) flake8 $(PACKAGE_NAME)
+	$(CONDA_ENV_RUN) nbqa isort                examples
+	$(CONDA_ENV_RUN) nbqa black                examples
+	$(CONDA_ENV_RUN) nbqa flake8 --ignore=E402 examples
 
 test:
 	$(CONDA_ENV_RUN) pytest -v --cov=$(PACKAGE_NAME) --cov-report=xml --color=yes $(PACKAGE_NAME)/tests/

diff --git a/devtools/envs/base.yaml b/devtools/envs/base.yaml
@@ -16,6 +16,11 @@ dependencies:
   - pytorch-cpu
   - pydantic
 
+  # Examples
+  - jupyter
+  - nbconvert
+  - nglview
+
   # Dev / Testing
   - rdkit
   - ambertools
@@ -27,12 +32,9 @@ dependencies:
   - pre-commit
   - isort
   - black
-  - black-jupyter
   - flake8
   - flake8-pyproject
-
-  - jupyter
-  - nbconvert
+  - nbqa
 
   - pytest
   - pytest-cov

diff --git a/examples/README.md b/examples/README.md
@@ -0,0 +1,7 @@
+# Examples
+
+This directory contains a number of examples of how to use `smirnoffee`. They currently include:
+
+* [Evaluating the energy of a molecule](compute-energy.ipynb)
+* [Minimizing the conformer of a molecule](conformer-minimization.ipynb)
+* [Computing the gradient of the energy w.r.t. force field parameters](parameter-gradients.ipynb)
diff --git a/examples/compute-energy.ipynb b/examples/compute-energy.ipynb
@@ -28,9 +28,9 @@
     }
    ],
    "source": [
-    "import torch\n",
     "import openff.toolkit\n",
     "import openff.units\n",
+    "import torch\n",
     "\n",
     "molecule = openff.toolkit.Molecule.from_smiles(\"CC(=O)NC1=CC=C(C=C1)O\")\n",
     "molecule.generate_conformers(\n",
@@ -163,16 +163,6 @@
     }
    },
    "id": "2ac831ea0b442177"
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "Here we have provided a single conformer, but multiple can be batched together by stacking them along the first axis for faster evaluation."
-   ],
-   "metadata": {
-    "collapsed": false
-   },
-   "id": "c668293cf08c444f"
   }
  ],
  "metadata": {

diff --git a/examples/conformer-minimization.ipynb b/examples/conformer-minimization.ipynb
@@ -0,0 +1,242 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "261b79c7042b8a6f",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "# Conformer Minimization\n",
+    "\n",
+    "This example will show how to optimize a conformer of paracetamol.\n",
+    "\n",
+    "Load in a paracetamol molecule, generate a conformer for it, and perturb the conformer to ensure it needs minimization."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "b081ee3aecf864ac",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-15T11:06:21.014128Z",
+     "start_time": "2023-09-15T11:06:17.723820Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import openff.toolkit\n",
+    "import openff.units\n",
+    "import torch\n",
+    "\n",
+    "molecule = openff.toolkit.Molecule.from_smiles(\"CC(=O)NC1=CC=C(C=C1)O\")\n",
+    "molecule.generate_conformers(n_conformers=1)\n",
+    "\n",
+    "conformer = torch.tensor(molecule.conformers[0].m_as(openff.units.unit.angstrom)) * 1.10\n",
+    "conformer.requires_grad = True"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f4168aec7a72494c",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "We specify that the gradient of the conformer is required so that we can optimize it using PyTorch.\n",
+    "\n",
+    "Parameterize the molecule using OpenFF Interchange and convert it into a PyTorch tensor representation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "8d00fd2dcf4c27cf",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-15T11:06:25.334698Z",
+     "start_time": "2023-09-15T11:06:21.017138Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "05e38334e84f40cca6dd56100c36085c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": []
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import openff.interchange\n",
+    "\n",
+    "interchange = openff.interchange.Interchange.from_smirnoff(\n",
+    "    openff.toolkit.ForceField(\"openff-2.1.0.offxml\"),\n",
+    "    molecule.to_topology(),\n",
+    ")\n",
+    "\n",
+    "import smirnoffee.ff\n",
+    "\n",
+    "force_field, [topology] = smirnoffee.ff.convert_interchange(interchange)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "792cb057cb419fa8",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "We can minimize the conformer using any of PyTorch's optimizers. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "facd656a27cf46a8",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-15T11:06:25.643421Z",
+     "start_time": "2023-09-15T11:06:25.336911Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 0: E=427.2271423339844 kJ / mol\n",
+      "Epoch 5: E=29.65744972229004 kJ / mol\n",
+      "Epoch 10: E=-76.69727325439453 kJ / mol\n",
+      "Epoch 15: E=-92.81060791015625 kJ / mol\n",
+      "Epoch 20: E=-127.06442260742188 kJ / mol\n",
+      "Epoch 25: E=-154.0148162841797 kJ / mol\n",
+      "Epoch 30: E=-161.15867614746094 kJ / mol\n",
+      "Epoch 35: E=-169.47312927246094 kJ / mol\n",
+      "Epoch 40: E=-176.0826416015625 kJ / mol\n",
+      "Epoch 45: E=-176.5312957763672 kJ / mol\n",
+      "Epoch 50: E=-177.3107147216797 kJ / mol\n",
+      "Epoch 55: E=-178.5596466064453 kJ / mol\n",
+      "Epoch 60: E=-179.07481384277344 kJ / mol\n",
+      "Epoch 65: E=-179.67147827148438 kJ / mol\n",
+      "Epoch 70: E=-180.067626953125 kJ / mol\n",
+      "Epoch 74: E=-180.26405334472656 kJ / mol\n"
+     ]
+    }
+   ],
+   "source": [
+    "import smirnoffee.potentials\n",
+    "\n",
+    "optimizer = torch.optim.Adam([conformer], lr=0.02)\n",
+    "\n",
+    "for epoch in range(75):\n",
+    "    energy = smirnoffee.potentials.compute_energy(\n",
+    "        topology.parameters, conformer, force_field\n",
+    "    )\n",
+    "    energy.backward()\n",
+    "\n",
+    "    optimizer.step()\n",
+    "    optimizer.zero_grad()\n",
+    "\n",
+    "    if epoch % 5 == 0 or epoch == 74:\n",
+    "        print(f\"Epoch {epoch}: E={energy.item()} kJ / mol\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "360d6eb9cf2b6cc4",
+   "metadata": {
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "source": [
+    "We can then re-store the optimized conformer back into the molecule. Here we add the conformer to the molecule's conformer list, but we could also replace the original conformer."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "eaec04c4039ca59b",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-09-15T11:06:25.662680Z",
+     "start_time": "2023-09-15T11:06:25.643943Z"
+    },
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7ad66e9cb155431b9140d01aef214979",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "NGLWidget(max_frame=1)"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "molecule.add_conformer(conformer.detach().numpy() * openff.units.unit.angstrom)\n",
+    "molecule.visualize(backend=\"nglview\")"
+   ]
+  }
+ ],
+ "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.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}