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)