OpenMM DMFF plugin was developed for OpenMM to incorporate the trained JAX model from 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
.
The integration of the saved TensorFlow module with the DMFF plugin is accomplished using cppflow and the OpenMM C++ interface.
To install the DMFF plugin, please refer to the instructions provided in the Readme.
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 producesenergy
andforces
tensors. Below is an example usage of the script:
# 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
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.
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:
# 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
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:
# Create an OpenMM System object (dmff_system) with saved DMFF model.
from OpenMMDMFFPlugin import DMFFModel
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