diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 54325ca77c..ab0ba46042 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -11,7 +11,7 @@ from deepmd.env import default_tf_session_config from deepmd.utils.network import embedding_net from deepmd.utils.tabulate import DeepTabulate - +from deepmd.utils.type_embed import embed_atom_type class DescrptSeA (): @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) @@ -101,6 +101,11 @@ def __init__ (self, self.davg = None self.compress = False self.place_holders = {} + nei_type = np.array([]) + for ii in range(self.ntypes): + nei_type = np.append(nei_type, ii * np.ones(self.sel_a[ii])) # like a mask + self.nei_type = tf.constant(nei_type, dtype = tf.int32) + avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) sub_graph = tf.Graph() @@ -214,7 +219,7 @@ def compute_input_stats (self, sumr2 = np.sum(sumr2, axis = 0) suma2 = np.sum(suma2, axis = 0) for type_i in range(self.ntypes) : - davgunit = [sumr[type_i]/sumn[type_i], 0, 0, 0] + davgunit = [sumr[type_i]/(sumn[type_i]+1e-15), 0, 0, 0] dstdunit = [self._compute_std(sumr2[type_i], sumr[type_i], sumn[type_i]), self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), @@ -440,11 +445,15 @@ def _pass_filter(self, reuse = None, suffix = '', trainable = True) : + if input_dict is not None: + type_embedding = input_dict.get('type_embedding', None) + else: + type_embedding = None start_index = 0 inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]]) output = [] output_qmat = [] - if not self.type_one_side: + if not self.type_one_side and type_embedding is None: for type_i in range(self.ntypes): inputs_i = tf.slice (inputs, [ 0, start_index* self.ndescrpt], @@ -460,7 +469,7 @@ def _pass_filter(self, inputs_i = inputs inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) type_i = -1 - layer, qmat = self._filter(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_all'+suffix, natoms=natoms, reuse=reuse, seed = self.seed, trainable = trainable, activation_fn = self.filter_activation_fn) + layer, qmat = self._filter(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_all'+suffix, natoms=natoms, reuse=reuse, seed = self.seed, trainable = trainable, activation_fn = self.filter_activation_fn, type_embedding=type_embedding) layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()]) qmat = tf.reshape(qmat, [tf.shape(inputs)[0], natoms[0] * self.get_dim_rot_mat_1() * 3]) output.append(layer) @@ -516,75 +525,152 @@ def _compute_dstats_sys_smth (self, def _compute_std (self,sumv2, sumv, sumn) : + if sumn == 0: + return 1e-2 val = np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn)) if np.abs(val) < 1e-2: val = 1e-2 return val - def _filter(self, - inputs, - type_input, - natoms, - activation_fn=tf.nn.tanh, - stddev=1.0, - bavg=0.0, - name='linear', - reuse=None, - seed=None, - trainable = True): + def _concat_type_embedding( + self, + xyz_scatter, + nframes, + natoms, + type_embedding, + ): + te_out_dim = type_embedding.get_shape().as_list()[-1] + nei_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(self.nei_type,dtype=tf.int32)) #nnei*nchnl + nei_embed = tf.tile(nei_embed,(nframes*natoms[0],1)) + nei_embed = tf.reshape(nei_embed,[-1,te_out_dim]) + embedding_input = tf.concat([xyz_scatter,nei_embed],1) + if not self.type_one_side: + atm_embed = embed_atom_type(self.ntypes, natoms, type_embedding) + atm_embed = tf.tile(atm_embed,(1,self.nnei)) + atm_embed = tf.reshape(atm_embed,[-1,te_out_dim]) + embedding_input = tf.concat([embedding_input,atm_embed],1) + return embedding_input + + + def _filter_lower( + self, + start_index, + incrs_index, + inputs, + nframes, + natoms, + type_embedding=None, + is_exclude = False, + activation_fn = None, + bavg = 0.0, + stddev = 1.0, + seed = None, + trainable = True, + suffix = '', + ): + """ + input env matrix, returns R.G + """ + outputs_size = [1] + self.filter_neuron + # cut-out inputs + # with natom x (nei_type_i x 4) + inputs_i = tf.slice (inputs, + [ 0, start_index* 4], + [-1, incrs_index* 4] ) + shape_i = inputs_i.get_shape().as_list() + # with (natom x nei_type_i) x 4 + inputs_reshape = tf.reshape(inputs_i, [-1, 4]) + # 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: + type_embedding = tf.cast(type_embedding, self.filter_precision) + 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') + # with (natom x nei_type_i) x out_size + if self.compress and (not is_exclude): + info = [self.lower, self.upper, self.upper * self.table_config[0], self.table_config[1], self.table_config[2], self.table_config[3]] + if self.type_one_side: + net = 'filter_-1_net_' + str(type_i) + else: + net = 'filter_' + str(type_input) + '_net_' + str(type_i) + return op_module.tabulate_fusion(self.table.data[net].astype(self.filter_np_precision), info, xyz_scatter, tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]), last_layer_size = outputs_size[-1]) + else: + if (not is_exclude): + xyz_scatter = embedding_net( + xyz_scatter, + self.filter_neuron, + self.filter_precision, + activation_fn = activation_fn, + resnet_dt = self.filter_resnet_dt, + name_suffix = suffix, + stddev = stddev, + bavg = bavg, + seed = seed, + trainable = trainable) + else: + w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=GLOBAL_TF_FLOAT_PRECISION) + xyz_scatter = tf.matmul(xyz_scatter, w) + # natom x nei_type_i x out_size + xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1])) + return tf.matmul(tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]), xyz_scatter, transpose_a = True) + + + def _filter( + self, + inputs, + type_input, + natoms, + type_embedding = None, + activation_fn=tf.nn.tanh, + stddev=1.0, + bavg=0.0, + name='linear', + reuse=None, + seed=None, + trainable = True): + nframes = tf.shape(tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]))[0] # natom x (nei x 4) shape = inputs.get_shape().as_list() outputs_size = [1] + self.filter_neuron outputs_size_2 = self.n_axis_neuron with tf.variable_scope(name, reuse=reuse): start_index = 0 - xyz_scatter_total = [] - for type_i in range(self.ntypes): - # cut-out inputs - # with natom x (nei_type_i x 4) - inputs_i = tf.slice (inputs, - [ 0, start_index* 4], - [-1, self.sel_a[type_i]* 4] ) - start_index += self.sel_a[type_i] - shape_i = inputs_i.get_shape().as_list() - # with (natom x nei_type_i) x 4 - inputs_reshape = tf.reshape(inputs_i, [-1, 4]) - # with (natom x nei_type_i) x 1 - xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0,0],[-1,1]),[-1,1]) - # with (natom x nei_type_i) x out_size - if self.compress and (type_input, type_i) not in self.exclude_types: - info = [self.lower, self.upper, self.upper * self.table_config[0], self.table_config[1], self.table_config[2], self.table_config[3]] - if self.type_one_side: - net = 'filter_-1_net_' + str(type_i) - else: - net = 'filter_' + str(type_input) + '_net_' + str(type_i) - if type_i == 0: - xyz_scatter_1 = op_module.tabulate_fusion(self.table.data[net].astype(self.filter_np_precision), info, xyz_scatter, tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]), last_layer_size = outputs_size[-1]) - else: - xyz_scatter_1 += op_module.tabulate_fusion(self.table.data[net].astype(self.filter_np_precision), info, xyz_scatter, tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]), last_layer_size = outputs_size[-1]) - else: - if (type_input, type_i) not in self.exclude_types: - xyz_scatter = embedding_net(xyz_scatter, - self.filter_neuron, - self.filter_precision, - activation_fn = activation_fn, - resnet_dt = self.filter_resnet_dt, - name_suffix = "_"+str(type_i), - stddev = stddev, - bavg = bavg, - seed = seed, - trainable = trainable) - else: - w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=GLOBAL_TF_FLOAT_PRECISION) - xyz_scatter = tf.matmul(xyz_scatter, w) - # natom x nei_type_i x out_size - xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1])) - # xyz_scatter_total.append(xyz_scatter) - if type_i == 0 : - xyz_scatter_1 = tf.matmul(tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]), xyz_scatter, transpose_a = True) - else : - xyz_scatter_1 += tf.matmul(tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]), xyz_scatter, transpose_a = True) + if type_embedding is None: + for type_i in range(self.ntypes): + ret = self._filter_lower( + start_index, self.sel_a[type_i], + inputs, + nframes, + natoms, + type_embedding = type_embedding, + is_exclude = (type_input, type_i) in self.exclude_types, + activation_fn = activation_fn, + stddev = stddev, + bavg = bavg, + seed = seed, + trainable = trainable, + suffix = "_"+str(type_i)) + if type_i == 0: + xyz_scatter_1 = ret + else: + xyz_scatter_1+= ret + start_index += self.sel_a[type_i] + else : + xyz_scatter_1 = self._filter_lower( + start_index, np.cumsum(self.sel_a)[-1], + inputs, + nframes, + natoms, + type_embedding = type_embedding, + is_exclude = False, + activation_fn = activation_fn, + stddev = stddev, + bavg = bavg, + seed = seed, + trainable = trainable) # natom x nei x outputs_size # xyz_scatter = tf.concat(xyz_scatter_total, axis=1) # natom x nei x 4 diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py index 5d60ffcc09..823a77a28c 100644 --- a/deepmd/fit/ener.py +++ b/deepmd/fit/ener.py @@ -8,6 +8,7 @@ from deepmd.utils.network import one_layer from deepmd.descriptor import DescrptLocFrame from deepmd.descriptor import DescrptSeA +from deepmd.utils.type_embed import embed_atom_type from deepmd.env import global_cvt_2_tf_float from deepmd.env import GLOBAL_TF_FLOAT_PRECISION @@ -208,6 +209,72 @@ def compute_input_stats(self, def _compute_std (self, sumv2, sumv, sumn) : return np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn)) + + def _build_lower( + self, + start_index, + natoms, + inputs, + fparam = None, + aparam = None, + bias_atom_e = 0.0, + suffix = '', + reuse = None + ): + # cut-out inputs + inputs_i = tf.slice (inputs, + [ 0, start_index* self.dim_descrpt], + [-1, natoms* self.dim_descrpt] ) + inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt]) + layer = inputs_i + if fparam is not None: + ext_fparam = tf.tile(fparam, [1, natoms]) + ext_fparam = tf.reshape(ext_fparam, [-1, self.numb_fparam]) + ext_fparam = tf.cast(ext_fparam,self.fitting_precision) + layer = tf.concat([layer, ext_fparam], axis = 1) + if aparam is not None: + ext_aparam = tf.slice(aparam, + [ 0, start_index * self.numb_aparam], + [-1, natoms * self.numb_aparam]) + ext_aparam = tf.reshape(ext_aparam, [-1, self.numb_aparam]) + ext_aparam = tf.cast(ext_aparam,self.fitting_precision) + layer = tf.concat([layer, ext_aparam], axis = 1) + + for ii in range(0,len(self.n_neuron)) : + if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : + layer+= one_layer( + layer, + self.n_neuron[ii], + name='layer_'+str(ii)+suffix, + reuse=reuse, + seed = self.seed, + use_timestep = self.resnet_dt, + activation_fn = self.fitting_activation_fn, + precision = self.fitting_precision, + trainable = self.trainable[ii]) + else : + layer = one_layer( + layer, + self.n_neuron[ii], + name='layer_'+str(ii)+suffix, + reuse=reuse, + seed = self.seed, + activation_fn = self.fitting_activation_fn, + precision = self.fitting_precision, + trainable = self.trainable[ii]) + final_layer = one_layer( + layer, + 1, + activation_fn = None, + bavg = bias_atom_e, + name='final_layer_'+suffix, + reuse=reuse, + seed = self.seed, + precision = self.fitting_precision, + trainable = self.trainable[-1]) + + return final_layer + def build (self, @@ -215,7 +282,7 @@ def build (self, natoms : tf.Tensor, input_dict : dict = {}, reuse : bool = None, - suffix : str = '' + suffix : str = '', ) -> tf.Tensor: """ Build the computational graph for fitting net @@ -279,12 +346,14 @@ def build (self, trainable = False, initializer = tf.constant_initializer(self.aparam_inv_std)) - start_index = 0 inputs = tf.cast(tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) + inputs_zero = tf.zeros_like(inputs, dtype=GLOBAL_TF_FLOAT_PRECISION) if bias_atom_e is not None : assert(len(bias_atom_e) == self.ntypes) + fparam = None + aparam = None if self.numb_fparam > 0 : fparam = input_dict['fparam'] fparam = tf.reshape(fparam, [-1, self.numb_fparam]) @@ -294,62 +363,61 @@ def build (self, aparam = tf.reshape(aparam, [-1, self.numb_aparam]) aparam = (aparam - t_aparam_avg) * t_aparam_istd aparam = tf.reshape(aparam, [-1, self.numb_aparam * natoms[0]]) + + if input_dict is not None: + type_embedding = input_dict.get('type_embedding', None) + else: + type_embedding = None + if type_embedding is not None: + atype_embed = embed_atom_type(self.ntypes, natoms, type_embedding) + else: + atype_embed = None - for type_i in range(self.ntypes): - # cut-out inputs - inputs_i = tf.slice (inputs, - [ 0, start_index* self.dim_descrpt], - [-1, natoms[2+type_i]* self.dim_descrpt] ) - inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt]) - layer = inputs_i - if self.numb_fparam > 0 : - ext_fparam = tf.tile(fparam, [1, natoms[2+type_i]]) - ext_fparam = tf.reshape(ext_fparam, [-1, self.numb_fparam]) - ext_fparam = tf.cast(ext_fparam,self.fitting_precision) - layer = tf.concat([layer, ext_fparam], axis = 1) - if self.numb_aparam > 0 : - ext_aparam = tf.slice(aparam, - [ 0, start_index * self.numb_aparam], - [-1, natoms[2+type_i] * self.numb_aparam]) - ext_aparam = tf.reshape(ext_aparam, [-1, self.numb_aparam]) - ext_aparam = tf.cast(ext_aparam,self.fitting_precision) - layer = tf.concat([layer, ext_aparam], axis = 1) - start_index += natoms[2+type_i] - - if bias_atom_e is None : - type_bias_ae = 0.0 - else : - type_bias_ae = bias_atom_e[type_i] - - for ii in range(0,len(self.n_neuron)) : - if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : - layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, trainable = self.trainable[ii]) + if atype_embed is None: + start_index = 0 + for type_i in range(self.ntypes): + if bias_atom_e is None : + type_bias_ae = 0.0 else : - layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, trainable = self.trainable[ii]) - final_layer = one_layer(layer, 1, activation_fn = None, bavg = type_bias_ae, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, trainable = self.trainable[-1]) - - if type_i < len(self.atom_ener) and self.atom_ener[type_i] is not None: - inputs_zero = tf.zeros_like(inputs_i, dtype=GLOBAL_TF_FLOAT_PRECISION) - layer = inputs_zero - if self.numb_fparam > 0 : - layer = tf.concat([layer, ext_fparam], axis = 1) - if self.numb_aparam > 0 : - layer = tf.concat([layer, ext_aparam], axis = 1) - for ii in range(0,len(self.n_neuron)) : - if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : - layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=True, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, trainable = self.trainable[ii]) - else : - layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=True, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, trainable = self.trainable[ii]) - zero_layer = one_layer(layer, 1, activation_fn = None, bavg = type_bias_ae, name='final_layer_type_'+str(type_i)+suffix, reuse=True, seed = self.seed, precision = self.fitting_precision, trainable = self.trainable[-1]) - final_layer += self.atom_ener[type_i] - zero_layer - - final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i]]) - - # concat the results - if type_i == 0: - outs = final_layer - else: - outs = tf.concat([outs, final_layer], axis = 1) + type_bias_ae = bias_atom_e[type_i] + final_layer = self._build_lower( + start_index, natoms[2+type_i], + inputs, fparam, aparam, + bias_atom_e=type_bias_ae, suffix='_type_'+str(type_i)+suffix, reuse=reuse + ) + # concat the results + if type_i < len(self.atom_ener) and self.atom_ener[type_i] is not None: + zero_layer = self._build_lower( + start_index, natoms[2+type_i], + inputs_zero, fparam, aparam, + bias_atom_e=type_bias_ae, suffix='_zero_type_'+str(type_i)+suffix, reuse=reuse + ) + final_layer += self.atom_ener[type_i] - zero_layer + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i]]) + # concat the results + if type_i == 0: + outs = final_layer + else: + outs = tf.concat([outs, final_layer], axis = 1) + start_index += natoms[2+type_i] + # with type embedding + else: + if len(self.atom_ener) > 0: + raise RuntimeError("setting atom_ener is not supported by type embedding") + atype_embed = tf.cast(atype_embed, self.fitting_precision) + type_shape = atype_embed.get_shape().as_list() + inputs = tf.concat( + [tf.reshape(inputs,[-1,self.dim_descrpt]),atype_embed], + axis=1 + ) + self.dim_descrpt = self.dim_descrpt + type_shape[1] + inputs = tf.cast(tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) + final_layer = self._build_lower( + 0, natoms[0], + inputs, fparam, aparam, + bias_atom_e=0.0, suffix=suffix, reuse=reuse + ) + outs = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[0]]) if self.tot_ener_zero: force_tot_ener = 0.0 @@ -363,6 +431,3 @@ def build (self, return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION) - - - diff --git a/deepmd/model/ener.py b/deepmd/model/ener.py index 3cc3d4bd8b..6d7230b6cd 100644 --- a/deepmd/model/ener.py +++ b/deepmd/model/ener.py @@ -15,6 +15,7 @@ def __init__ ( self, descrpt, fitting, + typeebd = None, type_map : List[str] = None, data_stat_nbatch : int = 10, data_stat_protect : float = 1e-2, @@ -55,6 +56,8 @@ def __init__ ( # fitting self.fitting = fitting self.numb_fparam = self.fitting.get_numb_fparam() + # type embedding + self.typeebd = typeebd # other inputs if type_map is None: self.type_map = [] @@ -138,6 +141,15 @@ def build (self, coord = tf.reshape (coord_, [-1, natoms[1] * 3]) atype = tf.reshape (atype_, [-1, natoms[1]]) + # type embedding if any + if self.typeebd is not None: + type_embedding = self.typeebd.build( + self.ntypes, + reuse = reuse, + suffix = suffix, + ) + input_dict['type_embedding'] = type_embedding + dout \ = self.descrpt.build(coord_, atype_, diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index 62934b0faa..ead6840051 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -21,6 +21,7 @@ from deepmd.loss import EnerStdLoss, EnerDipoleLoss, TensorLoss from deepmd.utils.learning_rate import LearningRateExp from deepmd.utils.neighbor_stat import NeighborStat +from deepmd.utils.type_embed import TypeEmbedNet from tensorflow.python.client import timeline from deepmd.env import op_module @@ -84,6 +85,7 @@ def _init_param(self, jdata): model_param = j_must_have(jdata, 'model') descrpt_param = j_must_have(model_param, 'descriptor') fitting_param = j_must_have(model_param, 'fitting_net') + typeebd_param = model_param.get('type_embedding', None) self.model_param = model_param self.descrpt_param = descrpt_param @@ -132,12 +134,25 @@ def _init_param(self, jdata): else : raise RuntimeError('unknow fitting type ' + fitting_type) + # type embedding + if typeebd_param is not None: + self.typeebd = TypeEmbedNet( + typeebd_param['neuron'], + typeebd_param['resnet_dt'], + typeebd_param['seed'], + typeebd_param['activation_function'], + typeebd_param['precision'] + ) + else: + self.typeebd = None + # init model # infer model type by fitting_type if fitting_type == 'ener': self.model = EnerModel( self.descrpt, self.fitting, + self.typeebd, model_param.get('type_map'), model_param.get('data_stat_nbatch', 10), model_param.get('data_stat_protect', 1e-2), diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a29fd3b7f2..a4583b1b40 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -18,6 +18,24 @@ def make_link(content, ref_key): else f'`{content} <#{ref_key}>`_' +def type_embedding_args(): + doc_neuron = 'Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.' + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' + doc_seed = 'Random seed for parameter initialization' + doc_activation_function = f'The activation function in the embedding net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())}' + doc_precision = f'The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())}' + doc_trainable = 'If the parameters in the embedding net are trainable' + + return [ + Argument("neuron", list, optional = True, default = [2, 4, 8], doc = doc_neuron), + Argument("activation_function", str, optional = True, default = 'tanh', doc = doc_activation_function), + Argument("resnet_dt", bool, optional = True, default = False, doc = doc_resnet_dt), + Argument("precision", str, optional = True, default = "float64", doc = doc_precision), + Argument("trainable", bool, optional = True, default = True, doc = doc_trainable), + Argument("seed", [int,None], optional = True, doc = doc_seed), + ] + + # --- Descriptor configurations: --- # def descrpt_local_frame_args (): doc_sel_a = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel_a[i]` gives the selected number of type-i neighbors. The full relative coordinates of the neighbors are used by the descriptor.' @@ -79,7 +97,7 @@ def descrpt_se_t_args(): doc_activation_function = f'The activation function in the embedding net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())}' doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' doc_precision = f'The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())}' - doc_trainable = 'If the parameters in the embedding net is trainable' + doc_trainable = 'If the parameters in the embedding net are trainable' doc_seed = 'Random seed for parameter initialization' doc_set_davg_zero = 'Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used' @@ -119,7 +137,7 @@ def descrpt_se_r_args(): doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' doc_type_one_side = 'Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets' doc_precision = f'The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())}' - doc_trainable = 'If the parameters in the embedding net is trainable' + doc_trainable = 'If the parameters in the embedding net are trainable' doc_seed = 'Random seed for parameter initialization' doc_exclude_types = 'The Excluded types' doc_set_davg_zero = 'Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used' @@ -303,10 +321,11 @@ def modifier_variant_type_args(): doc = doc_modifier_type) -def model_args (): +def model_args (): doc_type_map = 'A list of strings. Give the name to each type of atoms.' doc_data_stat_nbatch = 'The model determines the normalization from the statistics of the data. This key specifies the number of `frames` in each `system` used for statistics.' doc_data_stat_protect = 'Protect parameter for atomic energy regression.' + doc_type_embedding = "The type embedding." doc_descrpt = 'The descriptor of atomic environment.' doc_fitting = 'The fitting of physical properties.' doc_modifier = 'The modifier of model output.' @@ -323,6 +342,7 @@ def model_args (): Argument("smin_alpha", float, optional = True, doc = doc_smin_alpha), Argument("sw_rmin", float, optional = True, doc = doc_sw_rmin), Argument("sw_rmax", float, optional = True, doc = doc_sw_rmax), + Argument("type_embedding", dict, type_embedding_args(), [], doc = doc_type_embedding), 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), diff --git a/deepmd/utils/network.py b/deepmd/utils/network.py index ab32308bdd..a19f0e0d6f 100644 --- a/deepmd/utils/network.py +++ b/deepmd/utils/network.py @@ -90,7 +90,8 @@ def embedding_net(xx, trainable: boolean If the netowk is trainable """ - outputs_size = [1] + network_size + input_shape = xx.get_shape().as_list() + outputs_size = [input_shape[1]] + network_size for ii in range(1, len(outputs_size)): w = tf.get_variable('matrix_'+str(ii)+name_suffix, @@ -150,4 +151,4 @@ def variable_summaries(var: tf.Variable, name: str): tf.summary.scalar('stddev', stddev) tf.summary.scalar('max', tf.reduce_max(var)) tf.summary.scalar('min', tf.reduce_min(var)) - tf.summary.histogram('histogram', var) \ No newline at end of file + tf.summary.histogram('histogram', var) diff --git a/deepmd/utils/type_embed.py b/deepmd/utils/type_embed.py new file mode 100644 index 0000000000..55083c6f9d --- /dev/null +++ b/deepmd/utils/type_embed.py @@ -0,0 +1,137 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.utils.network import one_layer +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import op_module +from deepmd.env import default_tf_session_config +from deepmd.utils.network import embedding_net + +import math +from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, get_np_precision +from deepmd.utils.argcheck import list_to_doc +from deepmd.utils.tabulate import DeepTabulate + + +def embed_atom_type( + ntypes : int, + natoms : tf.Tensor, + type_embedding : tf.Tensor, +): + """ + Make the embedded type for the atoms in system. + The atoms are assumed to be sorted according to the type, + thus their types are described by a `tf.Tensor` natoms, see explanation below. + + Parameters + ---------- + ntypes: + Number of types. + natoms: + The number of atoms. This tensor has the length of Ntypes + 2 + natoms[0]: number of local atoms + natoms[1]: total number of atoms held by this processor + natoms[i]: 2 <= i < Ntypes+2, number of type i atoms + type_embedding: + The type embedding. + It has the shape of [ntypes, embedding_dim] + + Returns + ------- + atom_embedding + The embedded type of each atom. + It has the shape of [numb_atoms, embedding_dim] + """ + te_out_dim = type_embedding.get_shape().as_list()[-1] + atype = [] + for ii in range(ntypes): + atype.append(tf.tile([ii], [natoms[2+ii]])) + atype = tf.concat(atype, axis = 0) + atm_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(atype,dtype=tf.int32)) #(nf*natom)*nchnl + atm_embed = tf.reshape(atm_embed,[-1,te_out_dim]) + return atm_embed + + +class TypeEmbedNet(): + @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + def __init__( + self, + neuron: List[int]=[], + resnet_dt: bool = False, + activation_function: str = 'tanh', + precision: str = 'default', + trainable: bool = True, + seed: int = 1, + )->None: + """ + Constructor + Parameters + ---------- + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + resnet_dt + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + activation_function + The activation function in the embedding net. Supported options are {0} + precision + The precision of the embedding net parameters. Supported options are {1} + trainable + If the weights of embedding net are trainable. + seed + Random seed for initializing the network parameters. + """ + self.neuron = neuron + self.seed = seed + self.filter_resnet_dt = resnet_dt + self.filter_precision = get_precision(precision) + self.filter_activation_fn = get_activation_func(activation_function) + self.trainable = trainable + + + def build( + self, + ntypes: int, + reuse = None, + suffix = '', + ): + """ + Build the computational graph for the descriptor + + Parameters + ---------- + ntypes + Number of atom types. + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Returns + ------- + embedded_types + The computational graph for embedded types + """ + types = tf.convert_to_tensor( + [ii for ii in range(ntypes)], + dtype = tf.int32 + ) + ebd_type = tf.cast(tf.one_hot(tf.cast(types,dtype=tf.int32),int(ntypes)), self.filter_precision) + ebd_type = tf.reshape(ebd_type, [-1, ntypes]) + name = 'type_embed_net' + suffix + with tf.variable_scope(name, reuse=reuse): + ebd_type = embedding_net( + ebd_type, + self.neuron, + activation_fn = self.filter_activation_fn, + precision = self.filter_precision, + resnet_dt = self.filter_resnet_dt, + seed = self.seed, + trainable = self.trainable) + ebd_type = tf.reshape(ebd_type, [-1, self.neuron[-1]]) # nnei * neuron[-1] + self.ebd_type = tf.identity(ebd_type, name ='t_typeebd') + return self.ebd_type + + diff --git a/doc/train-input-auto.rst b/doc/train-input-auto.rst index edfb634654..7700615360 100644 --- a/doc/train-input-auto.rst +++ b/doc/train-input-auto.rst @@ -60,6 +60,62 @@ model: The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided. + .. _`model/type_embedding`: + + type_embedding: + | type: ``dict`` + | argument path: ``model/type_embedding`` + + The type embedding. + + .. _`model/type_embedding/neuron`: + + neuron: + | type: ``list``, optional, default: ``[2, 4, 8]`` + | argument path: ``model/type_embedding/neuron`` + + Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built. + + .. _`model/type_embedding/activation_function`: + + activation_function: + | type: ``str``, optional, default: ``tanh`` + | argument path: ``model/type_embedding/activation_function`` + + The activation function in the embedding net. Supported activation functions are "relu", "relu6", "softplus", "sigmoid", "tanh", "gelu". + + .. _`model/type_embedding/resnet_dt`: + + resnet_dt: + | type: ``bool``, optional, default: ``False`` + | argument path: ``model/type_embedding/resnet_dt`` + + Whether to use a "Timestep" in the skip connection + + .. _`model/type_embedding/precision`: + + precision: + | type: ``str``, optional, default: ``float64`` + | argument path: ``model/type_embedding/precision`` + + The precision of the embedding net parameters, supported options are "default", "float16", "float32", "float64". + + .. _`model/type_embedding/trainable`: + + trainable: + | type: ``bool``, optional, default: ``True`` + | argument path: ``model/type_embedding/trainable`` + + If the parameters in the embedding net are trainable + + .. _`model/type_embedding/seed`: + + seed: + | type: ``int`` | ``NoneType``, optional + | argument path: ``model/type_embedding/seed`` + + Random seed for parameter initialization + .. _`model/descriptor`: descriptor: @@ -341,7 +397,7 @@ model: | type: ``bool``, optional, default: ``True`` | argument path: ``model/descriptor[se_e2_r]/trainable`` - If the parameters in the embedding net is trainable + If the parameters in the embedding net are trainable .. _`model/descriptor[se_e2_r]/seed`: @@ -434,7 +490,7 @@ model: | type: ``bool``, optional, default: ``True`` | argument path: ``model/descriptor[se_e3]/trainable`` - If the parameters in the embedding net is trainable + If the parameters in the embedding net are trainable .. _`model/descriptor[se_e3]/seed`: @@ -1014,7 +1070,8 @@ loss: | type: ``dict``, optional | argument path: ``loss`` - The definition of loss function. The type of the loss depends on the type of the fitting. For fitting type `ener`, the prefactors before energy, force, virial and atomic energy losses may be provided. For fitting type `dipole`, `polar` and `global_polar`, the loss may be an empty `dict` or unset. + The definition of loss function. The loss type should be set to the fitting type or left unset. + \. Depending on the value of *type*, different sub args are accepted. @@ -1024,13 +1081,19 @@ loss: type: | type: ``str`` (flag key), default: ``ener`` | argument path: ``loss/type`` - | possible choices: |code:loss[ener]|_ + | possible choices: |code:loss[ener]|_, |code:loss[dipole]|_, |code:loss[polar]|_, |code:loss[global_polar]|_ - The type of the loss. + The type of the loss. The loss type should be set to the fitting type or left unset. \. .. |code:loss[ener]| replace:: ``ener`` .. _`code:loss[ener]`: `loss[ener]`_ + .. |code:loss[dipole]| replace:: ``dipole`` + .. _`code:loss[dipole]`: `loss[dipole]`_ + .. |code:loss[polar]| replace:: ``polar`` + .. _`code:loss[polar]`: `loss[polar]`_ + .. |code:loss[global_polar]| replace:: ``global_polar`` + .. _`code:loss[global_polar]`: `loss[global_polar]`_ .. |flag:loss/type| replace:: *type* .. _`flag:loss/type`: `loss/type`_ @@ -1113,6 +1176,69 @@ loss: If provided, relative force error will be used in the loss. The difference of force will be normalized by the magnitude of the force in the label with a shift given by `relative_f`, i.e. DF_i / ( || F || + relative_f ) with DF denoting the difference between prediction and label and || F || denoting the L2 norm of the label. + .. _`loss[dipole]`: + + When |flag:loss/type|_ is set to ``dipole``: + + .. _`loss[dipole]/pref_weight`: + + pref_weight: + | type: ``float`` | ``NoneType`` | ``int``, optional, default: ``None`` + | argument path: ``loss[dipole]/pref_weight`` + + The prefactor of the weight of global loss. It should be larger than or equal to 0. If not provided, training will be atomic mode, i.e. atomic label should be provided. + + .. _`loss[dipole]/pref_atomic_weight`: + + pref_atomic_weight: + | type: ``float`` | ``NoneType`` | ``int``, optional, default: ``None`` + | argument path: ``loss[dipole]/pref_atomic_weight`` + + The prefactor of the weight of atomic loss. It should be larger than or equal to 0. If it's not provided and global weight is provided, training will be global mode, i.e. global label should be provided. If both global and atomic weight are not provided, training will be atomic mode, i.e. atomic label should be provided. + + + .. _`loss[polar]`: + + When |flag:loss/type|_ is set to ``polar``: + + .. _`loss[polar]/pref_weight`: + + pref_weight: + | type: ``float`` | ``NoneType`` | ``int``, optional, default: ``None`` + | argument path: ``loss[polar]/pref_weight`` + + The prefactor of the weight of global loss. It should be larger than or equal to 0. If not provided, training will be atomic mode, i.e. atomic label should be provided. + + .. _`loss[polar]/pref_atomic_weight`: + + pref_atomic_weight: + | type: ``float`` | ``NoneType`` | ``int``, optional, default: ``None`` + | argument path: ``loss[polar]/pref_atomic_weight`` + + The prefactor of the weight of atomic loss. It should be larger than or equal to 0. If it's not provided and global weight is provided, training will be global mode, i.e. global label should be provided. If both global and atomic weight are not provided, training will be atomic mode, i.e. atomic label should be provided. + + + .. _`loss[global_polar]`: + + When |flag:loss/type|_ is set to ``global_polar``: + + .. _`loss[global_polar]/pref_weight`: + + pref_weight: + | type: ``float`` | ``NoneType`` | ``int``, optional, default: ``None`` + | argument path: ``loss[global_polar]/pref_weight`` + + The prefactor of the weight of global loss. It should be larger than or equal to 0. If it's not provided and atomic weight is provided, training will be atomic mode, i.e. atomic label should be provided. If both global and atomic weight are not provided, training will be global mode, i.e. global label should be provided. + + .. _`loss[global_polar]/pref_atomic_weight`: + + pref_atomic_weight: + | type: ``float`` | ``NoneType`` | ``int``, optional, default: ``None`` + | argument path: ``loss[global_polar]/pref_atomic_weight`` + + The prefactor of the weight of atomic loss. It should be larger than or equal to 0. If not provided, training will be global mode, i.e. global label should be provided. + + .. _`learning_rate`: learning_rate: diff --git a/examples/water/se_e2_a_tebd/input.json b/examples/water/se_e2_a_tebd/input.json new file mode 100644 index 0000000000..4f8b8698c7 --- /dev/null +++ b/examples/water/se_e2_a_tebd/input.json @@ -0,0 +1,72 @@ +{ + "_comment": " model parameters", + "model": { + "type_map": ["O", "H"], + "type_embedding":{ + "neuron": [2, 4, 8], + "resnet_dt": false, + "seed": 1 + }, + "descriptor" :{ + "type": "se_e2_a", + "sel": [46, 92], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [25, 50, 100], + "resnet_dt": false, + "axis_neuron": 16, + "type_one_side": true, + "seed": 1, + "_comment": " that's all" + }, + "fitting_net" : { + "neuron": [240, 240, 240], + "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": ["../data/data_0/", "../data/data_1/", "../data/data_2/"], + "batch_size": "auto", + "_comment": "that's all" + }, + "validation_data":{ + "systems": ["../data/data_3"], + "batch_size": 1, + "numb_btch": 3, + "_comment": "that's all" + }, + "numb_steps": 1000000, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 1000, + "_comment": "that's all" + }, + + "_comment": "that's all" +} + diff --git a/source/tests/common.py b/source/tests/common.py index 1b1f37c3f5..3bde871077 100644 --- a/source/tests/common.py +++ b/source/tests/common.py @@ -1,4 +1,4 @@ -import os, sys, dpdata +import os, sys, dpdata, shutil import numpy as np import pathlib @@ -22,6 +22,10 @@ def j_loader(filename): return dp_j_loader(tests_path/filename) +def del_data(): + if os.path.isdir('system'): + shutil.rmtree('system') + def gen_data() : tmpdata = Data(rand_pert = 0.1, seed = 1) sys = dpdata.LabeledSystem() diff --git a/source/tests/test_descrpt_se_a_type.py b/source/tests/test_descrpt_se_a_type.py new file mode 100644 index 0000000000..b62d93e514 --- /dev/null +++ b/source/tests/test_descrpt_se_a_type.py @@ -0,0 +1,240 @@ + +import dpdata,os,sys,unittest +import numpy as np +from deepmd.env import tf +import pickle +from common import Data,gen_data, j_loader + +from deepmd.utils.data_system import DataSystem +from deepmd.descriptor import DescrptSeA +from deepmd.fit import EnerFitting +from deepmd.model import EnerModel +from deepmd.common import j_must_have +from deepmd.utils.type_embed import embed_atom_type, TypeEmbedNet + +GLOBAL_ENER_FLOAT_PRECISION = tf.float64 +GLOBAL_TF_FLOAT_PRECISION = tf.float64 +GLOBAL_NP_FLOAT_PRECISION = np.float64 + +class TestModel(unittest.TestCase): + def setUp(self) : + gen_data() + + def test_descriptor_two_sides(self): + jfile = 'water_se_a_type.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') + sel = j_must_have (jdata['model']['descriptor'], 'sel') + ntypes=len(sel) + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt = None) + + test_data = data.get_test () + numb_test = 1 + + # set parameters + jdata['model']['descriptor']['neuron'] = [5, 5, 5] + jdata['model']['descriptor']['axis_neuron'] = 2 + typeebd_param = {'neuron' : [5, 5, 5], + 'resnet_dt': False, + 'seed': 1, + } + + # init models + typeebd = TypeEmbedNet( + neuron = typeebd_param['neuron'], + resnet_dt = typeebd_param['resnet_dt'], + seed = typeebd_param['seed'], + ) + + jdata['model']['descriptor'].pop('type', None) + descrpt = DescrptSeA(**jdata['model']['descriptor']) + + # 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']] + } + 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, [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 + + type_embedding = typeebd.build( + ntypes, + suffix = "_se_a_type_des_ebd_2sdies" + ) + + dout \ + = descrpt.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + {'type_embedding' : type_embedding}, + reuse = False, + suffix = "_se_a_type_des_2sides" + ) + + 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 = tf.Session() + sess.run(tf.global_variables_initializer()) + [model_dout] = sess.run([dout], + feed_dict = feed_dict_test) + model_dout = model_dout.reshape([-1]) + + ref_dout = [0.0005722682145569174,-0.00020202686217742682,-0.00020202686217742682,7.13250554992363e-05,-0.0014770058171250015,0.000521468690207748,-0.001143865186937176,0.0004038453384193948,0.0005617335409639567,-0.00019831394075147532, + 0.00048086740718842236,-0.0001693584775806112,-0.0001693584775806112,5.966987137476082e-05,-0.0012342029581315136,0.00043492340851472783,-0.0009566016612537016,0.00033706767041080107,0.00047065988464132244,-0.0001657950398095401, + 0.0003647849239740657,-0.00013744939018250384,-0.00013744939018250384,5.1825826955234744e-05,-0.00096004206555711,0.00036185565262332876,-0.0007267433909643961,0.0002738914365542745,0.00038019365906978136,-0.00014322754331896057, + 0.0004675256930823109,-0.00017634410399626168,-0.00017634410399626168,6.652672908755666e-05,-0.0012328062885292486,0.00046500213384094614,-0.0009328887521346069,0.0003518668613172834,0.0004877847509912577,-0.00018396318824508986, + 0.0005154794374703516,-0.00019422534512034776,-0.00019422534512034776,7.318151797939947e-05,-0.0013576642997136488,0.0005115548790018505,-0.0010275333676074971,0.00038716440070070385,0.0005376426714609369,-0.00020257810468163985, + 0.0004482204892297628,-0.00016887749501640607,-0.00016887749501640607,6.364643102775375e-05,-0.001181345877677835,0.0004452029242063362,-0.0008941636427724908,0.0003369586197174627,0.0004677878512312651,-0.00017625260641095753] + + places = 10 + for ii in range(model_dout.size) : + self.assertAlmostEqual(model_dout[ii], ref_dout[ii], places = places) + + + def test_descriptor_one_side(self): + jfile = 'water_se_a_type.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') + sel = j_must_have (jdata['model']['descriptor'], 'sel') + ntypes=len(sel) + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt = None) + + test_data = data.get_test () + numb_test = 1 + + # set parameters + jdata['model']['descriptor']['neuron'] = [5, 5, 5] + jdata['model']['descriptor']['axis_neuron'] = 2 + jdata['model']['descriptor']['type_one_side'] = True + typeebd_param = {'neuron' : [5, 5, 5], + 'resnet_dt': False, + 'seed': 1, + } + + # init models + typeebd = TypeEmbedNet( + neuron = typeebd_param['neuron'], + resnet_dt = typeebd_param['resnet_dt'], + seed = typeebd_param['seed'], + ) + + jdata['model']['descriptor'].pop('type', None) + descrpt = DescrptSeA(**jdata['model']['descriptor']) + + # 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']] + } + 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, [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 + + type_embedding = typeebd.build( + ntypes, + suffix = "_se_a_type_des_ebd_1side" + ) + + dout \ + = descrpt.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + {'type_embedding' : type_embedding}, + reuse = False, + suffix = "_se_a_type_des_1side" + ) + + 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 = tf.Session() + sess.run(tf.global_variables_initializer()) + [model_dout] = sess.run([dout], + feed_dict = feed_dict_test) + model_dout = model_dout.reshape([-1]) + + ref_dout = [0.0009704469114440277,0.0007136310372560243,0.0007136310372560243,0.000524968274824758,-0.0019790100690810016,-0.0014556100390424947,-0.001318691223889266,-0.0009698525512440269,0.001937780602605409, + 0.0014251755182315322,0.0008158935519461114,0.0005943870925895051,0.0005943870925895051,0.0004340263490412088,-0.0016539827195947239,-0.0012066241021841376,-0.0011042186455562336,-0.0008051343572505189, + 0.0016229491738044255,0.0011833923257801077,0.0006020440527161554,0.00047526899287409847,0.00047526899287409847,0.00037538142786805136,-0.0012811397377036637,-0.0010116898098710776,-0.0008465095301785942, + -0.0006683577463042215,0.0012459039620461505,0.0009836962283627838,0.00077088529431722,0.0006105807630364827,0.0006105807630364827,0.00048361458700877996,-0.0016444700616024337,-0.001302510079662288, + -0.0010856603485807576,-0.0008598975276238373,0.00159730642327918,0.001265146946434076,0.0008495806081447204,0.000671787466824433,0.000671787466824433,0.0005312928157964384,-0.0018105890543181475, + -0.001431844407277983,-0.0011956722392735362,-0.000945544277375045,0.0017590147511761475,0.0013910348287283414,0.0007393644735054756,0.0005850536182149991,0.0005850536182149991,0.0004631887654949332, + -0.0015760302086346792,-0.0012475134925387294,-0.001041074331192672,-0.0008239586048523492,0.0015319673563669856,0.0012124704278707746] + + places = 10 + for ii in range(model_dout.size) : + self.assertAlmostEqual(model_dout[ii], ref_dout[ii], places = places) + + + + diff --git a/source/tests/test_fitting_ener_type.py b/source/tests/test_fitting_ener_type.py new file mode 100644 index 0000000000..f7209be37c --- /dev/null +++ b/source/tests/test_fitting_ener_type.py @@ -0,0 +1,115 @@ + +import dpdata,os,sys,unittest +import numpy as np +from deepmd.env import tf +import pickle +from common import Data,gen_data, j_loader + +from deepmd.utils.data_system import DataSystem +from deepmd.descriptor import DescrptSeA +from deepmd.fit import EnerFitting +from deepmd.model import EnerModel +from deepmd.common import j_must_have + +GLOBAL_ENER_FLOAT_PRECISION = tf.float64 +GLOBAL_TF_FLOAT_PRECISION = tf.float64 +GLOBAL_NP_FLOAT_PRECISION = np.float64 + +class TestModel(unittest.TestCase): + def setUp(self) : + gen_data() + + def test_fitting(self): + jfile = 'water_se_a_type.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') + sel = j_must_have (jdata['model']['descriptor'], 'sel') + ntypes=len(sel) + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt = None) + + test_data = data.get_test () + numb_test = 1 + + jdata['model']['descriptor']['neuron'] = [5, 5, 5] + jdata['model']['descriptor']['axis_neuron'] = 2 + + jdata['model']['descriptor'].pop('type', None) + descrpt = DescrptSeA(**jdata['model']['descriptor']) + jdata['model']['fitting_net']['descrpt'] = descrpt + fitting = EnerFitting(**jdata['model']['fitting_net']) + + # 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']] + } + + + 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, [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 + + dout = np.array([0.0005722682145569174,-0.00020202686217742682,-0.00020202686217742682,7.13250554992363e-05,-0.0014770058171250015,0.000521468690207748,-0.001143865186937176,0.0004038453384193948,0.0005617335409639567,-0.00019831394075147532, + 0.00048086740718842236,-0.0001693584775806112,-0.0001693584775806112,5.966987137476082e-05,-0.0012342029581315136,0.00043492340851472783,-0.0009566016612537016,0.00033706767041080107,0.00047065988464132244,-0.0001657950398095401, + 0.0003647849239740657,-0.00013744939018250384,-0.00013744939018250384,5.1825826955234744e-05,-0.00096004206555711,0.00036185565262332876,-0.0007267433909643961,0.0002738914365542745,0.00038019365906978136,-0.00014322754331896057, + 0.0004675256930823109,-0.00017634410399626168,-0.00017634410399626168,6.652672908755666e-05,-0.0012328062885292486,0.00046500213384094614,-0.0009328887521346069,0.0003518668613172834,0.0004877847509912577,-0.00018396318824508986, + 0.0005154794374703516,-0.00019422534512034776,-0.00019422534512034776,7.318151797939947e-05,-0.0013576642997136488,0.0005115548790018505,-0.0010275333676074971,0.00038716440070070385,0.0005376426714609369,-0.00020257810468163985, + 0.0004482204892297628,-0.00016887749501640607,-0.00016887749501640607,6.364643102775375e-05,-0.001181345877677835,0.0004452029242063362,-0.0008941636427724908,0.0003369586197174627,0.0004677878512312651,-0.00017625260641095753]) + type_embedding = np.array([1.4916816460764615,0.2720153234707013,-2.4385153754181985,-1.8454294510880027,2.874575701113528,1.1225116575801295,0.4204818970813372,-2.3784087249787587,-1.5053748251050598,2.769329403073084]) + + dout= dout.reshape([-1,10]) + type_embedding = type_embedding.reshape([ntypes,-1]) + atom_ener = fitting.build(tf.convert_to_tensor(dout), + t_natoms, + {'type_embedding':tf.convert_to_tensor(type_embedding)}, + reuse = False, + suffix = "se_a_type_fit_") + + 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 = tf.Session() + sess.run(tf.global_variables_initializer()) + [pred_atom_ener] = sess.run([atom_ener], + feed_dict = feed_dict_test) + + pred_atom_ener=pred_atom_ener.reshape([-1]) + + ref_atom_ener =[10.121733946849165,10.121712105320634,10.143275419743475,10.143299785396826,10.143311150431957,10.143295201182019] + + places = 10 + + + for ii in range(pred_atom_ener.size) : + self.assertAlmostEqual(pred_atom_ener[ii], ref_atom_ener[ii], places = places) + + diff --git a/source/tests/test_model_se_a.py b/source/tests/test_model_se_a.py index 837d344f90..73ae4c7d1e 100644 --- a/source/tests/test_model_se_a.py +++ b/source/tests/test_model_se_a.py @@ -2,7 +2,7 @@ import dpdata,os,sys,unittest import numpy as np from deepmd.env import tf -from common import Data,gen_data, j_loader +from common import Data, gen_data, del_data, j_loader from deepmd.utils.data_system import DataSystem from deepmd.descriptor import DescrptSeA @@ -18,6 +18,104 @@ class TestModel(unittest.TestCase): def setUp(self) : gen_data() + def tearDown(self): + del_data() + + def test_model_atom_ener(self): + jfile = 'water_se_a.json' + jdata = j_loader(jfile) + set_atom_ener = [0.02, 0.01] + jdata['model']['fitting_net']['atom_ener'] = set_atom_ener + + sys = dpdata.LabeledSystem() + sys.data['atom_names'] = ['foo', 'bar'] + sys.data['coords'] = np.array([0, 0, 0, 0, 0, 0]) + sys.data['atom_types'] = [0] + sys.data['cells'] = np.array([np.eye(3) * 30, np.eye(3) * 30]) + nframes = 2 + natoms = 1 + sys.data['coords'] = sys.data['coords'].reshape([nframes,natoms,3]) + sys.data['cells'] = sys.data['cells'].reshape([nframes,3,3]) + sys.data['energies'] = np.zeros([nframes,1]) + sys.data['forces'] = np.zeros([nframes,natoms,3]) + sys.to_deepmd_npy('system', prec=np.float64) + + 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) + descrpt = DescrptSeA(**jdata['model']['descriptor']) + jdata['model']['fitting_net']['descrpt'] = descrpt + fitting = EnerFitting(**jdata['model']['fitting_net']) + model = EnerModel(descrpt, fitting) + + test_data['natoms_vec'] = [1, 1, 1, 0] + + 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_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_atom_ener_0", + reuse = False) + energy = model_pred['energy'] + force = model_pred['force'] + virial = model_pred['virial'] + + feed_dict_test = {t_prop_c: test_data['prop_c'], + t_energy: test_data['energy'] [:numb_test], + t_coord: np.reshape(test_data['coord'] [:numb_test, :], [-1]), + t_box: test_data['box'] [:numb_test, :], + t_type: np.reshape([0], [-1]), + t_natoms: [1, 1, 1, 0], + t_mesh: test_data['default_mesh'], + is_training: False + } + sess = tf.Session() + sess.run(tf.global_variables_initializer()) + [e, f, v] = sess.run([energy, force, virial], + feed_dict = feed_dict_test) + self.assertAlmostEqual(e[0], set_atom_ener[0], places = 10) + + feed_dict_test[t_type] = np.reshape([1], [-1]) + feed_dict_test[t_natoms] = [1, 1, 0, 1] + [e, f, v] = sess.run([energy, force, virial], + feed_dict = feed_dict_test) + self.assertAlmostEqual(e[0], set_atom_ener[1], places = 10) + + + def test_model(self): jfile = 'water_se_a.json' jdata = j_loader(jfile) diff --git a/source/tests/test_model_se_a_srtab.py b/source/tests/test_model_se_a_srtab.py index 7ef371ae3b..991bd9b7e5 100644 --- a/source/tests/test_model_se_a_srtab.py +++ b/source/tests/test_model_se_a_srtab.py @@ -58,6 +58,7 @@ def test_model(self): model = EnerModel( descrpt, fitting, + None, jdata['model'].get('type_map'), jdata['model'].get('data_stat_nbatch'), jdata['model'].get('data_stat_protect'), diff --git a/source/tests/test_model_se_a_type.py b/source/tests/test_model_se_a_type.py new file mode 100644 index 0000000000..ebb9a445cb --- /dev/null +++ b/source/tests/test_model_se_a_type.py @@ -0,0 +1,133 @@ +import dpdata,os,sys,unittest +import numpy as np +from deepmd.env import tf +import pickle +from common import Data,gen_data, j_loader + +from deepmd.utils.data_system import DataSystem +from deepmd.descriptor import DescrptSeA +from deepmd.fit import EnerFitting +from deepmd.model import EnerModel +from deepmd.utils.type_embed import TypeEmbedNet +from deepmd.common import j_must_have + +GLOBAL_ENER_FLOAT_PRECISION = tf.float64 +GLOBAL_TF_FLOAT_PRECISION = tf.float64 +GLOBAL_NP_FLOAT_PRECISION = np.float64 + +class TestModel(unittest.TestCase): + def setUp(self) : + gen_data() + + def test_model(self): + jfile = 'water_se_a_type.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) + descrpt = DescrptSeA(**jdata['model']['descriptor']) + jdata['model']['fitting_net']['descrpt'] = descrpt + fitting = EnerFitting(**jdata['model']['fitting_net']) + typeebd_param = jdata['model']['type_embedding'] + typeebd = TypeEmbedNet( + neuron = typeebd_param['neuron'], + resnet_dt = typeebd_param['resnet_dt'], + seed = typeebd_param['seed']) + 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 + inputs_dict = {} + + model_pred \ + = model.build (t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + inputs_dict, + suffix = "se_a_type", + 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 = tf.Session() + sess.run(tf.global_variables_initializer()) + [e, f, v] = sess.run([energy, force, virial], + feed_dict = feed_dict_test) + # print(sess.run(model.type_embedding)) + # np.savetxt('tmp.out', sess.run(descrpt.dout, feed_dict = feed_dict_test), fmt='%.10e') + # # print(sess.run(model.atype_embed, feed_dict = feed_dict_test)) + # print(sess.run(fitting.inputs, feed_dict = feed_dict_test)) + # print(sess.run(fitting.outs, feed_dict = feed_dict_test)) + # print(sess.run(fitting.atype_embed, feed_dict = feed_dict_test)) + + e = e.reshape([-1]) + f = f.reshape([-1]) + v = v.reshape([-1]) + np.savetxt('e.out', e.reshape([1, -1]), delimiter=',') + np.savetxt('f.out', f.reshape([1, -1]), delimiter=',') + np.savetxt('v.out', v.reshape([1, -1]), delimiter=',') + + refe = [6.049065170680415804e+01] + reff = [1.021832439441947293e-01,1.122650466359011306e-01,3.927874278714531091e-03,1.407089812207832635e-01,1.312473824343091400e-01,-1.228371057389851181e-02,-1.109672154547165501e-01,6.582735820731049070e-02,1.251568633647655391e-03,7.933758749748777428e-02,-1.831777072317984367e-01,-6.173090134630876760e-03,-2.703597126460742794e-01,4.817856571062521104e-02,1.491963457594796399e-02,5.909711543832503466e-02,-1.743406457563475287e-01,-1.642276779780762769e-03] + refv = [-6.932736357193732823e-01,1.453756052949563837e-01,2.138263139115256783e-02,1.453756052949564392e-01,-3.880901656480436612e-01,-7.782259726407755700e-03,2.138263139115256437e-02,-7.782259726407749628e-03,-1.225285973678705374e-03] + + refe = np.reshape(refe, [-1]) + reff = np.reshape(reff, [-1]) + refv = np.reshape(refv, [-1]) + + places = 10 + 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) diff --git a/source/tests/test_type_embed.py b/source/tests/test_type_embed.py new file mode 100644 index 0000000000..4773bf2c45 --- /dev/null +++ b/source/tests/test_type_embed.py @@ -0,0 +1,50 @@ +import dpdata,os,sys,unittest +import numpy as np +from deepmd.env import tf +from deepmd.utils.type_embed import embed_atom_type, TypeEmbedNet + +class TestTypeEbd(unittest.TestCase): + def test_embed_atom_type(self): + ntypes = 3 + natoms = tf.constant([5, 5, 3, 0, 2]) + type_embedding = tf.constant( + [ + [1, 2, 3], + [3, 2, 1], + [7, 7, 7], + ]) + expected_out = [[1,2,3], + [1,2,3], + [1,2,3], + [7,7,7], + [7,7,7]] + atom_embed = embed_atom_type(ntypes, natoms, type_embedding) + sess = tf.Session() + atom_embed = sess.run(atom_embed) + for ii in range(5): + for jj in range(3): + self.assertAlmostEqual( + atom_embed[ii][jj], expected_out[ii][jj], places=10) + + def test_type_embed_net(self): + ten = TypeEmbedNet([2, 4, 8]) + type_embedding = ten.build(2) + sess = tf.Session() + sess.run(tf.global_variables_initializer()) + type_embedding = sess.run(type_embedding) + + expected_out = [ + 1.429967002262267917e+00,-9.138175897677495163e-01,-3.799606588218059633e-01,-2.143157692726757046e-01,2.341138114260268743e+00,-1.568346043255314015e+00,8.917082000854256174e-01,-1.500356675378008209e+00, + 8.955885646123034061e-01,-5.835326470989941061e-01,-1.465708662924672057e+00,-4.052047884085572260e-01,1.367825594590430072e+00,-2.736204307656463497e-01,-4.044263041521370394e-01,-9.438057524881729998e-01 + ] + expected_out = np.reshape(expected_out, [2, 8]) + + # 2 types + self.assertEqual(type_embedding.shape[0], 2) + # size of embedded vec 8 + self.assertEqual(type_embedding.shape[1], 8) + # check value + for ii in range(2): + for jj in range(8): + self.assertAlmostEqual( + type_embedding[ii][jj], expected_out[ii][jj], places=10) diff --git a/source/tests/water_se_a_type.json b/source/tests/water_se_a_type.json new file mode 100644 index 0000000000..5fd5b27ab7 --- /dev/null +++ b/source/tests/water_se_a_type.json @@ -0,0 +1,61 @@ +{ + "_comment": " model parameters", + "model" : { + "type_embedding":{ + "neuron": [5,10,10], + "resnet_dt": false, + "seed": 1 + }, + "descriptor" :{ + "type": "se_a", + "sel": [46, 92], + "rcut_smth": 5.80, + "rcut": 6.00, + "neuron": [25, 50, 100], + "resnet_dt": false, + "type_one_side": false, + "axis_neuron": 16, + "seed": 1 + }, + "fitting_net" : { + "neuron": [240, 240, 240], + "resnet_dt": true, + "seed": 1 + } + }, + + + "_comment": " traing controls", + "systems": ["system"], + "set_prefix": "set", + "stop_batch": 1000000, + "batch_size": 1, + "start_lr": 0.005, + "decay_steps": 5000, + "decay_rate": 0.95, + + "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, + + "seed": 1, + + "_comment": " display and restart", + "_comment": " frequencies counted in batch", + "disp_file": "lcurve.out", + "disp_freq": 100, + "numb_test": 1, + "save_freq": 1000, + "save_ckpt": "model.ckpt", + "load_ckpt": "model.ckpt", + "disp_training": true, + "time_training": true, + "profiling": false, + "profiling_file": "timeline.json", + + "_comment": "that's all" +} +