From 348fa6a56d2aad5df96e1115e707f0fe12b29104 Mon Sep 17 00:00:00 2001 From: Ye Ding Date: Thu, 9 Nov 2023 11:46:20 +0800 Subject: [PATCH] Update the doc for openmm dmff plugin usage. (#146) * Add issue templates for feature request and bug-report * Add script for dmff model saving. * Remove issue template from devel branch. * debug workflow * remove debug * Update ut.yml. Install mdtraj by conda. * Add openmm dmff plugin to backend * Update .gitignore * Add test workflow for openmm dmff plugin * Update workflow for openmm dmff plugin test. * Update workflow for openmm dmff plugin test. * Update workflow for openmm dmff plugin test. * Update workflow for openmm dmff plugin test. * Update workflow for openmm dmff plugin test. * Specific setuptools version in workflow yaml * Update tensorflow version * Enable double precision in test omm dmff plugin workflow * update nve assert checking in omm dmff plugin workflow * Fix no axis type bug (#120) * add rules of local axis for NoAxisType * add test for NoAxisType * Update omm plugin and save_dmff2tf.py for aux compatible. * Fix test issue. * Update test_openmm_dmff_plugin.yml * Update the doc for openmm dmff plugin usage * Revision on the doc of DMFF plugin * Upload tf c directory. * Remove tf2.9.1_c.tar --------- Co-authored-by: plumbum082 <54125554+plumbum082@users.noreply.github.com> --- backend/openmm_dmff_plugin/README.md | 13 ++-- docs/user_guide/4.7OpenMMplugin.md | 110 +++++++++++++-------------- 2 files changed, 57 insertions(+), 66 deletions(-) diff --git a/backend/openmm_dmff_plugin/README.md b/backend/openmm_dmff_plugin/README.md index 07a8032e0..c059c5eb4 100644 --- a/backend/openmm_dmff_plugin/README.md +++ b/backend/openmm_dmff_plugin/README.md @@ -2,7 +2,7 @@ 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`. +To use it, you need to save your DMFF model with the script in `DMFF/backend/save_dmff2tf.py`. ## Installation @@ -15,10 +15,9 @@ 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. +Install the precompiled libtensorflow_cc library from conda. ```shell - -conda install -c deepmodeling libtensorflow_cc=2.9.1=cuda112h02da4e0_0 +conda install -y libtensorflow_cc=2.9.1 -c conda-forge ``` 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 @@ -37,9 +36,9 @@ mkdir ${CONDA_PREFIX}/include/cppflow cp -r include/cppflow ${CONDA_PREFIX}/include/ ``` -### Install the OpenMM DMFF plugin from source +### Install the OpenMM DMFF plugin from the source -Compile the plugin from source with following steps. +Compile the plugin from the source with the following steps. 1. Set up environment variables. ```shell export OPENMM_INSTALLED_DIR=$CONDA_PREFIX @@ -49,7 +48,7 @@ Compile the plugin from source with following steps. mkdir build && cd build ``` -2. Run `cmake` command with required parameters. +2. Run `cmake` command with the required parameters. ```shell cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR} make && make install diff --git a/docs/user_guide/4.7OpenMMplugin.md b/docs/user_guide/4.7OpenMMplugin.md index 07a8032e0..14c2d5192 100644 --- a/docs/user_guide/4.7OpenMMplugin.md +++ b/docs/user_guide/4.7OpenMMplugin.md @@ -1,75 +1,67 @@ -# OpenMM Plugin for DMFF +# OpenMM DMFF Plugin +## 1. Theory -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`. +OpenMM DMFF plugin was developed for [OpenMM](http://openmm.org) to incorporate the trained JAX model from [DMFF](https://github.com/deepmodeling/DMFF) as an independent Force class for molecular dynamics simulations. +To utilize this plugin, you need to save your DMFF model using the `DMFF/backend/save_dmff2tf.py` script. +The `save_dmff2tf.py` script converts the DMFF model to a TensorFlow module using the experimental feature of JAX called [`jax2tf`](https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md). +The integration of the saved TensorFlow module with the DMFF plugin is accomplished using [cppflow](https://github.com/serizba/cppflow) and the OpenMM C++ interface. +To install the DMFF plugin, please refer to the instructions provided in the [Readme](backend/openmm_dmff_plugin/README.md). -## Installation -### Create environment with conda -Install the python, openmm and cudatoolkit. -```shell +## 2. Save DMFF model with script `save_dmff2tf.py` -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/ +The `save_dmff2tf.py` script requires at least two input files to run: +- input_pdb: A .pdb file containing the initial coordinates and box information. +- xml_files: The force field parameters optimized by DMFF. +Additionally, the `output` argument is required to specify the directory for saving the DMFF model. +The saved model produces `energy` and `forces` tensors. +Below is an example usage of the script: +```python +# For saving classical DMFF model. +python backend/save_dmff2tf.py --input_pdb examples/classical/lig.pdb + --xml_files examples/classical/lig-prm.xml + --output /tmp/dmff_classical_lig ``` -Download `cppflow` and move the headers library to environment path. -```shell +If additional .xml files are needed for topology definition in input files, the `--bond_definitions_xml` argument can be used to specify the path to the .xml file defining bond connections. -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/ +Please note that when saving the DMFF model that uses the `ADMPPmeForce` module, the `--has_aux=True` argument is required. +In this case, the saved model produces `energy`, `forces`, and `U_ind` tensors. +Below is an example usage of the script: +```python +# For saving ADMP DMFF model. +python backend/save_dmff2tf.py --input_pdb examples/water_fullpol/water_dimer.pdb + --xml_files examples/water_fullpol/forcefield.xml + --bond_definitions_xml examples/water_fullpol/residues.xml + --output /tmp/dmff_admp_water_dimer ``` -### 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`. +## 3. Class `DMFFModel` +The OpenMM DMFF Plugin can be easily used through the `DMFFModel` class. +This class provides a convenient way to work with the plugin. +See the example code below: ```python - +# Create an OpenMM System object (dmff_system) with saved DMFF model. from OpenMMDMFFPlugin import DMFFModel -# Set up the dmff_system with the dmff_model. -dmff_model = DMFFModel(dp_model) +dmff_model = DMFFModel(dp_model_dir) dmff_model.setUnitTransformCoefficients(1, 1, 1) dmff_system = dmff_model.createSystem(topology) ``` + +The `DMFFModel` class includes three methods: + +Function `setUnitTransformCoefficients(coordCoefficient, forceCoefficient, energyCoefficient)`: +- In the OpenMM context, the units for coordinates, energy and forces are constrained to nm, kJ/mol and kJ/(mol * nm) respectively. +- You need to provide three coefficients if the saved DMFF model requires different input and output units for coordinates, energy, and forces. + +Function `setHasAux(has_aux=False)`: +- Use this method when running simulations with the `ADMPPmeForce` module in the saved DMFF model. + +Function `createSystem(topology)`: +- This method constructs an OpenMM System object, with the OpenMM Topology object being the input. + +Once the OpenMM System object is constructed, +further simulation settings with the DMFF model are similar to those with other force fields in OpenMM. +A complete script for simulation with this plugin can be found in [test_dmff_plugin_nve.py](/backend/openmm_dmff_plugin//python/OpenMMDMFFPlugin/tests/test_dmff_plugin_nve.py)