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 init_variables(self,
graph : tf.Graph,
graph_def : tf.GraphDef,
model_type : str = "original_model",
suffix : str = "",
) -> None:
"""
Init the embedding net variables with the given frozen model
Parameters
----------
graph : tf.Graph
The input frozen model graph
graph_def : tf.GraphDef
The input frozen model graph_def
model_type : str
the type of the model
suffix : str
suffix to name scope
"""
if model_type == 'original_model':
self.descrpt.init_variables(graph, graph_def, suffix=suffix)
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("original_model", name = 'model_type', dtype = tf.string)
elif model_type == 'compressed_model':
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("compressed_model", name = 'model_type', dtype = tf.string)
else:
raise RuntimeError("Unknown model type %s" % model_type)
def setUp(self, data, pbc=True, sess=None):
self.sess = sess
self.data = data
self.natoms = self.data.get_natoms()
self.ntypes = self.data.get_ntypes()
self.sel = [12, 24]
self.sel_a = [0, 0]
self.rcut_smth = 2.45
self.rcut = 10.0
self.nnei = np.cumsum(self.sel)[-1]
self.ndescrpt = self.nnei * 1
davg = np.zeros([self.ntypes, self.ndescrpt])
dstd = np.ones([self.ntypes, self.ndescrpt])
self.t_avg = tf.constant(davg.astype(GLOBAL_NP_FLOAT_PRECISION))
self.t_std = tf.constant(dstd.astype(GLOBAL_NP_FLOAT_PRECISION))
if pbc:
self.default_mesh = np.zeros(6, dtype=np.int32)
self.default_mesh[3] = 2
self.default_mesh[4] = 2
self.default_mesh[5] = 2
else:
self.default_mesh = np.array([], dtype=np.int32)
# make place holder
self.coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION,
[None, self.natoms[0] * 3],
name='t_coord')
self.box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, 9],
name='t_box')
self.type = tf.placeholder(tf.int32, [None, self.natoms[0]],
name="t_type")
self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms")
self.efield = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION,
[None, self.natoms[0] * 3],
name='t_efield')
def _init_from_frz_model(self):
# get the model type from the frozen model(self.run_opt.init_frz_model)
try:
t_model_type = get_tensor_by_name(self.run_opt.init_frz_model, 'model_type')
self.model_type = bytes.decode(t_model_type)
except GraphWithoutTensorError as e:
# throw runtime error if there's no frozen model
if not os.path.exists(self.run_opt.init_frz_model):
raise RuntimeError(
"The input frozen model %s (%s) does not exist! Please check the path of the frozen model. " % (self.run_opt.init_frz_model, os.path.abspath(self.run_opt.init_frz_model))
) from e
# throw runtime error if the frozen_model has no model type information...
else:
raise RuntimeError(
"The input frozen model: %s has no 'model_type' information, "
"which is not supported by the 'dp train init-frz-model' interface. " % self.run_opt.init_frz_model
) from e
if self.fitting_type != 'ener':
raise RuntimeError("The 'dp train init-frz-model' command only supports the 'ener' type fitting net currently!")
# self.frz_model will control the self.model to import the descriptor from the given frozen model instead of building from scratch...
# initialize fitting net with the given compressed frozen model
if self.model_type == 'original_model':
self.descrpt.init_variables(self.run_opt.init_frz_model)
self.fitting.init_variables(get_fitting_net_variables(self.run_opt.init_frz_model))
tf.constant("original_model", name = 'model_type', dtype = tf.string)
elif self.model_type == 'compressed_model':
self.frz_model = self.run_opt.init_frz_model
self.fitting.init_variables(get_fitting_net_variables(self.frz_model))
tf.constant("compressed_model", name = 'model_type', dtype = tf.string)
else:
raise RuntimeError("Unknown model type %s" % self.model_type)
def __init__(self, data, comp=0):
self.sess = tf.Session()
self.data = data
self.natoms = self.data.get_natoms()
self.ntypes = self.data.get_ntypes()
self.sel_a = [12, 24]
self.sel_r = [12, 24]
self.rcut_a = -1
self.rcut_r = 10.0
self.axis_rule = [0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0]
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
davg = np.zeros([self.ntypes, self.ndescrpt])
dstd = np.ones([self.ntypes, self.ndescrpt])
self.t_avg = tf.constant(davg.astype(global_np_float_precision))
self.t_std = tf.constant(dstd.astype(global_np_float_precision))
self.default_mesh = np.zeros(6, dtype=np.int32)
self.default_mesh[3] = 2
self.default_mesh[4] = 2
self.default_mesh[5] = 2
# make place holder
self.coord = tf.placeholder(global_tf_float_precision,
[None, self.natoms[0] * 3],
name='t_coord')
self.box = tf.placeholder(global_tf_float_precision, [None, 9],
name='t_box')
self.type = tf.placeholder(tf.int32, [None, self.natoms[0]],
name="t_type")
self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms")
def get_nbor_stat(jdata, rcut):
max_rcut = get_rcut(jdata)
type_map = get_type_map(jdata)
if type_map and len(type_map) == 0:
type_map = None
train_data = get_data(jdata["training"]["training_data"], max_rcut, type_map, None)
train_data.get_batch()
data_ntypes = train_data.get_ntypes()
if type_map is not None:
map_ntypes = len(type_map)
else:
map_ntypes = data_ntypes
ntypes = max([map_ntypes, data_ntypes])
neistat = NeighborStat(ntypes, rcut)
min_nbor_dist, max_nbor_size = neistat.get_stat(train_data)
# moved from traier.py as duplicated
# TODO: this is a simple fix but we should have a clear
# architecture to call neighbor stat
tf.constant(min_nbor_dist,
name = 'train_attr/min_nbor_dist',
dtype = GLOBAL_ENER_FLOAT_PRECISION)
tf.constant(max_nbor_size,
name = 'train_attr/max_nbor_size',
dtype = tf.int32)
return min_nbor_dist, max_nbor_size
def init_variables(
self,
graph: tf.Graph,
graph_def: tf.GraphDef,
model_type: str = "original_model",
suffix: str = "",
) -> None:
"""
Init the embedding net variables with the given frozen model
Parameters
----------
graph : tf.Graph
The input frozen model graph
graph_def : tf.GraphDef
The input frozen model graph_def
model_type : str
the type of the model
suffix : str
suffix to name scope
"""
# self.frz_model will control the self.model to import the descriptor from the given frozen model instead of building from scratch...
# initialize fitting net with the given compressed frozen model
if model_type == 'original_model':
self.descrpt.init_variables(graph, graph_def, suffix=suffix)
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("original_model", name='model_type', dtype=tf.string)
elif model_type == 'compressed_model':
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("compressed_model", name='model_type', dtype=tf.string)
else:
raise RuntimeError("Unknown model type %s" % model_type)
if self.typeebd is not None:
self.typeebd.init_variables(graph, graph_def, suffix=suffix)
def test_op_tanh(self):
w = tf.constant(
[[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1, 1.1, 1.2]],
dtype='double')
x = tf.constant([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9],
[1.0, 1.1, 1.2]],
dtype='double')
b = tf.constant([[0.1], [0.2], [0.3], [0.4]], dtype='double')
xbar = tf.matmul(x, w) + b
y = tf.nn.tanh(xbar)
dy = op_module.unaggregated_dy_dx_s(y, w, xbar, tf.constant(1))
dy_array = tf.Session().run(dy)
answer = np.array(
[[
8.008666403121351973e-02, 1.513925729426658651e-01,
2.134733287761668430e-01, 2.661983049806041501e-01
],
[
4.010658815015744061e-02, 6.306476628799793926e-02,
7.332167904608145881e-02, 7.494218676568849269e-02
],
[
1.561705624394135218e-02, 1.994112926507514427e-02,
1.887519955881525671e-02, 1.576442161040989692e-02
],
[
5.492686739421748753e-03, 5.754985286040992763e-03,
4.493113544969218158e-03, 3.107638130764600777e-03
]])
places = 18
np.testing.assert_almost_equal(dy_array, answer, places)
def test_op_gelu(self):
w = tf.constant(
[[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1, 1.1, 1.2]],
dtype='double')
x = tf.constant([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9],
[1.0, 1.1, 1.2]],
dtype='double')
b = tf.constant([[0.1], [0.2], [0.3], [0.4]], dtype='double')
xbar = tf.matmul(x, w) + b
y = gelu(xbar)
dy = op_module.unaggregated_dy_dx_s(y, w, xbar, tf.constant(2))
dy_array = tf.Session().run(dy)
answer = np.array(
[[
8.549286163555620821e-02, 1.782905778685600906e-01,
2.776474599997448833e-01, 3.827650237273348965e-01
],
[
1.089906023807040714e-01, 2.230820937721638697e-01,
3.381867859682909927e-01, 4.513008399758057232e-01
],
[
1.124254240556722684e-01, 2.209918074710395253e-01,
3.238894323148118759e-01, 4.220357318198978414e-01
],
[
1.072173273655498138e-01, 2.082159073100979807e-01,
3.059816075270163083e-01, 4.032981557798429595e-01
]])
places = 18
np.testing.assert_almost_equal(dy_array, answer, places)
def setUp(self, data, pbc=True):
self.sess = tf.Session()
self.data = data
self.natoms = self.data.get_natoms()
self.ntypes = self.data.get_ntypes()
self.sel = [12, 24]
self.sel_a = [0, 0]
self.rcut_smth = 2.45
self.rcut = 10.0
self.nnei = np.cumsum(self.sel)[-1]
self.ndescrpt = self.nnei * 1
davg = np.zeros([self.ntypes, self.ndescrpt])
dstd = np.ones([self.ntypes, self.ndescrpt])
self.t_avg = tf.constant(davg.astype(global_np_float_precision))
self.t_std = tf.constant(dstd.astype(global_np_float_precision))
if pbc:
self.default_mesh = np.zeros(6, dtype=np.int32)
self.default_mesh[3] = 2
self.default_mesh[4] = 2
self.default_mesh[5] = 2
else:
self.default_mesh = np.array([], dtype=np.int32)
# make place holder
self.coord = tf.placeholder(global_tf_float_precision,
[None, self.natoms[0] * 3],
name='t_coord')
self.box = tf.placeholder(global_tf_float_precision, [None, 9],
name='t_box')
self.type = tf.placeholder(tf.int32, [None, self.natoms[0]],
name="t_type")
self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms")
def setUp(self):
self.sess = self.test_session().__enter__()
self.natoms = [5, 5, 2, 3]
self.ntypes = 2
self.sel_a = [12, 24]
self.sel_r = [0, 0]
self.rcut_a = -1
self.rcut_r_smth = 2.45
self.rcut_r = 10.0
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
davg = np.zeros([self.ntypes, self.ndescrpt])
dstd = np.ones([self.ntypes, self.ndescrpt])
self.t_avg = tf.constant(davg.astype(GLOBAL_NP_FLOAT_PRECISION))
self.t_std = tf.constant(dstd.astype(GLOBAL_NP_FLOAT_PRECISION))
# no pbc
self.default_mesh = np.array([], dtype=np.int32)
# make place holder
self.coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION,
[None, self.natoms[0] * 3],
name='t_coord')
self.efield = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION,
[None, self.natoms[0] * 3],
name='t_efield')
self.box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, 9],
name='t_box')
self.type = tf.placeholder(tf.int32, [None, self.natoms[0]],
name="t_type")
self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms")
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_ : 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 build(self, data=None, stop_batch=0):
self.ntypes = self.model.get_ntypes()
self.stop_batch = stop_batch
if self.numb_fparam > 0:
log.info("training with %d frame parameter(s)" % self.numb_fparam)
else:
log.info("training without frame parameter")
if not self.is_compress:
# Usually, the type number of the model should be equal to that of the data
# However, nt_model > nt_data should be allowed, since users may only want to
# train using a dataset that only have some of elements
if self.ntypes < data.get_ntypes():
raise ValueError(
"The number of types of the training data is %d, but that of the "
"model is only %d. The latter must be no less than the former. "
"You may need to reset one or both of them. Usually, the former "
"is given by `model/type_map` in the training parameter (if set) "
"or the maximum number in the training data. The latter is given "
"by `model/descriptor/sel` in the training parameter." %
(data.get_ntypes(), self.ntypes))
self.type_map = data.get_type_map()
self.batch_size = data.get_batch_size()
self.model.data_stat(data)
# config the init_frz_model command
if self.run_opt.init_mode == 'init_from_frz_model':
self._init_from_frz_model()
# neighbor_stat is moved to train.py as duplicated
# TODO: this is a simple fix but we should have a clear
# architecture to call neighbor stat
else:
self.descrpt.enable_compression(
self.model_param['compress']["min_nbor_dist"],
self.model_param['compress']['model_file'],
self.model_param['compress']['table_config'][0],
self.model_param['compress']['table_config'][1],
self.model_param['compress']['table_config'][2],
self.model_param['compress']['table_config'][3])
self.fitting.init_variables(
get_fitting_net_variables(
self.model_param['compress']['model_file']))
if self.is_compress or self.model_type == 'compressed_model':
tf.constant("compressed_model", name='model_type', dtype=tf.string)
else:
tf.constant("original_model", name='model_type', dtype=tf.string)
self._build_lr()
self._build_network(data)
self._build_training()
def test_embed_atom_type(self):
ntypes = 3
natoms = tf.constant([5, 5, 3, 0, 2])
type_embedding = tf.constant([
[1, 2, 3],
[3, 2, 1],
[7, 7, 7],
])
expected_out = [[1, 2, 3], [1, 2, 3], [1, 2, 3], [7, 7, 7], [7, 7, 7]]
atom_embed = embed_atom_type(ntypes, natoms, type_embedding)
sess = self.test_session().__enter__()
atom_embed = sess.run(atom_embed)
np.testing.assert_almost_equal(atom_embed, expected_out, 10)
def __init__(self, jdata, descrpt):
# model param
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
args = ClassArg()\
.add('numb_fparam', int, default = 0)\
.add('numb_aparam', int, default = 0)\
.add('neuron', list, default = [120,120,120], alias = 'n_neuron')\
.add('resnet_dt', bool, default = True)\
.add('rcond', float, default = 1e-3) \
.add('seed', int) \
.add('atom_ener', list, default = [])\
.add("activation_function", str, default = "tanh")\
.add("precision", str, default = "default")\
.add("trainable", [list, bool], default = True)
class_data = args.parse(jdata)
self.numb_fparam = class_data['numb_fparam']
self.numb_aparam = class_data['numb_aparam']
self.n_neuron = class_data['neuron']
self.resnet_dt = class_data['resnet_dt']
self.rcond = class_data['rcond']
self.seed = class_data['seed']
self.fitting_activation_fn = get_activation_func(
class_data["activation_function"])
self.fitting_precision = get_precision(class_data['precision'])
self.trainable = class_data['trainable']
if type(self.trainable) is bool:
self.trainable = [self.trainable] * (len(self.n_neuron) + 1)
assert (len(self.trainable) == len(self.n_neuron) +
1), 'length of trainable should be that of n_neuron + 1'
self.atom_ener = []
for at, ae in enumerate(class_data['atom_ener']):
if ae is not None:
self.atom_ener.append(
tf.constant(ae,
global_tf_float_precision,
name="atom_%d_ener" % at))
else:
self.atom_ener.append(None)
self.useBN = False
self.bias_atom_e = None
# data requirement
if self.numb_fparam > 0:
add_data_requirement('fparam',
self.numb_fparam,
atomic=False,
must=True,
high_prec=False)
self.fparam_avg = None
self.fparam_std = None
self.fparam_inv_std = None
if self.numb_aparam > 0:
add_data_requirement('aparam',
self.numb_aparam,
atomic=True,
must=True,
high_prec=False)
self.aparam_avg = None
self.aparam_std = None
self.aparam_inv_std = None
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 test_nopbc_self_built_nlist(self):
tem, tem_deriv, trij, tnlist \
= op_module.prod_env_mat_a (
self.tcoord,
self.ttype,
self.tnatoms,
self.tbox,
tf.constant(np.zeros(0, dtype = np.int32)),
self.t_avg,
self.t_std,
rcut_a = -1,
rcut_r = self.rcut,
rcut_r_smth = self.rcut_smth,
sel_a = self.sel,
sel_r = [0, 0])
self.sess.run(tf.global_variables_initializer())
dem, dem_deriv, drij, dnlist = self.sess.run(
[tem, tem_deriv, trij, tnlist],
feed_dict={
self.tcoord: self.dcoord,
self.ttype: self.dtype,
self.tbox: self.dbox,
self.tnatoms: self.dnatoms
})
self.assertEqual(dem.shape, (self.nframes, self.nloc * self.ndescrpt))
self.assertEqual(dem_deriv.shape,
(self.nframes, self.nloc * self.ndescrpt * 3))
self.assertEqual(drij.shape, (self.nframes, self.nloc * self.nnei * 3))
self.assertEqual(dnlist.shape, (self.nframes, self.nloc * self.nnei))
for ff in range(self.nframes):
np.testing.assert_almost_equal(dem[ff], self.nopbc_expected_output,
5)
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 _compute_dstats_sys_smth (self,
data_coord,
data_box,
data_atype,
natoms_vec,
mesh) :
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():
descrpt, descrpt_deriv, rij, nlist \
= op_module.descrpt_se_a (tf.constant(data_coord),
tf.constant(data_atype),
tf.constant(natoms_vec, dtype = tf.int32),
tf.constant(data_box),
tf.constant(mesh),
tf.constant(avg_zero),
tf.constant(std_ones),
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.sess.run(tf.global_variables_initializer())
# sub_sess = tf.Session(graph = sub_graph,
# config=tf.ConfigProto(intra_op_parallelism_threads=self.run_opt.num_intra_threads,
# inter_op_parallelism_threads=self.run_opt.num_inter_threads
# ))
sub_sess = tf.Session(graph = sub_graph)
dd_all = sub_sess.run(descrpt)
sub_sess.close()
natoms = natoms_vec
dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]])
start_index = 0
sysr = []
sysa = []
sysn = []
sysr2 = []
sysa2 = []
for type_i in range(self.ntypes):
end_index = start_index + self.ndescrpt * natoms[2+type_i]
dd = dd_all[:, start_index:end_index]
dd = np.reshape(dd, [-1, self.ndescrpt])
start_index = end_index
# compute
dd = np.reshape (dd, [-1, 4])
ddr = dd[:,:1]
dda = dd[:,1:]
sumr = np.sum(ddr)
suma = np.sum(dda) / 3.
sumn = dd.shape[0]
sumr2 = np.sum(np.multiply(ddr, ddr))
suma2 = np.sum(np.multiply(dda, dda)) / 3.
sysr.append(sumr)
sysa.append(suma)
sysn.append(sumn)
sysr2.append(sumr2)
sysa2.append(suma2)
return sysr, sysr2, sysa, sysa2, sysn
def _make_data(self, xx, idx):
with self.sub_graph.as_default():
with self.sub_sess.as_default():
xx = tf.reshape(xx, [xx.size, -1])
for layer in range(self.layer_size):
if layer == 0:
xbar = tf.matmul(
xx, self.matrix["layer_" + str(layer + 1)]
[idx]) + self.bias["layer_" + str(layer + 1)][idx]
yy = self._layer_0(
xx, self.matrix["layer_" + str(layer + 1)][idx],
self.bias["layer_" + str(layer + 1)][idx])
dy = op_module.unaggregated_dy_dx_s(
yy, self.matrix["layer_" + str(layer + 1)][idx],
xbar, tf.constant(self.functype))
dy2 = op_module.unaggregated_dy2_dx_s(
yy, dy,
self.matrix["layer_" + str(layer + 1)][idx], xbar,
tf.constant(self.functype))
else:
ybar = tf.matmul(
yy, self.matrix["layer_" + str(layer + 1)]
[idx]) + self.bias["layer_" + str(layer + 1)][idx]
tt, zz = self._layer_1(
yy, self.matrix["layer_" + str(layer + 1)][idx],
self.bias["layer_" + str(layer + 1)][idx])
dz = op_module.unaggregated_dy_dx(
zz - tt,
self.matrix["layer_" + str(layer + 1)][idx], dy,
ybar, tf.constant(self.functype))
dy2 = op_module.unaggregated_dy2_dx(
zz - tt,
self.matrix["layer_" + str(layer + 1)][idx], dy,
dy2, ybar, tf.constant(self.functype))
dy = dz
yy = zz
vv = zz.eval()
dd = dy.eval()
d2 = dy2.eval()
return vv, dd, d2
def build_fv_graph(self) -> tf.Tensor:
"""
Build the computational graph for the force and virial inference.
"""
with tf.variable_scope('modifier_attr'):
t_mdl_name = tf.constant(self.model_name,
name='mdl_name',
dtype=tf.string)
t_modi_type = tf.constant(self.modifier_prefix,
name='type',
dtype=tf.string)
t_mdl_charge_map = tf.constant(' '.join(
[str(ii) for ii in self.model_charge_map]),
name='mdl_charge_map',
dtype=tf.string)
t_sys_charge_map = tf.constant(' '.join(
[str(ii) for ii in self.sys_charge_map]),
name='sys_charge_map',
dtype=tf.string)
t_ewald_h = tf.constant(self.ewald_h,
name='ewald_h',
dtype=tf.float64)
t_ewald_b = tf.constant(self.ewald_beta,
name='ewald_beta',
dtype=tf.float64)
with self.graph.as_default():
return self._build_fv_graph_inner()
def _compute_dstats_sys_nonsmth (self,
data_coord,
data_box,
data_atype,
natoms_vec,
mesh,
reuse = None) :
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():
descrpt, descrpt_deriv, rij, nlist, axis, rot_mat \
= op_module.descrpt (tf.constant(data_coord),
tf.constant(data_atype),
tf.constant(natoms_vec, dtype = tf.int32),
tf.constant(data_box),
tf.constant(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.sess.run(tf.global_variables_initializer())
# sub_sess = tf.Session(graph = sub_graph,
# config=tf.ConfigProto(intra_op_parallelism_threads=self.run_opt.num_intra_threads,
# inter_op_parallelism_threads=self.run_opt.num_inter_threads
# ))
sub_sess = tf.Session(graph = sub_graph)
dd_all = sub_sess.run(descrpt)
sub_sess.close()
natoms = natoms_vec
dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]])
start_index = 0
sysv = []
sysn = []
sysv2 = []
for type_i in range(self.ntypes):
end_index = start_index + self.ndescrpt * natoms[2+type_i]
dd = dd_all[:, start_index:end_index]
dd = np.reshape(dd, [-1, self.ndescrpt])
start_index = end_index
# compute
sumv = np.sum(dd, axis = 0)
sumn = dd.shape[0]
sumv2 = np.sum(np.multiply(dd,dd), axis = 0)
sysv.append(sumv)
sysn.append(sumn)
sysv2.append(sumv2)
return sysv, sysv2, sysn
def test_nopbc_self_built_nlist_deriv(self):
hh = 1e-4
tem, tem_deriv, trij, tnlist \
= op_module.prod_env_mat_a (
self.tcoord,
self.ttype,
self.tnatoms,
self.tbox,
tf.constant(np.zeros(0, dtype = np.int32)),
self.t_avg,
self.t_std,
rcut_a = -1,
rcut_r = self.rcut,
rcut_r_smth = self.rcut_smth,
sel_a = self.sel,
sel_r = [0, 0])
self.sess.run(tf.global_variables_initializer())
self.check_deriv_numerical_deriv(hh, tem, tem_deriv, trij, tnlist)
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_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_fv_graph(self):
with tf.variable_scope('modifier_attr') :
t_mdl_name = tf.constant(self.model_name,
name = 'mdl_name',
dtype = tf.string)
t_modi_type = tf.constant(self.modifier_prefix,
name = 'type',
dtype = tf.string)
t_mdl_charge_map = tf.constant(' '.join([str(ii) for ii in self.model_charge_map]),
name = 'mdl_charge_map',
dtype = tf.string)
t_sys_charge_map = tf.constant(' '.join([str(ii) for ii in self.sys_charge_map]),
name = 'sys_charge_map',
dtype = tf.string)
t_ewald_h = tf.constant(self.ewald_h,
name = 'ewald_h',
dtype = tf.float64)
t_ewald_b = tf.constant(self.ewald_beta,
name = 'ewald_beta',
dtype = tf.float64)
with self.graph.as_default():
return self._build_fv_graph_inner()
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 __init__(self, jdata):
args = ClassArg()\
.add('sel', list, must = True) \
.add('rcut', float, default = 6.0) \
.add('rcut_smth',float, default = 5.5) \
.add('neuron', list, default = [10, 20, 40]) \
.add('axis_neuron', int, default = 4, alias = 'n_axis_neuron') \
.add('resnet_dt',bool, default = False) \
.add('trainable',bool, default = True) \
.add('seed', int) \
.add('exclude_types', list, default = []) \
.add('set_davg_zero', bool, default = False) \
.add('activation_function', str, default = 'tanh') \
.add('precision', str, default = "default")
class_data = args.parse(jdata)
self.sel_a = class_data['sel']
self.rcut_r = class_data['rcut']
self.rcut_r_smth = class_data['rcut_smth']
self.filter_neuron = class_data['neuron']
self.n_axis_neuron = class_data['axis_neuron']
self.filter_resnet_dt = class_data['resnet_dt']
self.seed = class_data['seed']
self.trainable = class_data['trainable']
self.filter_activation_fn = get_activation_func(
class_data['activation_function'])
self.filter_precision = get_precision(class_data['precision'])
exclude_types = class_data['exclude_types']
self.exclude_types = set()
for tt in exclude_types:
assert (len(tt) == 2)
self.exclude_types.add((tt[0], tt[1]))
self.exclude_types.add((tt[1], tt[0]))
self.set_davg_zero = class_data['set_davg_zero']
# descrpt config
self.sel_r = [0 for ii in range(len(self.sel_a))]
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.useBN = False
self.dstd = None
self.davg = 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_sea_'
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 \
= op_module.descrpt_se_a(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,
rcut_r_smth = self.rcut_r_smth,
sel_a = self.sel_a,
sel_r = self.sel_r)
self.sub_sess = tf.Session(graph=sub_graph,
config=default_tf_session_config)
def setUp(self):
self.sess = self.test_session().__enter__()
self.inputs = tf.constant([0., 1., 2.], dtype=tf.float64)
self.ndata = 3
self.inputs = tf.reshape(self.inputs, [-1, 1])
self.places = 6
