diff --git a/deepmd/descriptor/__init__.py b/deepmd/descriptor/__init__.py index ab726d95c8..6ef6c50da5 100644 --- a/deepmd/descriptor/__init__.py +++ b/deepmd/descriptor/__init__.py @@ -14,6 +14,9 @@ from .se_a_ebd import ( DescrptSeAEbd, ) +from .se_a_ebd_v2 import ( + DescrptSeAEbdV2, +) from .se_a_ef import ( DescrptSeAEf, DescrptSeAEfLower, @@ -39,6 +42,7 @@ "DescrptHybrid", "DescrptLocFrame", "DescrptSeA", + "DescrptSeAEbdV2", "DescrptSeAEbd", "DescrptSeAEf", "DescrptSeAEfLower", diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index cceb72d4fb..8f0051cd4e 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -10,6 +10,7 @@ from deepmd.common import ( cast_precision, get_activation_func, + get_np_precision, get_precision, ) from deepmd.env import ( @@ -30,10 +31,17 @@ from deepmd.nvnmd.utils.config import ( nvnmd_cfg, ) +from deepmd.utils.compress import ( + get_extra_side_embedding_net_variable, + get_two_side_type_embedding, + get_type_embedding, + make_data, +) from deepmd.utils.errors import ( GraphWithoutTensorError, ) from deepmd.utils.graph import ( + get_pattern_nodes_from_graph_def, get_tensor_by_name_from_graph, ) from deepmd.utils.network import ( @@ -165,6 +173,7 @@ def __init__( uniform_seed: bool = False, multi_task: bool = False, spin: Optional[Spin] = None, + stripped_type_embedding: bool = False, **kwargs, ) -> None: """Constructor.""" @@ -185,6 +194,7 @@ def __init__( self.compress_activation_fn = get_activation_func(activation_function) self.filter_activation_fn = get_activation_func(activation_function) self.filter_precision = get_precision(precision) + self.filter_np_precision = get_np_precision(precision) self.exclude_types = set() for tt in exclude_types: assert len(tt) == 2 @@ -193,6 +203,9 @@ def __init__( self.set_davg_zero = set_davg_zero self.type_one_side = type_one_side self.spin = spin + self.stripped_type_embedding = stripped_type_embedding + self.extra_embeeding_net_variables = None + self.layer_size = len(neuron) # extend sel_a for spin system if self.spin is not None: @@ -463,6 +476,39 @@ def enable_compression( "The size of the next layer of the neural network must be twice the size of the previous layer." % ",".join([str(item) for item in self.filter_neuron]) ) + if self.stripped_type_embedding: + ret_two_side = get_pattern_nodes_from_graph_def( + graph_def, f"filter_type_all{suffix}/.+_two_side_ebd" + ) + ret_one_side = get_pattern_nodes_from_graph_def( + graph_def, f"filter_type_all{suffix}/.+_one_side_ebd" + ) + if len(ret_two_side) == 0 and len(ret_one_side) == 0: + raise RuntimeError( + "can not find variables of embedding net from graph_def, maybe it is not a compressible model." + ) + elif len(ret_one_side) != 0 and len(ret_two_side) != 0: + raise RuntimeError( + "both one side and two side embedding net varaibles are detected, it is a wrong model." + ) + elif len(ret_two_side) != 0: + self.final_type_embedding = get_two_side_type_embedding(self, graph) + self.matrix = get_extra_side_embedding_net_variable( + self, graph_def, "two_side", "matrix", suffix + ) + self.bias = get_extra_side_embedding_net_variable( + self, graph_def, "two_side", "bias", suffix + ) + self.extra_embedding = make_data(self, self.final_type_embedding) + else: + self.final_type_embedding = get_type_embedding(self, graph) + self.matrix = get_extra_side_embedding_net_variable( + self, graph_def, "one_side", "matrix", suffix + ) + self.bias = get_extra_side_embedding_net_variable( + self, graph_def, "one_side", "bias", suffix + ) + self.extra_embedding = make_data(self, self.final_type_embedding) self.compress = True self.table = DPTabulate( @@ -588,6 +634,7 @@ def build( coord = tf.reshape(coord_, [-1, natoms[1] * 3]) box = tf.reshape(box_, [-1, 9]) atype = tf.reshape(atype_, [-1, natoms[1]]) + self.atype = atype op_descriptor = ( build_op_descriptor() if nvnmd_cfg.enable else op_module.prod_env_mat_a @@ -606,6 +653,10 @@ def build( sel_a=self.sel_a, sel_r=self.sel_r, ) + nlist_t = tf.reshape(self.nlist + 1, [-1]) + atype_t = tf.concat([[self.ntypes], tf.reshape(self.atype, [-1])], axis=0) + self.nei_type_vec = tf.nn.embedding_lookup(atype_t, nlist_t) + # only used when tensorboard was set as true tf.summary.histogram("descrpt", self.descrpt) tf.summary.histogram("rij", self.rij) @@ -692,6 +743,8 @@ def _pass_filter( type_embedding = input_dict.get("type_embedding", None) else: type_embedding = None + if self.stripped_type_embedding and type_embedding is None: + raise RuntimeError("type_embedding is required for se_a_tebd_v2 model.") start_index = 0 inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]) output = [] @@ -901,13 +954,89 @@ def _filter_lower( # with (natom x nei_type_i) x 1 xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1]) if type_embedding is not None: - xyz_scatter = self._concat_type_embedding( - xyz_scatter, nframes, natoms, type_embedding - ) - if self.compress: - raise RuntimeError( - "compression of type embedded descriptor is not supported at the moment" + if self.stripped_type_embedding: + if self.type_one_side: + extra_embedding_index = self.nei_type_vec + else: + padding_ntypes = type_embedding.shape[0] + atype_expand = tf.reshape(self.atype, [-1, 1]) + idx_i = tf.tile(atype_expand * padding_ntypes, [1, self.nnei]) + idx_j = tf.reshape(self.nei_type_vec, [-1, self.nnei]) + idx = idx_i + idx_j + index_of_two_side = tf.reshape(idx, [-1]) + extra_embedding_index = index_of_two_side + + if not self.compress: + if self.type_one_side: + one_side_type_embedding_suffix = "_one_side_ebd" + net_output = embedding_net( + type_embedding, + self.filter_neuron, + self.filter_precision, + activation_fn=activation_fn, + resnet_dt=self.filter_resnet_dt, + name_suffix=one_side_type_embedding_suffix, + stddev=stddev, + bavg=bavg, + seed=self.seed, + trainable=trainable, + uniform_seed=self.uniform_seed, + initial_variables=self.extra_embeeding_net_variables, + mixed_prec=self.mixed_prec, + ) + net_output = tf.nn.embedding_lookup( + net_output, self.nei_type_vec + ) + else: + type_embedding_nei = tf.tile( + tf.reshape(type_embedding, [1, padding_ntypes, -1]), + [padding_ntypes, 1, 1], + ) # (ntypes) * ntypes * Y + type_embedding_center = tf.tile( + tf.reshape(type_embedding, [padding_ntypes, 1, -1]), + [1, padding_ntypes, 1], + ) # ntypes * (ntypes) * Y + two_side_type_embedding = tf.concat( + [type_embedding_nei, type_embedding_center], -1 + ) # ntypes * ntypes * (Y+Y) + two_side_type_embedding = tf.reshape( + two_side_type_embedding, + [-1, two_side_type_embedding.shape[-1]], + ) + + atype_expand = tf.reshape(self.atype, [-1, 1]) + idx_i = tf.tile(atype_expand * padding_ntypes, [1, self.nnei]) + idx_j = tf.reshape(self.nei_type_vec, [-1, self.nnei]) + idx = idx_i + idx_j + index_of_two_side = tf.reshape(idx, [-1]) + self.extra_embedding_index = index_of_two_side + + two_side_type_embedding_suffix = "_two_side_ebd" + net_output = embedding_net( + two_side_type_embedding, + self.filter_neuron, + self.filter_precision, + activation_fn=activation_fn, + resnet_dt=self.filter_resnet_dt, + name_suffix=two_side_type_embedding_suffix, + stddev=stddev, + bavg=bavg, + seed=self.seed, + trainable=trainable, + uniform_seed=self.uniform_seed, + initial_variables=self.extra_embeeding_net_variables, + mixed_prec=self.mixed_prec, + ) + net_output = tf.nn.embedding_lookup(net_output, idx) + net_output = tf.reshape(net_output, [-1, self.filter_neuron[-1]]) + else: + xyz_scatter = self._concat_type_embedding( + xyz_scatter, nframes, natoms, type_embedding ) + if self.compress: + raise RuntimeError( + "compression of type embedded descriptor is not supported when stripped_type_embedding == False" + ) # natom x 4 x outputs_size if nvnmd_cfg.enable: return filter_lower_R42GR( @@ -929,25 +1058,48 @@ def _filter_lower( self.embedding_net_variables, ) if self.compress and (not is_exclude): - if self.type_one_side: - net = "filter_-1_net_" + str(type_i) + if self.stripped_type_embedding: + net_output = tf.nn.embedding_lookup( + self.extra_embedding, extra_embedding_index + ) + net = "filter_net" + info = [ + self.lower[net], + self.upper[net], + self.upper[net] * self.table_config[0], + self.table_config[1], + self.table_config[2], + self.table_config[3], + ] + return op_module.tabulate_fusion_se_atten( + tf.cast(self.table.data[net], self.filter_precision), + info, + xyz_scatter, + tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), + net_output, + last_layer_size=outputs_size[-1], + is_sorted=False, + ) else: - net = "filter_" + str(type_input) + "_net_" + str(type_i) - info = [ - self.lower[net], - self.upper[net], - self.upper[net] * self.table_config[0], - self.table_config[1], - self.table_config[2], - self.table_config[3], - ] - return op_module.tabulate_fusion_se_a( - tf.cast(self.table.data[net], self.filter_precision), - info, - xyz_scatter, - tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), - last_layer_size=outputs_size[-1], - ) + if self.type_one_side: + net = "filter_-1_net_" + str(type_i) + else: + net = "filter_" + str(type_input) + "_net_" + str(type_i) + info = [ + self.lower[net], + self.upper[net], + self.upper[net] * self.table_config[0], + self.table_config[1], + self.table_config[2], + self.table_config[3], + ] + return op_module.tabulate_fusion_se_a( + tf.cast(self.table.data[net], self.filter_precision), + info, + xyz_scatter, + tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), + last_layer_size=outputs_size[-1], + ) else: if not is_exclude: # with (natom x nei_type_i) x out_size @@ -966,6 +1118,9 @@ def _filter_lower( initial_variables=self.embedding_net_variables, mixed_prec=self.mixed_prec, ) + + if self.stripped_type_embedding: + xyz_scatter = xyz_scatter * net_output + xyz_scatter if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift else: @@ -1179,3 +1334,10 @@ def init_variables( self.dstd = new_dstd if self.original_sel is None: self.original_sel = sel + + @property + def explicit_ntypes(self) -> bool: + """Explicit ntypes with type embedding.""" + if self.stripped_type_embedding: + return True + return False diff --git a/deepmd/descriptor/se_a_ebd_v2.py b/deepmd/descriptor/se_a_ebd_v2.py new file mode 100644 index 0000000000..c6e3cebc71 --- /dev/null +++ b/deepmd/descriptor/se_a_ebd_v2.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + List, + Optional, +) + +from deepmd.utils.spin import ( + Spin, +) + +from .descriptor import ( + Descriptor, +) +from .se_a import ( + DescrptSeA, +) + +log = logging.getLogger(__name__) + + +@Descriptor.register("se_a_tpe_v2") +@Descriptor.register("se_a_ebd_v2") +class DescrptSeAEbdV2(DescrptSeA): + r"""A compressible se_a_ebd model. + + This model is a warpper for DescriptorSeA, which set stripped_type_embedding=True. + """ + + def __init__( + self, + rcut: float, + rcut_smth: float, + sel: List[str], + neuron: List[int] = [24, 48, 96], + axis_neuron: int = 8, + resnet_dt: bool = False, + trainable: bool = True, + seed: Optional[int] = None, + type_one_side: bool = True, + exclude_types: List[List[int]] = [], + set_davg_zero: bool = False, + activation_function: str = "tanh", + precision: str = "default", + uniform_seed: bool = False, + multi_task: bool = False, + spin: Optional[Spin] = None, + **kwargs, + ) -> None: + DescrptSeA.__init__( + self, + rcut, + rcut_smth, + sel, + neuron=neuron, + axis_neuron=axis_neuron, + resnet_dt=resnet_dt, + trainable=trainable, + seed=seed, + type_one_side=type_one_side, + exclude_types=exclude_types, + set_davg_zero=set_davg_zero, + activation_function=activation_function, + precision=precision, + uniform_seed=uniform_seed, + multi_task=multi_task, + spin=spin, + stripped_type_embedding=True, + **kwargs, + ) diff --git a/deepmd/descriptor/se_a_mask.py b/deepmd/descriptor/se_a_mask.py index b9181fd6b0..780b34d294 100644 --- a/deepmd/descriptor/se_a_mask.py +++ b/deepmd/descriptor/se_a_mask.py @@ -128,6 +128,7 @@ def __init__( activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, + stripped_type_embedding: bool = False, **kwargs, ) -> None: """Constructor.""" @@ -159,6 +160,7 @@ def __init__( # numb of neighbors and numb of descrptors self.nnei_a = np.cumsum(self.sel_a)[-1] self.nnei = self.nnei_a + self.stripped_type_embedding = stripped_type_embedding self.ndescrpt_a = self.nnei_a * 4 self.ndescrpt = self.ndescrpt_a diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index c962952ec0..b0c65108e5 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -35,6 +35,11 @@ from deepmd.nvnmd.utils.config import ( nvnmd_cfg, ) +from deepmd.utils.compress import ( + get_extra_side_embedding_net_variable, + get_two_side_type_embedding, + make_data, +) from deepmd.utils.graph import ( get_attention_layer_variables_from_graph_def, get_pattern_nodes_from_graph_def, @@ -415,12 +420,14 @@ def enable_compression( min_nbor_dist, table_extrapolate, table_stride_1, table_stride_2 ) - self.final_type_embedding = self._get_two_side_type_embedding(graph) - self.matrix = self._get_two_side_embedding_net_variable( - graph_def, "matrix", suffix + self.final_type_embedding = get_two_side_type_embedding(self, graph) + self.matrix = get_extra_side_embedding_net_variable( + self, graph_def, "two_side", "matrix", suffix + ) + self.bias = get_extra_side_embedding_net_variable( + self, graph_def, "two_side", "bias", suffix ) - self.bias = self._get_two_side_embedding_net_variable(graph_def, "bias", suffix) - self.two_embd = self._make_data(self.final_type_embedding) + self.two_embd = make_data(self, self.final_type_embedding) self.davg = get_tensor_by_name_from_graph( graph, "descrpt_attr%s/t_avg" % suffix @@ -429,79 +436,6 @@ def enable_compression( graph, "descrpt_attr%s/t_std" % suffix ) - def _get_two_side_type_embedding(self, graph): - type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd") - type_embedding = type_embedding.astype(self.filter_np_precision) - type_embedding_shape = type_embedding.shape - type_embedding_nei = np.tile( - np.reshape(type_embedding, [1, type_embedding_shape[0], -1]), - [type_embedding_shape[0], 1, 1], - ) # (ntypes) * ntypes * Y - type_embedding_center = np.tile( - np.reshape(type_embedding, [type_embedding_shape[0], 1, -1]), - [1, type_embedding_shape[0], 1], - ) # ntypes * (ntypes) * Y - two_side_type_embedding = np.concatenate( - [type_embedding_nei, type_embedding_center], -1 - ) # ntypes * ntypes * (Y+Y) - two_side_type_embedding = np.reshape( - two_side_type_embedding, [-1, two_side_type_embedding.shape[-1]] - ) - return two_side_type_embedding - - def _get_two_side_embedding_net_variable(self, graph_def, varialbe_name, suffix): - ret = {} - for i in range(1, self.layer_size + 1): - target = get_pattern_nodes_from_graph_def( - graph_def, - f"filter_type_all{suffix}/{varialbe_name}_{i}_two_side_ebd", - ) - node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}_two_side_ebd"] - ret["layer_" + str(i)] = node - return ret - - def _layer_0(self, x, w, b): - return self.filter_activation_fn(tf.matmul(x, w) + b) - - def _layer_1(self, x, w, b): - t = tf.concat([x, x], axis=1) - return t, self.filter_activation_fn(tf.matmul(x, w) + b) + t - - def _make_data(self, xx): - with tf.Session() as sess: - for layer in range(self.layer_size): - if layer == 0: - if self.filter_neuron[0] == 1: - yy = ( - self._layer_0( - xx, - self.matrix["layer_" + str(layer + 1)], - self.bias["layer_" + str(layer + 1)], - ) - + xx - ) - elif self.filter_neuron[0] == 2: - tt, yy = self._layer_1( - xx, - self.matrix["layer_" + str(layer + 1)], - self.bias["layer_" + str(layer + 1)], - ) - else: - yy = self._layer_0( - xx, - self.matrix["layer_" + str(layer + 1)], - self.bias["layer_" + str(layer + 1)], - ) - else: - tt, zz = self._layer_1( - yy, - self.matrix["layer_" + str(layer + 1)], - self.bias["layer_" + str(layer + 1)], - ) - yy = zz - vv = sess.run(zz) - return vv - def build( self, coord_: tf.Tensor, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 7bd373b492..ae446ef348 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -432,6 +432,11 @@ def descrpt_se_atten_v2_args(): ] +@descrpt_args_plugin.register("se_a_ebd_v2", alias=["se_a_tpe_v2"]) +def descrpt_se_a_ebd_v2_args(): + return descrpt_se_a_args() + + @descrpt_args_plugin.register("se_a_mask") def descrpt_se_a_mask_args(): doc_sel = 'This parameter sets the number of selected neighbors for each type of atom. It can be:\n\n\ diff --git a/deepmd/utils/compress.py b/deepmd/utils/compress.py new file mode 100644 index 0000000000..c6e68dfe19 --- /dev/null +++ b/deepmd/utils/compress.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging + +import numpy as np + +from deepmd.env import ( + tf, +) +from deepmd.utils.graph import ( + get_pattern_nodes_from_graph_def, + get_tensor_by_name_from_graph, +) + +log = logging.getLogger(__name__) + + +def get_type_embedding(self, graph): + type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd") + type_embedding = type_embedding.astype(self.filter_np_precision) + return type_embedding + + +def get_two_side_type_embedding(self, graph): + type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd") + type_embedding = type_embedding.astype(self.filter_np_precision) + type_embedding_shape = type_embedding.shape + + type_embedding_nei = np.tile( + np.reshape(type_embedding, [1, type_embedding_shape[0], -1]), + [type_embedding_shape[0], 1, 1], + ) # (ntypes) * ntypes * Y + type_embedding_center = np.tile( + np.reshape(type_embedding, [type_embedding_shape[0], 1, -1]), + [1, type_embedding_shape[0], 1], + ) # ntypes * (ntypes) * Y + two_side_type_embedding = np.concatenate( + [type_embedding_nei, type_embedding_center], -1 + ) # ntypes * ntypes * (Y+Y) + two_side_type_embedding = np.reshape( + two_side_type_embedding, [-1, two_side_type_embedding.shape[-1]] + ) + return two_side_type_embedding + + +def get_extra_side_embedding_net_variable( + self, graph_def, type_side, varialbe_name, suffix +): + ret = {} + for i in range(1, self.layer_size + 1): + target = get_pattern_nodes_from_graph_def( + graph_def, + f"filter_type_all{suffix}/{varialbe_name}_{i}_{type_side}_ebd", + ) + node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}_{type_side}_ebd"] + ret["layer_" + str(i)] = node + return ret + + +def _layer_0(self, x, w, b): + return self.filter_activation_fn(tf.matmul(x, w) + b) + + +def _layer_1(self, x, w, b): + t = tf.concat([x, x], axis=1) + return t, self.filter_activation_fn(tf.matmul(x, w) + b) + t + + +def make_data(self, xx): + with tf.Session() as sess: + for layer in range(self.layer_size): + if layer == 0: + if self.filter_neuron[0] == 1: + yy = ( + _layer_0( + self, + xx, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + + xx + ) + elif self.filter_neuron[0] == 2: + tt, yy = _layer_1( + self, + xx, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + else: + yy = _layer_0( + self, + xx, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + else: + tt, zz = _layer_1( + self, + yy, + self.matrix["layer_" + str(layer + 1)], + self.bias["layer_" + str(layer + 1)], + ) + yy = zz + vv = sess.run(zz) + return vv diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index 883730b9d9..427887089a 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -176,7 +176,9 @@ def build( """ # tabulate range [lower, upper] with stride0 'stride0' lower, upper = self._get_env_mat_range(min_nbor_dist) - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( + self.descrpt, deepmd.descriptor.DescrptSeAEbdV2 + ): uu = np.max(upper) ll = np.min(lower) xx = np.arange(ll, uu, stride0, dtype=self.data_type) @@ -419,7 +421,9 @@ def _get_bias(self): bias = {} for layer in range(1, self.layer_size + 1): bias["layer_" + str(layer)] = [] - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( + self.descrpt, deepmd.descriptor.DescrptSeAEbdV2 + ): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/bias_{layer}" ] @@ -483,7 +487,9 @@ def _get_matrix(self): matrix = {} for layer in range(1, self.layer_size + 1): matrix["layer_" + str(layer)] = [] - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( + self.descrpt, deepmd.descriptor.DescrptSeAEbdV2 + ): node = self.embedding_net_nodes[ f"filter_type_all{self.suffix}/matrix_{layer}" ] @@ -687,7 +693,9 @@ def _spline5_switch(self, xx, rmin, rmax): def _get_layer_size(self): layer_size = 0 - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( + self.descrpt, deepmd.descriptor.DescrptSeAEbdV2 + ): layer_size = len(self.embedding_net_nodes) // 2 elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): layer_size = len(self.embedding_net_nodes) // ( @@ -737,7 +745,9 @@ def _all_excluded(self, ii: int) -> bool: def _get_table_size(self): table_size = 0 - if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten): + if isinstance(self.descrpt, deepmd.descriptor.DescrptSeAtten) or isinstance( + self.descrpt, deepmd.descriptor.DescrptSeAEbdV2 + ): table_size = 1 elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeA): table_size = self.ntypes * self.ntypes diff --git a/examples/water/se_e2_a_tebd/input.json b/examples/water/se_e2_a_tebd/input.json index 101c1a7a4f..1ec85a968c 100644 --- a/examples/water/se_e2_a_tebd/input.json +++ b/examples/water/se_e2_a_tebd/input.json @@ -16,7 +16,7 @@ "seed": 1 }, "descriptor": { - "type": "se_e2_a", + "type": "se_a_ebd_v2", "sel": [ 46, 92 diff --git a/source/tests/test_model_compression_se_a_ebd.py b/source/tests/test_model_compression_se_a_ebd.py new file mode 100644 index 0000000000..2a3163b062 --- /dev/null +++ b/source/tests/test_model_compression_se_a_ebd.py @@ -0,0 +1,587 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import subprocess as sp +import unittest + +import numpy as np + +# from deepmd.entrypoints.compress import compress +from common import ( + j_loader, + run_dp, + tests_path, +) + +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer import ( + DeepPot, +) + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 10 + + +def _file_delete(file): + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): + os.remove(file) + + +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b""): + if hasattr(line, "decode"): + line = line.decode("utf-8") + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode + + +def _init_models(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-se-e2-a-v2.pb") + compressed_model = str(tests_path / "dp-compressed-se-e2-a-v2.pb") + INPUT = str(tests_path / "input_se_a_ebd_v2.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + jdata["model"]["descriptor"] = {} + jdata["model"]["descriptor"]["type"] = "se_a_ebd_v2" + jdata["model"]["descriptor"]["type_one_side"] = False + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + with open("input_v2_compat.json.tempfile", "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def _init_models_exclude_types(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-exclude-types-se-e2-a-ebd-v2.pb") + compressed_model = str(tests_path / "dp-compressed-exclude-types-se-e2-a-ebd-v2.pb") + INPUT = str(tests_path / "input_se_a_ebd_v2.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["model"]["descriptor"] = {} + jdata["model"]["descriptor"]["type"] = "se_a_ebd_v2" + jdata["model"]["descriptor"]["exclude_types"] = [[0, 1]] + jdata["model"]["descriptor"]["type_one_side"] = False + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() +INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types() + + +class TestDeepPotAPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(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, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotANoPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(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 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotALargeBoxNoPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(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, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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(FROZEN_MODEL), + ) + water1 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(COMPRESSED_MODEL), + ) + ee0 = water0.get_potential_energy() + ff0 = water0.get_forces() + ee1 = water1.get_potential_energy() + ff1 = water1.get_forces() + nframes = 1 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + + +class TestDeepPotAPBCExcludeTypes(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL_ET) + self.dp_compressed = DeepPot(COMPRESSED_MODEL_ET) + 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, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + @classmethod + def tearDownClass(self): + _file_delete(INPUT_ET) + _file_delete(FROZEN_MODEL_ET) + _file_delete(COMPRESSED_MODEL_ET) + _file_delete("out.json") + _file_delete("compress.json") + _file_delete("checkpoint") + _file_delete("model.ckpt.meta") + _file_delete("model.ckpt.index") + _file_delete("model.ckpt.data-00000-of-00001") + _file_delete("model.ckpt-100.meta") + _file_delete("model.ckpt-100.index") + _file_delete("model.ckpt-100.data-00000-of-00001") + _file_delete("model-compression/checkpoint") + _file_delete("model-compression/model.ckpt.meta") + _file_delete("model-compression/model.ckpt.index") + _file_delete("model-compression/model.ckpt.data-00000-of-00001") + _file_delete("model-compression") + _file_delete("input_v2_compat.json") + _file_delete("lcurve.out") + + 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) diff --git a/source/tests/test_model_compression_se_a_ebd_type_one_side.py b/source/tests/test_model_compression_se_a_ebd_type_one_side.py new file mode 100644 index 0000000000..2f3d16b05f --- /dev/null +++ b/source/tests/test_model_compression_se_a_ebd_type_one_side.py @@ -0,0 +1,587 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import subprocess as sp +import unittest + +import numpy as np + +# from deepmd.entrypoints.compress import compress +from common import ( + j_loader, + run_dp, + tests_path, +) + +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer import ( + DeepPot, +) + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 10 + + +def _file_delete(file): + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): + os.remove(file) + + +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b""): + if hasattr(line, "decode"): + line = line.decode("utf-8") + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode + + +def _init_models(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-se-e2-a-v2-one-side.pb") + compressed_model = str(tests_path / "dp-compressed-se-e2-a-v2-one-side.pb") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + jdata["model"]["descriptor"] = {} + jdata["model"]["descriptor"]["type"] = "se_a_ebd_v2" + jdata["model"]["descriptor"]["type_one_side"] = True + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def _init_models_exclude_types(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-exclude-types-se-e2-a-v2-one-side.pb") + compressed_model = str( + tests_path / "dp-compressed-exclude-types-se-e2-a-v2-one-side.pb" + ) + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + jdata["model"]["descriptor"] = {} + jdata["model"]["descriptor"]["type"] = "se_a_ebd_v2" + jdata["model"]["descriptor"]["exclude_types"] = [[0, 1]] + jdata["model"]["descriptor"]["type_one_side"] = True + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + jdata["training"]["training_data"]["systems"] = data_file + jdata["training"]["validation_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() +INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types() + + +class TestDeepPotAPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(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, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotANoPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(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 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotALargeBoxNoPBC(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL) + self.dp_compressed = DeepPot(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, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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(FROZEN_MODEL), + ) + water1 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(COMPRESSED_MODEL), + ) + ee0 = water0.get_potential_energy() + ff0 = water0.get_forces() + ee1 = water1.get_potential_energy() + ff1 = water1.get_forces() + nframes = 1 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + + +class TestDeepPotAPBCExcludeTypes(unittest.TestCase): + @classmethod + def setUpClass(self): + self.dp_original = DeepPot(FROZEN_MODEL_ET) + self.dp_compressed = DeepPot(COMPRESSED_MODEL_ET) + 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, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + @classmethod + def tearDownClass(self): + _file_delete(INPUT_ET) + _file_delete(FROZEN_MODEL_ET) + _file_delete(COMPRESSED_MODEL_ET) + _file_delete("out.json") + _file_delete("compress.json") + _file_delete("checkpoint") + _file_delete("model.ckpt.meta") + _file_delete("model.ckpt.index") + _file_delete("model.ckpt.data-00000-of-00001") + _file_delete("model.ckpt-100.meta") + _file_delete("model.ckpt-100.index") + _file_delete("model.ckpt-100.data-00000-of-00001") + _file_delete("model-compression/checkpoint") + _file_delete("model-compression/model.ckpt.meta") + _file_delete("model-compression/model.ckpt.index") + _file_delete("model-compression/model.ckpt.data-00000-of-00001") + _file_delete("model-compression") + _file_delete("input_v2_compat.json") + _file_delete("lcurve.out") + + 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, 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 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) diff --git a/source/tests/test_model_se_a_ebd_v2.py b/source/tests/test_model_se_a_ebd_v2.py new file mode 100644 index 0000000000..71860890ce --- /dev/null +++ b/source/tests/test_model_se_a_ebd_v2.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +from common import ( + DataSystem, + gen_data, + j_loader, +) + +from deepmd.common import ( + j_must_have, +) +from deepmd.descriptor.se_a_ebd_v2 import ( + DescrptSeAEbdV2, +) +from deepmd.env import ( + tf, +) +from deepmd.fit import ( + EnerFitting, +) +from deepmd.model import ( + EnerModel, +) +from deepmd.utils.type_embed import ( + TypeEmbedNet, +) + +GLOBAL_ENER_FLOAT_PRECISION = tf.float64 +GLOBAL_TF_FLOAT_PRECISION = tf.float64 +GLOBAL_NP_FLOAT_PRECISION = np.float64 + + +class TestModel(tf.test.TestCase): + def setUp(self): + gen_data() + + def test_model(self): + jfile = "water_se_a_ebd.json" + jdata = j_loader(jfile) + + systems = j_must_have(jdata, "systems") + set_pfx = j_must_have(jdata, "set_prefix") + batch_size = j_must_have(jdata, "batch_size") + test_size = j_must_have(jdata, "numb_test") + batch_size = 1 + test_size = 1 + stop_batch = j_must_have(jdata, "stop_batch") + rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) + + test_data = data.get_test() + numb_test = 1 + + jdata["model"]["descriptor"].pop("type", None) + jdata["model"]["type_embedding"] = {} + jdata["model"]["type_embedding"]["neuron"] = [1] + jdata["model"]["type_embedding"]["resnet_dt"] = False + jdata["model"]["type_embedding"]["seed"] = 1 + typeebd_param = jdata["model"]["type_embedding"] + typeebd = TypeEmbedNet( + neuron=typeebd_param["neuron"], + activation_function=None, + resnet_dt=typeebd_param["resnet_dt"], + seed=typeebd_param["seed"], + uniform_seed=True, + padding=True, + ) + descrpt = DescrptSeAEbdV2( + **jdata["model"]["descriptor"], + ) + jdata["model"]["fitting_net"]["descrpt"] = descrpt + fitting = EnerFitting( + **jdata["model"]["fitting_net"], + ) + # fitting = EnerFitting(jdata['model']['fitting_net'], descrpt) + model = EnerModel(descrpt, fitting, typeebd) + + # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) + input_data = { + "coord": [test_data["coord"]], + "box": [test_data["box"]], + "type": [test_data["type"]], + "natoms_vec": [test_data["natoms_vec"]], + "default_mesh": [test_data["default_mesh"]], + } + model._compute_input_stat(input_data) + model.descrpt.bias_atom_e = data.compute_energy_shift() + + t_prop_c = tf.placeholder(tf.float32, [5], name="t_prop_c") + t_energy = tf.placeholder(GLOBAL_ENER_FLOAT_PRECISION, [None], name="t_energy") + t_force = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="t_force") + t_virial = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="t_virial") + t_atom_ener = tf.placeholder( + GLOBAL_TF_FLOAT_PRECISION, [None], name="t_atom_ener" + ) + t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_coord") + t_type = tf.placeholder(tf.int32, [None], name="i_type") + t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2], name="i_natoms") + t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, 9], name="i_box") + t_mesh = tf.placeholder(tf.int32, [None], name="i_mesh") + is_training = tf.placeholder(tf.bool) + t_fparam = None + + model_pred = model.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + t_fparam, + suffix="se_a_ebd_v2", + reuse=False, + ) + energy = model_pred["energy"] + force = model_pred["force"] + virial = model_pred["virial"] + atom_ener = model_pred["atom_ener"] + + feed_dict_test = { + t_prop_c: test_data["prop_c"], + t_energy: test_data["energy"][:numb_test], + t_force: np.reshape(test_data["force"][:numb_test, :], [-1]), + t_virial: np.reshape(test_data["virial"][:numb_test, :], [-1]), + t_atom_ener: np.reshape(test_data["atom_ener"][:numb_test, :], [-1]), + t_coord: np.reshape(test_data["coord"][:numb_test, :], [-1]), + t_box: test_data["box"][:numb_test, :], + t_type: np.reshape(test_data["type"][:numb_test, :], [-1]), + t_natoms: test_data["natoms_vec"], + t_mesh: test_data["default_mesh"], + is_training: False, + } + + sess = self.cached_session().__enter__() + sess.run(tf.global_variables_initializer()) + [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test) + + e = e.reshape([-1]) + f = f.reshape([-1]) + v = v.reshape([-1]) + + refe = [5.435394596262052014e-01] + reff = [ + 6.583728125594628944e-02, + 7.228993116083935744e-02, + 1.971543579114074483e-03, + 6.567474563776359853e-02, + 7.809421727465599983e-02, + -4.866958849094786890e-03, + -8.670511901715304004e-02, + 3.525374157021862048e-02, + 1.415748959800727487e-03, + 6.375813001810648473e-02, + -1.139053242798149790e-01, + -4.178593754384440744e-03, + -1.471737787218250215e-01, + 4.189712704724830872e-02, + 7.011731363309440038e-03, + 3.860874082716164030e-02, + -1.136296927731473005e-01, + -1.353471298745012206e-03, + ] + refv = [ + -4.243979601186427253e-01, + 1.097173849143971286e-01, + 1.227299373463585502e-02, + 1.097173849143970314e-01, + -2.462891443164323124e-01, + -5.711664180530139426e-03, + 1.227299373463585502e-02, + -5.711664180530143763e-03, + -6.217348853341628408e-04, + ] + refe = np.reshape(refe, [-1]) + reff = np.reshape(reff, [-1]) + refv = np.reshape(refv, [-1]) + + places = 6 + for ii in range(e.size): + self.assertAlmostEqual(e[ii], refe[ii], places=places) + for ii in range(f.size): + self.assertAlmostEqual(f[ii], reff[ii], places=places) + for ii in range(v.size): + self.assertAlmostEqual(v[ii], refv[ii], places=places)