def build(self, coord_, atype_, natoms, box, mesh, input_dict, suffix='', reuse=None): with tf.variable_scope('model_attr' + suffix, reuse=reuse): t_tmap = tf.constant(' '.join(self.type_map), name='tmap', dtype=tf.string) t_st = tf.constant(self.get_sel_type(), name='sel_type', dtype=tf.int32) t_mt = tf.constant(self.model_type, name='model_type', dtype=tf.string) coord = tf.reshape(coord_, [-1, natoms[1] * 3]) atype = tf.reshape(atype_, [-1, natoms[1]]) dout \ = self.descrpt.build(coord_, atype_, natoms, box, mesh, davg = self.davg, dstd = self.dstd, suffix = suffix, reuse = reuse) dout = tf.identity(dout, name='o_descriptor') rot_mat = self.descrpt.get_rot_mat() rot_mat = tf.identity(rot_mat, name='o_rot_mat') polar = self.fitting.build(dout, rot_mat, natoms, reuse=reuse, suffix=suffix) polar = tf.identity(polar, name='o_polar') return {'polar': polar}
def _identity_tensors(self, suffix : str = "") -> None: """Identify tensors which are expected to be stored and restored. Notes ----- These tensors will be indentitied: self.descrpt_reshape : o_rmat self.descrpt_deriv : o_rmat_deriv self.rij : o_rij self.nlist : o_nlist Thus, this method should be called during building the descriptor and after these tensors are initialized. Parameters ---------- suffix : str The suffix of the scope """ self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat' + suffix) self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv' + suffix) self.rij = tf.identity(self.rij, name = 'o_rij' + suffix) self.nlist = tf.identity(self.nlist, name = 'o_nlist' + suffix)
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, initial_variables = self.type_embedding_net_variables, uniform_seed = self.uniform_seed) 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
def build(self, coord_, atype_, natoms, box, mesh, input_dict, suffix='', reuse=None): 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) 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]]) dout \ = self.descrpt.build(coord_, atype_, natoms, box, mesh, davg = self.davg, dstd = self.dstd, suffix = suffix, reuse = reuse) dout = tf.identity(dout, name='o_descriptor') 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, input_dict, natoms, bias_atom_e=self.bias_atom_e, 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.tab_inter(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 return model_dict
def _filter_type_ext(self, inputs, natoms, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, trainable=True): # natom x (nei x 4) outputs_size = [1] + self.filter_neuron outputs_size_2 = self.n_axis_neuron with tf.variable_scope(name, reuse=reuse): start_index = 0 result_all = [] xyz_scatter_1_all = [] xyz_scatter_2_all = [] 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]) xyz_scatter = tf.reshape( tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1]) for ii in range(1, len(outputs_size)): w = tf.get_variable( 'matrix_' + str(ii) + '_' + str(type_i), [outputs_size[ii - 1], outputs_size[ii]], self.filter_precision, tf.random_normal_initializer( stddev=stddev / np.sqrt(outputs_size[ii] + outputs_size[ii - 1]), seed=seed), trainable=trainable) b = tf.get_variable('bias_' + str(ii) + '_' + str(type_i), [1, outputs_size[ii]], self.filter_precision, tf.random_normal_initializer( stddev=stddev, mean=bavg, seed=seed), trainable=trainable) if self.filter_resnet_dt: idt = tf.get_variable( 'idt_' + str(ii) + '_' + str(type_i), [1, outputs_size[ii]], self.filter_precision, tf.random_normal_initializer(stddev=0.001, mean=1.0, seed=seed), trainable=trainable) if outputs_size[ii] == outputs_size[ii - 1]: if self.filter_resnet_dt: xyz_scatter += activation_fn( tf.matmul(xyz_scatter, w) + b) * idt else: xyz_scatter += activation_fn( tf.matmul(xyz_scatter, w) + b) elif outputs_size[ii] == outputs_size[ii - 1] * 2: if self.filter_resnet_dt: xyz_scatter = tf.concat( [xyz_scatter, xyz_scatter], 1) + activation_fn( tf.matmul(xyz_scatter, w) + b) * idt else: xyz_scatter = tf.concat( [xyz_scatter, xyz_scatter], 1) + activation_fn( tf.matmul(xyz_scatter, w) + b) else: xyz_scatter = activation_fn( tf.matmul(xyz_scatter, w) + b) # natom x nei_type_i x out_size xyz_scatter = tf.reshape( xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1])) # natom x nei_type_i x 4 inputs_i_reshape = tf.reshape(inputs_i, [-1, shape_i[1] // 4, 4]) # natom x 4 x outputs_size xyz_scatter_1 = tf.matmul(inputs_i_reshape, xyz_scatter, transpose_a=True) xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape_i[1]) # natom x 4 x outputs_size_2 xyz_scatter_2 = tf.slice(xyz_scatter_1, [0, 0, 0], [-1, -1, outputs_size_2]) xyz_scatter_1_all.append(xyz_scatter_1) xyz_scatter_2_all.append(xyz_scatter_2) # for type_i in range(self.ntypes): # for type_j in range(type_i, self.ntypes): # # natom x outputs_size x outputs_size_2 # result = tf.matmul(xyz_scatter_1_all[type_i], xyz_scatter_2_all[type_j], transpose_a = True) # # natom x (outputs_size x outputs_size_2) # result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]]) # result_all.append(tf.identity(result)) xyz_scatter_2_coll = tf.concat(xyz_scatter_2_all, axis=2) for type_i in range(self.ntypes): # natom x outputs_size x (outputs_size_2 x ntypes) result = tf.matmul(xyz_scatter_1_all[type_i], xyz_scatter_2_coll, transpose_a=True) # natom x (outputs_size x outputs_size_2 x ntypes) result = tf.reshape( result, [-1, outputs_size_2 * self.ntypes * outputs_size[-1]]) result_all.append(tf.identity(result)) # natom x (ntypes x outputs_size x outputs_size_2 x ntypes) result_all = tf.concat(result_all, axis=1) return result_all
def build(self, coord_, atype_, natoms, box_, mesh, suffix='', reuse=None): davg = self.davg dstd = self.dstd with tf.variable_scope('descrpt_attr' + suffix, reuse=reuse): if davg is None: davg = np.zeros([self.ntypes, self.ndescrpt]) if dstd is None: dstd = np.ones([self.ntypes, self.ndescrpt]) t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), name='rcut', dtype=global_tf_float_precision) t_ntypes = tf.constant(self.ntypes, name='ntypes', dtype=tf.int32) t_ndescrpt = tf.constant(self.ndescrpt, name='ndescrpt', dtype=tf.int32) t_sel = tf.constant(self.sel_a, name='sel', dtype=tf.int32) self.t_avg = tf.get_variable( 't_avg', davg.shape, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(davg)) self.t_std = tf.get_variable( 't_std', dstd.shape, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(dstd)) coord = tf.reshape(coord_, [-1, natoms[1] * 3]) box = tf.reshape(box_, [-1, 9]) atype = tf.reshape(atype_, [-1, natoms[1]]) self.descrpt, self.descrpt_deriv, self.rij, self.nlist \ = op_module.descrpt_se_a (coord, atype, natoms, box, mesh, self.t_avg, self.t_std, rcut_a = self.rcut_a, rcut_r = self.rcut_r, rcut_r_smth = self.rcut_r_smth, sel_a = self.sel_a, sel_r = self.sel_r) self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) self.descrpt_reshape = tf.identity(self.descrpt_reshape, name='o_rmat') self.descrpt_deriv = tf.identity(self.descrpt_deriv, name='o_rmat_deriv') self.rij = tf.identity(self.rij, name='o_rij') self.nlist = tf.identity(self.nlist, name='o_nlist') self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, natoms, suffix=suffix, reuse=reuse, trainable=self.trainable) return self.dout
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, coord_, atype_, natoms, box, mesh, input_dict, frz_model = None, suffix = '', reuse = None): with tf.variable_scope('model_attr' + suffix, reuse = reuse) : t_tmap = tf.constant(' '.join(self.type_map), name = 'tmap', dtype = tf.string) t_st = tf.constant(self.get_sel_type(), name = 'sel_type', dtype = tf.int32) 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) t_od = tf.constant(self.get_out_size(), name = 'output_dim', dtype = tf.int32) natomsel = sum(natoms[2+type_i] for type_i in self.get_sel_type()) nout = self.get_out_size() 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]) rot_mat = self.descrpt.get_rot_mat() rot_mat = tf.identity(rot_mat, name = 'o_rot_mat'+suffix) output = self.fitting.build (dout, rot_mat, natoms, reuse = reuse, suffix = suffix) framesize = nout if "global" in self.model_type else natomsel * nout output = tf.reshape(output, [-1, framesize], name = 'o_' + self.model_type + suffix) model_dict = {self.model_type: output} if "global" not in self.model_type: gname = "global_"+self.model_type atom_out = tf.reshape(output, [-1, natomsel, nout]) global_out = tf.reduce_sum(atom_out, axis=1) global_out = tf.reshape(global_out, [-1, nout], name="o_" + gname + suffix) out_cpnts = tf.split(atom_out, nout, axis=-1) force_cpnts = [] virial_cpnts = [] atom_virial_cpnts = [] for out_i in out_cpnts: force_i, virial_i, atom_virial_i \ = self.descrpt.prod_force_virial(out_i, natoms) force_cpnts.append (tf.reshape(force_i, [-1, 3*natoms[1]])) virial_cpnts.append (tf.reshape(virial_i, [-1, 9])) atom_virial_cpnts.append(tf.reshape(atom_virial_i, [-1, 9*natoms[1]])) # [nframe x nout x (natom x 3)] force = tf.concat(force_cpnts, axis=1, name="o_force" + suffix) # [nframe x nout x 9] virial = tf.concat(virial_cpnts, axis=1, name="o_virial" + suffix) # [nframe x nout x (natom x 9)] atom_virial = tf.concat(atom_virial_cpnts, axis=1, name="o_atom_virial" + suffix) model_dict[gname] = global_out model_dict["force"] = force model_dict["virial"] = virial model_dict["atom_virial"] = atom_virial return model_dict
def build(self, coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, suffix='', reuse=None): 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]]) # 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, coord_, atype_, natoms, box_, mesh, input_dict, suffix='', reuse=None): efield = input_dict['efield'] davg = self.davg dstd = self.dstd with tf.variable_scope('descrpt_attr' + suffix, reuse=reuse): if davg is None: davg = np.zeros([self.ntypes, self.ndescrpt]) if dstd is None: dstd = np.ones([self.ntypes, self.ndescrpt]) t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), name='rcut', dtype=GLOBAL_TF_FLOAT_PRECISION) t_ntypes = tf.constant(self.ntypes, name='ntypes', dtype=tf.int32) t_ndescrpt = tf.constant(self.ndescrpt, name='ndescrpt', dtype=tf.int32) t_sel = tf.constant(self.sel_a, name='sel', dtype=tf.int32) self.t_avg = tf.get_variable( 't_avg', davg.shape, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(davg)) self.t_std = tf.get_variable( 't_std', dstd.shape, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(dstd)) coord = tf.reshape(coord_, [-1, natoms[1] * 3]) box = tf.reshape(box_, [-1, 9]) atype = tf.reshape(atype_, [-1, natoms[1]]) efield = tf.reshape(efield, [-1, 3]) efield = self._normalize_3d(efield) efield = tf.reshape(efield, [-1, natoms[0] * 3]) self.descrpt, self.descrpt_deriv, self.rij, self.nlist \ = self.op (coord, atype, natoms, box, mesh, efield, self.t_avg, self.t_std, rcut_a = self.rcut_a, rcut_r = self.rcut_r, rcut_r_smth = self.rcut_r_smth, sel_a = self.sel_a, sel_r = self.sel_r) self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) self.descrpt_reshape = tf.identity(self.descrpt_reshape, name='o_rmat') self.descrpt_deriv = tf.identity(self.descrpt_deriv, name='o_rmat_deriv') self.rij = tf.identity(self.rij, name='o_rij') self.nlist = tf.identity(self.nlist, name='o_nlist') # only used when tensorboard was set as true tf.summary.histogram('descrpt', self.descrpt) tf.summary.histogram('rij', self.rij) tf.summary.histogram('nlist', self.nlist) self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, atype, natoms, input_dict, suffix=suffix, reuse=reuse, trainable=self.trainable) tf.summary.histogram('embedding_net_output', self.dout) return self.dout
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 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 """ davg = self.davg dstd = self.dstd with tf.variable_scope('descrpt_attr' + suffix, reuse=reuse): if davg is None: davg = np.zeros([self.ntypes, self.ndescrpt]) if dstd is None: dstd = np.ones([self.ntypes, self.ndescrpt]) t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), name='rcut', dtype=GLOBAL_TF_FLOAT_PRECISION) t_ntypes = tf.constant(self.ntypes, name='ntypes', dtype=tf.int32) t_ndescrpt = tf.constant(self.ndescrpt, name='ndescrpt', dtype=tf.int32) t_sel = tf.constant(self.sel_a, name='sel', dtype=tf.int32) self.t_avg = tf.get_variable( 't_avg', davg.shape, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(davg)) self.t_std = tf.get_variable( 't_std', dstd.shape, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(dstd)) coord = tf.reshape(coord_, [-1, natoms[1] * 3]) box = tf.reshape(box_, [-1, 9]) atype = tf.reshape(atype_, [-1, natoms[1]]) self.descrpt, self.descrpt_deriv, self.rij, self.nlist \ = op_module.prod_env_mat_a (coord, atype, natoms, box, mesh, self.t_avg, self.t_std, rcut_a = self.rcut_a, rcut_r = self.rcut_r, rcut_r_smth = self.rcut_r_smth, sel_a = self.sel_a, sel_r = self.sel_r) self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) self.descrpt_reshape = tf.identity(self.descrpt_reshape, name='o_rmat') self.descrpt_deriv = tf.identity(self.descrpt_deriv, name='o_rmat_deriv') self.rij = tf.identity(self.rij, name='o_rij') self.nlist = tf.identity(self.nlist, name='o_nlist') self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, atype, natoms, input_dict, suffix=suffix, reuse=reuse, trainable=self.trainable) return self.dout