def build(self,
coord_,
atype_,
natoms,
box_,
mesh,
davg=None,
dstd=None,
suffix='',
reuse=None):
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)
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.dout, self.qmat = self._pass_filter(self.descrpt_reshape,
natoms,
suffix=suffix,
reuse=reuse,
trainable=self.trainable)
return self.dout
def _pass_filter(self,
inputs,
atype,
natoms,
input_dict,
reuse=None,
suffix='',
trainable=True):
start_index = 0
inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]])
output = []
output_qmat = []
inputs_i = inputs
inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
type_i = -1
layer, qmat = self._filter(inputs_i,
type_i,
name='filter_type_all' + suffix,
natoms=natoms,
reuse=reuse,
trainable=trainable,
activation_fn=self.filter_activation_fn)
layer = tf.reshape(
layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()])
# qmat = tf.reshape(qmat, [tf.shape(inputs)[0], natoms[0] * self.get_dim_rot_mat_1() * 3])
output.append(layer)
# output_qmat.append(qmat)
output = tf.concat(output, axis=1)
# output_qmat = tf.concat(output_qmat, axis = 1)
return output, None
def build(self, coord_, atype_, natoms, box, mesh, suffix='', reuse=None):
davg = self.davg
dstd = self.dstd
if davg is None:
davg = [
np.zeros([self.descrpt_a.ntypes, self.descrpt_a.ndescrpt]),
np.zeros([self.descrpt_r.ntypes, self.descrpt_r.ndescrpt])
]
if dstd is None:
dstd = [
np.ones([self.descrpt_a.ntypes, self.descrpt_a.ndescrpt]),
np.ones([self.descrpt_r.ntypes, self.descrpt_r.ndescrpt])
]
# dout
self.dout_a = self.descrpt_a.build(coord_,
atype_,
natoms,
box,
mesh,
suffix=suffix + '_a',
reuse=reuse)
self.dout_r = self.descrpt_r.build(coord_,
atype_,
natoms,
box,
mesh,
suffix=suffix,
reuse=reuse)
self.dout_a = tf.reshape(self.dout_a,
[-1, self.descrpt_a.get_dim_out()])
self.dout_r = tf.reshape(self.dout_r,
[-1, self.descrpt_r.get_dim_out()])
self.dout = tf.concat([self.dout_a, self.dout_r], axis=1)
self.dout = tf.reshape(self.dout, [-1, natoms[0] * self.get_dim_out()])
return self.dout
def _pass_filter(self,
inputs,
natoms,
reuse=None,
suffix='',
trainable=True):
start_index = 0
inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]])
shape = inputs.get_shape().as_list()
output = []
output_qmat = []
for type_i in range(self.ntypes):
inputs_i = tf.slice(inputs, [0, start_index * self.ndescrpt],
[-1, natoms[2 + type_i] * self.ndescrpt])
inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
layer, qmat = self._filter(inputs_i,
name='filter_type_' + str(type_i) +
suffix,
natoms=natoms,
reuse=reuse,
seed=self.seed,
trainable=trainable)
layer = tf.reshape(
layer,
[tf.shape(inputs)[0], natoms[2 + type_i] * self.get_dim_out()])
qmat = tf.reshape(qmat, [
tf.shape(inputs)[0],
natoms[2 + type_i] * self.get_dim_rot_mat_1() * 3
])
output.append(layer)
output_qmat.append(qmat)
start_index += natoms[2 + type_i]
output = tf.concat(output, axis=1)
output_qmat = tf.concat(output_qmat, axis=1)
return output, output_qmat
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 _pass_filter(self,
inputs,
natoms,
reuse=None,
suffix='',
trainable=True):
start_index = 0
inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]])
output = []
for type_i in range(self.ntypes):
inputs_i = tf.slice(inputs, [0, start_index * self.ndescrpt],
[-1, natoms[2 + type_i] * self.ndescrpt])
inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
layer = self._filter_r(tf.cast(inputs_i, self.filter_precision),
type_i,
name='filter_type_' + str(type_i) + suffix,
natoms=natoms,
reuse=reuse,
seed=self.seed,
trainable=trainable,
activation_fn=self.filter_activation_fn)
layer = tf.reshape(
layer,
[tf.shape(inputs)[0], natoms[2 + type_i] * self.get_dim_out()])
output.append(layer)
start_index += natoms[2 + type_i]
output = tf.concat(output, axis=1)
return output
def _pass_filter(self,
inputs,
atype,
natoms,
input_dict,
reuse=None,
suffix='',
trainable=True):
# nf x na x ndescrpt
# nf x na x (nnei x 4)
inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt])
layer, qmat = self._ebd_filter(tf.cast(inputs, self.filter_precision),
atype,
natoms,
input_dict,
name='filter_type_all' + suffix,
reuse=reuse,
seed=self.seed,
trainable=trainable,
activation_fn=self.filter_activation_fn)
output = tf.reshape(
layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()])
output_qmat = tf.reshape(
qmat,
[tf.shape(inputs)[0], natoms[0] * self.get_dim_rot_mat_1() * 3])
return output, output_qmat
def _type_embed(self,
atype,
ndim=1,
reuse=None,
suffix='',
trainable=True):
ebd_type = tf.cast(atype, self.filter_precision)
ebd_type = ebd_type / float(self.ntypes)
ebd_type = tf.reshape(ebd_type, [-1, ndim])
for ii in range(self.type_nlayer):
name = 'type_embed_layer_' + str(ii)
ebd_type = one_layer(ebd_type,
self.type_nchanl,
activation_fn=self.filter_activation_fn,
precision=self.filter_precision,
name=name,
reuse=reuse,
seed=self.seed + ii,
trainable=trainable)
name = 'type_embed_layer_' + str(self.type_nlayer)
ebd_type = one_layer(ebd_type,
self.type_nchanl,
activation_fn=None,
precision=self.filter_precision,
name=name,
reuse=reuse,
seed=self.seed + ii,
trainable=trainable)
ebd_type = tf.reshape(ebd_type, [tf.shape(atype)[0], self.type_nchanl])
return ebd_type
def _filter_r(self,
inputs,
type_input,
natoms,
activation_fn=tf.nn.tanh,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
trainable=True):
# natom x nei
outputs_size = [1] + self.filter_neuron
with tf.variable_scope(name, reuse=reuse):
start_index = 0
xyz_scatter_total = []
for type_i in range(self.ntypes):
# cut-out inputs
# with natom x nei_type_i
inputs_i = tf.slice(inputs, [0, start_index],
[-1, self.sel_r[type_i]])
start_index += self.sel_r[type_i]
shape_i = inputs_i.get_shape().as_list()
# with (natom x nei_type_i) x 1
xyz_scatter = tf.reshape(inputs_i, [-1, 1])
if (type_input, type_i) not in self.exclude_types:
xyz_scatter = embedding_net(
xyz_scatter,
self.filter_neuron,
self.filter_precision,
activation_fn=activation_fn,
resnet_dt=self.filter_resnet_dt,
name_suffix="_" + str(type_i),
stddev=stddev,
bavg=bavg,
seed=self.seed,
trainable=trainable,
uniform_seed=self.uniform_seed,
initial_variables=self.embedding_net_variables,
)
if (not self.uniform_seed) and (self.seed is not None):
self.seed += self.seed_shift
# natom x nei_type_i x out_size
xyz_scatter = tf.reshape(
xyz_scatter, (-1, shape_i[1], outputs_size[-1]))
else:
natom = tf.shape(inputs)[0]
xyz_scatter = tf.cast(
tf.fill((natom, shape_i[1], outputs_size[-1]), 0.),
GLOBAL_TF_FLOAT_PRECISION)
xyz_scatter_total.append(xyz_scatter)
# natom x nei x outputs_size
xyz_scatter = tf.concat(xyz_scatter_total, axis=1)
# natom x outputs_size
#
res_rescale = 1. / 5.
result = tf.reduce_mean(xyz_scatter, axis=1) * res_rescale
return result
def _net(self, inputs, name, reuse=False):
with tf.variable_scope(name, reuse=reuse):
net_w = tf.get_variable('net_w', [self.ndescrpt],
GLOBAL_TF_FLOAT_PRECISION,
tf.constant_initializer(self.net_w_i))
dot_v = tf.matmul(tf.reshape(inputs, [-1, self.ndescrpt]),
tf.reshape(net_w, [self.ndescrpt, 1]))
return tf.reshape(dot_v, [-1])
def _net(self, inputs, name, reuse=False):
with tf.variable_scope(name, reuse=reuse):
net_w = tf.get_variable('net_w', [self.ndescrpt],
global_tf_float_precision,
tf.constant_initializer(self.net_w_i))
dot_v = tf.matmul(tf.reshape(inputs, [-1, self.ndescrpt]),
tf.reshape(net_w, [self.ndescrpt, 1]))
return tf.reshape(dot_v, [-1])
def build(self, input_d, rot_mat, natoms, reuse=None, suffix=''):
inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]])
outs = self.polar_fitting.build(input_d, rot_mat, natoms, reuse,
suffix)
# nframes x natoms x 9
outs = tf.reshape(outs, [tf.shape(inputs)[0], -1, 9])
outs = tf.reduce_sum(outs, axis=1)
return tf.reshape(outs, [-1])
def _concat_type_embedding(
self,
xyz_scatter,
nframes,
natoms,
type_embedding,
):
'''Concatenate `type_embedding` of neighbors and `xyz_scatter`.
If not self.type_one_side, concatenate `type_embedding` of center atoms as well.
Parameters
----------
xyz_scatter:
shape is [nframes*natoms[0]*self.nnei, 1]
nframes:
shape is []
natoms:
shape is [1+1+self.ntypes]
type_embedding:
shape is [self.ntypes, Y] where Y=jdata['type_embedding']['neuron'][-1]
Returns
-------
embedding:
environment of each atom represented by embedding.
'''
te_out_dim = type_embedding.get_shape().as_list()[-1]
nei_embed = tf.nn.embedding_lookup(
type_embedding,
tf.cast(self.nei_type,
dtype=tf.int32)) # shape is [self.nnei, 1+te_out_dim]
nei_embed = tf.tile(
nei_embed,
(nframes * natoms[0],
1)) # shape is [nframes*natoms[0]*self.nnei, te_out_dim]
nei_embed = tf.reshape(nei_embed, [-1, te_out_dim])
embedding_input = tf.concat(
[xyz_scatter, nei_embed],
1) # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim]
if not self.type_one_side:
atm_embed = embed_atom_type(
self.ntypes, natoms,
type_embedding) # shape is [natoms[0], te_out_dim]
atm_embed = tf.tile(
atm_embed,
(nframes, self.nnei
)) # shape is [nframes*natoms[0], self.nnei*te_out_dim]
atm_embed = tf.reshape(
atm_embed,
[-1, te_out_dim
]) # shape is [nframes*natoms[0]*self.nnei, te_out_dim]
embedding_input = tf.concat(
[embedding_input, atm_embed], 1
) # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim+te_out_dim]
return embedding_input
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
"""
with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) :
t_rcut = tf.constant(self.get_rcut(),
name = 'rcut',
dtype = GLOBAL_TF_FLOAT_PRECISION)
t_ntypes = tf.constant(self.get_ntypes(),
name = 'ntypes',
dtype = tf.int32)
all_dout = []
for idx,ii in enumerate(self.descrpt_list):
dout = ii.build(coord_, atype_, natoms, box_, mesh, input_dict, suffix=suffix+f'_{idx}', reuse=reuse)
dout = tf.reshape(dout, [-1, ii.get_dim_out()])
all_dout.append(dout)
dout = tf.concat(all_dout, axis = 1)
dout = tf.reshape(dout, [-1, natoms[0] * self.get_dim_out()])
return dout
def one_layer(inputs,
outputs_size,
activation_fn=tf.nn.tanh,
precision=global_tf_float_precision,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
seed=None,
use_timestep=False,
trainable=True,
useBN=False):
with tf.variable_scope(name, reuse=reuse):
shape = inputs.get_shape().as_list()
w = tf.get_variable(
'matrix', [shape[1], outputs_size],
precision,
tf.random_normal_initializer(stddev=stddev /
np.sqrt(shape[1] + outputs_size),
seed=seed),
trainable=trainable)
b = tf.get_variable('bias', [outputs_size],
precision,
tf.random_normal_initializer(stddev=stddev,
mean=bavg,
seed=seed),
trainable=trainable)
hidden = tf.matmul(inputs, w) + b
if activation_fn != None and use_timestep:
idt = tf.get_variable('idt', [outputs_size],
precision,
tf.random_normal_initializer(stddev=0.001,
mean=0.1,
seed=seed),
trainable=trainable)
if activation_fn != None:
if useBN:
None
# hidden_bn = self._batch_norm(hidden, name=name+'_normalization', reuse=reuse)
# return activation_fn(hidden_bn)
else:
if use_timestep:
return tf.reshape(activation_fn(hidden),
[-1, outputs_size]) * idt
else:
return tf.reshape(activation_fn(hidden),
[-1, outputs_size])
else:
if useBN:
None
# return self._batch_norm(hidden, name=name+'_normalization', reuse=reuse)
else:
return hidden
def _embedding_net(self,
inputs,
natoms,
filter_neuron,
activation_fn=tf.nn.tanh,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
seed=None,
trainable=True):
'''
inputs: nf x na x (nei x 4)
outputs: nf x na x nei x output_size
'''
# natom x (nei x 4)
inputs = tf.reshape(inputs, [-1, self.ndescrpt])
shape = inputs.get_shape().as_list()
outputs_size = [1] + filter_neuron
with tf.variable_scope(name, reuse=reuse):
xyz_scatter_total = []
# with natom x (nei x 4)
inputs_i = inputs
shape_i = inputs_i.get_shape().as_list()
# with (natom x nei) x 4
inputs_reshape = tf.reshape(inputs_i, [-1, 4])
# with (natom x nei) x 1
xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]),
[-1, 1])
# with (natom x nei) x out_size
xyz_scatter = embedding_net(xyz_scatter,
self.filter_neuron,
self.filter_precision,
activation_fn=activation_fn,
resnet_dt=self.filter_resnet_dt,
stddev=stddev,
bavg=bavg,
seed=seed,
trainable=trainable)
# natom x nei x out_size
xyz_scatter = tf.reshape(xyz_scatter,
(-1, shape_i[1] // 4, outputs_size[-1]))
xyz_scatter_total.append(xyz_scatter)
# natom x nei x outputs_size
xyz_scatter = tf.concat(xyz_scatter_total, axis=1)
# nf x natom x nei x outputs_size
xyz_scatter = tf.reshape(
xyz_scatter,
[tf.shape(inputs)[0], natoms[0], self.nnei, outputs_size[-1]])
return xyz_scatter
def _pass_filter(self,
inputs,
natoms,
reuse=None,
suffix='',
trainable=True):
start_index = 0
inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]])
output = []
if not (self.type_one_side and len(self.exclude_types) == 0):
for type_i in range(self.ntypes):
inputs_i = tf.slice(inputs, [0, start_index * self.ndescrpt],
[-1, natoms[2 + type_i] * self.ndescrpt])
inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
if self.type_one_side:
# reuse NN parameters for all types to support type_one_side along with exclude_types
reuse = tf.AUTO_REUSE
filter_name = 'filter_type_all' + suffix
else:
filter_name = 'filter_type_' + str(type_i) + suffix
layer = self._filter_r(inputs_i,
type_i,
name=filter_name,
natoms=natoms,
reuse=reuse,
trainable=trainable,
activation_fn=self.filter_activation_fn)
layer = tf.reshape(layer, [
tf.shape(inputs)[0],
natoms[2 + type_i] * self.get_dim_out()
])
output.append(layer)
start_index += natoms[2 + type_i]
else:
inputs_i = inputs
inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
type_i = -1
layer = self._filter_r(inputs_i,
type_i,
name='filter_type_all' + suffix,
natoms=natoms,
reuse=reuse,
trainable=trainable,
activation_fn=self.filter_activation_fn)
layer = tf.reshape(
layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()])
output.append(layer)
output = tf.concat(output, axis=1)
return output
def comp_ef(self, dcoord, dbox, dtype, tnatoms, name, reuse=None):
dout = self.descrpt.build(dcoord,
dtype,
tnatoms,
dbox,
self.default_mesh, {"efield": self.efield},
suffix=name,
reuse=reuse)
inputs_reshape = tf.reshape(dout, [-1, self.descrpt.get_dim_out()])
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)
force, virial, av = self.descrpt.prod_force_virial(
atom_ener_reshape, tnatoms)
return energy, force, virial
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 build (self,
coord_,
atype_,
natoms,
box,
mesh,
davg,
dstd,
suffix = '',
reuse = None):
# dout
self.dout_a = self.descrpt_a.build(coord_, atype_, natoms, box, mesh, davg[0], dstd[0], suffix=suffix+'_a', reuse=reuse)
self.dout_r = self.descrpt_r.build(coord_, atype_, natoms, box, mesh, davg[1], dstd[1], suffix=suffix+'_r', reuse=reuse)
self.dout_a = tf.reshape(self.dout_a, [-1, self.descrpt_a.get_dim_out()])
self.dout_r = tf.reshape(self.dout_r, [-1, self.descrpt_r.get_dim_out()])
self.dout = tf.concat([self.dout_a, self.dout_r], axis = 1)
self.dout = tf.reshape(self.dout, [-1, natoms[0] * self.get_dim_out()])
return self.dout
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 _pass_filter(self,
inputs,
natoms,
reuse=None,
suffix='',
trainable=True):
start_index = 0
inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]])
output = []
output_qmat = []
if not self.type_one_side:
for type_i in range(self.ntypes):
inputs_i = tf.slice(inputs, [0, start_index * self.ndescrpt],
[-1, natoms[2 + type_i] * self.ndescrpt])
inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
layer, qmat = self._filter(
tf.cast(inputs_i, self.filter_precision),
type_i,
name='filter_type_' + str(type_i) + suffix,
natoms=natoms,
reuse=reuse,
seed=self.seed,
trainable=trainable,
activation_fn=self.filter_activation_fn)
layer = tf.reshape(layer, [
tf.shape(inputs)[0],
natoms[2 + type_i] * self.get_dim_out()
])
qmat = tf.reshape(qmat, [
tf.shape(inputs)[0],
natoms[2 + type_i] * self.get_dim_rot_mat_1() * 3
])
output.append(layer)
output_qmat.append(qmat)
start_index += natoms[2 + type_i]
else:
inputs_i = inputs
inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt])
type_i = -1
layer, qmat = self._filter(tf.cast(inputs_i,
self.filter_precision),
type_i,
name='filter_type_all' + suffix,
natoms=natoms,
reuse=reuse,
seed=self.seed,
trainable=trainable,
activation_fn=self.filter_activation_fn)
layer = tf.reshape(
layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()])
qmat = tf.reshape(qmat, [
tf.shape(inputs)[0], natoms[0] * self.get_dim_rot_mat_1() * 3
])
output.append(layer)
output_qmat.append(qmat)
output = tf.concat(output, axis=1)
output_qmat = tf.concat(output_qmat, axis=1)
return output, output_qmat
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 build(self,
input_d,
rot_mat,
natoms,
reuse=None,
suffix='') -> tf.Tensor:
"""
Build the computational graph for fitting net
Parameters
----------
input_d
The input descriptor
rot_mat
The rotation matrix from the descriptor.
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
reuse
The weights in the networks should be reused when get the variable.
suffix
Name suffix to identify this descriptor
Returns
-------
polar
The system polarizability
"""
inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]])
outs = self.polar_fitting.build(input_d, rot_mat, natoms, reuse,
suffix)
# nframes x natoms x 9
outs = tf.reshape(outs, [tf.shape(inputs)[0], -1, 9])
outs = tf.reduce_sum(outs, axis=1)
tf.summary.histogram('fitting_net_output', outs)
return tf.reshape(outs, [-1])
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 build_efv(self, dcoord, dbox, dtype, tnatoms, name, op, reuse=None):
efield = tf.reshape(self.efield, [-1, 3])
efield = self._normalize_3d(efield)
efield = tf.reshape(efield, [-1, tnatoms[0] * 3])
if op != op_module.prod_env_mat_a:
descrpt = DescrptSeAEfLower(
op, **{
'sel': self.sel_a,
'rcut': 6,
'rcut_smth': 5.5,
'seed': 1,
'uniform_seed': True
})
else:
descrpt = DescrptSeA(
**{
'sel': self.sel_a,
'rcut': 6,
'rcut_smth': 0.5,
'seed': 1,
'uniform_seed': True
})
dout = descrpt.build(dcoord,
dtype,
tnatoms,
dbox,
tf.constant(self.default_mesh),
{'efield': efield},
suffix=name,
reuse=reuse)
dout = tf.reshape(dout, [-1, descrpt.get_dim_out()])
atom_ener = tf.reduce_sum(dout, axis=1)
atom_ener_reshape = tf.reshape(atom_ener, [-1, self.natoms[0]])
energy = tf.reduce_sum(atom_ener_reshape, axis=1)
force, virial, atom_vir \
= descrpt.prod_force_virial (atom_ener, tnatoms)
return energy, force, virial, atom_ener, atom_vir
def build (self,
input_d,
rot_mat,
natoms,
reuse = None,
suffix = '') :
start_index = 0
inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision)
rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]])
count = 0
outs_list = []
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice (inputs,
[ 0, start_index* self.dim_descrpt],
[-1, natoms[2+type_i]* self.dim_descrpt] )
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
rot_mat_i = tf.slice (rot_mat,
[ 0, start_index* 9],
[-1, natoms[2+type_i]* 9] )
rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3])
start_index += natoms[2+type_i]
if not type_i in self.sel_type :
continue
layer = inputs_i
for ii in range(0,len(self.n_neuron)) :
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
else :
layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
# (nframes x natoms) x 9
final_layer = one_layer(layer, 9, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, final_layer = True)
# (nframes x natoms) x 3 x 3
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 3, 3])
# (nframes x natoms) x 3 x 3
final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1])
# (nframes x natoms) x 3 x 3(coord)
final_layer = tf.matmul(final_layer, rot_mat_i)
# (nframes x natoms) x 3(coord) x 3(coord)
final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True)
# nframes x natoms x 3 x 3
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3])
# concat the results
outs_list.append(final_layer)
count += 1
outs = tf.concat(outs_list, axis = 1)
tf.summary.histogram('fitting_net_output', outs)
return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
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(self, learning_rate, natoms, model_dict, label_dict, suffix):
coord = model_dict['coord']
energy = model_dict['energy']
atom_ener = model_dict['atom_ener']
nframes = tf.shape(atom_ener)[0]
natoms = tf.shape(atom_ener)[1]
# build energy dipole
atom_ener0 = atom_ener - tf.reshape(
tf.tile(
tf.reshape(energy / global_cvt_2_ener_float(natoms), [-1, 1]),
[1, natoms]), [nframes, natoms])
coord = tf.reshape(coord, [nframes, natoms, 3])
atom_ener0 = tf.reshape(atom_ener0, [nframes, 1, natoms])
ener_dipole = tf.matmul(atom_ener0, coord)
ener_dipole = tf.reshape(ener_dipole, [nframes, 3])
energy_hat = label_dict['energy']
ener_dipole_hat = label_dict['energy_dipole']
find_energy = label_dict['find_energy']
find_ener_dipole = label_dict['find_energy_dipole']
l2_ener_loss = tf.reduce_mean(tf.square(energy - energy_hat),
name='l2_' + suffix)
ener_dipole_reshape = tf.reshape(ener_dipole, [-1])
ener_dipole_hat_reshape = tf.reshape(ener_dipole_hat, [-1])
l2_ener_dipole_loss = tf.reduce_mean(
tf.square(ener_dipole_reshape - ener_dipole_hat_reshape),
name='l2_' + suffix)
# atom_norm_ener = 1./ global_cvt_2_ener_float(natoms[0])
atom_norm_ener = 1. / global_cvt_2_ener_float(natoms)
pref_e = global_cvt_2_ener_float(
find_energy * (self.limit_pref_e +
(self.start_pref_e - self.limit_pref_e) *
learning_rate / self.starter_learning_rate))
pref_ed = global_cvt_2_tf_float(
find_ener_dipole * (self.limit_pref_ed +
(self.start_pref_ed - self.limit_pref_ed) *
learning_rate / self.starter_learning_rate))
l2_loss = 0
more_loss = {}
l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
l2_loss += global_cvt_2_ener_float(pref_ed * l2_ener_dipole_loss)
more_loss['l2_ener_loss'] = l2_ener_loss
more_loss['l2_ener_dipole_loss'] = l2_ener_dipole_loss
self.l2_l = l2_loss
self.l2_more = more_loss
return l2_loss, more_loss
def build (self,
input_d,
rot_mat,
natoms,
reuse = None,
suffix = '') :
start_index = 0
inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision)
rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]])
count = 0
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice (inputs,
[ 0, start_index* self.dim_descrpt],
[-1, natoms[2+type_i]* self.dim_descrpt] )
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
rot_mat_i = tf.slice (rot_mat,
[ 0, start_index* 9],
[-1, natoms[2+type_i]* 9] )
rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3])
start_index += natoms[2+type_i]
if not type_i in self.sel_type :
continue
layer = inputs_i
for ii in range(0,len(self.n_neuron)) :
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed)
else :
layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
# (nframes x natoms) x (nwfc x 3)
final_layer = one_layer(layer, self.wfc_numb * 3, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed)
if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
# (nframes x natoms) x nwfc(wc) x 3(coord_local)
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.wfc_numb, 3])
# (nframes x natoms) x nwfc(wc) x 3(coord)
final_layer = tf.matmul(final_layer, rot_mat_i)
# nframes x natoms x nwfc(wc) x 3(coord_local)
final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], self.wfc_numb, 3])
# concat the results
if count == 0:
outs = final_layer
else:
outs = tf.concat([outs, final_layer], axis = 1)
count += 1
tf.summary.histogram('fitting_net_output', outs)
return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
