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 <[email protected]>
deepmodeling · Nov 9, 2023 · 348fa6a · 348fa6a
1 parent e2cbf0a
commit 348fa6a
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 66 deletions.
diff --git 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
@@ -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)