Add openmm-dmff plugin to the backend of DMFF project

Add openmm-dmff plugin and related script to freeze dmff potential and run it with OpenMM MD engine.

.github/workflows/test_openmm_dmff_plugin.yml

@@ -0,0 +1,50 @@
+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_omm -y python=${{ matrix.python-version }} numpy openmm=7.7 -c conda-forge
+          conda activate dmff_omm
+          conda install -y libtensorflow_cc=2.9.1 -c conda-forge
+          pip install setuptools==59.5.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_omm
+          cd ${GITHUB_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} -DUSE_HIGH_PRECISION=ON
+          make && make install 
+          make PythonInstall
+      - name: Run Tests
+        run: |
+          source $CONDA/bin/activate dmff_omm
+          python -m OpenMMDMFFPlugin.tests.test_dmff_plugin_nve -n 100

.gitignore

@@ -784,4 +784,4 @@ FodyWeavers.xsd
 */_version.py
 
 # hmtff cache
-*.hmtff/
+*.hmtff/

backend/openmm_dmff_plugin/.gitignore

@@ -0,0 +1,6 @@
+**/build
+**/.vscode
+**/__pycache__
+*.dcd
+**/output
+**/.devcontainer

backend/openmm_dmff_plugin/CMakeLists.txt

@@ -0,0 +1,149 @@
+#---------------------------------------------------
+# OpenMM DMFF Plugin
+#----------------------------------------------------
+
+CMAKE_MINIMUM_REQUIRED(VERSION 3.5)
+SET(CMAKE_VERBOSE_MAKEFILE OFF)
+
+project(OpenMMDMFF)
+
+# We need to know where OpenMM is installed so we can access the headers and libraries.
+SET(OPENMM_DIR "/usr/local/openmm" CACHE PATH "Where OpenMM is installed")
+INCLUDE_DIRECTORIES("${OPENMM_DIR}/include")
+LINK_DIRECTORIES("${OPENMM_DIR}/lib" "${OPENMM_DIR}/lib/plugins")
+
+
+# Include the cppflow
+SET(CPPFLOW_DIR "/usr/local/cppflow/" CACHE PATH "Where cppflow is installed")
+INCLUDE_DIRECTORIES("${CPPFLOW_DIR}/include")
+LINK_DIRECTORIES("${CPPFLOW_DIR}/lib")
+
+# Include the tensorflow
+SET(TENSORFLOW_DIR "/usr/local/tensorflow/" CACHE PATH "Where libtensorflow is installed")
+INCLUDE_DIRECTORIES("${TENSORFLOW_DIR}/include")
+LINK_DIRECTORIES("${TENSORFLOW_DIR}/lib")
+
+# Specify the C++ version we are building for.
+SET (CMAKE_CXX_STANDARD 17)
+
+set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-ignored-attributes -Wl,--allow-multiple-definition  -Wl,--no-as-needed")
+
+# Set if high precision (double) is used. By default, it is off.
+OPTION(USE_HIGH_PRECISION "Use high precision (double) for calculations" OFF)
+
+IF(USE_HIGH_PRECISION)
+    message(STATUS "USE_HIGH_PRECISION is enabled")
+    ADD_DEFINITIONS(-DUSE_HIGH_PRECISION)
+endif ()
+
+if (USE_HIGH_PRECISION)
+    set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DHIGH_PRECISION")
+else ()
+    set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS}")
+endif ()
+
+# Select where to install
+IF(${CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT})
+    IF(WIN32)
+        SET(CMAKE_INSTALL_PREFIX "$ENV{ProgramFiles}/OpenMM" CACHE PATH "Where to install the plugin" FORCE)
+    ELSE(WIN32)
+        SET(CMAKE_INSTALL_PREFIX "${OPENMM_DIR}" CACHE PATH "Where to install the plugin" FORCE)
+    ENDIF(WIN32)
+ENDIF(${CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT})
+
+# Put all the tests and libraries in a single output directory.
+IF(NOT EXECUTABLE_OUTPUT_PATH)
+  SET(EXECUTABLE_OUTPUT_PATH ${PROJECT_BINARY_DIR}
+      CACHE INTERNAL "Single output directory for building all executables.")
+ENDIF()
+IF(NOT LIBRARY_OUTPUT_PATH)
+  SET(LIBRARY_OUTPUT_PATH ${PROJECT_BINARY_DIR}
+      CACHE INTERNAL "Single output directory for building all libraries.")
+ENDIF()
+SET(${PROJECT_NAME}_EXECUTABLE_DIR ${EXECUTABLE_OUTPUT_PATH}/${CMAKE_CFG_INTDIR})
+SET(${PROJECT_NAME}_LIBRARY_DIR    ${LIBRARY_OUTPUT_PATH}/${CMAKE_CFG_INTDIR})
+
+# The source is organized into subdirectories, but we handle them all from
+# this CMakeLists file rather than letting CMake visit them as SUBDIRS.
+SET(DMFF_PLUGIN_SOURCE_SUBDIRS openmmapi serialization)
+
+# Set the library name
+SET(DMFF_LIBRARY_NAME OpenMMDMFF)
+SET(SHARED_DMFF_TARGET ${DMFF_LIBRARY_NAME})
+
+# These are all the places to search for header files that are to be part of the API.
+SET(API_INCLUDE_DIRS "openmmapi/include" "openmmapi/include/internal")
+
+# Locate header files.
+SET(API_INCLUDE_FILES)
+FOREACH(dir ${API_INCLUDE_DIRS})
+    FILE(GLOB fullpaths ${dir}/*.h)
+    SET(API_INCLUDE_FILES ${API_INCLUDE_FILES} ${fullpaths})
+ENDFOREACH(dir)
+
+# Collect source files
+SET(SOURCE_FILES) # empty
+SET(SOURCE_INCLUDE_FILES)
+FOREACH(subdir ${DMFF_PLUGIN_SOURCE_SUBDIRS})
+    FILE(GLOB src_files  ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.cpp)
+    FILE(GLOB incl_files ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.h)
+    SET(SOURCE_FILES         ${SOURCE_FILES}         ${src_files})   #append
+    SET(SOURCE_INCLUDE_FILES ${SOURCE_INCLUDE_FILES} ${incl_files})
+
+    ## Make sure we find these locally before looking in OpenMM/include if
+    ## OpenMM was previously installed there.
+    INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/include)
+ENDFOREACH(subdir)
+
+# Create the library.
+
+ADD_LIBRARY(${SHARED_DMFF_TARGET} SHARED ${SOURCE_FILES} ${SOURCE_INCLUDE_FILES} ${API_INCLUDE_FILES})
+SET_TARGET_PROPERTIES(${SHARED_DMFF_TARGET}
+    PROPERTIES COMPILE_FLAGS "-DDMFF_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}"
+    LINK_FLAGS "${EXTRA_COMPILE_FLAGS}")
+TARGET_LINK_LIBRARIES(${SHARED_DMFF_TARGET} OpenMM)
+TARGET_LINK_LIBRARIES(${SHARED_DMFF_TARGET} tensorflow_cc)
+
+INSTALL_TARGETS(/lib RUNTIME_DIRECTORY /lib ${SHARED_DMFF_TARGET})
+
+# install headers
+FILE(GLOB API_ONLY_INCLUDE_FILES "openmmapi/include/*.h")
+INSTALL (FILES ${API_ONLY_INCLUDE_FILES} DESTINATION include)
+FILE(GLOB API_ONLY_INCLUDE_FILES_INTERNAL "openmmapi/include/internal/*.h")
+INSTALL (FILES ${API_ONLY_INCLUDE_FILES_INTERNAL} DESTINATION include/internal)
+
+# Enable testing
+# Set no testing here.
+ENABLE_TESTING()
+ADD_SUBDIRECTORY(serialization/tests)
+
+# Use reference platform default.
+ADD_SUBDIRECTORY(platforms/reference)
+
+FIND_PACKAGE(CUDA QUIET)
+IF(CUDA_FOUND)
+    message(STATUS "CUDA found, building CUDA implementation")
+    SET(PLUGIN_BUILD_CUDA_LIB ON CACHE BOOL "Build implementation for CUDA: ON")
+ELSE(CUDA_FOUND)
+    message(STATUS "CUDA not found, not building CUDA implementation")
+    SET(PLUGIN_BUILD_CUDA_LIB OFF CACHE BOOL "Build implementation for CUDA: OFF")
+ENDIF(CUDA_FOUND)
+IF(PLUGIN_BUILD_CUDA_LIB)
+    ADD_SUBDIRECTORY(platforms/cuda)
+ENDIF(PLUGIN_BUILD_CUDA_LIB)
+
+
+# Build the Python API
+FIND_PROGRAM(PYTHON_EXECUTABLE python)
+FIND_PROGRAM(SWIG_EXECUTABLE swig)
+IF(PYTHON_EXECUTABLE AND SWIG_EXECUTABLE)
+    message(STATUS "Python and SWIG found, building Python wrappers")
+    SET(PLUGIN_BUILD_PYTHON_WRAPPERS ON CACHE BOOL "Build wrappers for Python")
+ELSE(PYTHON_EXECUTABLE AND SWIG_EXECUTABLE)
+    message(FATAL_ERROR "Python or SWIG not found, not building Python wrappers")
+    SET(PLUGIN_BUILD_PYTHON_WRAPPERS OFF CACHE BOOL "Build wrappers for Python")
+ENDIF(PYTHON_EXECUTABLE AND SWIG_EXECUTABLE)
+IF(PLUGIN_BUILD_PYTHON_WRAPPERS)
+    ADD_SUBDIRECTORY(python)
+ENDIF(PLUGIN_BUILD_PYTHON_WRAPPERS)
+

backend/openmm_dmff_plugin/README.md

@@ -0,0 +1,75 @@
+# OpenMM Plugin for DMFF
+
+
+This is a plugin for [OpenMM](http://openmm.org) that used the trained JAX model by [DMFF](https://github.com/deepmodeling/DMFF) as an independent Force class for dynamics.
+To use it, you need to save you DMFF model with the script in `DMFF/backend/save_dmff2tf.py`.
+
+## Installation
+
+### Create environment with conda
+Install the python, openmm and cudatoolkit.
+```shell
+
+mkdir omm_dmff_working_dir && cd omm_dmff_working_dir
+conda create -n dmff_omm -c conda-forge python=3.9 openmm cudatoolkit=11.6
+conda activate dmff_omm
+```
+### Download `libtensorflow_cc` and install `cppflow` package
+Install the precompiled libtensorflow_cc library from deepmodeling channel.
+```shell
+
+conda install -c deepmodeling libtensorflow_cc=2.9.1=cuda112h02da4e0_0
+```
+Download the tensorflow sources file. Copy the `c` direcotry in source code to installed header files of tensorflow library, since it's needed by package `cppflow`.
+```shell
+
+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/
+```
+Download `cppflow` and move the headers library to environment path.
+```shell
+
+git clone https://github.com/serizba/cppflow.git
+cd cppflow
+git apply DMFF/backend/openmm_dmff_plugin/tests/cppflow_empty_constructor.patch
+mkdir ${CONDA_PREFIX}/include/cppflow
+cp -r include/cppflow ${CONDA_PREFIX}/include/
+```
+
+### Install the OpenMM DMFF plugin from source 
+
+Compile the plugin from source with following steps.
+1. Set up environment variables.
+   ```shell
+   export OPENMM_INSTALLED_DIR=$CONDA_PREFIX
+   export CPPFLOW_INSTALLED_DIR=$CONDA_PREFIX
+   export LIBTENSORFLOW_INSTALLED_DIR=$CONDA_PREFIX
+   cd DMFF/backend/openmm_dmff_plugin/
+   mkdir build && cd build
+   ```
+
+2. Run `cmake` command with required parameters.
+   ```shell
+   cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR}
+   make && make install
+   make PythonInstall
+   ```
+
+3. Test the plugin in Python interface, reference platform.
+   ```shell
+   python -m OpenMMDMFFPlugin.tests.test_dmff_plugin_nve -n 100
+   python -m OpenMMDMFFPlugin.tests.test_dmff_plugin_nvt -n 100 --platform CUDA
+   ```
+## Usage
+Add the following lines to your Python script to use the plugin.
+More details can refer to the script in `python/OpenMMDMFFPlugin/tests/test_dmff_plugin_nve.py`.
+
+```python
+
+from OpenMMDMFFPlugin import DMFFModel
+# Set up the dmff_system with the dmff_model.    
+dmff_model = DMFFModel(dp_model)
+dmff_model.setUnitTransformCoefficients(1, 1, 1)
+dmff_system = dmff_model.createSystem(topology)
+```