[Feature Request] MD Engine Support (OpenMM) for trained model with DMFF

[dingye18]

Summary

Implement the communication of OpenMM and DMFF trained model in c++ interface.
Enhance the production running efficiency with popular MD engine.

Motivation

Instead of inheriting from the existing interaction forms in MD engine support force fields, implementation of the MD Engine support for calling DMFF models directly have three advantages.

• Enhanced the flexibility of the DMFF function forms in training. No need to consider the limitations of the MD engine supported force fields forms.
• Enhanced the simulation efficiency with the neighbor list informations that fetched from MD engine.
• Simplifies the use threshold for users after using DMFF to train the model.

Suggested Solutions

Refer to openmm_deepmd_plugin for openmm_dmff_plugin architecture.
I will create an openmm_dmff_plugin repo with my own account first.
After the implementation of openmm_dmff_plugin, it will be transfer to deepmodeling account and added into the DMFF as an submodule.

The main problem is the neighbor list information fetching from OpenMM in CUDA platform.
Reference platform called an serial function with O(n) complexity to construct the neighbor list, which is easy to be called in plugin.
However, CUDA platform split the atoms into different block region and construct the interacted neighor blocks parallelized.
How to fetch the information of neighbor list from CUDA context without loss of accuracy is required to be solved rigorously.

Any suggestions and comments on this solution are welcomed !

Further Information, Files, and Links

OpenMM v8.0.0 made some changes for the compatible of openmm-tensorflow/openmm-torch. If these changes are suitable to the development of openmm_dmff_plugin is unknown to me until now.

[dingye18]

A simple version of OpenMM DMFF plugin is implemented.
Both reference and CUDA platforms are included.
But the CUDA platform was not tested. (Not recommended for now).
Reference platform has been tested.

To test the CUDA platform, a frozen JAX graph with GPU support is needed.
Besides, the numerical type for coordinates, box size, forces and energy should be double in the frozen JAX graph for the compatibility of OpenMM context.

[dingye18]

Test in DMFF CUDA plugin is passed. How to fetch the neighbor list information from openmm cuda context is still a problem.

[WangXinyan940]

Thanks to @dingye18 , the OpenMM-DMFF plugin has been implemented as one of the most important features in our 1.0.0 version. This issue would be happily closed.