Add support for NequIP models

README.md

@@ -8,6 +8,9 @@ and [ANI-2x](https://pubs.acs.org/doi/full/10.1021/acs.jctc.0c00121), using the
 suitable for small molecules involving a limited set of elements and no charges.  Future releases will add new potential
 functions that support a much wider range of molecules.
 
+Additionally [NequIP](https://www.nature.com/articles/s41467-022-29939-5) potential functions can be used using the implementation from https://github.com/mir-group/nequip , see the [README](examples/nequip/README.md) in `/examples/nequip` for more information.
+
+
 ### Installation
 
 OpenMM-ML can be installed with conda or mamba.

examples/nequip/README.md

@@ -0,0 +1,47 @@
+# NequIP models in OpenMM-ML
+
+This directory contains examples for running simulations using a NequIP potential.
+
+## Installation
+
+
+first install openmm-torch from conda-forge:
+
+```
+conda install -c conda-forge openmm-torch
+```
+
+Then install NequIP development branch and this version of openmm-ml using pip
+
+```
+pip install git+https://github.com/mir-group/nequip@develop
+pip install git+https://github.com/sef43/openmm-ml@nequip
+```
+
+## Usage
+
+Once you have a deployed trained NequIP model you can use it as the potential in OpenMM-ML:
+
+```python
+from openmmml import MLPotential
+
+# create a System with NequIP MLP
+
+# need to specify the unit conversion factors from the NequIP model units to OpenMM units.
+# e.g.:
+# distance: model is in Angstrom, OpenMM is in nanometers
+A_to_nm = 0.1
+# energy: model is in kcal/mol, OpenMM is in kJ/mol
+kcal_to_kJ_per_mol = 4.184
+
+potential = MLPotential('nequip', model_path='example_model_deployed.pth',
+                        distance_to_nm=A_to_nm,
+                        energy_to_kJ_per_mol=kcal_to_kJ_per_mol)
+
+system = potential.createSystem(topology)
+```
+
+## Example
+
+### run_nequip.ipynb
+Runs a simulation using the model created by NequIP example [config/example.yaml](https://github.com/mir-group/nequip/blob/main/configs/example.yaml). It is available as a python script: [`run_nequip.py`](run_nequip.py) and a Jupyter notebook [`run_nequip.ipynb`](run_nequip.ipynb) which can be run on Colab.

examples/nequip/run_nequip.ipynb

@@ -0,0 +1,186 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Example: using OpenMM to run a simulation with a NequIP ML potential\n",
+    "\n",
+    "You can run this example directly in your browser: [![Open On Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/sef43/openmm-ml/blob/nequip/examples/nequip/run_nequip.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Install Conda\n",
+    "\n",
+    "[Conda](https://docs.conda.io/) is a package and environment manager. On Google Colab, Conda is installed with [conda-colab](https://github.com/jaimergp/condacolab). On your computer, you should follow these [installation instructions](https://docs.conda.io/projects/conda/en/latest/user-guide/install/index.html).\n",
+    "\n",
+    "⚠️ Do not use conda-colab on your computer!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install -q condacolab\n",
+    "import condacolab\n",
+    "condacolab.install_mambaforge() # use mamba on colab because it is faster than conda"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Install software\n",
+    "\n",
+    "First install everything we can from [conda-forge](https://conda-forge.org/).\n",
+    "Then use pip.\n",
+    "\n",
+    "⚠️ The installation might take up to 10 min!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#https://github.com/openmm/openmm-torch/issues/88\n",
+    "%env CONDA_OVERRIDE_CUDA=12.0\n",
+    "!mamba install -c conda-forge openmm-torch pytorch=*=cuda*\n",
+    "\n",
+    "!pip install git+https://github.com/mir-group/nequip@develop\n",
+    "!pip install git+https://github.com/sef43/openmm-ml@nequip"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Get the files we need to run the example"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!wget https://raw.githubusercontent.com/sef43/openmm-ml/nequip/examples/nequip/toluene.pdb\n",
+    "!wget https://raw.githubusercontent.com/sef43/openmm-ml/nequip/examples/nequip/example_model_deployed.pth"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run simulation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "#\"Step\",\"Potential Energy (kJ/mole)\",\"Temperature (K)\"\n",
+      "100,-710452.5,220.47627134902876\n",
+      "200,-710464.5625,204.22321067994645\n",
+      "300,-710484.25,312.68608178697696\n",
+      "400,-710462.375,250.56526717177863\n",
+      "500,-710464.25,289.43386954186155\n",
+      "600,-710448.9375,219.63591412548993\n",
+      "700,-710446.0,261.7770563640068\n",
+      "800,-710466.75,454.97388216259844\n",
+      "900,-710454.0,401.38716825310536\n",
+      "1000,-710460.5,429.09933310001867\n",
+      "-710513.4375 kJ/mol\n"
+     ]
+    }
+   ],
+   "source": [
+    "import openmm\n",
+    "import openmm.app as app\n",
+    "import openmm.unit as unit\n",
+    "from openmmml import MLPotential\n",
+    "from sys import stdout\n",
+    "\n",
+    "\"\"\"\n",
+    "Uses a deployed trained NequIP model from the basic example on https://github.com/mir-group/nequip\n",
+    ">>> nequip-train configs/example.yaml\n",
+    ">>> nequip-deploy build --train-dir path/to/training/session/ example_model_deployed.pth\n",
+    "\"\"\"\n",
+    "\n",
+    "# load toluene structure\n",
+    "pdb = app.PDBFile('toluene.pdb')\n",
+    "\n",
+    "# create a System with NequIP MLP\n",
+    "\n",
+    "# need to specify the unit conversion factors from the NequIP model units to OpenMM units.\n",
+    "# distance: model is in Angstrom, OpenMM is in nanometers\n",
+    "A_to_nm = 0.1\n",
+    "# energy: model is in kcal/mol, OpenMM is in kJ/mol\n",
+    "kcal_to_kJ_per_mol = 4.184\n",
+    "\n",
+    "potential = MLPotential('nequip', model_path='example_model_deployed.pth',\n",
+    "                        distance_to_nm=A_to_nm,\n",
+    "                        energy_to_kJ_per_mol=kcal_to_kJ_per_mol)\n",
+    "\n",
+    "system = potential.createSystem(pdb.topology)\n",
+    "\n",
+    "# run langevin dynamics at 300K for 1000 steps\n",
+    "integrator = openmm.LangevinIntegrator(300*unit.kelvin, 10.0/unit.picoseconds, 1.0*unit.femtosecond)\n",
+    "simulation = app.Simulation(pdb.topology, system, integrator)\n",
+    "simulation.context.setPositions(pdb.positions)\n",
+    "simulation.reporters.append(app.PDBReporter('output.pdb', 100))\n",
+    "simulation.reporters.append(app.StateDataReporter(stdout, 100, step=True,\n",
+    "        potentialEnergy=True, temperature=True))\n",
+    "\n",
+    "simulation.step(1000)\n",
+    "\n",
+    "# Minimize the energy\n",
+    "simulation.minimizeEnergy()\n",
+    "energy = simulation.context.getState(getEnergy=True).getPotentialEnergy()\n",
+    "print(energy)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "neqmm",
+   "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"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "413bac2e6b7bbc56e28392734551129ace8c921fbbb9e9273c3258497d1da304"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

examples/nequip/run_nequip.py

@@ -0,0 +1,43 @@
+import openmm
+import openmm.app as app
+import openmm.unit as unit
+from openmmml import MLPotential
+from sys import stdout
+
+"""
+Uses a deployed trained NequIP model from the basic example on https://github.com/mir-group/nequip
+>>> nequip-train configs/example.yaml
+>>> nequip-deploy build --train-dir path/to/training/session/ example_model_deployed.pth
+"""
+
+# load toluene structure
+pdb = app.PDBFile('toluene.pdb')
+
+# create a System with NequIP MLP
+
+# need to specify the unit conversion factors from the NequIP model units to OpenMM units.
+# distance: model is in Angstrom, OpenMM is in nanometers
+A_to_nm = 0.1
+# energy: model is in kcal/mol, OpenMM is in kJ/mol
+kcal_to_kJ_per_mol = 4.184
+
+potential = MLPotential('nequip', model_path='example_model_deployed.pth',
+                        distance_to_nm=A_to_nm,
+                        energy_to_kJ_per_mol=kcal_to_kJ_per_mol)
+
+system = potential.createSystem(pdb.topology)
+
+# run langevin dynamics at 300K for 1000 steps
+integrator = openmm.LangevinIntegrator(300*unit.kelvin, 10.0/unit.picoseconds, 1.0*unit.femtosecond)
+simulation = app.Simulation(pdb.topology, system, integrator)
+simulation.context.setPositions(pdb.positions)
+simulation.reporters.append(app.PDBReporter('output.pdb', 100))
+simulation.reporters.append(app.StateDataReporter(stdout, 100, step=True,
+        potentialEnergy=True, temperature=True, speed=True))
+
+simulation.step(1000)
+
+# Minimize the energy
+simulation.minimizeEnergy()
+energy=simulation.context.getState(getEnergy=True).getPotentialEnergy()
+print(energy)

examples/nequip/toluene.pdb

@@ -0,0 +1,23 @@
+HETATM    1  C1  UNL     1       2.199  -0.143   0.062  1.00  0.00           C  
+HETATM    2  C2  UNL     1       0.713  -0.073   0.011  1.00  0.00           C  
+HETATM    3  C3  UNL     1       0.076   1.155   0.101  1.00  0.00           C  
+HETATM    4  C4  UNL     1      -1.298   1.288   0.060  1.00  0.00           C  
+HETATM    5  C5  UNL     1      -2.077   0.159  -0.076  1.00  0.00           C  
+HETATM    6  C6  UNL     1      -1.445  -1.066  -0.167  1.00  0.00           C  
+HETATM    7  C7  UNL     1      -0.070  -1.200  -0.126  1.00  0.00           C  
+HETATM    8  H1  UNL     1       2.463  -0.598   1.048  1.00  0.00           H  
+HETATM    9  H2  UNL     1       2.621   0.858   0.021  1.00  0.00           H  
+HETATM   10  H3  UNL     1       2.604  -0.841  -0.701  1.00  0.00           H  
+HETATM   11  H4  UNL     1       0.727   2.035   0.209  1.00  0.00           H  
+HETATM   12  H5  UNL     1      -1.731   2.295   0.138  1.00  0.00           H  
+HETATM   13  H6  UNL     1      -3.158   0.259  -0.109  1.00  0.00           H  
+HETATM   14  H7  UNL     1      -2.051  -1.972  -0.275  1.00  0.00           H  
+HETATM   15  H8  UNL     1       0.429  -2.158  -0.196  1.00  0.00           H  
+CONECT    1    2    8    9   10
+CONECT    2    3    3    7
+CONECT    3    4   11
+CONECT    4    5    5   12
+CONECT    5    6   13
+CONECT    6    7    7   14
+CONECT    7   15
+END

openmmml/models/__init__.py

@@ -1 +1 @@
-from . import anipotential
+from . import anipotential, nequippotential, utils