def _pass_filter(self, inputs, natoms, reuse=None, suffix='', trainable=True): start_index = 0 inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]]) shape = inputs.get_shape().as_list() output = [] output_qmat = [] for type_i in range(self.ntypes): inputs_i = tf.slice(inputs, [0, start_index * self.ndescrpt], [-1, natoms[2 + type_i] * self.ndescrpt]) inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) layer, qmat = self._filter(inputs_i, name='filter_type_' + str(type_i) + suffix, natoms=natoms, reuse=reuse, seed=self.seed, trainable=trainable) layer = tf.reshape( layer, [tf.shape(inputs)[0], natoms[2 + type_i] * self.get_dim_out()]) qmat = tf.reshape(qmat, [ tf.shape(inputs)[0], natoms[2 + type_i] * self.get_dim_rot_mat_1() * 3 ]) output.append(layer) output_qmat.append(qmat) start_index += natoms[2 + type_i] output = tf.concat(output, axis=1) output_qmat = tf.concat(output_qmat, axis=1) return output, output_qmat
def _pass_filter(self, inputs, atype, natoms, input_dict, reuse=None, suffix='', trainable=True): # nf x na x ndescrpt # nf x na x (nnei x 4) inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]) layer, qmat = self._ebd_filter(tf.cast(inputs, self.filter_precision), atype, natoms, input_dict, name='filter_type_all' + suffix, reuse=reuse, seed=self.seed, trainable=trainable, activation_fn=self.filter_activation_fn) output = tf.reshape( layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()]) output_qmat = tf.reshape( qmat, [tf.shape(inputs)[0], natoms[0] * self.get_dim_rot_mat_1() * 3]) return output, output_qmat
def _pass_filter(self, inputs, natoms, reuse = None, suffix = '', trainable = True) : start_index = 0 inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]]) output = [] if not self.type_one_side: for type_i in range(self.ntypes): inputs_i = tf.slice (inputs, [ 0, start_index* self.ndescrpt], [-1, natoms[2+type_i]* self.ndescrpt] ) inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) layer = self._filter_r(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_'+str(type_i)+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn) layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[2+type_i] * self.get_dim_out()]) output.append(layer) start_index += natoms[2+type_i] else : inputs_i = inputs inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) type_i = -1 layer = self._filter_r(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_all'+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn) layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()]) output.append(layer) output = tf.concat(output, axis = 1) return output
def build (self, input_d, rot_mat, natoms, reuse = None, suffix = '') : start_index = 0 inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]]) count = 0 outs_list = [] 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]) rot_mat_i = tf.slice (rot_mat, [ 0, start_index* 9], [-1, natoms[2+type_i]* 9] ) rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3]) start_index += natoms[2+type_i] if not type_i in self.sel_type : continue layer = inputs_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) 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) # (nframes x natoms) x 9 final_layer = one_layer(layer, 9, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, final_layer = True) # (nframes x natoms) x 3 x 3 final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 3, 3]) # (nframes x natoms) x 3 x 3 final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1]) # (nframes x natoms) x 3 x 3(coord) final_layer = tf.matmul(final_layer, rot_mat_i) # (nframes x natoms) x 3(coord) x 3(coord) final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True) # nframes x natoms x 3 x 3 final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3]) # concat the results outs_list.append(final_layer) count += 1 outs = tf.concat(outs_list, axis = 1) tf.summary.histogram('fitting_net_output', outs) return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
def build(self, learning_rate, natoms, model_dict, label_dict, suffix): coord = model_dict['coord'] energy = model_dict['energy'] atom_ener = model_dict['atom_ener'] nframes = tf.shape(atom_ener)[0] natoms = tf.shape(atom_ener)[1] # build energy dipole atom_ener0 = atom_ener - tf.reshape( tf.tile( tf.reshape(energy / global_cvt_2_ener_float(natoms), [-1, 1]), [1, natoms]), [nframes, natoms]) coord = tf.reshape(coord, [nframes, natoms, 3]) atom_ener0 = tf.reshape(atom_ener0, [nframes, 1, natoms]) ener_dipole = tf.matmul(atom_ener0, coord) ener_dipole = tf.reshape(ener_dipole, [nframes, 3]) energy_hat = label_dict['energy'] ener_dipole_hat = label_dict['energy_dipole'] find_energy = label_dict['find_energy'] find_ener_dipole = label_dict['find_energy_dipole'] l2_ener_loss = tf.reduce_mean(tf.square(energy - energy_hat), name='l2_' + suffix) ener_dipole_reshape = tf.reshape(ener_dipole, [-1]) ener_dipole_hat_reshape = tf.reshape(ener_dipole_hat, [-1]) l2_ener_dipole_loss = tf.reduce_mean( tf.square(ener_dipole_reshape - ener_dipole_hat_reshape), name='l2_' + suffix) # atom_norm_ener = 1./ global_cvt_2_ener_float(natoms[0]) atom_norm_ener = 1. / global_cvt_2_ener_float(natoms) pref_e = global_cvt_2_ener_float( find_energy * (self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * learning_rate / self.starter_learning_rate)) pref_ed = global_cvt_2_tf_float( find_ener_dipole * (self.limit_pref_ed + (self.start_pref_ed - self.limit_pref_ed) * learning_rate / self.starter_learning_rate)) l2_loss = 0 more_loss = {} l2_loss += atom_norm_ener * (pref_e * l2_ener_loss) l2_loss += global_cvt_2_ener_float(pref_ed * l2_ener_dipole_loss) more_loss['l2_ener_loss'] = l2_ener_loss more_loss['l2_ener_dipole_loss'] = l2_ener_dipole_loss self.l2_l = l2_loss self.l2_more = more_loss return l2_loss, more_loss
def build (self, input_d, rot_mat, natoms, reuse = None, suffix = '') : start_index = 0 inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]]) count = 0 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]) rot_mat_i = tf.slice (rot_mat, [ 0, start_index* 9], [-1, natoms[2+type_i]* 9] ) rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3]) start_index += natoms[2+type_i] if not type_i in self.sel_type : continue layer = inputs_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, uniform_seed = self.uniform_seed) 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, uniform_seed = self.uniform_seed) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x (nwfc x 3) final_layer = one_layer(layer, self.wfc_numb * 3, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x nwfc(wc) x 3(coord_local) final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.wfc_numb, 3]) # (nframes x natoms) x nwfc(wc) x 3(coord) final_layer = tf.matmul(final_layer, rot_mat_i) # nframes x natoms x nwfc(wc) x 3(coord_local) final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], self.wfc_numb, 3]) # concat the results if count == 0: outs = final_layer else: outs = tf.concat([outs, final_layer], axis = 1) count += 1 tf.summary.histogram('fitting_net_output', outs) return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
def _pass_filter(self, inputs, natoms, reuse=None, suffix='', trainable=True): start_index = 0 inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]]) output = [] if not (self.type_one_side and len(self.exclude_types) == 0): for type_i in range(self.ntypes): inputs_i = tf.slice(inputs, [0, start_index * self.ndescrpt], [-1, natoms[2 + type_i] * self.ndescrpt]) inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) if self.type_one_side: # reuse NN parameters for all types to support type_one_side along with exclude_types reuse = tf.AUTO_REUSE filter_name = 'filter_type_all' + suffix else: filter_name = 'filter_type_' + str(type_i) + suffix layer = self._filter_r(inputs_i, type_i, name=filter_name, natoms=natoms, reuse=reuse, trainable=trainable, activation_fn=self.filter_activation_fn) layer = tf.reshape(layer, [ tf.shape(inputs)[0], natoms[2 + type_i] * self.get_dim_out() ]) output.append(layer) start_index += natoms[2 + type_i] else: inputs_i = inputs inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) type_i = -1 layer = self._filter_r(inputs_i, type_i, name='filter_type_all' + suffix, natoms=natoms, reuse=reuse, trainable=trainable, activation_fn=self.filter_activation_fn) layer = tf.reshape( layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()]) output.append(layer) output = tf.concat(output, axis=1) return output
def _type_embed(self, atype, ndim=1, reuse=None, suffix='', trainable=True): ebd_type = tf.cast(atype, self.filter_precision) ebd_type = ebd_type / float(self.ntypes) ebd_type = tf.reshape(ebd_type, [-1, ndim]) for ii in range(self.type_nlayer): name = 'type_embed_layer_' + str(ii) ebd_type = one_layer(ebd_type, self.type_nchanl, activation_fn=self.filter_activation_fn, precision=self.filter_precision, name=name, reuse=reuse, seed=self.seed + ii, trainable=trainable) name = 'type_embed_layer_' + str(self.type_nlayer) ebd_type = one_layer(ebd_type, self.type_nchanl, activation_fn=None, precision=self.filter_precision, name=name, reuse=reuse, seed=self.seed + ii, trainable=trainable) ebd_type = tf.reshape(ebd_type, [tf.shape(atype)[0], self.type_nchanl]) return ebd_type
def _pass_filter(self, inputs, atype, natoms, input_dict, reuse=None, suffix='', trainable=True): start_index = 0 inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]]) output = [] output_qmat = [] inputs_i = inputs inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) type_i = -1 layer, qmat = self._filter(inputs_i, type_i, name='filter_type_all' + suffix, natoms=natoms, reuse=reuse, trainable=trainable, activation_fn=self.filter_activation_fn) 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) # output_qmat.append(qmat) output = tf.concat(output, axis=1) # output_qmat = tf.concat(output_qmat, axis = 1) return output, None
def _filter_r(self, inputs, type_input, natoms, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, trainable=True): # natom x nei outputs_size = [1] + self.filter_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 inputs_i = tf.slice(inputs, [0, start_index], [-1, self.sel_r[type_i]]) start_index += self.sel_r[type_i] shape_i = inputs_i.get_shape().as_list() # with (natom x nei_type_i) x 1 xyz_scatter = tf.reshape(inputs_i, [-1, 1]) 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=self.seed, trainable=trainable, uniform_seed=self.uniform_seed, initial_variables=self.embedding_net_variables, ) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # natom x nei_type_i x out_size xyz_scatter = tf.reshape( xyz_scatter, (-1, shape_i[1], outputs_size[-1])) else: natom = tf.shape(inputs)[0] xyz_scatter = tf.cast( tf.fill((natom, shape_i[1], outputs_size[-1]), 0.), GLOBAL_TF_FLOAT_PRECISION) xyz_scatter_total.append(xyz_scatter) # natom x nei x outputs_size xyz_scatter = tf.concat(xyz_scatter_total, axis=1) # natom x outputs_size # res_rescale = 1. / 5. result = tf.reduce_mean(xyz_scatter, axis=1) * res_rescale return result
def build(self, input_d, rot_mat, natoms, reuse=None, suffix=''): inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]) outs = self.polar_fitting.build(input_d, rot_mat, natoms, reuse, suffix) # nframes x natoms x 9 outs = tf.reshape(outs, [tf.shape(inputs)[0], -1, 9]) outs = tf.reduce_sum(outs, axis=1) return tf.reshape(outs, [-1])
def _type_embedding_net_one_side(self, mat_g, atype, natoms, name='', reuse=None, seed=None, trainable=True): outputs_size = self.filter_neuron[-1] nframes = tf.shape(mat_g)[0] # (nf x natom x nei) x (outputs_size x chnl x chnl) mat_g = tf.reshape(mat_g, [nframes * natoms[0] * self.nnei, outputs_size]) mat_g = one_layer(mat_g, outputs_size * self.type_nchanl, activation_fn=None, precision=self.filter_precision, name=name + '_amplify', reuse=reuse, seed=self.seed, trainable=trainable) # nf x natom x nei x outputs_size x chnl mat_g = tf.reshape( mat_g, [nframes, natoms[0], self.nnei, outputs_size, self.type_nchanl]) # nf x natom x outputs_size x nei x chnl mat_g = tf.transpose(mat_g, perm=[0, 1, 3, 2, 4]) # nf x natom x outputs_size x (nei x chnl) mat_g = tf.reshape( mat_g, [nframes, natoms[0], outputs_size, self.nnei * self.type_nchanl]) # nei x nchnl ebd_nei_type = self._type_embed(self.nei_type, reuse=reuse, trainable=True, suffix='') # (nei x nchnl) ebd_nei_type = tf.reshape(ebd_nei_type, [self.nnei * self.type_nchanl]) # nf x natom x outputs_size x (nei x chnl) mat_g = tf.multiply(mat_g, ebd_nei_type) # nf x natom x outputs_size x nei x chnl mat_g = tf.reshape( mat_g, [nframes, natoms[0], outputs_size, self.nnei, self.type_nchanl]) # nf x natom x outputs_size x nei mat_g = tf.reduce_mean(mat_g, axis=4) # nf x natom x nei x outputs_size mat_g = tf.transpose(mat_g, perm=[0, 1, 3, 2]) # (nf x natom) x nei x outputs_size mat_g = tf.reshape(mat_g, [nframes * natoms[0], self.nnei, outputs_size]) return mat_g
def _embedding_net(self, inputs, natoms, filter_neuron, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, trainable=True): ''' inputs: nf x na x (nei x 4) outputs: nf x na x nei x output_size ''' # natom x (nei x 4) inputs = tf.reshape(inputs, [-1, self.ndescrpt]) shape = inputs.get_shape().as_list() outputs_size = [1] + filter_neuron with tf.variable_scope(name, reuse=reuse): xyz_scatter_total = [] # with natom x (nei x 4) inputs_i = inputs shape_i = inputs_i.get_shape().as_list() # with (natom x nei) x 4 inputs_reshape = tf.reshape(inputs_i, [-1, 4]) # with (natom x nei) x 1 xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1]) # with (natom x nei) x out_size xyz_scatter = embedding_net(xyz_scatter, self.filter_neuron, self.filter_precision, activation_fn=activation_fn, resnet_dt=self.filter_resnet_dt, stddev=stddev, bavg=bavg, seed=seed, trainable=trainable) # natom x nei x out_size xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1])) xyz_scatter_total.append(xyz_scatter) # natom x nei x outputs_size xyz_scatter = tf.concat(xyz_scatter_total, axis=1) # nf x natom x nei x outputs_size xyz_scatter = tf.reshape( xyz_scatter, [tf.shape(inputs)[0], natoms[0], self.nnei, outputs_size[-1]]) return xyz_scatter
def _enrich(self, dipole, dof=3): coll = [] sel_start_idx = 0 for type_i in range(self.ntypes): if type_i in self.sel_type: di = tf.slice(dipole, [0, sel_start_idx * dof], [-1, self.t_natoms[2 + type_i] * dof]) sel_start_idx += self.t_natoms[2 + type_i] else: di = tf.zeros( [tf.shape(dipole)[0], self.t_natoms[2 + type_i] * dof], dtype=GLOBAL_TF_FLOAT_PRECISION) coll.append(di) return tf.concat(coll, axis=1)
def _enrich(self, dipole, dof = 3): coll = [] sel_start_idx = 0 for type_i in range(self.ntypes): if type_i in self.sel_type: di = tf.slice(dipole, [ 0, sel_start_idx * dof], [-1, self.t_natoms[2+type_i] * dof]) sel_start_idx += self.t_natoms[2+type_i] else: di = tf.zeros([tf.shape(dipole)[0], self.t_natoms[2+type_i] * dof], dtype = global_tf_float_precision) coll.append(di) return tf.concat(coll, axis = 1)
def build(self, input_d, rot_mat, natoms, reuse=None, suffix='') -> tf.Tensor: """ Build the computational graph for fitting net Parameters ---------- input_d The input descriptor rot_mat The rotation matrix from the descriptor. 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 reuse The weights in the networks should be reused when get the variable. suffix Name suffix to identify this descriptor Returns ------- polar The system polarizability """ inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]) outs = self.polar_fitting.build(input_d, rot_mat, natoms, reuse, suffix) # nframes x natoms x 9 outs = tf.reshape(outs, [tf.shape(inputs)[0], -1, 9]) outs = tf.reduce_sum(outs, axis=1) tf.summary.histogram('fitting_net_output', outs) return tf.reshape(outs, [-1])
def _build_fv_graph_inner(self): self.t_ef = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_ef') nf = 10 nfxnas = 64 * nf nfxna = 192 * nf nf = -1 nfxnas = -1 nfxna = -1 self.t_box_reshape = tf.reshape(self.t_box, [-1, 9]) t_nframes = tf.shape(self.t_box_reshape)[0] # (nframes x natoms_sel) x 1 x 3 self.t_ef_reshape = tf.reshape(self.t_ef, [nfxnas, 1, 3]) # (nframes x natoms) x ndescrpt self.descrpt = self.graph.get_tensor_by_name( os.path.join(self.modifier_prefix, 'o_rmat:0')) self.descrpt_deriv = self.graph.get_tensor_by_name( os.path.join(self.modifier_prefix, 'o_rmat_deriv:0')) self.nlist = self.graph.get_tensor_by_name( os.path.join(self.modifier_prefix, 'o_nlist:0')) self.rij = self.graph.get_tensor_by_name( os.path.join(self.modifier_prefix, 'o_rij:0')) # self.descrpt_reshape = tf.reshape(self.descrpt, [nf, 192 * self.ndescrpt]) # self.descrpt_deriv = tf.reshape(self.descrpt_deriv, [nf, 192 * self.ndescrpt * 3]) # nframes x (natoms_sel x 3) self.t_tensor_reshpe = tf.reshape(self.t_tensor, [t_nframes, -1]) # nframes x (natoms x 3) self.t_tensor_reshpe = self._enrich(self.t_tensor_reshpe, dof=3) # (nframes x natoms) x 3 self.t_tensor_reshpe = tf.reshape(self.t_tensor_reshpe, [nfxna, 3]) # (nframes x natoms) x 1 self.t_dipole_x = tf.slice(self.t_tensor_reshpe, [0, 0], [nfxna, 1]) self.t_dipole_y = tf.slice(self.t_tensor_reshpe, [0, 1], [nfxna, 1]) self.t_dipole_z = tf.slice(self.t_tensor_reshpe, [0, 2], [nfxna, 1]) self.t_dipole_z = tf.reshape(self.t_dipole_z, [nfxna, 1]) # (nframes x natoms) x ndescrpt [self.t_dipole_x_d] = tf.gradients(self.t_dipole_x, self.descrpt) [self.t_dipole_y_d] = tf.gradients(self.t_dipole_y, self.descrpt) [self.t_dipole_z_d] = tf.gradients(self.t_dipole_z, self.descrpt) # nframes x (natoms x ndescrpt) self.t_dipole_x_d = tf.reshape(self.t_dipole_x_d, [-1, self.t_natoms[0] * self.ndescrpt]) self.t_dipole_y_d = tf.reshape(self.t_dipole_y_d, [-1, self.t_natoms[0] * self.ndescrpt]) self.t_dipole_z_d = tf.reshape(self.t_dipole_z_d, [-1, self.t_natoms[0] * self.ndescrpt]) # nframes x (natoms_sel x ndescrpt) self.t_dipole_x_d = self._slice_descrpt_deriv(self.t_dipole_x_d) self.t_dipole_y_d = self._slice_descrpt_deriv(self.t_dipole_y_d) self.t_dipole_z_d = self._slice_descrpt_deriv(self.t_dipole_z_d) # (nframes x natoms_sel) x ndescrpt self.t_dipole_x_d = tf.reshape(self.t_dipole_x_d, [nfxnas, self.ndescrpt]) self.t_dipole_y_d = tf.reshape(self.t_dipole_y_d, [nfxnas, self.ndescrpt]) self.t_dipole_z_d = tf.reshape(self.t_dipole_z_d, [nfxnas, self.ndescrpt]) # (nframes x natoms_sel) x 3 x ndescrpt self.t_dipole_d = tf.concat( [self.t_dipole_x_d, self.t_dipole_y_d, self.t_dipole_z_d], axis=1) self.t_dipole_d = tf.reshape(self.t_dipole_d, [nfxnas, 3 * self.ndescrpt]) # (nframes x natoms_sel) x 3 x ndescrpt self.t_dipole_d = tf.reshape(self.t_dipole_d, [-1, 3, self.ndescrpt]) # (nframes x natoms_sel) x 1 x ndescrpt self.t_ef_d = tf.matmul(self.t_ef_reshape, self.t_dipole_d) # nframes x (natoms_sel x ndescrpt) self.t_ef_d = tf.reshape(self.t_ef_d, [t_nframes, -1]) # nframes x (natoms x ndescrpt) self.t_ef_d = self._enrich(self.t_ef_d, dof=self.ndescrpt) self.t_ef_d = tf.reshape(self.t_ef_d, [nf, self.t_natoms[0] * self.ndescrpt]) # t_ef_d is force (with -1), prod_forc takes deriv, so we need the opposite self.t_ef_d_oppo = -self.t_ef_d force = op_module.prod_force_se_a(self.t_ef_d_oppo, self.descrpt_deriv, self.nlist, self.t_natoms, n_a_sel=self.nnei_a, n_r_sel=self.nnei_r) virial, atom_virial \ = op_module.prod_virial_se_a (self.t_ef_d_oppo, self.descrpt_deriv, self.rij, self.nlist, self.t_natoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) force = tf.identity(force, name='o_dm_force') virial = tf.identity(virial, name='o_dm_virial') atom_virial = tf.identity(atom_virial, name='o_dm_av') return force, virial, atom_virial
def build( self, inputs: tf.Tensor, natoms: tf.Tensor, input_dict: dict = None, reuse: bool = None, suffix: str = '', ) -> tf.Tensor: """ Build the computational graph for fitting net Parameters ---------- inputs The input descriptor input_dict Additional dict for inputs. if numb_fparam > 0, should have input_dict['fparam'] if numb_aparam > 0, should have input_dict['aparam'] 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 reuse The weights in the networks should be reused when get the variable. suffix Name suffix to identify this descriptor Returns ------- ener The system energy """ if input_dict is None: input_dict = {} bias_atom_e = self.bias_atom_e if self.numb_fparam > 0 and (self.fparam_avg is None or self.fparam_inv_std is None): raise RuntimeError( 'No data stat result. one should do data statisitic, before build' ) if self.numb_aparam > 0 and (self.aparam_avg is None or self.aparam_inv_std is None): raise RuntimeError( 'No data stat result. one should do data statisitic, before build' ) with tf.variable_scope('fitting_attr' + suffix, reuse=reuse): t_dfparam = tf.constant(self.numb_fparam, name='dfparam', dtype=tf.int32) t_daparam = tf.constant(self.numb_aparam, name='daparam', dtype=tf.int32) if self.numb_fparam > 0: t_fparam_avg = tf.get_variable( 't_fparam_avg', self.numb_fparam, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(self.fparam_avg)) t_fparam_istd = tf.get_variable( 't_fparam_istd', self.numb_fparam, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(self.fparam_inv_std)) if self.numb_aparam > 0: t_aparam_avg = tf.get_variable( 't_aparam_avg', self.numb_aparam, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(self.aparam_avg)) t_aparam_istd = tf.get_variable( 't_aparam_istd', self.numb_aparam, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(self.aparam_inv_std)) inputs = tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]]) if len(self.atom_ener): # only for atom_ener nframes = input_dict.get('nframes') if nframes is not None: # like inputs, but we don't want to add a dependency on inputs inputs_zero = tf.zeros((nframes, self.dim_descrpt * natoms[0]), dtype=self.fitting_precision) else: inputs_zero = tf.zeros_like(inputs, dtype=self.fitting_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]) fparam = (fparam - t_fparam_avg) * t_fparam_istd if self.numb_aparam > 0: aparam = input_dict['aparam'] 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]]) type_embedding = input_dict.get('type_embedding', None) if type_embedding is not None: atype_embed = embed_atom_type(self.ntypes, natoms, type_embedding) atype_embed = tf.tile(atype_embed, [tf.shape(inputs)[0], 1]) else: atype_embed = None if atype_embed is None: start_index = 0 outs_list = [] for type_i in range(self.ntypes): if bias_atom_e is None: type_bias_ae = 0.0 else: 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='_type_' + str(type_i) + suffix, reuse=True) final_layer += self.atom_ener[type_i] - zero_layer final_layer = tf.reshape( final_layer, [tf.shape(inputs)[0], natoms[2 + type_i]]) outs_list.append(final_layer) start_index += natoms[2 + type_i] # concat the results # concat once may be faster than multiple concat outs = tf.concat(outs_list, axis=1) # 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.reshape(inputs, [-1, self.dim_descrpt * natoms[0]]) 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]]) # add atom energy bias; TF will broadcast to all batches # tf.repeat is avaiable in TF>=2.1 or TF 1.15 _TF_VERSION = Version(TF_VERSION) if (Version('1.15') <= _TF_VERSION < Version('2') or _TF_VERSION >= Version('2.1')) and self.bias_atom_e is not None: outs += tf.repeat( tf.Variable(self.bias_atom_e, dtype=self.fitting_precision, trainable=False, name="bias_atom_ei"), natoms[2:]) if self.tot_ener_zero: force_tot_ener = 0.0 outs = tf.reshape(outs, [-1, natoms[0]]) outs_mean = tf.reshape(tf.reduce_mean(outs, axis=1), [-1, 1]) outs_mean = outs_mean - tf.ones_like( outs_mean, dtype=GLOBAL_TF_FLOAT_PRECISION) * ( force_tot_ener / global_cvt_2_tf_float(natoms[0])) outs = outs - outs_mean outs = tf.reshape(outs, [-1]) tf.summary.histogram('fitting_net_output', outs) return tf.reshape(outs, [-1])
def _ebd_filter(self, inputs, atype, natoms, input_dict, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, trainable=True): outputs_size = self.filter_neuron[-1] outputs_size_2 = self.n_axis_neuron # nf x natom x (nei x 4) nframes = tf.shape(inputs)[0] shape = tf.reshape(inputs, [-1, self.ndescrpt]).get_shape().as_list() # nf x natom x nei x outputs_size mat_g = self._embedding_net(inputs, natoms, self.filter_neuron, activation_fn=activation_fn, stddev=stddev, bavg=bavg, name=name, reuse=reuse, seed=seed, trainable=trainable) # nf x natom x nei x outputs_size mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size]) # (nf x natom) x nei x outputs_size if self.type_one_side: if self.numb_aparam > 0: aparam = input_dict['aparam'] xyz_scatter \ = self._type_embedding_net_one_side_aparam(mat_g, atype, natoms, aparam, name = name, reuse = reuse, seed = seed, trainable = trainable) else: xyz_scatter \ = self._type_embedding_net_one_side(mat_g, atype, natoms, name = name, reuse = reuse, seed = seed, trainable = trainable) else: xyz_scatter \ = self._type_embedding_net_two_sides(mat_g, atype, natoms, name = name, reuse = reuse, seed = seed, trainable = trainable) # natom x nei x 4 inputs_reshape = tf.reshape(inputs, [-1, shape[1] // 4, 4]) # natom x 4 x outputs_size xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a=True) xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape[1]) # natom x 4 x outputs_size_2 xyz_scatter_2 = tf.slice(xyz_scatter_1, [0, 0, 0], [-1, -1, outputs_size_2]) # # natom x 3 x outputs_size_2 # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1]) # natom x 3 x outputs_size_1 qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1]) # natom x outputs_size_2 x 3 qmat = tf.transpose(qmat, perm=[0, 2, 1]) # natom x outputs_size x outputs_size_2 result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True) # natom x (outputs_size x outputs_size_2) result = tf.reshape(result, [-1, outputs_size_2 * outputs_size]) return result, qmat
def build(self, learning_rate, natoms, model_dict, label_dict, suffix): polar_hat = label_dict[self.label_name] atomic_polar_hat = label_dict["atomic_" + self.label_name] polar = tf.reshape(model_dict[self.tensor_name], [-1]) find_global = label_dict['find_' + self.label_name] find_atomic = label_dict['find_atomic_' + self.label_name] # YHT: added for global / local dipole combination l2_loss = global_cvt_2_tf_float(0.0) more_loss = { "local_loss": global_cvt_2_tf_float(0.0), "global_loss": global_cvt_2_tf_float(0.0) } if self.local_weight > 0.0: local_loss = global_cvt_2_tf_float(find_atomic) * tf.reduce_mean( tf.square(self.scale * (polar - atomic_polar_hat)), name='l2_' + suffix) more_loss['local_loss'] = local_loss l2_loss += self.local_weight * local_loss self.l2_loss_local_summary = tf.summary.scalar( 'l2_local_loss', tf.sqrt(more_loss['local_loss'])) if self.global_weight > 0.0: # Need global loss atoms = 0 if self.type_sel is not None: for w in self.type_sel: atoms += natoms[2 + w] else: atoms = natoms[0] nframes = tf.shape(polar)[0] // self.tensor_size // atoms # get global results global_polar = tf.reshape( tf.reduce_sum(tf.reshape(polar, [nframes, -1, self.tensor_size]), axis=1), [-1]) #if self.atomic: # If label is local, however # global_polar_hat = tf.reshape(tf.reduce_sum(tf.reshape( # polar_hat, [nframes, -1, self.tensor_size]), axis=1),[-1]) #else: # global_polar_hat = polar_hat global_loss = global_cvt_2_tf_float(find_global) * tf.reduce_mean( tf.square(self.scale * (global_polar - polar_hat)), name='l2_' + suffix) more_loss['global_loss'] = global_loss self.l2_loss_global_summary = tf.summary.scalar( 'l2_global_loss', tf.sqrt(more_loss['global_loss']) / global_cvt_2_tf_float(atoms)) # YWolfeee: should only consider atoms with dipole, i.e. atoms # atom_norm = 1./ global_cvt_2_tf_float(natoms[0]) atom_norm = 1. / global_cvt_2_tf_float(atoms) global_loss *= atom_norm l2_loss += self.global_weight * global_loss self.l2_more = more_loss self.l2_l = l2_loss self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss)) return l2_loss, more_loss
def _pass_filter(self, inputs, atype, natoms, input_dict, 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 and type_embedding is None: for type_i in range(self.ntypes): inputs_i = tf.slice(inputs, [0, start_index * self.ndescrpt], [-1, natoms[2 + type_i] * self.ndescrpt]) inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) layer, qmat = self._filter( tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_' + str(type_i) + suffix, natoms=natoms, reuse=reuse, trainable=trainable, activation_fn=self.filter_activation_fn) layer = tf.reshape(layer, [ tf.shape(inputs)[0], natoms[2 + type_i] * self.get_dim_out() ]) qmat = tf.reshape(qmat, [ tf.shape(inputs)[0], natoms[2 + type_i] * self.get_dim_rot_mat_1() * 3 ]) output.append(layer) output_qmat.append(qmat) start_index += natoms[2 + type_i] else: 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, 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) output_qmat.append(qmat) output = tf.concat(output, axis=1) output_qmat = tf.concat(output_qmat, axis=1) return output, output_qmat
def build (self, coord_, atype_, natoms, box, mesh, input_dict, frz_model = None, suffix = '', reuse = None): if input_dict is None: input_dict = {} with tf.variable_scope('model_attr' + suffix, reuse = reuse) : t_tmap = tf.constant(' '.join(self.type_map), name = 'tmap', dtype = tf.string) t_mt = tf.constant(self.model_type, name = 'model_type', dtype = tf.string) t_ver = tf.constant(MODEL_VERSION, name = 'model_version', dtype = tf.string) if self.srtab is not None : tab_info, tab_data = self.srtab.get() self.tab_info = tf.get_variable('t_tab_info', tab_info.shape, dtype = tf.float64, trainable = False, initializer = tf.constant_initializer(tab_info, dtype = tf.float64)) self.tab_data = tf.get_variable('t_tab_data', tab_data.shape, dtype = tf.float64, trainable = False, initializer = tf.constant_initializer(tab_data, dtype = tf.float64)) coord = tf.reshape (coord_, [-1, natoms[1] * 3]) atype = tf.reshape (atype_, [-1, natoms[1]]) input_dict['nframes'] = tf.shape(coord)[0] # 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 if frz_model == None: dout \ = self.descrpt.build(coord_, atype_, natoms, box, mesh, input_dict, suffix = suffix, reuse = reuse) dout = tf.identity(dout, name='o_descriptor') else: tf.constant(self.rcut, name = 'descrpt_attr/rcut', dtype = GLOBAL_TF_FLOAT_PRECISION) tf.constant(self.ntypes, name = 'descrpt_attr/ntypes', dtype = tf.int32) feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh) return_elements = [*self.descrpt.get_tensor_names(), 'o_descriptor:0'] imported_tensors \ = self._import_graph_def_from_frz_model(frz_model, feed_dict, return_elements) dout = imported_tensors[-1] self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1]) if self.srtab is not None : nlist, rij, sel_a, sel_r = self.descrpt.get_nlist() nnei_a = np.cumsum(sel_a)[-1] nnei_r = np.cumsum(sel_r)[-1] atom_ener = self.fitting.build (dout, natoms, input_dict, reuse = reuse, suffix = suffix) if self.srtab is not None : sw_lambda, sw_deriv \ = op_module.soft_min_switch(atype, rij, nlist, natoms, sel_a = sel_a, sel_r = sel_r, alpha = self.smin_alpha, rmin = self.sw_rmin, rmax = self.sw_rmax) inv_sw_lambda = 1.0 - sw_lambda # NOTICE: # atom energy is not scaled, # force and virial are scaled tab_atom_ener, tab_force, tab_atom_virial \ = op_module.pair_tab(self.tab_info, self.tab_data, atype, rij, nlist, natoms, sw_lambda, sel_a = sel_a, sel_r = sel_r) energy_diff = tab_atom_ener - tf.reshape(atom_ener, [-1, natoms[0]]) tab_atom_ener = tf.reshape(sw_lambda, [-1]) * tf.reshape(tab_atom_ener, [-1]) atom_ener = tf.reshape(inv_sw_lambda, [-1]) * atom_ener energy_raw = tab_atom_ener + atom_ener else : energy_raw = atom_ener energy_raw = tf.reshape(energy_raw, [-1, natoms[0]], name = 'o_atom_energy'+suffix) energy = tf.reduce_sum(global_cvt_2_ener_float(energy_raw), axis=1, name='o_energy'+suffix) force, virial, atom_virial \ = self.descrpt.prod_force_virial (atom_ener, natoms) if self.srtab is not None : sw_force \ = op_module.soft_min_force(energy_diff, sw_deriv, nlist, natoms, n_a_sel = nnei_a, n_r_sel = nnei_r) force = force + sw_force + tab_force force = tf.reshape (force, [-1, 3 * natoms[1]], name = "o_force"+suffix) if self.srtab is not None : sw_virial, sw_atom_virial \ = op_module.soft_min_virial (energy_diff, sw_deriv, rij, nlist, natoms, n_a_sel = nnei_a, n_r_sel = nnei_r) atom_virial = atom_virial + sw_atom_virial + tab_atom_virial virial = virial + sw_virial \ + tf.reduce_sum(tf.reshape(tab_atom_virial, [-1, natoms[1], 9]), axis = 1) virial = tf.reshape (virial, [-1, 9], name = "o_virial"+suffix) atom_virial = tf.reshape (atom_virial, [-1, 9 * natoms[1]], name = "o_atom_virial"+suffix) model_dict = {} model_dict['energy'] = energy model_dict['force'] = force model_dict['virial'] = virial model_dict['atom_ener'] = energy_raw model_dict['atom_virial'] = atom_virial model_dict['coord'] = coord model_dict['atype'] = atype return model_dict
def build(self, input_d, rot_mat, natoms, reuse=None, suffix=''): start_index = 0 inputs = tf.cast( tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]]) count = 0 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]) rot_mat_i = tf.slice(rot_mat, [0, start_index * self.dim_rot_mat], [-1, natoms[2 + type_i] * self.dim_rot_mat]) rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) start_index += natoms[2 + type_i] if not type_i in self.sel_type: continue layer = inputs_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) 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) # (nframes x natoms) x naxis final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn=None, name='final_layer_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed, precision=self.fitting_precision) # (nframes x natoms) x 1 * naxis final_layer = tf.reshape(final_layer, [ tf.shape(inputs)[0] * natoms[2 + type_i], 1, self.dim_rot_mat_1 ]) # (nframes x natoms) x 1 x 3(coord) final_layer = tf.matmul(final_layer, rot_mat_i) # nframes x natoms x 3 final_layer = tf.reshape( final_layer, [tf.shape(inputs)[0], natoms[2 + type_i], 3]) # concat the results if count == 0: outs = final_layer else: outs = tf.concat([outs, final_layer], axis=1) count += 1 return tf.cast(tf.reshape(outs, [-1]), global_tf_float_precision)
def build(self, inputs, input_dict, natoms, reuse=None, suffix=''): bias_atom_e = self.bias_atom_e if self.numb_fparam > 0 and (self.fparam_avg is None or self.fparam_inv_std is None): raise RuntimeError( 'No data stat result. one should do data statisitic, before build' ) if self.numb_aparam > 0 and (self.aparam_avg is None or self.aparam_inv_std is None): raise RuntimeError( 'No data stat result. one should do data statisitic, before build' ) with tf.variable_scope('fitting_attr' + suffix, reuse=reuse): t_dfparam = tf.constant(self.numb_fparam, name='dfparam', dtype=tf.int32) t_daparam = tf.constant(self.numb_aparam, name='daparam', dtype=tf.int32) if self.numb_fparam > 0: t_fparam_avg = tf.get_variable( 't_fparam_avg', self.numb_fparam, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(self.fparam_avg)) t_fparam_istd = tf.get_variable( 't_fparam_istd', self.numb_fparam, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(self.fparam_inv_std)) if self.numb_aparam > 0: t_aparam_avg = tf.get_variable( 't_aparam_avg', self.numb_aparam, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(self.aparam_avg)) t_aparam_istd = tf.get_variable( 't_aparam_istd', self.numb_aparam, dtype=global_tf_float_precision, 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) if bias_atom_e is not None: assert (len(bias_atom_e) == self.ntypes) if self.numb_fparam > 0: fparam = input_dict['fparam'] fparam = tf.reshape(fparam, [-1, self.numb_fparam]) fparam = (fparam - t_fparam_avg) * t_fparam_istd if self.numb_aparam > 0: aparam = input_dict['aparam'] 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]]) 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]) 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]) 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]) 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) return tf.cast(tf.reshape(outs, [-1]), global_tf_float_precision)
def build(self, input_d, rot_mat, natoms, reuse=None, suffix=''): start_index = 0 inputs = tf.cast( tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]]) count = 0 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]) rot_mat_i = tf.slice(rot_mat, [0, start_index * self.dim_rot_mat], [-1, natoms[2 + type_i] * self.dim_rot_mat]) rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) start_index += natoms[2 + type_i] if not type_i in self.sel_type: continue layer = inputs_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) 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) if self.fit_diag: bavg = np.zeros(self.dim_rot_mat_1) # bavg[0] = self.avgeig[0] # bavg[1] = self.avgeig[1] # bavg[2] = self.avgeig[2] # (nframes x natoms) x naxis final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn=None, name='final_layer_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed, bavg=bavg, precision=self.fitting_precision) # (nframes x natoms) x naxis final_layer = tf.reshape(final_layer, [ tf.shape(inputs)[0] * natoms[2 + type_i], self.dim_rot_mat_1 ]) # (nframes x natoms) x naxis x naxis final_layer = tf.matrix_diag(final_layer) else: bavg = np.zeros(self.dim_rot_mat_1 * self.dim_rot_mat_1) # bavg[0*self.dim_rot_mat_1+0] = self.avgeig[0] # bavg[1*self.dim_rot_mat_1+1] = self.avgeig[1] # bavg[2*self.dim_rot_mat_1+2] = self.avgeig[2] # (nframes x natoms) x (naxis x naxis) final_layer = one_layer( layer, self.dim_rot_mat_1 * self.dim_rot_mat_1, activation_fn=None, name='final_layer_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed, bavg=bavg, precision=self.fitting_precision) # (nframes x natoms) x naxis x naxis final_layer = tf.reshape(final_layer, [ tf.shape(inputs)[0] * natoms[2 + type_i], self.dim_rot_mat_1, self.dim_rot_mat_1 ]) # (nframes x natoms) x naxis x naxis final_layer = final_layer + tf.transpose(final_layer, perm=[0, 2, 1]) # (nframes x natoms) x naxis x 3(coord) final_layer = tf.matmul(final_layer, rot_mat_i) # (nframes x natoms) x 3(coord) x 3(coord) final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a=True) # nframes x natoms x 3 x 3 final_layer = tf.reshape( final_layer, [tf.shape(inputs)[0], natoms[2 + type_i], 3, 3]) # shift and scale sel_type_idx = self.sel_type.index(type_i) final_layer = final_layer * self.scale[sel_type_idx] final_layer = final_layer + self.diag_shift[sel_type_idx] * tf.eye( 3, batch_shape=[tf.shape(inputs)[0], natoms[2 + type_i]], dtype=global_tf_float_precision) # concat the results if count == 0: outs = final_layer else: outs = tf.concat([outs, final_layer], axis=1) count += 1 return tf.cast(tf.reshape(outs, [-1]), global_tf_float_precision)
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, 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 all_excluded = all([(type_input, type_i) in self.exclude_types for type_i in range(self.ntypes)]) if all_excluded: # all types are excluded so result and qmat should be zeros # we can safaly return a zero matrix... # See also https://stackoverflow.com/a/34725458/9567349 # result: natom x outputs_size x outputs_size_2 # qmat: natom x outputs_size x 3 natom = tf.shape(inputs)[0] result = tf.cast( tf.fill((natom, outputs_size_2, outputs_size[-1]), 0.), GLOBAL_TF_FLOAT_PRECISION) qmat = tf.cast(tf.fill((natom, outputs_size[-1], 3), 0.), GLOBAL_TF_FLOAT_PRECISION) return result, qmat with tf.variable_scope(name, reuse=reuse): start_index = 0 type_i = 0 # natom x 4 x outputs_size if type_embedding is None: rets = [] for type_i in range(self.ntypes): ret = self._filter_lower(type_i, type_input, 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, trainable=trainable, suffix="_" + str(type_i)) if (type_input, type_i) not in self.exclude_types: # add zero is meaningless; skip rets.append(ret) start_index += self.sel_a[type_i] # faster to use accumulate_n than multiple add xyz_scatter_1 = tf.accumulate_n(rets) else: xyz_scatter_1 = self._filter_lower( type_i, type_input, 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, trainable=trainable) # natom x nei x outputs_size # xyz_scatter = tf.concat(xyz_scatter_total, axis=1) # natom x nei x 4 # inputs_reshape = tf.reshape(inputs, [-1, shape[1]//4, 4]) # natom x 4 x outputs_size # xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a = True) if self.original_sel is None: # shape[1] = nnei * 4 nnei = shape[1] / 4 else: nnei = tf.cast( tf.Variable(np.sum(self.original_sel), dtype=tf.int32, trainable=False, name="nnei"), self.filter_precision) xyz_scatter_1 = xyz_scatter_1 / nnei # natom x 4 x outputs_size_2 xyz_scatter_2 = tf.slice(xyz_scatter_1, [0, 0, 0], [-1, -1, outputs_size_2]) # # natom x 3 x outputs_size_2 # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1]) # natom x 3 x outputs_size_1 qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1]) # natom x outputs_size_1 x 3 qmat = tf.transpose(qmat, perm=[0, 2, 1]) # natom x outputs_size x outputs_size_2 result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True) # natom x (outputs_size x outputs_size_2) result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]]) return result, qmat
def _filter_lower( self, type_i, type_input, start_index, incrs_index, inputs, nframes, natoms, type_embedding=None, is_exclude=False, activation_fn=None, bavg=0.0, stddev=1.0, 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() natom = tf.shape(inputs_i)[0] # 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: 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' ) # natom x 4 x outputs_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_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 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=self.seed, trainable=trainable, uniform_seed=self.uniform_seed, initial_variables=self.embedding_net_variables, mixed_prec=self.mixed_prec) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift else: # we can safely return the final xyz_scatter filled with zero directly return tf.cast(tf.fill((natom, 4, outputs_size[-1]), 0.), self.filter_precision) # natom x nei_type_i x out_size xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1])) # When using tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]) below # [588 24] -> [588 6 4] correct # but if sel is zero # [588 0] -> [147 0 4] incorrect; the correct one is [588 0 4] # So we need to explicitly assign the shape to tf.shape(inputs_i)[0] instead of -1 # natom x 4 x outputs_size return tf.matmul(tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), xyz_scatter, transpose_a=True)
def build (self, input_d : tf.Tensor, rot_mat : tf.Tensor, natoms : tf.Tensor, reuse : bool = None, suffix : str = '') : """ Build the computational graph for fitting net Parameters ---------- input_d The input descriptor rot_mat The rotation matrix from the descriptor. 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 reuse The weights in the networks should be reused when get the variable. suffix Name suffix to identify this descriptor Returns ------- atomic_polar The atomic polarizability """ start_index = 0 inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]) rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]]) count = 0 outs_list = [] 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]) rot_mat_i = tf.slice (rot_mat, [ 0, start_index* self.dim_rot_mat], [-1, natoms[2+type_i]* self.dim_rot_mat] ) rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) start_index += natoms[2+type_i] if not type_i in self.sel_type : continue layer = inputs_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, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec) 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, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift if self.fit_diag : bavg = np.zeros(self.dim_rot_mat_1) # bavg[0] = self.avgeig[0] # bavg[1] = self.avgeig[1] # bavg[2] = self.avgeig[2] # (nframes x natoms) x naxis final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x naxis final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1]) # (nframes x natoms) x naxis x naxis final_layer = tf.matrix_diag(final_layer) else : bavg = np.zeros(self.dim_rot_mat_1*self.dim_rot_mat_1) # bavg[0*self.dim_rot_mat_1+0] = self.avgeig[0] # bavg[1*self.dim_rot_mat_1+1] = self.avgeig[1] # bavg[2*self.dim_rot_mat_1+2] = self.avgeig[2] # (nframes x natoms) x (naxis x naxis) final_layer = one_layer(layer, self.dim_rot_mat_1*self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x naxis x naxis final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1, self.dim_rot_mat_1]) # (nframes x natoms) x naxis x naxis final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1]) # (nframes x natoms) x naxis x 3(coord) final_layer = tf.matmul(final_layer, rot_mat_i) # (nframes x natoms) x 3(coord) x 3(coord) final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True) # nframes x natoms x 3 x 3 final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3]) # shift and scale sel_type_idx = self.sel_type.index(type_i) final_layer = final_layer * self.scale[sel_type_idx] final_layer = final_layer + self.constant_matrix[sel_type_idx] * tf.eye(3, batch_shape=[tf.shape(inputs)[0], natoms[2+type_i]], dtype = GLOBAL_TF_FLOAT_PRECISION) # concat the results outs_list.append(final_layer) count += 1 outs = tf.concat(outs_list, axis = 1) tf.summary.histogram('fitting_net_output', outs) return tf.reshape(outs, [-1])
def build(self, input_d: tf.Tensor, rot_mat: tf.Tensor, natoms: tf.Tensor, reuse: bool = None, suffix: str = '') -> tf.Tensor: """ Build the computational graph for fitting net Parameters ---------- input_d The input descriptor rot_mat The rotation matrix from the descriptor. 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 reuse The weights in the networks should be reused when get the variable. suffix Name suffix to identify this descriptor Returns ------- dipole The atomic dipole. """ start_index = 0 inputs = tf.cast( tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]]) count = 0 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]) rot_mat_i = tf.slice(rot_mat, [0, start_index * self.dim_rot_mat], [-1, natoms[2 + type_i] * self.dim_rot_mat]) rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) start_index += natoms[2 + type_i] if not type_i in self.sel_type: continue layer = inputs_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, uniform_seed=self.uniform_seed) 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, uniform_seed=self.uniform_seed) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x naxis final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn=None, name='final_layer_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed, precision=self.fitting_precision, uniform_seed=self.uniform_seed) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x 1 * naxis final_layer = tf.reshape(final_layer, [ tf.shape(inputs)[0] * natoms[2 + type_i], 1, self.dim_rot_mat_1 ]) # (nframes x natoms) x 1 x 3(coord) final_layer = tf.matmul(final_layer, rot_mat_i) # nframes x natoms x 3 final_layer = tf.reshape( final_layer, [tf.shape(inputs)[0], natoms[2 + type_i], 3]) # concat the results if count == 0: outs = final_layer else: outs = tf.concat([outs, final_layer], axis=1) count += 1 tf.summary.histogram('fitting_net_output', outs) return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
def build(self, coord_: tf.Tensor, atype_: tf.Tensor, natoms: tf.Tensor, box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, reuse: bool = None, suffix: str = '') -> tf.Tensor: """ Build the computational graph for the descriptor Parameters ---------- coord_ The coordinate of atoms atype_ The type of atoms 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 mesh For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed. input_dict Dictionary for additional inputs. Should have 'efield'. reuse The weights in the networks should be reused when get the variable. suffix Name suffix to identify this descriptor Returns ------- descriptor The output descriptor """ self.dout_vert = self.descrpt_vert.build(coord_, atype_, natoms, box_, mesh, input_dict) self.dout_para = self.descrpt_para.build(coord_, atype_, natoms, box_, mesh, input_dict, reuse=True) coord = tf.reshape(coord_, [-1, natoms[1] * 3]) nframes = tf.shape(coord)[0] self.dout_vert = tf.reshape( self.dout_vert, [nframes * natoms[0], self.descrpt_vert.get_dim_out()]) self.dout_para = tf.reshape( self.dout_para, [nframes * natoms[0], self.descrpt_para.get_dim_out()]) self.dout = tf.concat([self.dout_vert, self.dout_para], axis=1) self.dout = tf.reshape(self.dout, [nframes, natoms[0] * self.get_dim_out()]) self.qmat = self.descrpt_vert.qmat + self.descrpt_para.qmat tf.summary.histogram('embedding_net_output', self.dout) return self.dout