Add test workflow for openmm dmff plugin

.github/workflows/test_openmm_dmff_plugin.yml

@@ -0,0 +1,51 @@
+name: OpenMM DMFF Plugin tests.
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.9]
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install Dependencies
+        run: |
+          source $CONDA/bin/activate
+          conda create -n dmff -y python=${{ matrix.python-version }} numpy openmm==7.7.0 mdtraj=1.9.7 pytest rdkit biopandas openbabel -c conda-forge
+          conda activate dmff
+          pip install --upgrade pip
+          pip install jax==0.3.15 jaxlib==0.3.15 jax_md==0.2.0 pymbar==4.0.1 chex==0.1.4 dm-haiku==0.0.7 tqdm
+          conda install -c deepmodeling libtensorflow_cc=2.9.1=cuda112h02da4e0_0
+          mkdir /tmp/omm_dmff_working_dir
+          cd /tmp/omm_dmff_working_dir
+          wget https://github.com/tensorflow/tensorflow/archive/refs/tags/v2.9.1.tar.gz
+          tar -xvf v2.9.1.tar.gz
+          cp -r tensorflow-2.9.1/tensorflow/c ${CONDA_PREFIX}/include/tensorflow/
+          git clone https://github.com/serizba/cppflow.git
+          cd cppflow
+          git apply ${GITHUB_WORKSPACE}/backend/openmm_dmff_plugin/tests/cppflow_empty_constructor.patch
+          mkdir ${CONDA_PREFIX}/include/cppflow
+          cp -r include/cppflow ${CONDA_PREFIX}/include/
+      - name: Install OpenMM DMFF Plugin
+        run: |
+          source $CONDA/bin/activate dmff
+          cd ${GITHU_WORKSPACE}/backend/openmm_dmff_plugin
+          mkdir build && cd build
+          export OPENMM_INSTALLED_DIR=$CONDA_PREFIX
+          export CPPFLOW_INSTALLED_DIR=$CONDA_PREFIX
+          export LIBTENSORFLOW_INSTALLED_DIR=$CONDA_PREFIX
+          cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR}
+          make && make install 
+          make PythonInstall
+      - name: Run Tests
+        run: |
+          source $CONDA/bin/activate dmff
+          python -m OpenMMDMFFPlugin.tests.test_dmff_plugin_nve -n 100

backend/save_dmff2tf.py

@@ -3,11 +3,9 @@
 import jax
 import jax.numpy as jnp
 from jax.experimental import jax2tf
-# The model is saved in double precision by default. 
-# Since forces accuracy in double precision is needed in molecular dynamics simulations, 
-# we need to enable double precision in JAX.
-from jax import config
-config.update("jax_enable_x64", True)
+# The model is saved in float32 precision by default. 
+#from jax import config
+#config.update("jax_enable_x64", True)
 import openmm.app as app
 import openmm.unit as unit
 import tensorflow as tf
@@ -71,6 +69,6 @@ def potential_engrad(positions, box, pairs):
     )
     dmff_model = tf.Module()
     dmff_model.f = tf.function(f_tf, autograph=False,
-                            input_signature=[tf.TensorSpec(shape=[natoms,3], dtype=tf.float64), tf.TensorSpec(shape=[3,3], dtype=tf.float64), tf.TensorSpec(shape=tf.TensorShape([None, 2]), dtype=tf.int32)])
+                            input_signature=[tf.TensorSpec(shape=[natoms,3], dtype=tf.float32), tf.TensorSpec(shape=[3,3], dtype=tf.float32), tf.TensorSpec(shape=tf.TensorShape([None, 2]), dtype=tf.int32)])
 
     tf.saved_model.save(dmff_model, output_dir, options=tf.saved_model.SaveOptions(experimental_custom_gradients=True))