def test_modify(self):
dcm = DipoleChargeModifier(
os.path.join(modifier_datapath, "dipole.pb"), [-1, -3],
[1, 1, 1, 1, 1], 1, 0.25)
ve0, vf0, vv0 = dcm.eval(self.coords0, self.box0, self.atom_types0)
ve1, vf1, vv1 = dcm.eval(self.coords1, self.box1, self.atom_types1)
vf01 = vf0[:, self.idx_map, :]
for ii in range(self.nframes):
self.assertAlmostEqual(ve0[ii],
ve1[ii],
msg='energy %d should match' % ii)
for ii in range(self.nframes):
for jj in range(9):
self.assertAlmostEqual(vv0[ii][jj],
vv1[ii][jj],
msg='virial [%d,%d] should match' %
(ii, jj))
for ii in range(self.nframes):
for jj in range(self.natoms):
for dd in range(3):
self.assertAlmostEqual(
vf01[ii][jj][dd],
vf1[ii][jj][dd],
msg="force [%d,%d,%d] dose not match" % (ii, jj, dd))
def _do_work(jdata, run_opt):
# init the model
model = NNPTrainer(jdata, run_opt=run_opt)
rcut = model.model.get_rcut()
type_map = model.model.get_type_map()
# init params and run options
assert ('training' in jdata)
systems = j_must_have(jdata['training'], 'systems')
if type(systems) == str:
systems = expand_sys_str(systems)
set_pfx = j_must_have(jdata['training'], 'set_prefix')
seed = None
if 'seed' in jdata['training'].keys(): seed = jdata['training']['seed']
if seed is not None:
seed = seed % (2**32)
np.random.seed(seed)
batch_size = j_must_have(jdata['training'], 'batch_size')
test_size = j_must_have(jdata['training'], 'numb_test')
stop_batch = j_must_have(jdata['training'], 'stop_batch')
sys_probs = jdata['training'].get('sys_probs')
auto_prob_style = jdata['training'].get('auto_prob_style', 'prob_sys_size')
if len(type_map) == 0:
# empty type_map
ipt_type_map = None
else:
ipt_type_map = type_map
# data modifier
modifier = None
modi_data = jdata['model'].get("modifier", None)
if modi_data is not None:
if modi_data['type'] == 'dipole_charge':
modifier = DipoleChargeModifier(modi_data['model_name'],
modi_data['model_charge_map'],
modi_data['sys_charge_map'],
modi_data['ewald_h'],
modi_data['ewald_beta'])
else:
raise RuntimeError('unknown modifier type ' +
str(modi_data['type']))
# init data
data = DeepmdDataSystem(systems,
batch_size,
test_size,
rcut,
set_prefix=set_pfx,
type_map=ipt_type_map,
modifier=modifier)
data.print_summary(run_opt,
sys_probs=sys_probs,
auto_prob_style=auto_prob_style)
data.add_dict(data_requirement)
# build the model with stats from the first system
model.build(data, stop_batch)
# train the model with the provided systems in a cyclic way
start_time = time.time()
model.train(data)
end_time = time.time()
run_opt.message("finished training\nwall time: %.3f s" %
(end_time - start_time))
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 _test_fv (self):
dcm = DipoleChargeModifier(os.path.join(modifier_datapath, "dipole.pb"),
[-8],
[6, 1],
1,
0.25)
data = Data()
coord, box, atype = data.get_data()
atype = atype[0]
ve, vf, vv = dcm.eval(coord, box, atype)
hh = global_default_fv_hh
hh=1e-4
places = global_default_places
places=1
nframes = coord.shape[0]
ndof = coord.shape[1]
natoms = ndof // 3
vf = np.reshape(vf, [nframes, -1])
for ii in range(ndof):
coordp = np.copy(coord)
coordm = np.copy(coord)
coordp[:,ii] += hh
coordm[:,ii] -= hh
ep, _, __ = dcm.eval(coordp, box, atype, eval_fv = False)
em, _, __ = dcm.eval(coordm, box, atype, eval_fv = False)
num_f = -(ep - em) / (2.*hh)
for ff in range(nframes):
self.assertAlmostEqual(vf[ff,ii], num_f[ff],
places = places,
msg = 'frame %d dof %d does not match' % (ff, ii))
box3 = np.reshape(box, [nframes, 3,3])
rbox3 = np.linalg.inv(box3)
coord3 = np.reshape(coord, [nframes, natoms, 3])
rcoord3 = np.matmul(coord3, rbox3)
num_deriv = np.zeros([nframes,3,3])
for ii in range(3):
for jj in range(3):
box3p = np.copy(box3)
box3m = np.copy(box3)
box3p[:,ii,jj] = box3[:,ii,jj] + hh
box3m[:,ii,jj] = box3[:,ii,jj] - hh
boxp = np.reshape(box3p, [-1,9])
boxm = np.reshape(box3m, [-1,9])
coord3p = np.matmul(rcoord3, box3p)
coord3m = np.matmul(rcoord3, box3m)
coordp = np.reshape(coord3p, [nframes,-1])
coordm = np.reshape(coord3m, [nframes,-1])
ep, _, __ = dcm.eval(coordp, boxp, atype, eval_fv = False)
em, _, __ = dcm.eval(coordm, boxm, atype, eval_fv = False)
num_deriv[:,ii,jj] = -(ep - em) / (2.*hh)
# box3t = np.transpose(box3, [0,2,1])
# t_esti = np.matmul(num_deriv, box3t)
num_deriv = np.transpose(num_deriv, [0,2,1])
t_esti = np.matmul(num_deriv, box3)
print(t_esti, '\n', vv.reshape([-1, 3, 3]))
for ff in range(nframes):
for ii in range(3):
for jj in range(3):
self.assertAlmostEqual(t_esti[ff][ii][jj], vv[ff,ii*3+jj],
places = places,
msg = "frame %d virial component [%d,%d] failed" % (ff, ii, jj))