diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index ab0ba46042..aa8b88145c 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -555,6 +555,8 @@ def _concat_type_embedding( def _filter_lower( self, + type_i, + type_input, start_index, incrs_index, inputs, @@ -638,9 +640,11 @@ def _filter( outputs_size_2 = self.n_axis_neuron with tf.variable_scope(name, reuse=reuse): start_index = 0 + type_i = 0 if type_embedding is None: for type_i in range(self.ntypes): ret = self._filter_lower( + type_i, type_input, start_index, self.sel_a[type_i], inputs, nframes, @@ -660,6 +664,7 @@ def _filter( start_index += self.sel_a[type_i] else : xyz_scatter_1 = self._filter_lower( + type_i, type_input, start_index, np.cumsum(self.sel_a)[-1], inputs, nframes, diff --git a/deepmd/entrypoints/compress.py b/deepmd/entrypoints/compress.py index acfeae5dee..2193f0c1f9 100644 --- a/deepmd/entrypoints/compress.py +++ b/deepmd/entrypoints/compress.py @@ -65,8 +65,8 @@ def compress( jdata = j_loader(INPUT) if "model" not in jdata.keys(): jdata = convert_input_v0_v1(jdata, warning=True, dump="input_v1_compat.json") - jdata = normalize(jdata) jdata["model"]["compress"] = {} + jdata["model"]["compress"]["type"] = 'se_e2_a' jdata["model"]["compress"]["compress"] = True jdata["model"]["compress"]["model_file"] = input jdata["model"]["compress"]["table_config"] = [ @@ -75,11 +75,15 @@ def compress( 10 * stride, int(frequency), ] + # be careful here, if one want to refine the model + jdata["training"]["numb_steps"] = jdata["training"]["save_freq"] + jdata = normalize(jdata) + # check the descriptor info of the input file assert ( - jdata["model"]["descriptor"]["type"] == "se_a" - ), "Model compression error: descriptor type must be se_a!" + jdata["model"]["descriptor"]["type"] == "se_a" or jdata["model"]["descriptor"]["type"] == "se_e2_a" + ), "Model compression error: descriptor type must be se_a or se_e2_a!" assert ( jdata["model"]["descriptor"]["resnet_dt"] is False ), "Model compression error: descriptor resnet_dt must be false!" @@ -87,8 +91,6 @@ def compress( # stage 1: training or refining the model with tabulation log.info("\n\n") log.info("stage 1: train or refine the model with tabulation") - # be careful here, if one want to refine the model - jdata["training"]["stop_batch"] = jdata["training"]["save_freq"] control_file = "compress.json" with open(control_file, "w") as fp: json.dump(jdata, fp, indent=4) diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py index 823a77a28c..b6892df3e3 100644 --- a/deepmd/fit/ener.py +++ b/deepmd/fit/ener.py @@ -267,7 +267,7 @@ def _build_lower( 1, activation_fn = None, bavg = bias_atom_e, - name='final_layer_'+suffix, + name='final_layer'+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index cc8274dea4..ceafe89193 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -320,6 +320,29 @@ def modifier_variant_type_args(): optional = False, doc = doc_modifier_type) +# --- model compression configurations: --- # +def model_compression(): + doc_compress = "The name of the frozen model file." + doc_model_file = f"The input model file, which will be compressed by the DeePMD-kit." + doc_table_config = f"The arguments of model compression, including extrapolate(scale of model extrapolation), stride(uniform stride of tabulation's first and second table), and frequency(frequency of tabulation overflow check)." + + return [ + Argument("compress", bool, optional = False, default = True, doc = doc_compress), + Argument("model_file", str, optional = False, default = 'frozen_model.pb', doc = doc_model_file), + Argument("table_config", list, optional = False, default = [5, 0.01, 0.1, -1], doc = doc_table_config), + ] + +# --- model compression configurations: --- # +def model_compression_type_args(): + doc_compress_type = "The type of model compression, which should be consistent with the descriptor type." + + return Variant("type", [ + Argument("se_e2_a", dict, model_compression(), alias = ['se_a']) + ], + optional = True, + default_tag = 'se_e2_a', + doc = doc_compress_type) + def model_args (): doc_type_map = 'A list of strings. Give the name to each type of atoms.' @@ -333,6 +356,7 @@ def model_args (): doc_smin_alpha = 'The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when `use_srtab` is provided.' doc_sw_rmin = 'The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided.' doc_sw_rmax = 'The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided.' + doc_compress_config = 'Model compression configurations' ca = Argument("model", dict, [Argument("type_map", list, optional = True, doc = doc_type_map), @@ -346,6 +370,7 @@ def model_args (): Argument("descriptor", dict, [], [descrpt_variant_type_args()], doc = doc_descrpt), Argument("fitting_net", dict, [], [fitting_variant_type_args()], doc = doc_fitting), Argument("modifier", dict, [], [modifier_variant_type_args()], optional = True, doc = doc_modifier), + Argument("compress", dict, [], [model_compression_type_args()], optional = True, doc = doc_compress_config) ]) # print(ca.gen_doc()) return ca diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index ad13365f55..1590d459dd 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -41,13 +41,17 @@ def __init__(self, self.sub_graph, self.sub_graph_def = self._load_sub_graph() self.sub_sess = tf.Session(graph = self.sub_graph) - self.sel_a = self.graph.get_operation_by_name('DescrptSeA').get_attr('sel_a') + try: + self.sel_a = self.graph.get_operation_by_name('ProdEnvMatA').get_attr('sel_a') + self.descrpt = self.graph.get_operation_by_name ('ProdEnvMatA') + except Exception: + self.sel_a = self.graph.get_operation_by_name('DescrptSeA').get_attr('sel_a') + self.descrpt = self.graph.get_operation_by_name ('DescrptSeA') self.ntypes = self._get_tensor_value(self.graph.get_tensor_by_name ('descrpt_attr/ntypes:0')) - self.davg = self._get_tensor_value(self.graph.get_tensor_by_name ('descrpt_attr/t_avg:0')) self.dstd = self._get_tensor_value(self.graph.get_tensor_by_name ('descrpt_attr/t_std:0')) - self.descrpt = self.graph.get_operation_by_name ('DescrptSeA') + self.rcut = self.descrpt.get_attr('rcut_r') self.rcut_smth = self.descrpt.get_attr('rcut_r_smth') diff --git a/source/tests/model_compression/data/set.000/box.npy b/source/tests/model_compression/data/set.000/box.npy new file mode 100644 index 0000000000..1f5de3f920 Binary files /dev/null and b/source/tests/model_compression/data/set.000/box.npy differ diff --git a/source/tests/model_compression/data/set.000/coord.npy b/source/tests/model_compression/data/set.000/coord.npy new file mode 100644 index 0000000000..3ab4571007 Binary files /dev/null and b/source/tests/model_compression/data/set.000/coord.npy differ diff --git a/source/tests/model_compression/data/set.000/energy.npy b/source/tests/model_compression/data/set.000/energy.npy new file mode 100644 index 0000000000..35899bf0f1 Binary files /dev/null and b/source/tests/model_compression/data/set.000/energy.npy differ diff --git a/source/tests/model_compression/data/set.000/force.npy b/source/tests/model_compression/data/set.000/force.npy new file mode 100644 index 0000000000..a1f696c186 Binary files /dev/null and b/source/tests/model_compression/data/set.000/force.npy differ diff --git a/source/tests/model_compression/data/type.raw b/source/tests/model_compression/data/type.raw new file mode 100644 index 0000000000..97e8fdfcf8 --- /dev/null +++ b/source/tests/model_compression/data/type.raw @@ -0,0 +1,192 @@ +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 diff --git a/source/tests/model_compression/data/type_map.raw b/source/tests/model_compression/data/type_map.raw new file mode 100644 index 0000000000..e900768b1d --- /dev/null +++ b/source/tests/model_compression/data/type_map.raw @@ -0,0 +1,2 @@ +O +H diff --git a/source/tests/model_compression/input.json b/source/tests/model_compression/input.json new file mode 100644 index 0000000000..599311e784 --- /dev/null +++ b/source/tests/model_compression/input.json @@ -0,0 +1,66 @@ +{ + "_comment": " model parameters", + "model": { +"type_map": ["O", "H"], +"descriptor" :{ + "type": "se_e2_a", + "sel": [46, 92], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [4, 8, 16], + "resnet_dt": false, + "axis_neuron": 16, + "seed": 1, + "_comment": " that's all" +}, +"fitting_net" : { + "neuron": [20, 20, 20], + "resnet_dt": true, + "seed": 1, + "_comment": " that's all" +}, +"_comment": " that's all" + }, + + "learning_rate" :{ +"type": "exp", +"decay_steps": 5000, +"start_lr": 0.001, +"stop_lr": 3.51e-8, +"_comment": "that's all" + }, + + "loss" :{ +"type": "ener", +"start_pref_e": 0.02, +"limit_pref_e": 1, +"start_pref_f": 1000, +"limit_pref_f": 1, +"start_pref_v": 0, +"limit_pref_v": 0, +"_comment": " that's all" + }, + + "training" : { +"training_data": { + "systems": ["model_compression/data"], + "batch_size": "auto", + "_comment": "that's all" +}, +"validation_data":{ + "systems": ["model_compression/data"], + "batch_size": 1, + "numb_btch": 3, + "_comment": "that's all" +}, +"numb_steps": 100, +"seed": 10, +"disp_file": "lcurve.out", +"disp_freq": 100, +"save_freq": 100, +"_comment": "that's all" + }, + + "_comment": "that's all" +} + diff --git a/source/tests/test_model_compression.py b/source/tests/test_model_compression.py new file mode 100644 index 0000000000..ac3529caa1 --- /dev/null +++ b/source/tests/test_model_compression.py @@ -0,0 +1,380 @@ +import os,sys,platform,shutil,dpdata,json +import numpy as np +import unittest + +from deepmd.infer import DeepPot +from deepmd.env import MODEL_VERSION +# from deepmd.entrypoints.compress import compress +from common import j_loader, tests_path + +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +if GLOBAL_NP_FLOAT_PRECISION == np.float32 : + default_places = 4 +else : + default_places = 6 + +def file_delete(file) : + if os.path.exists(file): + os.remove(file) + +class TestDeepPotAPBC(unittest.TestCase) : + def setUp(self): + self.data_file = str(tests_path / os.path.join("model_compression", "data")) + self.frozen_model = str(tests_path / "dp-original.pb") + self.compressed_model = str(tests_path / "dp-compressed.pb") + self.INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["training"]["training_data"]["systems"] = self.data_file + jdata["training"]["validation_data"]["systems"] = self.data_file + with open(self.INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = os.system("dp train " + self.INPUT) + ret = os.system("dp freeze -o " + self.frozen_model) + ret = os.system("dp compress " + self.INPUT + " -i " + self.frozen_model + " -o " + self.compressed_model) + assert(ret == 0), "Model compression error!" + + self.dp_original = DeepPot(self.frozen_model) + self.dp_compressed = DeepPot(self.compressed_model) + self.coords = np.array([12.83, 2.56, 2.18, + 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, + 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, + 4.27, 3.22, 1.56]) + self.atype = [0, 1, 1, 0, 1, 1] + self.box = np.array([13., 0., 0., 0., 13., 0., 0., 0., 13.]) + + def tearDown(self): + file_delete(self.INPUT) + file_delete(self.frozen_model) + file_delete(self.compressed_model) + file_delete("out.json") + file_delete("compress.json") + file_delete("checkpoint") + file_delete("lcurve.out") + file_delete("model.ckpt.meta") + file_delete("model.ckpt.index") + file_delete("model.ckpt.data-00000-of-00001") + + def test_attrs(self): + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places = default_places) + self.assertEqual(self.dp_original.get_type_map(), ['O', 'H']) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_compressed.get_rcut(), 6.0, places = default_places) + self.assertEqual(self.dp_compressed.get_type_map(), ['O', 'H']) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self): + ee0, ff0, vv0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = False) + ee1, ff1, vv1 = self.dp_compressed.eval(self.coords, self.box, self.atype, atomic = False) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes,1)) + self.assertEqual(ff0.shape, (nframes,natoms,3)) + self.assertEqual(vv0.shape, (nframes,9)) + self.assertEqual(ee1.shape, (nframes,1)) + self.assertEqual(ff1.shape, (nframes,natoms,3)) + self.assertEqual(vv1.shape, (nframes,9)) + # check values + for ii in range(ff0.size): + self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) + for ii in range(nframes): + self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) + for ii in range(nframes, 9): + self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + + def test_1frame_atm(self): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = True) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(self.coords, self.box, self.atype, atomic = True) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes,1)) + self.assertEqual(ff0.shape, (nframes,natoms,3)) + self.assertEqual(vv0.shape, (nframes,9)) + self.assertEqual(ae0.shape, (nframes,natoms,1)) + self.assertEqual(av0.shape, (nframes,natoms,9)) + self.assertEqual(ee1.shape, (nframes,1)) + self.assertEqual(ff1.shape, (nframes,natoms,3)) + self.assertEqual(vv1.shape, (nframes,9)) + self.assertEqual(ae1.shape, (nframes,natoms,1)) + self.assertEqual(av1.shape, (nframes,natoms,9)) + # check values + for ii in range(ff0.size): + self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) + for ii in range(ae0.size): + self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) + for ii in range(av0.size): + self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) + for ii in range(nframes): + self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) + for ii in range(nframes, 9): + self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + + def test_2frame_atm(self): + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(coords2, box2, self.atype, atomic = True) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(coords2, box2, self.atype, atomic = True) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes,1)) + self.assertEqual(ff0.shape, (nframes,natoms,3)) + self.assertEqual(vv0.shape, (nframes,9)) + self.assertEqual(ae0.shape, (nframes,natoms,1)) + self.assertEqual(av0.shape, (nframes,natoms,9)) + self.assertEqual(ee1.shape, (nframes,1)) + self.assertEqual(ff1.shape, (nframes,natoms,3)) + self.assertEqual(vv1.shape, (nframes,9)) + self.assertEqual(ae1.shape, (nframes,natoms,1)) + self.assertEqual(av1.shape, (nframes,natoms,9)) + + # check values + for ii in range(ff0.size): + self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) + for ii in range(ae0.size): + self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) + for ii in range(av0.size): + self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) + for ii in range(nframes): + self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) + for ii in range(nframes, 9): + self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + + +class TestDeepPotANoPBC(unittest.TestCase) : + def setUp(self): + self.data_file = str(tests_path / os.path.join("model_compression", "data")) + self.frozen_model = str(tests_path / "dp-original.pb") + self.compressed_model = str(tests_path / "dp-compressed.pb") + self.INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["training"]["training_data"]["systems"] = self.data_file + jdata["training"]["validation_data"]["systems"] = self.data_file + with open(self.INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = os.system("dp train " + self.INPUT) + ret = os.system("dp freeze -o " + self.frozen_model) + ret = os.system("dp compress " + self.INPUT + " -i " + self.frozen_model + " -o " + self.compressed_model) + assert(ret == 0), "Model compression error!" + + self.dp_original = DeepPot(self.frozen_model) + self.dp_compressed = DeepPot(self.compressed_model) + self.coords = np.array([12.83, 2.56, 2.18, + 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, + 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, + 4.27, 3.22, 1.56]) + self.atype = [0, 1, 1, 0, 1, 1] + self.box = None + + def tearDown(self): + file_delete(self.INPUT) + file_delete(self.frozen_model) + file_delete(self.compressed_model) + file_delete("out.json") + file_delete("compress.json") + file_delete("checkpoint") + file_delete("lcurve.out") + file_delete("model.ckpt.meta") + file_delete("model.ckpt.index") + file_delete("model.ckpt.data-00000-of-00001") + + def test_1frame(self): + ee0, ff0, vv0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = False) + ee1, ff1, vv1 = self.dp_compressed.eval(self.coords, self.box, self.atype, atomic = False) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes,1)) + self.assertEqual(ff0.shape, (nframes,natoms,3)) + self.assertEqual(vv0.shape, (nframes,9)) + self.assertEqual(ee1.shape, (nframes,1)) + self.assertEqual(ff1.shape, (nframes,natoms,3)) + self.assertEqual(vv1.shape, (nframes,9)) + # check values + for ii in range(ff0.size): + self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) + for ii in range(nframes): + self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) + for ii in range(nframes, 9): + self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + + def test_1frame_atm(self): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = True) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(self.coords, self.box, self.atype, atomic = True) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes,1)) + self.assertEqual(ff0.shape, (nframes,natoms,3)) + self.assertEqual(vv0.shape, (nframes,9)) + self.assertEqual(ae0.shape, (nframes,natoms,1)) + self.assertEqual(av0.shape, (nframes,natoms,9)) + self.assertEqual(ee1.shape, (nframes,1)) + self.assertEqual(ff1.shape, (nframes,natoms,3)) + self.assertEqual(vv1.shape, (nframes,9)) + self.assertEqual(ae1.shape, (nframes,natoms,1)) + self.assertEqual(av1.shape, (nframes,natoms,9)) + # check values + for ii in range(ff0.size): + self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) + for ii in range(ae0.size): + self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) + for ii in range(av0.size): + self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) + for ii in range(nframes): + self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) + for ii in range(nframes, 9): + self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + + def test_2frame_atm(self): + coords2 = np.concatenate((self.coords, self.coords)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(coords2, self.box, self.atype, atomic = True) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(coords2, self.box, self.atype, atomic = True) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes,1)) + self.assertEqual(ff0.shape, (nframes,natoms,3)) + self.assertEqual(vv0.shape, (nframes,9)) + self.assertEqual(ae0.shape, (nframes,natoms,1)) + self.assertEqual(av0.shape, (nframes,natoms,9)) + self.assertEqual(ee1.shape, (nframes,1)) + self.assertEqual(ff1.shape, (nframes,natoms,3)) + self.assertEqual(vv1.shape, (nframes,9)) + self.assertEqual(ae1.shape, (nframes,natoms,1)) + self.assertEqual(av1.shape, (nframes,natoms,9)) + + # check values + for ii in range(ff0.size): + self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) + for ii in range(ae0.size): + self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) + for ii in range(av0.size): + self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) + for ii in range(nframes): + self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) + for ii in range(nframes, 9): + self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + + +class TestDeepPotALargeBoxNoPBC(unittest.TestCase) : + def setUp(self): + self.data_file = str(tests_path / os.path.join("model_compression", "data")) + self.frozen_model = str(tests_path / "dp-original.pb") + self.compressed_model = str(tests_path / "dp-compressed.pb") + self.INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["training"]["training_data"]["systems"] = self.data_file + jdata["training"]["validation_data"]["systems"] = self.data_file + with open(self.INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = os.system("dp train " + self.INPUT) + ret = os.system("dp freeze -o " + self.frozen_model) + ret = os.system("dp compress " + self.INPUT + " -i " + self.frozen_model + " -o " + self.compressed_model) + assert(ret == 0), "Model compression error!" + + self.dp_original = DeepPot(self.frozen_model) + self.dp_compressed = DeepPot(self.compressed_model) + self.coords = np.array([12.83, 2.56, 2.18, + 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, + 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, + 4.27, 3.22, 1.56]) + self.atype = [0, 1, 1, 0, 1, 1] + self.box = np.array([19., 0., 0., 0., 13., 0., 0., 0., 13.]) + + def tearDown(self): + file_delete(self.INPUT) + file_delete(self.frozen_model) + file_delete(self.compressed_model) + file_delete("out.json") + file_delete("compress.json") + file_delete("checkpoint") + file_delete("lcurve.out") + file_delete("model.ckpt.meta") + file_delete("model.ckpt.index") + file_delete("model.ckpt.data-00000-of-00001") + + def test_1frame(self): + ee0, ff0, vv0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = False) + ee1, ff1, vv1 = self.dp_compressed.eval(self.coords, self.box, self.atype, atomic = False) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes,1)) + self.assertEqual(ff0.shape, (nframes,natoms,3)) + self.assertEqual(vv0.shape, (nframes,9)) + self.assertEqual(ee1.shape, (nframes,1)) + self.assertEqual(ff1.shape, (nframes,natoms,3)) + self.assertEqual(vv1.shape, (nframes,9)) + # check values + for ii in range(ff0.size): + self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) + for ii in range(nframes): + self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) + for ii in range(nframes, 9): + self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + + def test_1frame_atm(self): + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = True) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval(self.coords, self.box, self.atype, atomic = True) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes,1)) + self.assertEqual(ff0.shape, (nframes,natoms,3)) + self.assertEqual(vv0.shape, (nframes,9)) + self.assertEqual(ae0.shape, (nframes,natoms,1)) + self.assertEqual(av0.shape, (nframes,natoms,9)) + self.assertEqual(ee1.shape, (nframes,1)) + self.assertEqual(ff1.shape, (nframes,natoms,3)) + self.assertEqual(vv1.shape, (nframes,9)) + self.assertEqual(ae1.shape, (nframes,natoms,1)) + self.assertEqual(av1.shape, (nframes,natoms,9)) + # check values + for ii in range(ff0.size): + self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) + for ii in range(ae0.size): + self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) + for ii in range(av0.size): + self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) + for ii in range(nframes): + self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) + for ii in range(nframes, 9): + self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + + def test_ase(self): + from ase import Atoms + from deepmd.calculator import DP + water0 = Atoms('OHHOHH', + positions=self.coords.reshape((-1,3)), + cell=self.box.reshape((3,3)), + calculator=DP(self.frozen_model)) + water1 = Atoms('OHHOHH', + positions=self.coords.reshape((-1,3)), + cell=self.box.reshape((3,3)), + calculator=DP(self.compressed_model)) + ee0 = water0.get_potential_energy() + ff0 = water0.get_forces() + ee1 = water1.get_potential_energy() + ff1 = water1.get_forces() + nframes = 1 + for ii in range(ff0.size): + self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) + for ii in range(nframes): + self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) +