def test_py_interface(self):
hh = 1e-4
places = 4
sess = tf.Session()
t_energy, t_force, t_virial \
= op_module.ewald_recp(self.coord, self.charge, self.nloc, self.box,
ewald_h = self.ewald_h,
ewald_beta = self.ewald_beta)
[e, f, v] = sess.run([t_energy, t_force, t_virial],
feed_dict = {
self.coord: self.dcoord.reshape([-1]),
self.charge: self.dcharge.reshape([-1]),
self.box: self.dbox.reshape([-1]),
self.nloc: [self.natoms],
})
er = EwaldRecp(self.ewald_h, self.ewald_beta)
e1, f1, v1 = er.eval(self.dcoord, self.dcharge, self.dbox)
for ff in range(self.nframes):
self.assertAlmostEqual(e[ff], e1[ff],
places = places,
msg = "frame %d energy failed" % (ff))
for idx in range(self.natoms):
for dd in range(3):
self.assertAlmostEqual(f[ff, idx*3+dd], f1[ff,idx*3+dd],
places = places,
msg = "frame %d force component [%d,%d] failed" % (ff, idx, dd))
for d0 in range(3):
for d1 in range(3):
self.assertAlmostEqual(v[ff, d0*3+d1], v[ff,d0*3+d1],
places = places,
msg = "frame %d virial component [%d,%d] failed" % (ff, d0, d1))
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 get_tensor_by_name_from_graph(graph: tf.Graph,
tensor_name: str) -> tf.Tensor:
"""
Load tensor value from the given tf.Graph object
Parameters
----------
graph : tf.Graph
The input TensorFlow graph
tensor_name : str
Indicates which tensor which will be loaded from the frozen model
Returns
-------
tf.Tensor
The tensor which was loaded from the frozen model
Raises
------
GraphWithoutTensorError
Whether the tensor_name is within the frozen model
"""
try:
tensor = graph.get_tensor_by_name(tensor_name + ':0')
except KeyError as e:
raise GraphWithoutTensorError() from e
with tf.Session(graph=graph) as sess:
tensor = run_sess(sess, tensor)
return tensor
def __init__(self, model_file):
model_file = model_file
graph = self.load_graph(model_file)
t_mt = graph.get_tensor_by_name('load/model_attr/model_type:0')
sess = tf.Session(graph=graph)
[mt] = sess.run([t_mt], feed_dict={})
self.model_type = mt.decode('utf-8')
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 _init_sess_serial(self) :
self.sess = tf.Session(
config=tf.ConfigProto(intra_op_parallelism_threads=self.run_opt.num_intra_threads,
inter_op_parallelism_threads=self.run_opt.num_inter_threads
))
self.saver = tf.train.Saver()
saver = self.saver
if self.run_opt.init_mode == 'init_from_scratch' :
self._message("initialize model from scratch")
init_op = tf.global_variables_initializer()
self.sess.run(init_op)
fp = open(self.disp_file, "w")
fp.close ()
elif self.run_opt.init_mode == 'init_from_model' :
self._message("initialize from model %s" % self.run_opt.init_model)
init_op = tf.global_variables_initializer()
self.sess.run(init_op)
saver.restore (self.sess, self.run_opt.init_model)
self.sess.run(self.global_step.assign(0))
fp = open(self.disp_file, "w")
fp.close ()
elif self.run_opt.init_mode == 'restart' :
self._message("restart from model %s" % self.run_opt.restart)
init_op = tf.global_variables_initializer()
self.sess.run(init_op)
saver.restore (self.sess, self.run_opt.restart)
else :
raise RuntimeError ("unkown init mode")
def __init__(self,
model_file,
variable_name,
variable_dof) :
DeepEval.__init__(self, model_file)
self.model_file = model_file
self.graph = self.load_graph (self.model_file)
self.variable_name = variable_name
self.variable_dof = variable_dof
# checkout input/output tensors from graph
self.t_ntypes = self.graph.get_tensor_by_name ('load/descrpt_attr/ntypes:0')
self.t_rcut = self.graph.get_tensor_by_name ('load/descrpt_attr/rcut:0')
self.t_tmap = self.graph.get_tensor_by_name ('load/model_attr/tmap:0')
self.t_sel_type= self.graph.get_tensor_by_name ('load/model_attr/sel_type:0')
# inputs
self.t_coord = self.graph.get_tensor_by_name ('load/t_coord:0')
self.t_type = self.graph.get_tensor_by_name ('load/t_type:0')
self.t_natoms = self.graph.get_tensor_by_name ('load/t_natoms:0')
self.t_box = self.graph.get_tensor_by_name ('load/t_box:0')
self.t_mesh = self.graph.get_tensor_by_name ('load/t_mesh:0')
# outputs
self.t_tensor = self.graph.get_tensor_by_name ('load/o_%s:0' % self.variable_name)
# start a tf session associated to the graph
self.sess = tf.Session (graph = self.graph)
[self.ntypes, self.rcut, self.tmap, self.tselt] = self.sess.run([self.t_ntypes, self.t_rcut, self.t_tmap, self.t_sel_type])
self.tmap = self.tmap.decode('UTF-8').split()
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 __init__(self, data):
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))
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 __init__(self, model_file):
self.model_file = model_file
self.graph = self.load_graph(self.model_file)
# checkout input/output tensors from graph
self.t_ntypes = self.graph.get_tensor_by_name(
'load/descrpt_attr/ntypes:0')
self.t_rcut = self.graph.get_tensor_by_name('load/descrpt_attr/rcut:0')
self.t_dfparam = self.graph.get_tensor_by_name(
'load/fitting_attr/dfparam:0')
self.t_tmap = self.graph.get_tensor_by_name('load/model_attr/tmap:0')
# inputs
self.t_coord = self.graph.get_tensor_by_name('load/t_coord:0')
self.t_type = self.graph.get_tensor_by_name('load/t_type:0')
self.t_natoms = self.graph.get_tensor_by_name('load/t_natoms:0')
self.t_box = self.graph.get_tensor_by_name('load/t_box:0')
self.t_mesh = self.graph.get_tensor_by_name('load/t_mesh:0')
# outputs
self.t_energy = self.graph.get_tensor_by_name('load/o_energy:0')
self.t_force = self.graph.get_tensor_by_name('load/o_force:0')
self.t_virial = self.graph.get_tensor_by_name('load/o_virial:0')
self.t_ae = self.graph.get_tensor_by_name('load/o_atom_energy:0')
self.t_av = self.graph.get_tensor_by_name('load/o_atom_virial:0')
self.t_fparam = None
# check if the graph has fparam
for op in self.graph.get_operations():
if op.name == 'load/t_fparam':
self.t_fparam = self.graph.get_tensor_by_name(
'load/t_fparam:0')
self.has_fparam = self.t_fparam is not None
# start a tf session associated to the graph
self.sess = tf.Session(graph=self.graph)
[self.ntypes, self.rcut, self.dfparam, self.tmap] = self.sess.run(
[self.t_ntypes, self.t_rcut, self.t_dfparam, self.t_tmap])
self.tmap = self.tmap.decode('UTF-8').split()
def __init__(self, ntypes: int, rcut: float) -> None:
"""
Constructor
"""
self.rcut = rcut
self.ntypes = ntypes
self.place_holders = {}
sub_graph = tf.Graph()
with sub_graph.as_default():
for ii in ['coord', 'box']:
self.place_holders[ii] = tf.placeholder(
GLOBAL_NP_FLOAT_PRECISION, [None, None], name='t_' + ii)
self.place_holders['type'] = tf.placeholder(tf.int32, [None, None],
name='t_type')
self.place_holders['natoms_vec'] = tf.placeholder(
tf.int32, [self.ntypes + 2], name='t_natoms')
self.place_holders['default_mesh'] = tf.placeholder(tf.int32,
[None],
name='t_mesh')
self._max_nbor_size, self._min_nbor_dist \
= op_module.neighbor_stat(self.place_holders['coord'],
self.place_holders['type'],
self.place_holders['natoms_vec'],
self.place_holders['box'],
self.place_holders['default_mesh'],
rcut = self.rcut)
self.sub_sess = tf.Session(graph=sub_graph,
config=default_tf_session_config)
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 __init__(self,
model_file,
default_tf_graph = False) :
DeepEval.__init__(self, model_file, default_tf_graph = default_tf_graph)
# self.model_file = model_file
# self.graph = self.load_graph (self.model_file)
# checkout input/output tensors from graph
self.t_ntypes = self.graph.get_tensor_by_name ('load/descrpt_attr/ntypes:0')
self.t_rcut = self.graph.get_tensor_by_name ('load/descrpt_attr/rcut:0')
self.t_dfparam= self.graph.get_tensor_by_name ('load/fitting_attr/dfparam:0')
self.t_daparam= self.graph.get_tensor_by_name ('load/fitting_attr/daparam:0')
self.t_tmap = self.graph.get_tensor_by_name ('load/model_attr/tmap:0')
# inputs
self.t_coord = self.graph.get_tensor_by_name ('load/t_coord:0')
self.t_type = self.graph.get_tensor_by_name ('load/t_type:0')
self.t_natoms = self.graph.get_tensor_by_name ('load/t_natoms:0')
self.t_box = self.graph.get_tensor_by_name ('load/t_box:0')
self.t_mesh = self.graph.get_tensor_by_name ('load/t_mesh:0')
# outputs
self.t_energy = self.graph.get_tensor_by_name ('load/o_energy:0')
self.t_force = self.graph.get_tensor_by_name ('load/o_force:0')
self.t_virial = self.graph.get_tensor_by_name ('load/o_virial:0')
self.t_ae = self.graph.get_tensor_by_name ('load/o_atom_energy:0')
self.t_av = self.graph.get_tensor_by_name ('load/o_atom_virial:0')
self.t_fparam = None
self.t_aparam = None
# check if the graph has fparam
for op in self.graph.get_operations():
if op.name == 'load/t_fparam' :
self.t_fparam = self.graph.get_tensor_by_name ('load/t_fparam:0')
self.has_fparam = self.t_fparam is not None
# check if the graph has aparam
for op in self.graph.get_operations():
if op.name == 'load/t_aparam' :
self.t_aparam = self.graph.get_tensor_by_name ('load/t_aparam:0')
self.has_aparam = self.t_aparam is not None
# start a tf session associated to the graph
self.sess = tf.Session (graph = self.graph, config=default_tf_session_config)
[self.ntypes, self.rcut, self.dfparam, self.daparam, self.tmap] = self.sess.run([self.t_ntypes, self.t_rcut, self.t_dfparam, self.t_daparam, self.t_tmap])
self.tmap = self.tmap.decode('UTF-8').split()
# setup modifier
try:
t_modifier_type = self.graph.get_tensor_by_name('load/modifier_attr/type:0')
self.modifier_type = self.sess.run(t_modifier_type).decode('UTF-8')
except ValueError:
self.modifier_type = None
except KeyError:
self.modifier_type = None
if self.modifier_type == 'dipole_charge':
t_mdl_name = self.graph.get_tensor_by_name('load/modifier_attr/mdl_name:0')
t_mdl_charge_map = self.graph.get_tensor_by_name('load/modifier_attr/mdl_charge_map:0')
t_sys_charge_map = self.graph.get_tensor_by_name('load/modifier_attr/sys_charge_map:0')
t_ewald_h = self.graph.get_tensor_by_name('load/modifier_attr/ewald_h:0')
t_ewald_beta = self.graph.get_tensor_by_name('load/modifier_attr/ewald_beta:0')
[mdl_name, mdl_charge_map, sys_charge_map, ewald_h, ewald_beta] = self.sess.run([t_mdl_name, t_mdl_charge_map, t_sys_charge_map, t_ewald_h, t_ewald_beta])
mdl_charge_map = [int(ii) for ii in mdl_charge_map.decode('UTF-8').split()]
sys_charge_map = [int(ii) for ii in sys_charge_map.decode('UTF-8').split()]
self.dm = DipoleChargeModifier(mdl_name, mdl_charge_map, sys_charge_map, ewald_h = ewald_h, ewald_beta = ewald_beta)
def __init__(self, model_file, load_prefix='load', default_tf_graph=False):
self.graph = self._load_graph(model_file,
prefix=load_prefix,
default_tf_graph=default_tf_graph)
t_mt = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 'model_attr/model_type:0'))
sess = tf.Session(graph=self.graph, config=default_tf_session_config)
[mt] = sess.run([t_mt], feed_dict={})
self.model_type = mt.decode('utf-8')
def model_type(self) -> str:
"""Get type of model.
:type:str
"""
if not self._model_type:
t_mt = self._get_tensor("model_attr/model_type:0")
sess = tf.Session(graph=self.graph,
config=default_tf_session_config)
[mt] = run_sess(sess, [t_mt], feed_dict={})
self._model_type = mt.decode("utf-8")
return self._model_type
def freeze_graph(model_folder, output, output_node_names=None):
# We retrieve our checkpoint fullpath
checkpoint = tf.train.get_checkpoint_state(model_folder)
input_checkpoint = checkpoint.model_checkpoint_path
# We precise the file fullname of our freezed graph
absolute_model_folder = "/".join(input_checkpoint.split('/')[:-1])
output_graph = absolute_model_folder + "/" + output
# Before exporting our graph, we need to precise what is our output node
# This is how TF decides what part of the Graph he has to keep and what part it can dump
# NOTE: this variable is plural, because you can have multiple output nodes
# output_node_names = "energy_test,force_test,virial_test,t_rcut"
# We clear devices to allow TensorFlow to control on which device it will load operations
clear_devices = True
# We import the meta graph and retrieve a Saver
saver = tf.train.import_meta_graph(input_checkpoint + '.meta',
clear_devices=clear_devices)
# We retrieve the protobuf graph definition
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
nodes = [n.name for n in input_graph_def.node]
# We start a session and restore the graph weights
with tf.Session() as sess:
saver.restore(sess, input_checkpoint)
model_type = sess.run('model_attr/model_type:0',
feed_dict={}).decode('utf-8')
if 'modifier_attr/type' in nodes:
modifier_type = sess.run('modifier_attr/type:0',
feed_dict={}).decode('utf-8')
else:
modifier_type = None
if output_node_names is None:
output_node_names = _make_node_names(model_type, modifier_type)
print('The following nodes will be frozen: %s' % output_node_names)
# We use a built-in TF helper to export variables to constants
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess, # The session is used to retrieve the weights
input_graph_def, # The graph_def is used to retrieve the nodes
output_node_names.split(
","
) # The output node names are used to select the usefull nodes
)
# Finally we serialize and dump the output graph to the filesystem
with tf.gfile.GFile(output_graph, "wb") as f:
f.write(output_graph_def.SerializeToString())
print("%d ops in the final graph." % len(output_graph_def.node))
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 _init_session(self):
config = get_tf_session_config()
device, idx = self.run_opt.my_device.split(":", 1)
if device == "gpu":
config.gpu_options.visible_device_list = idx
self.sess = tf.Session(config=config)
# Initializes or restore global variables
init_op = tf.global_variables_initializer()
if self.run_opt.is_chief:
self.saver = tf.train.Saver(save_relative_paths=True)
if self.run_opt.init_mode == 'init_from_scratch' :
log.info("initialize model from scratch")
run_sess(self.sess, init_op)
if not self.is_compress:
fp = open(self.disp_file, "w")
fp.close ()
elif self.run_opt.init_mode == 'init_from_model' :
log.info("initialize from model %s" % self.run_opt.init_model)
run_sess(self.sess, init_op)
self.saver.restore (self.sess, self.run_opt.init_model)
run_sess(self.sess, self.global_step.assign(0))
fp = open(self.disp_file, "w")
fp.close ()
elif self.run_opt.init_mode == 'restart' :
log.info("restart from model %s" % self.run_opt.restart)
run_sess(self.sess, init_op)
self.saver.restore (self.sess, self.run_opt.restart)
elif self.run_opt.init_mode == 'init_from_frz_model' :
log.info("initialize training from the frozen model")
run_sess(self.sess, init_op)
fp = open(self.disp_file, "w")
fp.close ()
else :
raise RuntimeError ("unkown init mode")
else:
run_sess(self.sess, init_op)
self.saver = None
# Ensure variable consistency among tasks when training starts
if self.run_opt.is_distrib:
bcast_op = self.run_opt._HVD.broadcast_global_variables(0)
if self.run_opt.is_chief:
log.info('broadcast global variables to other tasks')
else:
log.info('receive global variables from task#0')
run_sess(self.sess, bcast_op)
def setUp(self, data):
self.sess = tf.Session()
self.data = data
self.natoms = self.data.get_natoms()
self.ntypes = self.data.get_ntypes()
param_a = {
'sel': [12, 24],
'rcut': 4,
'rcut_smth': 3.5,
'neuron': [5, 10, 20],
'seed': 1,
}
param_r = {
'sel': [20, 40],
'rcut': 6,
'rcut_smth': 6.5,
'neuron': [10, 20, 40],
'seed': 1,
}
param = {'a': param_a, 'r': param_r}
self.descrpt = DescrptSeAR(param)
self.ndescrpt = self.descrpt.get_dim_out()
# davg = np.zeros ([self.ntypes, self.ndescrpt])
# dstd = np.ones ([self.ntypes, self.ndescrpt])
# self.t_avg = tf.constant(davg.astype(np.float64))
# self.t_std = tf.constant(dstd.astype(np.float64))
avg_a = np.zeros([self.ntypes, self.descrpt.descrpt_a.ndescrpt])
std_a = np.ones([self.ntypes, self.descrpt.descrpt_a.ndescrpt])
avg_r = np.zeros([self.ntypes, self.descrpt.descrpt_r.ndescrpt])
std_r = np.ones([self.ntypes, self.descrpt.descrpt_r.ndescrpt])
self.avg = [avg_a, avg_r]
self.std = [std_a, std_r]
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 __init__(self, hh, beta):
self.hh = hh
self.beta = beta
with tf.Graph().as_default() as graph:
# place holders
self.t_nloc = tf.placeholder(tf.int32, [1], name="t_nloc")
self.t_coord = tf.placeholder(global_tf_float_precision, [None],
name='t_coord')
self.t_charge = tf.placeholder(global_tf_float_precision, [None],
name='t_charge')
self.t_box = tf.placeholder(global_tf_float_precision, [None],
name='t_box')
# output
self.t_energy, self.t_force, self.t_virial \
= op_module.ewald_recp(self.t_coord, self.t_charge, self.t_nloc, self.t_box,
ewald_h = self.hh,
ewald_beta = self.beta)
self.sess = tf.Session(graph=graph, config=default_tf_session_config)
def _setUp(self):
args = Args()
run_opt = RunOptions(args, False)
with open (args.INPUT, 'r') as fp:
jdata = json.load (fp)
# init model
model = NNPTrainer (jdata, run_opt = run_opt)
rcut = model.model.get_rcut()
# init data system
systems = j_must_have(jdata['training'], 'systems')
set_pfx = j_must_have(jdata['training'], 'set_prefix')
batch_size = j_must_have(jdata['training'], 'batch_size')
test_size = j_must_have(jdata['training'], 'numb_test')
data = DeepmdDataSystem(systems,
batch_size,
test_size,
rcut,
set_prefix=set_pfx)
data.add_dict(data_requirement)
# clear the default graph
tf.reset_default_graph()
# build the model with stats from the first system
model.build (data)
# freeze the graph
with tf.Session() as sess:
init_op = tf.global_variables_initializer()
sess.run(init_op)
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
nodes = "o_dipole,o_rmat,o_rmat_deriv,o_nlist,o_rij,descrpt_attr/rcut,descrpt_attr/ntypes,descrpt_attr/sel,descrpt_attr/ndescrpt,model_attr/tmap,model_attr/sel_type,model_attr/model_type"
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess,
input_graph_def,
nodes.split(",")
)
output_graph = os.path.join(modifier_datapath, 'dipole.pb')
with tf.gfile.GFile(output_graph, "wb") as f:
f.write(output_graph_def.SerializeToString())
def __init__(self,
model_file,
variable_name,
variable_dof,
load_prefix='load',
default_tf_graph=False):
DeepEval.__init__(self,
model_file,
load_prefix=load_prefix,
default_tf_graph=default_tf_graph)
# self.model_file = model_file
# self.graph = self.load_graph (self.model_file)
self.variable_name = variable_name
self.variable_dof = variable_dof
# checkout input/output tensors from graph
self.t_ntypes = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 'descrpt_attr/ntypes:0'))
self.t_rcut = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 'descrpt_attr/rcut:0'))
self.t_tmap = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 'model_attr/tmap:0'))
self.t_sel_type = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 'model_attr/sel_type:0'))
# inputs
self.t_coord = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 't_coord:0'))
self.t_type = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 't_type:0'))
self.t_natoms = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 't_natoms:0'))
self.t_box = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 't_box:0'))
self.t_mesh = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 't_mesh:0'))
# outputs
self.t_tensor = self.graph.get_tensor_by_name(
os.path.join(load_prefix, 'o_%s:0' % self.variable_name))
# start a tf session associated to the graph
self.sess = tf.Session(graph=self.graph,
config=default_tf_session_config)
[self.ntypes, self.rcut, self.tmap, self.tselt] = self.sess.run(
[self.t_ntypes, self.t_rcut, self.t_tmap, self.t_sel_type])
self.tmap = self.tmap.decode('UTF-8').split()
def __init__(self,
model_file: "Path",
load_prefix: str = 'load',
default_tf_graph: bool = False) -> None:
"""Constructor"""
DeepEval.__init__(self,
model_file,
load_prefix=load_prefix,
default_tf_graph=default_tf_graph)
# check model type
model_type = self.tensors["t_tensor"][2:-2]
assert self.model_type == model_type, \
f"expect {model_type} model but got {self.model_type}"
# now load tensors to object attributes
for attr_name, tensor_name in self.tensors.items():
self._get_tensor(tensor_name, attr_name)
# load optional tensors if possible
optional_tensors = {
"t_global_tensor": f"o_global_{model_type}:0",
"t_force": "o_force:0",
"t_virial": "o_virial:0",
"t_atom_virial": "o_atom_virial:0"
}
try:
# first make sure these tensor all exists (but do not modify self attr)
for attr_name, tensor_name in optional_tensors.items():
self._get_tensor(tensor_name)
# then put those into self.attrs
for attr_name, tensor_name in optional_tensors.items():
self._get_tensor(tensor_name, attr_name)
except KeyError:
self._support_gfv = False
else:
self.tensors.update(optional_tensors)
self._support_gfv = True
# start a tf session associated to the graph
self.sess = tf.Session(graph=self.graph,
config=default_tf_session_config)
self._run_default_sess()
self.tmap = self.tmap.decode('UTF-8').split()
def model_version(self) -> str:
"""Get version of model.
Returns
-------
str
version of model
"""
if not self._model_version:
try:
t_mt = self._get_tensor("model_attr/model_version:0")
except KeyError:
# For deepmd-kit version 0.x - 1.x, set model version to 0.0
self._model_version = "0.0"
else:
sess = tf.Session(graph=self.graph,
config=default_tf_session_config)
[mt] = run_sess(sess, [t_mt], feed_dict={})
self._model_version = mt.decode("utf-8")
return self._model_version
def test_force(self):
hh = 1e-4
places = 4
sess = tf.Session()
t_energy, t_force, t_virial \
= op_module.ewald_recp(self.coord, self.charge, self.nloc, self.box,
ewald_h = self.ewald_h,
ewald_beta = self.ewald_beta)
[force] = sess.run([t_force],
feed_dict = {
self.coord: self.dcoord.reshape([-1]),
self.charge: self.dcharge.reshape([-1]),
self.box: self.dbox.reshape([-1]),
self.nloc: [self.natoms],
})
for idx in range(self.natoms):
for dd in range(3):
dcoordp = np.copy(self.dcoord)
dcoordm = np.copy(self.dcoord)
dcoordp[:,idx*3+dd] = self.dcoord[:,idx*3+dd] + hh
dcoordm[:,idx*3+dd] = self.dcoord[:,idx*3+dd] - hh
energyp = sess.run([t_energy],
feed_dict = {
self.coord: dcoordp.reshape([-1]),
self.charge: self.dcharge.reshape([-1]),
self.box: self.dbox.reshape([-1]),
self.nloc: [self.natoms],
})
energym = sess.run([t_energy],
feed_dict = {
self.coord: dcoordm.reshape([-1]),
self.charge: self.dcharge.reshape([-1]),
self.box: self.dbox.reshape([-1]),
self.nloc: [self.natoms],
})
c_force = -(energyp[0] - energym[0]) / (2*hh)
for ff in range(self.nframes):
self.assertAlmostEqual(c_force[ff], force[ff,idx*3+dd],
places = places,
msg = "frame %d force component [%d,%d] failed" % (ff, idx, dd))
def setUp(self, data, comp=0, pbc=True):
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))
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 __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 test_model(self):
jfile = 'wfc.json'
with open(jfile) as fp:
jdata = json.load(fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata['model']['descriptor'], 'rcut')
data = DataSystem(systems,
set_pfx,
batch_size,
test_size,
rcut,
run_opt=None)
test_data = data.get_test()
numb_test = 1
descrpt = DescrptLocFrame(jdata['model']['descriptor'])
fitting = WFCFitting(jdata['model']['fitting_net'], descrpt)
model = WFCModel(jdata['model'], descrpt, fitting)
input_data = {
'coord': [test_data['coord']],
'box': [test_data['box']],
'type': [test_data['type']],
'natoms_vec': [test_data['natoms_vec']],
'default_mesh': [test_data['default_mesh']],
'fparam': [test_data['fparam']],
}
model._compute_dstats(input_data)
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
t_energy = tf.placeholder(global_ener_float_precision, [None],
name='t_energy')
t_force = tf.placeholder(global_tf_float_precision, [None],
name='t_force')
t_virial = tf.placeholder(global_tf_float_precision, [None],
name='t_virial')
t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
name='t_atom_ener')
t_coord = tf.placeholder(global_tf_float_precision, [None],
name='i_coord')
t_type = tf.placeholder(tf.int32, [None], name='i_type')
t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
name='i_natoms')
t_box = tf.placeholder(global_tf_float_precision, [None, 9],
name='i_box')
t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
is_training = tf.placeholder(tf.bool)
t_fparam = None
model_pred \
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
t_fparam,
suffix = "wfc",
reuse = False)
wfc = model_pred['wfc']
feed_dict_test = {
t_prop_c: test_data['prop_c'],
t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
t_box: test_data['box'][:numb_test, :],
t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
is_training: False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[p] = sess.run([wfc], feed_dict=feed_dict_test)
p = p.reshape([-1])
refp = [
-9.105016838228578990e-01, 7.196284362034099935e-01,
-9.548516928185298014e-02, 2.764615027095288724e+00,
2.661319598995644520e-01, 7.579512949131941846e-02,
-2.107409067376114997e+00, -1.299080016614967414e-01,
-5.962778584850070285e-01, 2.913899917663253514e-01,
-1.226917174638697094e+00, 1.829523069930876655e+00,
1.015704024959750873e+00, -1.792333611099589386e-01,
5.032898080485321834e-01, 1.808561721292949453e-01,
2.468863482075112081e+00, -2.566442546384765100e-01,
-1.467453783795173994e-01, -1.822963931552128658e+00,
5.843600156865462747e-01, -1.493875280832117403e+00,
1.693322352814763398e-01, -1.877325443995481624e+00
]
places = 6
for ii in range(p.size):
self.assertAlmostEqual(p[ii], refp[ii], places=places)
def test_model(self):
jfile = 'water.json'
with open(jfile) as fp:
jdata = json.load(fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata['model']['descriptor'], 'rcut')
data = DataSystem(systems,
set_pfx,
batch_size,
test_size,
rcut,
run_opt=None)
test_data = data.get_test()
numb_test = 1
descrpt = DescrptLocFrame(jdata['model']['descriptor'])
fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
model = Model(jdata['model'], descrpt, fitting)
# model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']])
input_data = {
'coord': [test_data['coord']],
'box': [test_data['box']],
'type': [test_data['type']],
'natoms_vec': [test_data['natoms_vec']],
'default_mesh': [test_data['default_mesh']]
}
model._compute_input_stat(input_data)
model.fitting.bias_atom_e = data.compute_energy_shift()
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
t_energy = tf.placeholder(global_ener_float_precision, [None],
name='t_energy')
t_force = tf.placeholder(global_tf_float_precision, [None],
name='t_force')
t_virial = tf.placeholder(global_tf_float_precision, [None],
name='t_virial')
t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
name='t_atom_ener')
t_coord = tf.placeholder(global_tf_float_precision, [None],
name='i_coord')
t_type = tf.placeholder(tf.int32, [None], name='i_type')
t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
name='i_natoms')
t_box = tf.placeholder(global_tf_float_precision, [None, 9],
name='i_box')
t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
is_training = tf.placeholder(tf.bool)
t_fparam = None
model_pred \
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
t_fparam,
suffix = "loc_frame",
reuse = False)
energy = model_pred['energy']
force = model_pred['force']
virial = model_pred['virial']
atom_ener = model_pred['atom_ener']
feed_dict_test = {
t_prop_c: test_data['prop_c'],
t_energy: test_data['energy'][:numb_test],
t_force: np.reshape(test_data['force'][:numb_test, :], [-1]),
t_virial: np.reshape(test_data['virial'][:numb_test, :], [-1]),
t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :],
[-1]),
t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
t_box: test_data['box'][:numb_test, :],
t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
is_training: False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test)
e = e.reshape([-1])
f = f.reshape([-1])
v = v.reshape([-1])
refe = [1.165945032784766511e+01]
reff = [
2.356319331246305437e-01, 1.772322096063349284e-01,
1.455439548950788684e-02, 1.968599426000810226e-01,
2.648214484898352983e-01, 7.595232354012236564e-02,
-2.121321856338151401e-01, -2.463886119018566037e-03,
-2.075636300914874069e-02, -9.360310077571798101e-03,
-1.751965198776750943e-01, -2.046405309983102827e-02,
-1.990194093283037535e-01, -1.828347741191920298e-02,
-6.916374506995154325e-02, -1.197997068502068031e-02,
-2.461097746875573200e-01, 1.987744214930105627e-02
]
refv = [
-4.998509978510510265e-01, -1.966169437179327711e-02,
1.136130543869883977e-02, -1.966169437179334650e-02,
-4.575353297894450555e-01, -2.668666556859019493e-03,
1.136130543869887100e-02, -2.668666556859039876e-03,
2.455466940358383508e-03
]
refe = np.reshape(refe, [-1])
reff = np.reshape(reff, [-1])
refv = np.reshape(refv, [-1])
places = 10
for ii in range(e.size):
self.assertAlmostEqual(e[ii], refe[ii], places=places)
for ii in range(f.size):
self.assertAlmostEqual(f[ii], reff[ii], places=places)
for ii in range(v.size):
self.assertAlmostEqual(v[ii], refv[ii], places=places)
def test_model(self):
jfile = 'water_se_a.json'
with open(jfile) as fp:
jdata = json.load(fp)
run_opt = RunOptions(None)
systems = j_must_have(jdata, 'systems')
set_pfx = j_must_have(jdata, 'set_prefix')
batch_size = j_must_have(jdata, 'batch_size')
test_size = j_must_have(jdata, 'numb_test')
batch_size = 1
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata['model']['descriptor'], 'rcut')
data = DataSystem(systems,
set_pfx,
batch_size,
test_size,
rcut,
run_opt=None)
test_data = data.get_test()
numb_test = 1
descrpt = DescrptSeA(jdata['model']['descriptor'])
fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
model = Model(jdata['model'], descrpt, fitting)
# model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']])
input_data = {
'coord': [test_data['coord']],
'box': [test_data['box']],
'type': [test_data['type']],
'natoms_vec': [test_data['natoms_vec']],
'default_mesh': [test_data['default_mesh']]
}
model._compute_dstats(input_data)
model.bias_atom_e = data.compute_energy_shift()
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
t_energy = tf.placeholder(global_ener_float_precision, [None],
name='t_energy')
t_force = tf.placeholder(global_tf_float_precision, [None],
name='t_force')
t_virial = tf.placeholder(global_tf_float_precision, [None],
name='t_virial')
t_atom_ener = tf.placeholder(global_tf_float_precision, [None],
name='t_atom_ener')
t_coord = tf.placeholder(global_tf_float_precision, [None],
name='i_coord')
t_type = tf.placeholder(tf.int32, [None], name='i_type')
t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2],
name='i_natoms')
t_box = tf.placeholder(global_tf_float_precision, [None, 9],
name='i_box')
t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh')
is_training = tf.placeholder(tf.bool)
t_fparam = None
model_pred \
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
t_fparam,
suffix = "se_a",
reuse = False)
energy = model_pred['energy']
force = model_pred['force']
virial = model_pred['virial']
atom_ener = model_pred['atom_ener']
feed_dict_test = {
t_prop_c: test_data['prop_c'],
t_energy: test_data['energy'][:numb_test],
t_force: np.reshape(test_data['force'][:numb_test, :], [-1]),
t_virial: np.reshape(test_data['virial'][:numb_test, :], [-1]),
t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :],
[-1]),
t_coord: np.reshape(test_data['coord'][:numb_test, :], [-1]),
t_box: test_data['box'][:numb_test, :],
t_type: np.reshape(test_data['type'][:numb_test, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
is_training: False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test)
e = e.reshape([-1])
f = f.reshape([-1])
v = v.reshape([-1])
refe = [6.135449167779321300e+01]
reff = [
7.799691562262310585e-02, 9.423098804815030483e-02,
3.790560997388224204e-03, 1.432522403799846578e-01,
1.148392791403983204e-01, -1.321871172563671148e-02,
-7.318966526325138000e-02, 6.516069212737778116e-02,
5.406418483320515412e-04, 5.870713761026503247e-02,
-1.605402669549013672e-01, -5.089516979826595386e-03,
-2.554593467731766654e-01, 3.092063507347833987e-02,
1.510355029451411479e-02, 4.869271842355533952e-02,
-1.446113274345035005e-01, -1.126524434771078789e-03
]
refv = [
-6.076776685178300053e-01, 1.103174323630009418e-01,
1.984250991380156690e-02, 1.103174323630009557e-01,
-3.319759402259439551e-01, -6.007404107650986258e-03,
1.984250991380157036e-02, -6.007404107650981921e-03,
-1.200076017439753642e-03
]
refe = np.reshape(refe, [-1])
reff = np.reshape(reff, [-1])
refv = np.reshape(refv, [-1])
places = 10
for ii in range(e.size):
self.assertAlmostEqual(e[ii], refe[ii], places=places)
for ii in range(f.size):
self.assertAlmostEqual(f[ii], reff[ii], places=places)
for ii in range(v.size):
self.assertAlmostEqual(v[ii], refv[ii], places=places)
