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
