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 __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 _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))
class DeepPot (DeepEval) :
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 get_ntypes(self) :
return self.ntypes
def get_rcut(self) :
return self.rcut
def get_dim_fparam(self) :
return self.dfparam
def get_dim_aparam(self) :
return self.daparam
def get_type_map(self):
return self.tmap
def eval(self,
coords,
cells,
atom_types,
fparam = None,
aparam = None,
atomic = False) :
if atomic :
if self.modifier_type is not None:
raise RuntimeError('modifier does not support atomic modification')
return self.eval_inner(coords, cells, atom_types, fparam = fparam, aparam = aparam, atomic = atomic)
else :
e, f, v = self.eval_inner(coords, cells, atom_types, fparam = fparam, aparam = aparam, atomic = atomic)
if self.modifier_type is not None:
me, mf, mv = self.dm.eval(coords, cells, atom_types)
e += me.reshape(e.shape)
f += mf.reshape(f.shape)
v += mv.reshape(v.shape)
return e, f, v
def eval_inner(self,
coords,
cells,
atom_types,
fparam = None,
aparam = None,
atomic = False) :
# standarize the shape of inputs
atom_types = np.array(atom_types, dtype = int).reshape([-1])
natoms = atom_types.size
coords = np.reshape(np.array(coords), [-1, natoms * 3])
nframes = coords.shape[0]
if cells is None:
pbc = False
# make cells to work around the requirement of pbc
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
if self.has_fparam :
assert(fparam is not None)
fparam = np.array(fparam)
if self.has_aparam :
assert(aparam is not None)
aparam = np.array(aparam)
# reshape the inputs
if self.has_fparam :
fdim = self.get_dim_fparam()
if fparam.size == nframes * fdim :
fparam = np.reshape(fparam, [nframes, fdim])
elif fparam.size == fdim :
fparam = np.tile(fparam.reshape([-1]), [nframes, 1])
else :
raise RuntimeError('got wrong size of frame param, should be either %d x %d or %d' % (nframes, fdim, fdim))
if self.has_aparam :
fdim = self.get_dim_aparam()
if aparam.size == nframes * natoms * fdim:
aparam = np.reshape(aparam, [nframes, natoms * fdim])
elif aparam.size == natoms * fdim :
aparam = np.tile(aparam.reshape([-1]), [nframes, 1])
elif aparam.size == fdim :
aparam = np.tile(aparam.reshape([-1]), [nframes, natoms])
else :
raise RuntimeError('got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d' % (nframes, natoms, fdim, natoms, fdim, fdim))
# sort inputs
coords, atom_types, imap = self.sort_input(coords, atom_types)
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert(natoms_vec[0] == natoms)
# evaluate
energy = []
force = []
virial = []
ae = []
av = []
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type ] = atom_types
t_out = [self.t_energy,
self.t_force,
self.t_virial]
if atomic :
t_out += [self.t_ae,
self.t_av]
for ii in range(nframes) :
feed_dict_test[self.t_coord] = np.reshape(coords[ii:ii+1, :], [-1])
feed_dict_test[self.t_box ] = np.reshape(cells [ii:ii+1, :], [-1])
if pbc:
feed_dict_test[self.t_mesh ] = make_default_mesh(cells[ii:ii+1, :])
else:
feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
if self.has_fparam:
feed_dict_test[self.t_fparam] = np.reshape(fparam[ii:ii+1, :], [-1])
if self.has_aparam:
feed_dict_test[self.t_aparam] = np.reshape(aparam[ii:ii+1, :], [-1])
v_out = self.sess.run (t_out, feed_dict = feed_dict_test)
energy.append(v_out[0])
force .append(v_out[1])
virial.append(v_out[2])
if atomic:
ae.append(v_out[3])
av.append(v_out[4])
# reverse map of the outputs
force = self.reverse_map(np.reshape(force, [nframes,-1,3]), imap)
if atomic :
ae = self.reverse_map(np.reshape(ae, [nframes,-1,1]), imap)
av = self.reverse_map(np.reshape(av, [nframes,-1,9]), imap)
energy = np.reshape(energy, [nframes, 1])
force = np.reshape(force, [nframes, natoms, 3])
virial = np.reshape(virial, [nframes, 9])
if atomic:
ae = np.reshape(ae, [nframes, natoms, 1])
av = np.reshape(av, [nframes, natoms, 9])
return energy, force, virial, ae, av
else :
return energy, force, virial
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))