diff --git a/deepmd/infer/model_devi.py b/deepmd/infer/model_devi.py index bc80ac78f6..0274384188 100644 --- a/deepmd/infer/model_devi.py +++ b/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) diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 24042444c8..8442f84156 100644 --- a/deepmd/utils/data.py +++ b/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 diff --git a/source/tests/common.py b/source/tests/common.py index e1788bb942..12723e52fc 100644 --- a/source/tests/common.py +++ b/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) diff --git a/source/tests/test_model_devi.py b/source/tests/test_model_devi.py index 3ce594bb8c..b1c5ec8ead 100644 --- a/source/tests/test_model_devi.py +++ b/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()