def test_prod_virial(self): tvirial, tatom_virial \ = op_module.prod_virial_se_a( self.tnet_deriv, self.tem_deriv, self.trij, self.tnlist, self.tnatoms, n_a_sel=self.nnei, n_r_sel=0) self.sess.run(tf.global_variables_initializer()) dvirial, datom_virial = self.sess.run( [tvirial, tatom_virial], feed_dict={ self.tnet_deriv: self.dnet_deriv, self.tem_deriv: self.dem_deriv, self.trij: self.drij, self.tnlist: self.dnlist, self.tnatoms: self.dnatoms }) self.assertEqual(dvirial.shape, (self.nframes, 9)) self.assertEqual(datom_virial.shape, (self.nframes, self.nall * 9)) for ff in range(self.nframes): np.testing.assert_almost_equal(dvirial[ff], self.expected_virial, 5) np.testing.assert_almost_equal(datom_virial[ff], self.expected_atom_virial, 5)
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_reshape) tf.summary.histogram('net_derivative', net_deriv) net_deriv_reshape = tf.reshape(net_deriv, [-1, natoms[0] * self.ndescrpt]) force \ = op_module.prod_force_se_a (net_deriv_reshape, self.descrpt_deriv, self.nlist, natoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) virial, atom_virial \ = op_module.prod_virial_se_a (net_deriv_reshape, self.descrpt_deriv, self.rij, self.nlist, 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 comp_ef (self, dcoord, dbox, dtype, tnatoms, name, reuse = None) : descrpt, descrpt_deriv, rij, nlist \ = op_module.descrpt_se_a_ef_para (dcoord, dtype, tnatoms, dbox, tf.constant(self.default_mesh), self.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) 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_a (net_deriv_reshape, descrpt_deriv, nlist, tnatoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) 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) return energy, force, 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_a (net_deriv_reshape, self.descrpt_deriv, self.nlist, natoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) virial, atom_virial \ = op_module.prod_virial_se_a (net_deriv_reshape, self.descrpt_deriv, self.rij, self.nlist, 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.prod_env_mat_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.pair_tab( 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