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