From 4e112335a6c4d09557c8220bc3dd6242e2cf595c Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 24 Jan 2024 10:11:11 -0500 Subject: [PATCH] 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 --- README.md | 1 + backend/dynamic_metadata.py | 2 +- deepmd_utils/driver.py | 73 ++++++++++++++++++++++++++++ doc/third-party/dpdata.md | 12 +++++ doc/third-party/index.md | 1 + doc/third-party/index.rst | 1 + doc/third-party/out-of-deepmd-kit.md | 13 ----- pyproject.toml | 3 ++ source/tests/test_deeppot_a.py | 52 ++++++++++++++++++++ 9 files changed, 144 insertions(+), 14 deletions(-) create mode 100644 deepmd_utils/driver.py create mode 100644 doc/third-party/dpdata.md diff --git a/README.md b/README.md index 81fdead098..27c8dab4bc 100644 --- a/README.md +++ b/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) diff --git a/backend/dynamic_metadata.py b/backend/dynamic_metadata.py index ab955c3cf8..210c04235e 100644 --- a/backend/dynamic_metadata.py +++ b/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", diff --git a/deepmd_utils/driver.py b/deepmd_utils/driver.py new file mode 100644 index 0000000000..b9e70f15e4 --- /dev/null +++ b/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 diff --git a/doc/third-party/dpdata.md b/doc/third-party/dpdata.md new file mode 100644 index 0000000000..05e0f6fb40 --- /dev/null +++ b/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. diff --git a/doc/third-party/index.md b/doc/third-party/index.md index 235337974c..419f1fbb5c 100644 --- a/doc/third-party/index.md +++ b/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) diff --git a/doc/third-party/index.rst b/doc/third-party/index.rst index f88a477fc7..cd0726a4bb 100644 --- a/doc/third-party/index.rst +++ b/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 diff --git a/doc/third-party/out-of-deepmd-kit.md b/doc/third-party/out-of-deepmd-kit.md index 71dc9adb23..2ebed6fb46 100644 --- a/doc/third-party/out-of-deepmd-kit.md +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index 550fbc4b54..e5515984fb 100644 --- a/pyproject.toml +++ b/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" diff --git a/source/tests/test_deeppot_a.py b/source/tests/test_deeppot_a.py index c229b4302c..ba3655b7f9 100644 --- a/source/tests/test_deeppot_a.py +++ b/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):