def __init__(self,
rcut: float,
sel_a : List[int],
sel_r : List[int],
axis_rule : List[int]
) -> None:
"""
Constructor
"""
# args = ClassArg()\
# .add('sel_a', list, must = True) \
# .add('sel_r', list, must = True) \
# .add('rcut', float, default = 6.0) \
# .add('axis_rule',list, must = True)
# class_data = args.parse(jdata)
self.sel_a = sel_a
self.sel_r = sel_r
self.axis_rule = axis_rule
self.rcut_r = rcut
# ntypes and rcut_a === -1
self.ntypes = len(self.sel_a)
assert(self.ntypes == len(self.sel_r))
self.rcut_a = -1
# numb of neighbors and numb of descrptors
self.nnei_a = np.cumsum(self.sel_a)[-1]
self.nnei_r = np.cumsum(self.sel_r)[-1]
self.nnei = self.nnei_a + self.nnei_r
self.ndescrpt_a = self.nnei_a * 4
self.ndescrpt_r = self.nnei_r * 1
self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
self.davg = None
self.dstd = None
self.place_holders = {}
avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
sub_graph = tf.Graph()
with sub_graph.as_default():
name_pfx = 'd_lf_'
for ii in ['coord', 'box']:
self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii)
self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type')
self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms')
self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh')
self.stat_descrpt, descrpt_deriv, rij, nlist, axis, rot_mat \
= op_module.descrpt (self.place_holders['coord'],
self.place_holders['type'],
self.place_holders['natoms_vec'],
self.place_holders['box'],
self.place_holders['default_mesh'],
tf.constant(avg_zero),
tf.constant(std_ones),
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.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config)
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)
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, self.axis, self.rot_mat \
= op_module.descrpt (coord,
atype,
natoms,
box,
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.descrpt = tf.reshape(self.descrpt, [-1, self.ndescrpt])
return self.descrpt
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,
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)
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, self.axis, self.rot_mat \
= op_module.descrpt (coord,
atype,
natoms,
box,
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.descrpt = tf.reshape(self.descrpt, [-1, self.ndescrpt])
tf.summary.histogram('descrpt', self.descrpt)
tf.summary.histogram('rij', self.rij)
tf.summary.histogram('nlist', self.nlist)
return self.descrpt
def comp_interpl_ef(self, dcoord, dbox, dtype, tnatoms, name, reuse=None):
descrpt, descrpt_deriv, rij, nlist, axis \
= op_module.descrpt (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,
sel_a = self.sel_a,
sel_r = self.sel_r,
axis_rule = self.axis_rule)
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(net_deriv_reshape,
descrpt_deriv,
nlist,
axis,
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(net_deriv_reshape,
descrpt_deriv,
rij,
nlist,
axis,
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
