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
