support fparam/aparam in dp model-devi (#2665)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

deepmd/infer/model_devi.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
+    Optional,
     Tuple,
 )
 
@@ -145,6 +146,8 @@ def calc_model_devi(
     fname=None,
     frequency=1,
     mixed_type=False,
+    fparam: Optional[np.ndarray] = None,
+    aparam: Optional[np.ndarray] = None,
 ):
     """Python interface to calculate model deviation.
 
@@ -164,6 +167,10 @@ def calc_model_devi(
         Steps between frames (if the system is given by molecular dynamics engine), default 1
     mixed_type : bool
         Whether the input atype is in mixed_type format or not
+    fparam : numpy.ndarray
+        frame specific parameters
+    aparam : numpy.ndarray
+        atomic specific parameters
 
     Returns
     -------
@@ -191,6 +198,8 @@ def calc_model_devi(
             coord,
             box,
             atype,
+            fparam=fparam,
+            aparam=aparam,
             mixed_type=mixed_type,
         )
         energies.append(ret[0] / natom)
@@ -248,9 +257,28 @@ def make_model_devi(
     if len(all_sys) == 0:
         raise RuntimeError("Did not find valid system")
     devis_coll = []
+
+    first_dp = dp_models[0]
+
     for system in all_sys:
         # create data-system
         dp_data = DeepmdData(system, set_prefix, shuffle_test=False, type_map=tmap)
+        if first_dp.get_dim_fparam() > 0:
+            dp_data.add(
+                "fparam",
+                first_dp.get_dim_fparam(),
+                atomic=False,
+                must=True,
+                high_prec=False,
+            )
+        if first_dp.get_dim_aparam() > 0:
+            dp_data.add(
+                "aparam",
+                first_dp.get_dim_aparam(),
+                atomic=True,
+                must=True,
+                high_prec=False,
+            )
         mixed_type = dp_data.mixed_type
 
         data_sets = [dp_data._load_set(set_name) for set_name in dp_data.dirs]
@@ -265,7 +293,23 @@ def make_model_devi(
                 atype = data["type"][0]
             if not dp_data.pbc:
                 box = None
-            devi = calc_model_devi(coord, box, atype, dp_models, mixed_type=mixed_type)
+            if first_dp.get_dim_fparam() > 0:
+                fparam = data["fparam"]
+            else:
+                fparam = None
+            if first_dp.get_dim_aparam() > 0:
+                aparam = data["aparam"]
+            else:
+                aparam = None
+            devi = calc_model_devi(
+                coord,
+                box,
+                atype,
+                dp_models,
+                mixed_type=mixed_type,
+                fparam=fparam,
+                aparam=aparam,
+            )
             nframes_tot += coord.shape[0]
             devis.append(devi)
         devis = np.vstack(devis)

deepmd/utils/data.py

@@ -82,7 +82,7 @@ def __init__(
         self.pbc = self._check_pbc(root)
         # enforce type_map if necessary
         self.enforce_type_map = False
-        if type_map is not None and self.type_map is not None:
+        if type_map is not None and self.type_map is not None and len(type_map):
             if not self.mixed_type:
                 atom_type_ = [
                     type_map.index(self.type_map[ii]) for ii in self.atom_type

source/tests/common.py

@@ -41,8 +41,8 @@ def del_data():
         shutil.rmtree("system_mixed_type")
 
 
-def gen_data_type_specific(nframes=1):
-    tmpdata = Data(rand_pert=0.1, seed=1, nframes=nframes)
+def gen_data_type_specific(nframes=1, dim_fparam=2):
+    tmpdata = Data(rand_pert=0.1, seed=1, nframes=nframes, dim_fparam=dim_fparam)
     sys = dpdata.LabeledSystem()
     sys.data["atom_names"] = ["foo", "bar"]
     sys.data["coords"] = tmpdata.coord
@@ -56,11 +56,14 @@ def gen_data_type_specific(nframes=1):
     sys.data["forces"] = np.zeros([nframes, natoms, 3])
     sys.to_deepmd_npy("system", prec=np.float64)
     np.save("system/set.000/fparam.npy", tmpdata.fparam)
-    np.save("system/set.000/aparam.npy", tmpdata.aparam.reshape([nframes, natoms, 2]))
+    np.save(
+        "system/set.000/aparam.npy",
+        tmpdata.aparam.reshape([nframes, natoms, dim_fparam]),
+    )
 
 
-def gen_data_mixed_type(nframes=1):
-    tmpdata = Data(rand_pert=0.1, seed=1, nframes=nframes)
+def gen_data_mixed_type(nframes=1, dim_fparam=2):
+    tmpdata = Data(rand_pert=0.1, seed=1, nframes=nframes, dim_fparam=dim_fparam)
     sys = dpdata.LabeledSystem()
     real_type_map = ["foo", "bar"]
     sys.data["atom_names"] = ["X"]
@@ -82,12 +85,12 @@ def gen_data_mixed_type(nframes=1):
     np.save("system_mixed_type/set.000/fparam.npy", tmpdata.fparam)
     np.save(
         "system_mixed_type/set.000/aparam.npy",
-        tmpdata.aparam.reshape([nframes, natoms, 2]),
+        tmpdata.aparam.reshape([nframes, natoms, dim_fparam]),
     )
 
 
-def gen_data_virtual_type(nframes=1, nghost=4):
-    tmpdata = Data(rand_pert=0.1, seed=1, nframes=nframes)
+def gen_data_virtual_type(nframes=1, nghost=4, dim_fparam=2):
+    tmpdata = Data(rand_pert=0.1, seed=1, nframes=nframes, dim_fparam=dim_fparam)
     sys = dpdata.LabeledSystem()
     real_type_map = ["foo", "bar"]
     sys.data["atom_names"] = ["X"]
@@ -129,25 +132,25 @@ def gen_data_virtual_type(nframes=1, nghost=4):
         "system_mixed_type/set.000/aparam.npy",
         np.concatenate(
             [
-                tmpdata.aparam.reshape([nframes, natoms, 2]),
-                np.zeros([nframes, nghost, 2]),
+                tmpdata.aparam.reshape([nframes, natoms, dim_fparam]),
+                np.zeros([nframes, nghost, dim_fparam]),
             ],
             axis=1,
         ),
     )
 
 
-def gen_data(nframes=1, mixed_type=False, virtual_type=False):
+def gen_data(nframes=1, mixed_type=False, virtual_type=False, dim_fparam=2):
     if not mixed_type:
-        gen_data_type_specific(nframes)
+        gen_data_type_specific(nframes, dim_fparam=dim_fparam)
     elif virtual_type:
-        gen_data_virtual_type(nframes)
+        gen_data_virtual_type(nframes, dim_fparam=dim_fparam)
     else:
-        gen_data_mixed_type(nframes)
+        gen_data_mixed_type(nframes, dim_fparam=dim_fparam)
 
 
 class Data:
-    def __init__(self, rand_pert=0.1, seed=1, box_scale=20, nframes=1):
+    def __init__(self, rand_pert=0.1, seed=1, box_scale=20, nframes=1, dim_fparam=2):
         coord = [
             [0.0, 0.0, 0.1],
             [1.1, 0.0, 0.1],
@@ -161,7 +164,7 @@ def __init__(self, rand_pert=0.1, seed=1, box_scale=20, nframes=1):
         self.coord = self._copy_nframes(self.coord)
         dp_random.seed(seed)
         self.coord += rand_pert * dp_random.random(self.coord.shape)
-        self.fparam = np.array([[0.1, 0.2]])
+        self.fparam = ((np.arange(dim_fparam) + 1) * 0.1).reshape(1, dim_fparam)
         self.aparam = np.tile(self.fparam, [1, 6])
         self.fparam = self._copy_nframes(self.fparam)
         self.aparam = self._copy_nframes(self.aparam)

source/tests/test_model_devi.py

@@ -91,3 +91,75 @@ def tearDown(self):
             os.remove(pb)
         os.remove(self.output)
         del_data()
+
+
+class TestMakeModelDeviFparamAparam(unittest.TestCase):
+    """Ensure dp model_devi accepts fparam and aparam."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.pbtxts = [
+            os.path.join(tests_path, "infer/fparam_aparam.pbtxt"),
+        ]
+        cls.graph_dirs = [pbtxt.replace("pbtxt", "pb") for pbtxt in cls.pbtxts]
+        for pbtxt, pb in zip(cls.pbtxts, cls.graph_dirs):
+            convert_pbtxt_to_pb(pbtxt, pb)
+        cls.graphs = [DeepPotential(pb) for pb in cls.graph_dirs]
+
+    @classmethod
+    def tearDownClass(cls):
+        for pb in cls.graph_dirs:
+            os.remove(pb)
+        cls.graphs = None
+
+    def setUp(self):
+        gen_data(dim_fparam=1)
+        self.data_dir = "system"
+        coord = np.load(os.path.join(self.data_dir, "set.000/coord.npy"))
+        box = np.load(os.path.join(self.data_dir, "set.000/box.npy"))
+        atype_ = np.loadtxt(os.path.join(self.data_dir, "type.raw"))
+        self.atype = np.zeros_like(atype_)
+        np.savetxt(os.path.join(self.data_dir, "type.raw"), self.atype)
+        self.coord = np.vstack([coord, coord])
+        self.box = np.vstack([box, box])
+        self.freq = 10
+
+        self.output = os.path.join(tests_path, "model_devi.out")
+        self.expect = np.zeros(8)
+        nframes = self.box.size // 9
+        self.fparam = np.repeat([0.25852028], nframes).reshape((nframes, 1))
+        self.aparam = np.repeat(self.fparam, self.atype.size).reshape(
+            (nframes, self.atype.size, 1)
+        )
+
+    def test_calc_model_devi(self):
+        model_devi = calc_model_devi(
+            self.coord,
+            None,
+            self.atype,
+            self.graphs,
+            frequency=self.freq,
+            fname=self.output,
+            fparam=self.fparam,
+            aparam=self.aparam,
+        )
+        self.assertAlmostEqual(model_devi[0][0], 0)
+        self.assertAlmostEqual(model_devi[1][0], self.freq)
+        np.testing.assert_almost_equal(model_devi[0][1:8], self.expect[1:8], 6)
+        np.testing.assert_almost_equal(model_devi[0][1:8], model_devi[1][1:8], 6)
+        self.assertTrue(os.path.isfile(self.output))
+
+    def test_make_model_devi(self):
+        make_model_devi(
+            models=self.graph_dirs,
+            system=self.data_dir,
+            set_prefix="set",
+            output=self.output,
+            frequency=self.freq,
+        )
+        x = np.loadtxt(self.output)
+        np.testing.assert_allclose(x, self.expect, 6)
+
+    def tearDown(self):
+        os.remove(self.output)
+        del_data()