def _init_sess_serial(self):
self.sess = tf.Session(config=default_tf_session_config)
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 virial_dw_test(inter,
testCase,
places=global_default_places,
hh=global_default_dw_hh,
suffix=''):
dcoord, dbox, dtype = inter.data.get_data()
feed_dict_test0 = {
inter.coord: dcoord,
inter.box: dbox,
inter.type: dtype,
inter.tnatoms: inter.natoms
}
w0 = np.ones(inter.ndescrpt)
inter.net_w_i = np.copy(w0)
t_ll, t_dw = inter.comp_v_dw(inter.coord,
inter.box,
inter.type,
inter.tnatoms,
name="v_dw_test_0" + suffix)
inter.sess.run(tf.global_variables_initializer())
ll_0 = inter.sess.run(t_ll, feed_dict=feed_dict_test0)
dw_0 = inter.sess.run(t_dw, feed_dict=feed_dict_test0)
absolut_e = []
relativ_e = []
test_list = range(inter.ndescrpt)
ntest = 3
if inter.sel_a[0] != 0:
test_list = np.concatenate((np.arange(0, ntest),
np.arange(inter.sel_a[0] * 4,
inter.sel_a[0] * 4 + ntest)))
else:
test_list = np.arange(0, ntest)
for ii in test_list:
inter.net_w_i = np.copy(w0)
inter.net_w_i[ii] += hh
t_ll, t_dw = inter.comp_v_dw(inter.coord,
inter.box,
inter.type,
inter.tnatoms,
name="v_dw_test_" + str(ii * 2 + 1) +
suffix)
inter.sess.run(tf.global_variables_initializer())
ll_1 = inter.sess.run(t_ll, feed_dict=feed_dict_test0)
inter.net_w_i[ii] -= 2. * hh
t_ll, t_dw = inter.comp_v_dw(inter.coord,
inter.box,
inter.type,
inter.tnatoms,
name="v_dw_test_" + str(ii * 2 + 2) +
suffix)
inter.sess.run(tf.global_variables_initializer())
ll_2 = inter.sess.run(t_ll, feed_dict=feed_dict_test0)
num_v = (ll_1 - ll_2) / (2. * hh)
ana_v = dw_0[ii]
testCase.assertAlmostEqual(num_v, ana_v, places=places)
def test_prod_virial(self):
tvirial, tatom_virial \
= op_module.prod_virial_se_a(
self.tnet_deriv,
self.tem_deriv,
self.trij,
self.tnlist,
self.tnatoms,
n_a_sel=self.nnei,
n_r_sel=0)
self.sess.run(tf.global_variables_initializer())
dvirial, datom_virial = self.sess.run(
[tvirial, tatom_virial],
feed_dict={
self.tnet_deriv: self.dnet_deriv,
self.tem_deriv: self.dem_deriv,
self.trij: self.drij,
self.tnlist: self.dnlist,
self.tnatoms: self.dnatoms
})
self.assertEqual(dvirial.shape, (self.nframes, 9))
self.assertEqual(datom_virial.shape, (self.nframes, self.nall * 9))
for ff in range(self.nframes):
np.testing.assert_almost_equal(dvirial[ff], self.expected_virial,
5)
np.testing.assert_almost_equal(datom_virial[ff],
self.expected_atom_virial, 5)
def test_prod_virial_grad(self):
tgrad_net \
= op_grads_module.prod_virial_se_a_grad(
self.tgrad,
self.tnet_deriv,
self.tem_deriv,
self.trij,
self.tnlist,
self.tnatoms,
n_a_sel=self.nnei,
n_r_sel=0)
self.sess.run(tf.global_variables_initializer())
dgrad_net = self.sess.run(tgrad_net,
feed_dict={
self.tgrad: self.dgrad,
self.tnet_deriv: self.dnet_deriv,
self.tem_deriv: self.dem_deriv,
self.trij: self.drij,
self.tnlist: self.dnlist,
self.tnatoms: self.dnatoms
})
self.assertEqual(dgrad_net.shape,
(self.nframes, self.nloc * self.ndescrpt))
for ff in range(self.nframes):
np.testing.assert_almost_equal(dgrad_net[ff],
self.expected_grad_net, 5)
def test_type_embed_net(self):
ten = TypeEmbedNet([2, 4, 8], seed=1, uniform_seed=True)
type_embedding = ten.build(2)
sess = self.test_session().__enter__()
sess.run(tf.global_variables_initializer())
type_embedding = sess.run(type_embedding)
expected_out = [
1.429967002262267917e+00, -9.138175897677495163e-01,
-3.799606588218059633e-01, -2.143157692726757046e-01,
2.341138114260268743e+00, -1.568346043255314015e+00,
8.917082000854256174e-01, -1.500356675378008209e+00,
8.955885646123034061e-01, -5.835326470989941061e-01,
-1.465708662924672057e+00, -4.052047884085572260e-01,
1.367825594590430072e+00, -2.736204307656463497e-01,
-4.044263041521370394e-01, -9.438057524881729998e-01
]
expected_out = np.reshape(expected_out, [2, 8])
# 2 types
self.assertEqual(type_embedding.shape[0], 2)
# size of embedded vec 8
self.assertEqual(type_embedding.shape[1], 8)
# check value
np.testing.assert_almost_equal(type_embedding, expected_out, 10)
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 test_prod_force_parallel(self):
forces = []
for ii in range(4):
tforce \
= op_module.parallel_prod_force_se_a(
self.tnet_deriv,
self.tem_deriv,
self.tnlist,
self.tnatoms,
n_a_sel=self.nnei,
n_r_sel=0,
parallel=True,
start_frac = ii/4,
end_frac = (ii+1)/4,
)
forces.append(tforce)
tforce = tf.add_n(forces)
self.sess.run(tf.global_variables_initializer())
dforce = self.sess.run(tforce,
feed_dict={
self.tnet_deriv: self.dnet_deriv,
self.tem_deriv: self.dem_deriv,
self.tnlist: self.dnlist,
self.tnatoms: self.dnatoms
})
self.assertEqual(dforce.shape, (self.nframes, self.nall * 3))
for ff in range(self.nframes):
np.testing.assert_almost_equal(dforce[ff], self.expected_force, 5)
def test_pbc_small_box(self):
data0 = Data()
data1 = Data(box_scale=2)
inter0 = Inter()
inter1 = Inter()
inter0.setUp(data0, pbc=True)
inter1.setUp(data1, pbc=False)
inter0.net_w_i = np.copy(np.ones(inter0.ndescrpt))
inter1.net_w_i = np.copy(np.ones(inter1.ndescrpt))
t_energy0, t_force0, t_virial0 \
= inter0.comp_ef (inter0.coord, inter0.box, inter0.type, inter0.tnatoms, name = "test_lf_pbc_sbox_true")
t_energy1, t_force1, t_virial1 \
= inter1.comp_ef (inter1.coord, inter1.box, inter1.type, inter1.tnatoms, name = "test_lf_pbc_sbox_false")
inter0.sess.run(tf.global_variables_initializer())
inter1.sess.run(tf.global_variables_initializer())
dcoord, dbox, dtype = data0.get_data()
[e0, f0, v0] = inter0.sess.run(
[t_energy0, t_force0, t_virial0],
feed_dict={
inter0.coord: dcoord,
inter0.box: dbox,
inter0.type: dtype,
inter0.tnatoms: inter0.natoms
})
dcoord, dbox, dtype = data1.get_data()
[e1, f1, v1] = inter1.sess.run(
[t_energy1, t_force1, t_virial1],
feed_dict={
inter1.coord: dcoord,
inter1.box: dbox,
inter1.type: dtype,
inter1.tnatoms: inter1.natoms
})
self.assertAlmostEqual(e0[0], e1[0])
for ii in range(f0[0].size):
# print(ii)
self.assertAlmostEqual(f0[0][ii], f1[0][ii])
for ii in range(v0[0].size):
# print(ii)
self.assertAlmostEqual(v0[0][ii], v1[0][ii])
def force_test(inter,
testCase,
places=global_default_places,
hh=global_default_fv_hh,
suffix=''):
# set weights
w0 = np.ones(inter.ndescrpt)
inter.net_w_i = np.copy(w0)
# make network
t_energy, t_force, t_virial \
= inter.comp_ef (inter.coord, inter.box, inter.type, inter.tnatoms, name = "test_f" + suffix)
inter.sess.run(tf.global_variables_initializer())
# get data
dcoord, dbox, dtype = inter.data.get_data()
# cmp e0, f0
[energy, force] = inter.sess.run(
[t_energy, t_force],
feed_dict={
inter.coord: dcoord,
inter.box: dbox,
inter.type: dtype,
inter.tnatoms: inter.natoms
})
# dim force
sel_idx = np.arange(inter.natoms[0])
for idx in sel_idx:
for dd in range(3):
dcoordp = np.copy(dcoord)
dcoordm = np.copy(dcoord)
dcoordp[0, idx * 3 + dd] = dcoord[0, idx * 3 + dd] + hh
dcoordm[0, idx * 3 + dd] = dcoord[0, idx * 3 + dd] - hh
[enerp] = inter.sess.run(
[t_energy],
feed_dict={
inter.coord: dcoordp,
inter.box: dbox,
inter.type: dtype,
inter.tnatoms: inter.natoms
})
[enerm] = inter.sess.run(
[t_energy],
feed_dict={
inter.coord: dcoordm,
inter.box: dbox,
inter.type: dtype,
inter.tnatoms: inter.natoms
})
c_force = -(enerp[0] - enerm[0]) / (2 * hh)
testCase.assertAlmostEqual(c_force,
force[0, idx * 3 + dd],
places=places,
msg="force component [%d,%d] failed" %
(idx, dd))
def test_pbc(self):
data = Data()
inter0 = Inter()
inter1 = Inter()
inter0.setUp(data, pbc=True, sess=self.test_session().__enter__())
inter1.setUp(data, pbc=False, sess=self.test_session().__enter__())
inter0.net_w_i = np.copy(np.ones(inter0.ndescrpt))
inter1.net_w_i = np.copy(np.ones(inter1.ndescrpt))
t_energy0, t_force0, t_virial0 \
= inter0.comp_ef (inter0.coord, inter0.box, inter0.type, inter0.tnatoms, name = "test_ser_pbc_true")
t_energy1, t_force1, t_virial1 \
= inter1.comp_ef (inter1.coord, inter1.box, inter1.type, inter1.tnatoms, name = "test_ser_pbc_false")
inter0.sess.run(tf.global_variables_initializer())
inter1.sess.run(tf.global_variables_initializer())
dcoord, dbox, dtype = data.get_data()
[e0, f0, v0] = inter0.sess.run(
[t_energy0, t_force0, t_virial0],
feed_dict={
inter0.coord: dcoord,
inter0.box: dbox,
inter0.type: dtype,
inter0.tnatoms: inter0.natoms
})
[e1, f1, v1] = inter1.sess.run(
[t_energy1, t_force1, t_virial1],
feed_dict={
inter1.coord: dcoord,
inter1.box: dbox,
inter1.type: dtype,
inter1.tnatoms: inter1.natoms
})
self.assertAlmostEqual(e0[0], e1[0])
np.testing.assert_almost_equal(f0[0], f1[0])
np.testing.assert_almost_equal(v0[0], v1[0])
def test_enlarger_net(self):
network_size = [3, 4]
out = embedding_net(self.inputs,
network_size,
tf.float64,
name_suffix='enlarger_net',
seed=1,
uniform_seed=True)
self.sess.run(tf.global_variables_initializer())
myout = self.sess.run(out)
refout = [[-0.1482171, -0.14177827, -0.76181204, 0.21266767],
[-0.27800543, -0.08974353, -0.78784335, 0.3485518],
[-0.36744368, -0.06285603, -0.80749876, 0.4347974]]
np.testing.assert_almost_equal(refout, myout, self.places)
def test_enlarger_net_1(self):
network_size = [4, 4]
out = embedding_net(self.inputs,
network_size,
tf.float64,
name_suffix='enlarger_net_1',
seed=1,
uniform_seed=True)
self.sess.run(tf.global_variables_initializer())
myout = self.sess.run(out)
refout = [[0.10842905, -0.61623145, -1.46738788, -0.01921788],
[0.09376136, -0.75526936, -1.64995884, 0.01076112],
[0.1033177, -0.8911794, -1.75530172, 0.00653156]]
np.testing.assert_almost_equal(refout, myout, self.places)
def test_enlarger_net_2(self):
network_size = [2, 4]
out = embedding_net(self.inputs,
network_size,
tf.float64,
name_suffix='enlarger_net_2',
seed=1,
uniform_seed=True)
self.sess.run(tf.global_variables_initializer())
myout = self.sess.run(out)
refout = [[0.24023149, -0.66311811, -0.50951819, -0.36873654],
[2.00858313, -0.05971232, 0.52272395, -0.12604478],
[3.39365063, 0.63492697, 1.5780069, 0.46445682]]
np.testing.assert_almost_equal(refout, myout, self.places)
def test_enlarger_net_1_idt(self):
network_size = [4, 4]
out = embedding_net(self.inputs,
network_size,
tf.float64,
name_suffix='enlarger_net_1_idt',
resnet_dt=True,
seed=1,
uniform_seed=True)
self.sess.run(tf.global_variables_initializer())
myout = self.sess.run(out)
refout = [[0.10839754, -0.6161336, -1.46673253, -0.01927138],
[0.09370214, -0.75516888, -1.64927868, 0.01067603],
[0.10323835, -0.89107102, -1.75460243, 0.00642493]]
np.testing.assert_almost_equal(refout, myout, self.places)
def test_enlarger_net_2(self):
network_size = [2, 4]
out = embedding_net(self.inputs,
network_size,
tf.float64,
name_suffix='enlarger_net_2_idt',
resnet_dt=True,
seed=1,
uniform_seed=True)
self.sess.run(tf.global_variables_initializer())
myout = self.sess.run(out)
refout = [[0.2403889, -0.66290763, -0.50883586, -0.36869913],
[2.00891479, -0.05936574, 0.52351633, -0.12579749],
[3.3940202, 0.63538459, 1.57887697, 0.46486689]]
np.testing.assert_almost_equal(refout, myout, self.places)
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):
run_opt = RunOptions(restart=None,
init_model=None,
log_path=None,
log_level=30,
mpi_log="master")
jdata = j_loader(INPUT)
# init model
model = DPTrainer(jdata, run_opt=run_opt)
rcut = model.model.get_rcut()
# init data system
systems = j_must_have(jdata['training'], 'systems')
#systems[0] = tests_path / systems[0]
systems = [tests_path / ii for ii in 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 self.test_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,model_attr/output_dim,model_attr/model_version"
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess, input_graph_def, nodes.split(","))
output_graph = str(tests_path /
os.path.join(modifier_datapath, 'dipole.pb'))
with tf.gfile.GFile(output_graph, "wb") as f:
f.write(output_graph_def.SerializeToString())
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 _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 test_prod_force(self):
tforce \
= op_module.prod_force_se_a(
self.tnet_deriv,
self.tem_deriv,
self.tnlist,
self.tnatoms,
n_a_sel=self.nnei,
n_r_sel=0)
self.sess.run(tf.global_variables_initializer())
dforce = self.sess.run(tforce,
feed_dict={
self.tnet_deriv: self.dnet_deriv,
self.tem_deriv: self.dem_deriv,
self.tnlist: self.dnlist,
self.tnatoms: self.dnatoms
})
self.assertEqual(dforce.shape, (self.nframes, self.nall * 3))
for ff in range(self.nframes):
np.testing.assert_almost_equal(dforce[ff], self.expected_force, 5)
def virial_test(inter,
testCase,
places=global_default_places,
hh=global_default_fv_hh,
suffix=''):
# set weights
w0 = np.ones(inter.ndescrpt)
inter.net_w_i = np.copy(w0)
# make network
t_energy, t_force, t_virial \
= inter.comp_ef (inter.coord, inter.box, inter.type, inter.tnatoms, name = "test_v" + suffix)
inter.sess.run(tf.global_variables_initializer())
# get data
dcoord, dbox, dtype = inter.data.get_test_box_data(hh)
# cmp e, f, v
[energy, force, virial] \
= inter.sess.run ([t_energy, t_force, t_virial],
feed_dict = {
inter.coord: dcoord,
inter.box: dbox,
inter.type: dtype,
inter.tnatoms: inter.natoms}
)
ana_vir = virial[0].reshape([3, 3])
num_vir = np.zeros([3, 3])
for ii in range(3):
for jj in range(3):
ep = energy[1 + (ii * 3 + jj) * 2 + 0]
em = energy[1 + (ii * 3 + jj) * 2 + 1]
num_vir[ii][jj] = -(ep - em) / (2. * hh)
num_vir = np.transpose(num_vir, [1, 0])
box3 = dbox[0].reshape([3, 3])
num_vir = np.matmul(num_vir, box3)
for ii in range(3):
for jj in range(3):
testCase.assertAlmostEqual(ana_vir[ii][jj],
num_vir[ii][jj],
places=places,
msg='virial component %d %d ' %
(ii, jj))
def virial_test(inter,
testCase,
places=global_default_places,
hh=global_default_fv_hh,
suffix=''):
# set weights
w0 = np.ones(inter.ndescrpt)
inter.net_w_i = np.copy(w0)
# make network
t_energy, t_force, t_virial \
= inter.comp_ef (inter.coord, inter.box, inter.type, inter.tnatoms, name = "test_v" + suffix)
inter.sess.run(tf.global_variables_initializer())
# get data
dcoord, dbox, dtype = inter.data.get_test_box_data(hh)
# cmp e, f, v
[energy, force, virial] \
= inter.sess.run ([t_energy, t_force, t_virial],
feed_dict = {
inter.coord: dcoord,
inter.box: dbox,
inter.type: dtype,
inter.tnatoms: inter.natoms}
)
# check
ana_vir3 = (virial[0][0] + virial[0][4] + virial[0][8]) / 3. / comp_vol(
dbox[0])
num_vir3 = -(energy[1] - energy[2]) / (comp_vol(dbox[1]) -
comp_vol(dbox[2]))
testCase.assertAlmostEqual(ana_vir3, num_vir3, places=places)
vir_idx = [0, 4, 8]
for dd in range(3):
ana_v = (virial[0][vir_idx[dd]] / comp_vol(dbox[0]))
idx = 2 * (dd + 1) + 1
num_v = (-(energy[idx] - energy[idx + 1]) /
(comp_vol(dbox[idx]) - comp_vol(dbox[idx + 1])))
testCase.assertAlmostEqual(ana_v, num_v, places=places)
def test_model(self):
jfile = 'water_se_r.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 = DescrptSeR(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.descrpt.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_r",
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.152085988309423925e+01]
reff = [-1.714443151616400110e-04,-1.315836609370952051e-04,-5.584120460897444674e-06,-7.197863450669731334e-05,-1.384609799994930676e-04,8.856091902774708468e-06,1.120578238869146797e-04,-7.428703645877488470e-05,9.370560731488587317e-07,-1.048347129617610465e-04,1.977876923815685781e-04,7.522050342771599598e-06,2.361772659657814205e-04,-5.774651813388292487e-05,-1.233143271630744828e-05,2.257277740226381951e-08,2.042905031476775584e-04,6.003548585097267914e-07]
refv = [1.035180911513190792e-03,-1.118982949050497126e-04,-2.383287813436022850e-05,-1.118982949050497126e-04,4.362023915782403281e-04,8.119543218224559240e-06,-2.383287813436022850e-05,8.119543218224559240e-06,1.201142938802945237e-06]
refe = np.reshape(refe, [-1])
reff = np.reshape(reff, [-1])
refv = np.reshape(refv, [-1])
places = 6
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 = '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_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)
def test_descriptor_two_sides(self):
jfile = 'water_se_a_type.json'
jdata = j_loader(jfile)
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 = 2
test_size = 1
stop_batch = j_must_have(jdata, 'stop_batch')
rcut = j_must_have(jdata['model']['descriptor'], 'rcut')
sel = j_must_have(jdata['model']['descriptor'], 'sel')
ntypes = len(sel)
data = DataSystem(systems,
set_pfx,
batch_size,
test_size,
rcut,
run_opt=None)
test_data = data.get_test()
numb_test = 1
# set parameters
jdata['model']['descriptor']['neuron'] = [5, 5, 5]
jdata['model']['descriptor']['axis_neuron'] = 2
typeebd_param = {
'neuron': [5, 5, 5],
'resnet_dt': False,
'seed': 1,
}
# init models
typeebd = TypeEmbedNet(neuron=typeebd_param['neuron'],
resnet_dt=typeebd_param['resnet_dt'],
seed=typeebd_param['seed'],
uniform_seed=True)
jdata['model']['descriptor'].pop('type', None)
descrpt = DescrptSeA(**jdata['model']['descriptor'], uniform_seed=True)
# 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']]
}
descrpt.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, [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
type_embedding = typeebd.build(ntypes,
suffix="_se_a_type_des_ebd_2sdies")
dout \
= descrpt.build(
t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
{'type_embedding' : type_embedding},
reuse = False,
suffix = "_se_a_type_des_2sides"
)
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 = self.test_session().__enter__()
sess.run(tf.global_variables_initializer())
[model_dout] = sess.run([dout], feed_dict=feed_dict_test)
model_dout = model_dout.reshape([-1])
ref_dout = [
0.0005722682145569174, -0.00020202686217742682,
-0.00020202686217742682, 7.13250554992363e-05,
-0.0014770058171250015, 0.000521468690207748,
-0.001143865186937176, 0.0004038453384193948,
0.0005617335409639567, -0.00019831394075147532,
0.00048086740718842236, -0.0001693584775806112,
-0.0001693584775806112, 5.966987137476082e-05,
-0.0012342029581315136, 0.00043492340851472783,
-0.0009566016612537016, 0.00033706767041080107,
0.00047065988464132244, -0.0001657950398095401,
0.0003647849239740657, -0.00013744939018250384,
-0.00013744939018250384, 5.1825826955234744e-05,
-0.00096004206555711, 0.00036185565262332876,
-0.0007267433909643961, 0.0002738914365542745,
0.00038019365906978136, -0.00014322754331896057,
0.0004675256930823109, -0.00017634410399626168,
-0.00017634410399626168, 6.652672908755666e-05,
-0.0012328062885292486, 0.00046500213384094614,
-0.0009328887521346069, 0.0003518668613172834,
0.0004877847509912577, -0.00018396318824508986,
0.0005154794374703516, -0.00019422534512034776,
-0.00019422534512034776, 7.318151797939947e-05,
-0.0013576642997136488, 0.0005115548790018505,
-0.0010275333676074971, 0.00038716440070070385,
0.0005376426714609369, -0.00020257810468163985,
0.0004482204892297628, -0.00016887749501640607,
-0.00016887749501640607, 6.364643102775375e-05,
-0.001181345877677835, 0.0004452029242063362,
-0.0008941636427724908, 0.0003369586197174627,
0.0004677878512312651, -0.00017625260641095753
]
places = 10
np.testing.assert_almost_equal(model_dout, ref_dout, places)
def test_model(self):
jfile = 'polar_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 = PolarFittingSeA(jdata['model']['fitting_net'], descrpt)
model = PolarModel(jdata['model'], descrpt, fitting)
model._compute_dstats([test_data['coord']], [test_data['box']],
[test_data['type']], [test_data['natoms_vec']],
[test_data['default_mesh']])
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 = "polar_se_a",
reuse = False)
polar = model_pred['polar']
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([polar], feed_dict=feed_dict_test)
p = p.reshape([-1])
refp = [
3.39695248e+01, 2.16564043e+01, 8.18501479e-01, 2.16564043e+01,
1.38211789e+01, 5.22775159e-01, 8.18501479e-01, 5.22775159e-01,
1.97847218e-02, 8.08467431e-01, 3.42081126e+00, -2.01072261e-01,
3.42081126e+00, 1.54924596e+01, -9.06153697e-01, -2.01072261e-01,
-9.06153697e-01, 5.30193262e-02
]
places = 6
for ii in range(p.size):
self.assertAlmostEqual(p[ii], refp[ii], places=places)
def test_rot_field_corot(self, suffix=''):
suffix = '_field_corot'
t_p_e, t_p_f, t_p_f, t_p_ae, t_p_av \
= self.build_efv (self.coord, self.box, self.type, self.tnatoms, name = "test_rot" + suffix, op = op_module.descrpt_se_a_ef_para)
t_v_e, t_v_f, t_v_f, t_v_ae, t_v_av \
= self.build_efv (self.coord, self.box, self.type, self.tnatoms, name = "test_rot" + suffix, op = op_module.descrpt_se_a_ef_vert, reuse = True)
t_e, t_f, t_f, t_ae, t_av \
= self.build_efv (self.coord, self.box, self.type, self.tnatoms, name = "test_rot" + suffix, op = op_module.prod_env_mat_a, reuse = True)
self.sess.run(tf.global_variables_initializer())
np.random.seed(0)
# make test data
nframes = 2
dcoord, dbox, dtype, defield = self.make_test_data(nframes)
[p_ae0] = self.sess.run(
[t_p_ae],
feed_dict={
self.coord: dcoord,
self.box: dbox,
self.type: dtype,
self.efield: defield,
self.tnatoms: self.natoms
})
[v_ae0] = self.sess.run(
[t_v_ae],
feed_dict={
self.coord: dcoord,
self.box: dbox,
self.type: dtype,
self.efield: defield,
self.tnatoms: self.natoms
})
[ae0] = self.sess.run(
[t_ae],
feed_dict={
self.coord: dcoord,
self.box: dbox,
self.type: dtype,
self.efield: defield,
self.tnatoms: self.natoms
})
# print(f0)
theta = 45. / 180. * np.pi
rr0 = self.rotate_mat(defield[0][0:3], theta)
rr1 = self.rotate_mat(defield[1][0:3], theta)
dcoord_ = np.copy(dcoord).reshape([nframes, -1, 3])
dcoord0 = np.matmul(dcoord_[0], rr0)
dcoord1 = np.matmul(dcoord_[1], rr1)
new_dcoord = np.concatenate([dcoord0, dcoord1],
axis=0).reshape([nframes, -1])
defield_ = np.copy(defield).reshape([nframes, -1, 3])
defield0 = np.matmul(defield_[0], rr0)
defield1 = np.matmul(defield_[1], rr1)
new_defield = np.concatenate([defield0, defield1],
axis=0).reshape([nframes, -1])
[p_ae1] = self.sess.run(
[t_p_ae],
feed_dict={
self.coord: new_dcoord,
self.box: dbox,
self.type: dtype,
self.efield: new_defield,
self.tnatoms: self.natoms
})
[v_ae1] = self.sess.run(
[t_v_ae],
feed_dict={
self.coord: new_dcoord,
self.box: dbox,
self.type: dtype,
self.efield: new_defield,
self.tnatoms: self.natoms
})
[ae1] = self.sess.run(
[t_ae],
feed_dict={
self.coord: new_dcoord,
self.box: dbox,
self.type: dtype,
self.efield: new_defield,
self.tnatoms: self.natoms
})
np.testing.assert_almost_equal(p_ae0, p_ae1)
np.testing.assert_almost_equal(v_ae0, v_ae1)
np.testing.assert_almost_equal(ae0, ae1)
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_rot_axis(self, suffix=''):
suffix = '_axis'
t_p_e, t_p_f, t_p_f, t_p_ae, t_p_av \
= self.build_efv (self.coord, self.box, self.type, self.tnatoms, name = "test_rot" + suffix, op = op_module.descrpt_se_a_ef_para)
t_v_e, t_v_f, t_v_f, t_v_ae, t_v_av \
= self.build_efv (self.coord, self.box, self.type, self.tnatoms, name = "test_rot" + suffix, op = op_module.descrpt_se_a_ef_vert, reuse = True)
t_e, t_f, t_f, t_ae, t_av \
= self.build_efv (self.coord, self.box, self.type, self.tnatoms, name = "test_rot" + suffix, op = op_module.prod_env_mat_a, reuse = True)
self.sess.run(tf.global_variables_initializer())
np.random.seed(0)
# make test data
nframes = 2
dcoord, dbox, dtype, defield = self.make_test_data(nframes)
[p_ae0] = self.sess.run(
[t_p_ae],
feed_dict={
self.coord: dcoord,
self.box: dbox,
self.type: dtype,
self.efield: defield,
self.tnatoms: self.natoms
})
[v_ae0] = self.sess.run(
[t_v_ae],
feed_dict={
self.coord: dcoord,
self.box: dbox,
self.type: dtype,
self.efield: defield,
self.tnatoms: self.natoms
})
[ae0] = self.sess.run(
[t_ae],
feed_dict={
self.coord: dcoord,
self.box: dbox,
self.type: dtype,
self.efield: defield,
self.tnatoms: self.natoms
})
# print(p_ae0)
# print(v_ae0)
# print(ae0)
for kk in range(0, self.natoms[0]):
# print(f0)
theta = 45. / 180. * np.pi
rr0 = self.rotate_mat(defield[0][kk * 3:kk * 3 + 3], theta)
# rr0 = self.rotate_mat([0, 0, 1], theta)
rr1 = self.rotate_mat(defield[1][kk * 3:kk * 3 + 3], theta)
# print(rr0, np.matmul(rr0, rr0.T), np.matmul(rr0.T, rr0))
# print(rr1)
dcoord_ = np.copy(dcoord).reshape([nframes, -1, 3])
dcoord0 = np.matmul(dcoord_[0], rr0)
dcoord1 = np.matmul(dcoord_[1], rr1)
new_dcoord = np.concatenate([dcoord0, dcoord1],
axis=0).reshape([nframes, -1])
defield_ = np.copy(defield).reshape([nframes, -1, 3])
defield0 = np.matmul(defield_[0], rr0)
defield1 = np.matmul(defield_[1], rr1)
new_defield = np.concatenate([defield0, defield1],
axis=0).reshape([nframes, -1])
[p_ae1] = self.sess.run(
[t_p_ae],
feed_dict={
self.coord: new_dcoord,
self.box: dbox,
self.type: dtype,
self.efield: defield,
self.tnatoms: self.natoms
})
[v_ae1] = self.sess.run(
[t_v_ae],
feed_dict={
self.coord: new_dcoord,
self.box: dbox,
self.type: dtype,
self.efield: defield,
self.tnatoms: self.natoms
})
[ae1] = self.sess.run(
[t_ae],
feed_dict={
self.coord: new_dcoord,
self.box: dbox,
self.type: dtype,
self.efield: defield,
self.tnatoms: self.natoms
})
for ii in range(0, self.natoms[0]):
for jj in range(0, self.natoms[0]):
diff = dcoord[0][3 * jj:3 * jj +
3] - dcoord[0][3 * ii:3 * ii + 3]
dot = np.dot(dcoord[0][3 * jj:3 * jj + 3],
dcoord[0][3 * ii:3 * ii + 3])
diff1 = new_dcoord[0][3 * jj:3 * jj +
3] - new_dcoord[0][3 * ii:3 * ii + 3]
dot1 = np.dot(new_dcoord[0][3 * jj:3 * jj + 3],
new_dcoord[0][3 * ii:3 * ii + 3])
assert (
np.abs(np.linalg.norm(diff) - np.linalg.norm(diff1)) <
1e-15)
assert (np.abs(dot - dot1) < 1e-15)
for ii in range(1, self.natoms[0]):
if ii == kk:
self.assertAlmostEqual(p_ae0[ii], p_ae1[ii])
self.assertAlmostEqual(v_ae0[ii], v_ae1[ii])
else:
self.assertNotAlmostEqual(p_ae0[ii], p_ae1[ii])
self.assertNotAlmostEqual(v_ae0[ii], v_ae1[ii])
