Adding OpenMM as an optional MD engine (#55)

Added experimental interface to OpenMM. Currently works only for MACE and runs in the local directory. * Initial OpenMM MD engine * Guard against incorrect usage * Added OpenMM usage and warning to README
duartegroup · Oct 30, 2023 · 3ba53a5 · 3ba53a5
1 parent 41fbb2c
commit 3ba53a5
Show file tree

Hide file tree

Showing 11 changed files with 533 additions and 22 deletions.
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
@@ -1,5 +1,5 @@
 ---
-name: Test with pytest
+name: pytest
 
 on:
   push:
@@ -8,7 +8,7 @@ on:
 
 jobs:
   test:
-    name: Env (${{ matrix.python-version }}, ${{ matrix.os }})
+    name: GAP (${{ matrix.python-version }}, ${{ matrix.os }})
     runs-on: ${{ matrix.os }}
 
     strategy:
@@ -37,9 +37,6 @@ jobs:
     - name: Basic install
       run: ./install_gap.sh
 
-    - name: List installed packages
-      run: micromamba list
-
     - name: Test basic install
       run: pytest
 
@@ -72,8 +69,5 @@ jobs:
     - name: MACE install
       run: ./install_mace.sh
 
-    - name: List installed packages
-      run: micromamba list
-
     - name: Test MACE install
       run: pytest
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -12,6 +12,7 @@ repos:
         args: ['--maxkb=500', '--enforce-all']
         exclude: mlptrain/sampling/tests/data.zip
       - id: check-yaml
+      - id: check-toml
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.

diff --git a/README.md b/README.md
@@ -24,14 +24,27 @@ Each model is installed into individual conda environment:
 # Install ACE
 ./install_ace.sh
 
-#Install MACE
+# Install MACE
 ./install_mace.sh 
 ```
 
 ### Notes
 
 - Units are: distance (Å), energy (eV), force (eV Å$`^{-1}`$), time (fs)
 
+## Using with OpenMM (Experimental!)
+
+The OpenMM backend only works with MACE at the moment. The necessary dependencies are installed automatically via conda:
+
+```console
+./install_mace.sh
+```
+
+You should now be able to run `water_openmm.py` in `./examples` or run the jupyter notebook on Google Colab [`water_openmm_colab.ipynb`](./examples/water_openmm_colab.ipynb).
+
+You can use OpenMM during active learning by passing the keyword argument `md_program="OpenMM"` to the `al_train` method.
+You can run MD with OpenMM using `mlptrain.md_openmm.run_mlp_md_openmm()`
+
 ## For developers
 
 We are happy to accept pull requests from users. Please first fork mlp-train repository. We use `pre-commit`, `Ruff` and `pytest` to check the code. Your PR needs to pass through these checks before is accepted. `Pre-commit` is installed as one the dependecies. To use it in your repository, run the following command in the mlp-train folder:

diff --git a/environment_mace.yml b/environment_mace.yml
@@ -20,7 +20,8 @@ dependencies:
   - scipy
   - xtb
   # MACE dependencies
-  - pytorch
+  - pytorch=2.0
+  - openmm-torch=1.1
   - torchvision
   - torchaudio
   - torch-ema
@@ -34,4 +35,6 @@ dependencies:
     # so needs to be installed together.
     - e3nn
     - mace@git+https://github.com/ACEsuit/[email protected]
+    # The upstream PR to openmm-ml has not been merged yet: https://github.com/openmm/openmm-ml/pull/61)
+    - openmmml@git+https://github.com/sef43/openmm-ml@mace
     - ase@git+https://gitlab.com/ase/ase@f2615a6e9a # for Plumed
diff --git a/examples/water_openmm.py b/examples/water_openmm.py
@@ -0,0 +1,30 @@
+import mlptrain as mlt
+
+mlt.Config.n_cores = 1
+
+
+if __name__ == '__main__':
+    system = mlt.System(mlt.Molecule('water.xyz'), box=None)
+
+    mace = mlt.potentials.MACE('water', system=system)
+
+    mace.al_train(
+        method_name='xtb',
+        temp=500,
+        max_active_iters=2,
+        max_active_time=50,
+        n_configs_iter=3,
+        md_program='OpenMM',
+    )
+
+    # Run some dynamics with the potential
+    trajectory = mlt.md_openmm.run_mlp_md_openmm(
+        configuration=system.random_configuration(),
+        mlp=mace,
+        fs=2000,
+        temp=300,
+        dt=0.5,
+        interval=100,
+    )
+
+    trajectory.save(filename='water_trajectory.xyz')
diff --git a/examples/water_openmm_colab.ipynb b/examples/water_openmm_colab.ipynb
@@ -0,0 +1,123 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "98up6lvBEmaE"
+   },
+   "source": [
+    "# MLP-Train Water with OpenMM on Colab\n",
+    "\n",
+    "You can run this notebook in Google Colab: [![Open On Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/sef43/mlp-train/blob/main/examples/water_openmm_colab.ipynb)\n",
+    "\n",
+    "MLP-Train basic example that runs on Colab and uses OpenMM as the MD engine."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 1000
+    },
+    "id": "Cb_3edQEEyxx",
+    "outputId": "b02b5b81-db06-49e1-81d2-780f0e72079b"
+   },
+   "outputs": [],
+   "source": [
+    "# Installation on Colab\n",
+    "!pip install -q condacolab\n",
+    "import condacolab\n",
+    "\n",
+    "condacolab.install_mambaforge()\n",
+    "# https://github.com/openmm/openmm-torch/issues/88\n",
+    "%env CONDA_OVERRIDE_CUDA=12.0\n",
+    "!mamba install -c conda-forge openmm-torch=1.1 pytorch=2.0 xtb cython rdkit\n",
+    "\n",
+    "# Only have python 3.10 available on Colab so cannot install autodE from\n",
+    "# conda-forge as there are no py3.10 builds for autode=v1.1.*\n",
+    "!git clone --branch v1.1.3 https://github.com/duartegroup/autodE.git\n",
+    "!cd autodE && python setup.py install\n",
+    "!cd ..\n",
+    "!pip install ./autodE\n",
+    "!pip install git+https://github.com/rosswhitfield/ase.git\n",
+    "!pip install git+https://github.com/ACEsuit/mace.git\n",
+    "!pip install git+https://github.com/sef43/openmm-ml@mace\n",
+    "!pip install git+https://github.com/sef43/mlp-train"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "jInwylB5aU3x",
+    "outputId": "efa7baaf-5ee2-469a-e131-a2616dde5890"
+   },
+   "outputs": [],
+   "source": [
+    "# Get the example input\n",
+    "!wget https://raw.githubusercontent.com/sef43/mlp-train/main/examples/water.xyz"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "rdjXObniFynd",
+    "outputId": "1ae4d16c-f2db-4f23-d42c-7115091fc3ad"
+   },
+   "outputs": [],
+   "source": [
+    "import mlptrain as mlt\n",
+    "\n",
+    "mlt.Config.n_cores = 1\n",
+    "\n",
+    "system = mlt.System(mlt.Molecule('water.xyz'), box=None)\n",
+    "\n",
+    "mace = mlt.potentials.MACE('water', system=system)\n",
+    "\n",
+    "mace.al_train(\n",
+    "    method_name='xtb',\n",
+    "    temp=500,\n",
+    "    max_active_iters=2,\n",
+    "    max_active_time=50,\n",
+    "    n_configs_iter=3,\n",
+    "    md_program='OpenMM',\n",
+    ")\n",
+    "\n",
+    "# Run some dynamics with the potential\n",
+    "trajectory = mlt.md_openmm.run_mlp_md_openmm(\n",
+    "    configuration=system.random_configuration(),\n",
+    "    mlp=mace,\n",
+    "    fs=2000,\n",
+    "    temp=300,\n",
+    "    dt=0.5,\n",
+    "    interval=100,\n",
+    ")\n",
+    "\n",
+    "trajectory.save(filename='water_trajectory.xyz')"
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/mlptrain/__init__.py b/mlptrain/__init__.py
@@ -3,7 +3,7 @@
 from mlptrain.molecule import Molecule
 from mlptrain.system import System
 from mlptrain.box import Box
-from mlptrain.sampling import md, UmbrellaSampling, Metadynamics
+from mlptrain.sampling import md, md_openmm, UmbrellaSampling, Metadynamics
 from mlptrain.sampling import Bias, PlumedBias, PlumedCalculator
 from mlptrain.sampling.plumed import plot_cv_versus_time, plot_cv1_and_cv2
 from mlptrain.utils import convert_ase_time, convert_ase_energy
@@ -41,6 +41,7 @@
            'convert_ase_time',
            'convert_ase_energy',
            'md',
+           'md_openmm',
            'loss',
            'selection',
            'potentials'

diff --git a/mlptrain/config.py b/mlptrain/config.py
@@ -24,6 +24,13 @@ class _ConfigClass:
                      'train_fraction': 0.9}
 
     # MACE params
+
+    try:
+        import torch
+        mace_device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    except ImportError:
+        mace_device = 'cpu'
+
     mace_params = {'valid_fraction': 0.1,
                    'config_type_weights': '{"Default":1.0}',
                    'model': 'MACE',
@@ -34,7 +41,7 @@ class _ConfigClass:
                    'batch_size': 10,
                    'r_max': 5,
                    'correlation': 3,
-                   'device': 'cuda',
+                   'device': mace_device ,
                    'calc_device': 'cpu',
                    'error_table': 'TotalMAE',
                    'swa': True,