add dpdata driver (#3174)

Add a dpdata driver via the plugin mechanism (override that in the
dpdata package) so it can benefit from the multiple-backend DeepPot.
Currently, the driver in the dpdata package has to support both v1 and
v2 for backward compatibility. When shipped within the deepmd-kit
package, it only needs to support the current deepmd-kit version.

---------

Signed-off-by: Jinzhe Zeng <[email protected]>

README.md

@@ -134,6 +134,7 @@ A full [document](doc/train/train-input-auto.rst) on options in the training inp
     - [C++ interface](doc/inference/cxx.md)
     - [Node.js interface](doc/inference/nodejs.md)
 - [Integrate with third-party packages](doc/third-party/index.rst)
+    - [Use deep potential with dpdata](doc/third-party/dpdata.md)
     - [Use deep potential with ASE](doc/third-party/ase.md)
     - [Run MD with LAMMPS](doc/third-party/lammps-command.md)
     - [Run path-integral MD with i-PI](doc/third-party/ipi.md)

backend/dynamic_metadata.py

@@ -33,7 +33,7 @@ def dynamic_metadata(
     elif field == "optional-dependencies":
         return {
             "test": [
-                "dpdata>=0.1.9",
+                "dpdata>=0.2.7",
                 "ase",
                 "pytest",
                 "pytest-cov",

deepmd_utils/driver.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""dpdata driver."""
+# Derived from https://github.com/deepmodeling/dpdata/blob/18a0ed5ebced8b1f6887038883d46f31ae9990a4/dpdata/plugins/deepmd.py#L361-L443
+# under LGPL-3.0-or-later license.
+# The original deepmd driver maintained in the dpdata package will be overriden.
+# The class in the dpdata package needs to handle different situations for v1 and v2 interface,
+# which is too complex with the development of deepmd-kit.
+# So, it will be a good idea to ship it with DeePMD-kit itself.
+import dpdata
+from dpdata.utils import (
+    sort_atom_names,
+)
+
+
+@dpdata.driver.Driver.register("dp")
+@dpdata.driver.Driver.register("deepmd")
+@dpdata.driver.Driver.register("deepmd-kit")
+class DPDriver(dpdata.driver.Driver):
+    """DeePMD-kit driver.
+
+    Parameters
+    ----------
+    dp : deepmd.DeepPot or str
+        The deepmd-kit potential class or the filename of the model.
+
+    Examples
+    --------
+    >>> DPDriver("frozen_model.pb")
+    """
+
+    def __init__(self, dp: str) -> None:
+        from deepmd_utils.infer.deep_pot import (
+            DeepPot,
+        )
+
+        if not isinstance(dp, DeepPot):
+            self.dp = DeepPot(dp, auto_batch_size=True)
+        else:
+            self.dp = dp
+
+    def label(self, data: dict) -> dict:
+        """Label a system data by deepmd-kit. Returns new data with energy, forces, and virials.
+
+        Parameters
+        ----------
+        data : dict
+            data with coordinates and atom types
+
+        Returns
+        -------
+        dict
+            labeled data with energies and forces
+        """
+        nframes = data["coords"].shape[0]
+        natoms = data["coords"].shape[1]
+        type_map = self.dp.get_type_map()
+        # important: dpdata type_map may not be the same as the model type_map
+        # note: while we want to change the type_map when feeding to DeepPot,
+        # we don't want to change the type_map in the returned data
+        sorted_data = sort_atom_names(data.copy(), type_map=type_map)
+        atype = sorted_data["atom_types"]
+
+        coord = data["coords"].reshape((nframes, natoms * 3))
+        if "nopbc" not in data:
+            cell = data["cells"].reshape((nframes, 9))
+        else:
+            cell = None
+        e, f, v = self.dp.eval(coord, cell, atype)
+        data = data.copy()
+        data["energies"] = e.reshape((nframes,))
+        data["forces"] = f.reshape((nframes, natoms, 3))
+        data["virials"] = v.reshape((nframes, 3, 3))
+        return data

doc/third-party/dpdata.md

@@ -0,0 +1,12 @@
+# Use deep potential with dpdata
+
+DeePMD-kit provides a driver for [dpdata](https://github.com/deepmodeling/dpdata) >=0.2.7 via the plugin mechanism, making it possible to call the `predict` method for `System` class:
+
+```py
+import dpdata
+
+dsys = dpdata.LabeledSystem("OUTCAR")
+dp_sys = dsys.predict("frozen_model_compressed.pb", driver="dp")
+```
+
+By inferring with the DP model `frozen_model_compressed.pb`, dpdata will generate a new labeled system `dp_sys` with inferred energies, forces, and virials.

doc/third-party/index.md

@@ -2,6 +2,7 @@
 
 Note that the model for inference is required to be compatible with the DeePMD-kit package. See [Model compatibility](../troubleshooting/model-compatability.html) for details.
 
+- [Use deep potential with dpdata](dpdata.md)
 - [Use deep potential with ASE](ase.md)
 - [Run MD with LAMMPS](lammps-command.md)
 - [Run path-integral MD with i-PI](ipi.md)

doc/third-party/index.rst

@@ -6,6 +6,7 @@ Note that the model for inference is required to be compatible with the DeePMD-k
 .. toctree::
    :maxdepth: 1
 
+   dpdata
    ase
    lammps-command
    ipi

doc/third-party/out-of-deepmd-kit.md

@@ -2,19 +2,6 @@
 
 The codes of the following interfaces are not a part of the DeePMD-kit package and maintained by other repositories. We list these interfaces here for user convenience.
 
-## dpdata
-
-[dpdata](https://github.com/deepmodeling/dpdata) provides the `predict` method for `System` class:
-
-```py
-import dpdata
-
-dsys = dpdata.LabeledSystem("OUTCAR")
-dp_sys = dsys.predict("frozen_model_compressed.pb")
-```
-
-By inferring with the DP model `frozen_model_compressed.pb`, dpdata will generate a new labeled system `dp_sys` with inferred energies, forces, and virials.
-
 ## OpenMM plugin for DeePMD-kit
 
 An [OpenMM](https://github.com/openmm/openmm) plugin is provided from [JingHuangLab/openmm_deepmd_plugin](https://github.com/JingHuangLab/openmm_deepmd_plugin), written by the [Huang Lab](http://www.compbiophysics.org/) at Westlake University.

pyproject.toml

@@ -55,6 +55,9 @@ deepmd = "deepmd.lmp:get_op_dir"
 [project.entry-points."dpgui"]
 "DeePMD-kit" = "deepmd_utils.utils.argcheck:gen_args"
 
+[project.entry-points."dpdata.plugins"]
+deepmd_driver = "deepmd_utils.driver:DPDriver"
+
 [project.urls]
 Homepage = "https://github.com/deepmodeling/deepmd-kit"
 documentation = "https://docs.deepmodeling.com/projects/deepmd"

source/tests/test_deeppot_a.py

@@ -4,6 +4,7 @@
 import unittest
 
 import ase.neighborlist
+import dpdata
 import numpy as np
 from common import (
     run_dp,
@@ -518,6 +519,32 @@ def test_2frame_atm(self):
         expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis=1)
         np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places)
 
+    def test_dpdata_driver(self):
+        nframes = 1
+        system = dpdata.System(
+            data={
+                "coords": self.coords.reshape((nframes, -1, 3)),
+                "cells": np.zeros((nframes, 3, 3)),
+                "atom_types": np.array(self.atype),
+                "orig": np.zeros((3,)),
+                "atom_names": ["O", "H"],
+                "atom_numbs": [2, 4],
+                "nopbc": True,
+            }
+        )
+        system_predicted = system.predict(self.dp, driver="dp")
+        np.testing.assert_almost_equal(
+            system_predicted["forces"].ravel(), self.expected_f.ravel(), default_places
+        )
+        expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1)
+        np.testing.assert_almost_equal(
+            system_predicted["energies"].ravel(), expected_se.ravel(), default_places
+        )
+        expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1)
+        np.testing.assert_almost_equal(
+            system_predicted["virials"].ravel(), expected_sv.ravel(), default_places
+        )
+
 
 class TestDeepPotALargeBoxNoPBC(unittest.TestCase):
     @classmethod
@@ -716,6 +743,31 @@ def test_ase(self):
         expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1)
         np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places)
 
+    def test_dpdata_driver(self):
+        nframes = 1
+        system = dpdata.System(
+            data={
+                "coords": self.coords.reshape((nframes, -1, 3)),
+                "cells": self.box.reshape((nframes, 3, 3)),
+                "atom_types": np.array(self.atype),
+                "orig": np.zeros((3,)),
+                "atom_names": ["O", "H"],
+                "atom_numbs": [2, 4],
+            }
+        )
+        system_predicted = system.predict("deeppot.pb", driver="dp")
+        np.testing.assert_almost_equal(
+            system_predicted["forces"].ravel(), self.expected_f.ravel(), default_places
+        )
+        expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis=1)
+        np.testing.assert_almost_equal(
+            system_predicted["energies"].ravel(), expected_se.ravel(), default_places
+        )
+        expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis=1)
+        np.testing.assert_almost_equal(
+            system_predicted["virials"].ravel(), expected_sv.ravel(), default_places
+        )
+
 
 class TestModelConvert(unittest.TestCase):
     def setUp(self):