def comp_ef(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): descrpt, descrpt_deriv, rij, nlist \ = op_module.descrpt_se_r (dcoord, dtype, tnatoms, dbox, tf.constant(self.default_mesh), self.t_avg, self.t_std, rcut = self.rcut, rcut_smth = self.rcut_smth, sel = self.sel) inputs_reshape = tf.reshape(descrpt, [-1, self.ndescrpt]) atom_ener = self._net(inputs_reshape, name, reuse=reuse) atom_ener_reshape = tf.reshape(atom_ener, [-1, self.natoms[0]]) energy = tf.reduce_sum(atom_ener_reshape, axis=1) net_deriv_ = tf.gradients(atom_ener, inputs_reshape) net_deriv = net_deriv_[0] net_deriv_reshape = tf.reshape(net_deriv, [-1, self.natoms[0] * self.ndescrpt]) force = op_module.prod_force_se_r(net_deriv_reshape, descrpt_deriv, nlist, tnatoms) virial, atom_vir = op_module.prod_virial_se_r(net_deriv_reshape, descrpt_deriv, rij, nlist, tnatoms) return energy, force, virial
def comp_v_dw(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): energy, force, virial = self.comp_ef(dcoord, dbox, dtype, tnatoms, name, reuse) with tf.variable_scope(name, reuse=True): net_w = tf.get_variable('net_w', [self.ndescrpt], GLOBAL_TF_FLOAT_PRECISION, tf.constant_initializer(self.net_w_i)) v_mag = tf.reduce_sum(virial) v_mag_dw = tf.gradients(v_mag, net_w) assert (len(v_mag_dw) == 1), "length of dw is wrong" return v_mag, v_mag_dw[0]
def comp_f_dw(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): energy, force, virial = self.comp_ef(dcoord, dbox, dtype, tnatoms, name, reuse) with tf.variable_scope(name, reuse=True): net_w = tf.get_variable('net_w', [self.ndescrpt], global_tf_float_precision, tf.constant_initializer(self.net_w_i)) f_mag = tf.reduce_sum(tf.nn.tanh(force)) f_mag_dw = tf.gradients(f_mag, net_w) assert (len(f_mag_dw) == 1), "length of dw is wrong" return f_mag, f_mag_dw[0]
def prod_force_virial( self, atom_ener: tf.Tensor, natoms: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: """ Compute force and virial Parameters ---------- atom_ener The atomic energy 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 Returns ------- force The force on atoms virial The total virial atom_virial The atomic virial """ [net_deriv] = tf.gradients(atom_ener, self.descrpt) tf.summary.histogram('net_derivative', net_deriv) net_deriv_reshape = tf.reshape(net_deriv, [-1, natoms[0] * self.ndescrpt]) force = op_module.prod_force(net_deriv_reshape, self.descrpt_deriv, self.nlist, self.axis, natoms, n_a_sel=self.nnei_a, n_r_sel=self.nnei_r) virial, atom_virial \ = op_module.prod_virial (net_deriv_reshape, self.descrpt_deriv, self.rij, self.nlist, self.axis, natoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) tf.summary.histogram('force', force) tf.summary.histogram('virial', virial) tf.summary.histogram('atom_virial', atom_virial) return force, virial, atom_virial
def prod_force_virial(self, atom_ener, natoms) : [net_deriv] = tf.gradients (atom_ener, self.descrpt_reshape) net_deriv_reshape = tf.reshape (net_deriv, [-1, natoms[0] * self.ndescrpt]) force \ = op_module.prod_force_se_r (net_deriv_reshape, self.descrpt_deriv, self.nlist, natoms) virial, atom_virial \ = op_module.prod_virial_se_r (net_deriv_reshape, self.descrpt_deriv, self.rij, self.nlist, natoms) return force, virial, atom_virial
def _build_training(self): trainable_variables = tf.trainable_variables() optimizer = tf.train.AdamOptimizer(learning_rate = self.learning_rate) if self.run_opt.is_distrib : optimizer = tf.train.SyncReplicasOptimizer( optimizer, replicas_to_aggregate = self.run_opt.cluster_spec.num_tasks("worker"), total_num_replicas = self.run_opt.cluster_spec.num_tasks("worker"), name = "sync_replicas") self.sync_replicas_hook = optimizer.make_session_run_hook(self.run_opt.is_chief) grads = tf.gradients(self.l2_l, trainable_variables) apply_op = optimizer.apply_gradients (zip (grads, trainable_variables), global_step=self.global_step, name='train_step') train_ops = [apply_op] + self._extra_train_ops self.train_op = tf.group(*train_ops) self._message("built training")
def comp_ef(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): t_default_mesh = tf.constant(self.default_mesh) descrpt, descrpt_deriv, rij, nlist, axis, rot_mat \ = op_module.descrpt (dcoord, dtype, tnatoms, dbox, t_default_mesh, self.t_avg, self.t_std, rcut_a = self.rcut_a, rcut_r = self.rcut_r, sel_a = self.sel_a, sel_r = self.sel_r, axis_rule = self.axis_rule) self.axis = axis self.nlist = nlist self.descrpt = descrpt inputs_reshape = tf.reshape(descrpt, [-1, self.ndescrpt]) atom_ener = self._net(inputs_reshape, name, reuse=reuse) atom_ener_reshape = tf.reshape(atom_ener, [-1, self.natoms[0]]) energy = tf.reduce_sum(atom_ener_reshape, axis=1) net_deriv_ = tf.gradients(atom_ener, inputs_reshape) net_deriv = net_deriv_[0] net_deriv_reshape = tf.reshape(net_deriv, [-1, self.natoms[0] * self.ndescrpt]) force = op_module.prod_force(net_deriv_reshape, descrpt_deriv, nlist, axis, tnatoms, n_a_sel=self.nnei_a, n_r_sel=self.nnei_r) virial, atom_vir = op_module.prod_virial(net_deriv_reshape, descrpt_deriv, rij, nlist, axis, tnatoms, n_a_sel=self.nnei_a, n_r_sel=self.nnei_r) return energy, force, virial
def prod_force_virial(self, atom_ener, natoms) : [net_deriv] = tf.gradients (atom_ener, self.descrpt) net_deriv_reshape = tf.reshape (net_deriv, [-1, natoms[0] * self.ndescrpt]) force = op_module.prod_force (net_deriv_reshape, self.descrpt_deriv, self.nlist, self.axis, natoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) virial, atom_virial \ = op_module.prod_virial (net_deriv_reshape, self.descrpt_deriv, self.rij, self.nlist, self.axis, natoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) return force, virial, atom_virial
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 comp_ef(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): descrpt, descrpt_deriv, rij, nlist \ = op_module.descrpt_se_a (dcoord, dtype, tnatoms, dbox, tf.constant(self.default_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) inputs_reshape = tf.reshape(descrpt, [-1, self.ndescrpt]) atom_ener = self._net(inputs_reshape, name, reuse=reuse) sw_lambda, sw_deriv \ = op_module.soft_min_switch(dtype, rij, nlist, tnatoms, sel_a = self.sel_a, sel_r = self.sel_r, alpha = self.smin_alpha, rmin = self.sw_rmin, rmax = self.sw_rmax) inv_sw_lambda = 1.0 - sw_lambda tab_atom_ener, tab_force, tab_atom_virial \ = op_module.tab_inter(self.tab_info, self.tab_data, dtype, rij, nlist, tnatoms, sw_lambda, sel_a = self.sel_a, sel_r = self.sel_r) energy_diff = tab_atom_ener - tf.reshape(atom_ener, [-1, self.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 energy_raw = tf.reshape(energy_raw, [-1, self.natoms[0]]) energy = tf.reduce_sum(energy_raw, axis=1) net_deriv_ = tf.gradients(atom_ener, inputs_reshape) net_deriv = net_deriv_[0] net_deriv_reshape = tf.reshape(net_deriv, [-1, self.natoms[0] * self.ndescrpt]) force = op_module.prod_force_se_a(net_deriv_reshape, descrpt_deriv, nlist, tnatoms, n_a_sel=self.nnei_a, n_r_sel=self.nnei_r) sw_force \ = op_module.soft_min_force(energy_diff, sw_deriv, nlist, tnatoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) force = force + sw_force + tab_force virial, atom_vir = op_module.prod_virial_se_a(net_deriv_reshape, descrpt_deriv, rij, nlist, tnatoms, n_a_sel=self.nnei_a, n_r_sel=self.nnei_r) sw_virial, sw_atom_virial \ = op_module.soft_min_virial (energy_diff, sw_deriv, rij, nlist, tnatoms, n_a_sel = self.nnei_a, n_r_sel = self.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, self.natoms[1], 9]), axis = 1) return energy, force, virial