def predict(self, dp):
"""
Predict energies and forces by deepmd-kit.
Parameters
----------
dp : deepmd.DeepPot or str
The deepmd-kit potential class or the filename of the model.
Returns
-------
labeled_sys : LabeledSystem
The labeled system.
"""
try:
# DP 1.x
import deepmd.DeepPot as DeepPot
except ModuleNotFoundError:
# DP 2.x
from deepmd.infer import DeepPot
if not isinstance(dp, DeepPot):
dp = DeepPot(dp)
type_map = dp.get_type_map()
ori_sys = self.copy()
ori_sys.sort_atom_names(type_map=type_map)
atype = ori_sys['atom_types']
labeled_sys = LabeledSystem()
if 'auto_batch_size' not in DeepPot.__init__.__code__.co_varnames:
for ss in self:
coord = ss['coords'].reshape((1, ss.get_natoms()*3))
if not ss.nopbc:
cell = ss['cells'].reshape((1, 9))
else:
cell = None
e, f, v = dp.eval(coord, cell, atype)
data = ss.data
data['energies'] = e.reshape((1, 1))
data['forces'] = f.reshape((1, ss.get_natoms(), 3))
data['virials'] = v.reshape((1, 3, 3))
this_sys = LabeledSystem.from_dict({'data': data})
labeled_sys.append(this_sys)
else:
# since v2.0.2, auto batch size is supported
coord = self.data['coords'].reshape((self.get_nframes(), self.get_natoms()*3))
if not self.nopbc:
cell = self.data['cells'].reshape((self.get_nframes(), 9))
else:
cell = None
e, f, v = dp.eval(coord, cell, atype)
data = self.data.copy()
data['energies'] = e.reshape((self.get_nframes(), 1))
data['forces'] = f.reshape((self.get_nframes(), self.get_natoms(), 3))
data['virials'] = v.reshape((self.get_nframes(), 3, 3))
labeled_sys = LabeledSystem.from_dict({'data': data})
return labeled_sys
class DP_fparam(Calculator):
"""
modify from DP class of deepmd
the vanilla one does not support fparam
"""
name = "DP"
implemented_properties = ["energy", "forces", "stress"]
def __init__(self, model, label="DP", fparam=0, **kwargs):
Calculator.__init__(self, label=label, **kwargs)
self.dp = DeepPot(model)
self.type_dict = dict(
zip(self.dp.get_type_map(), range(self.dp.get_ntypes())))
self.fparam = [fparam]
def calculate(self,
atoms=None,
properties=["energy", "forces", "stress"],
system_changes=all_changes):
coord = atoms.get_positions().reshape([1, -1])
cell = atoms.get_cell().reshape([1, -1])
symbols = atoms.get_chemical_symbols()
atype = [self.type_dict[k] for k in symbols]
e, f, v = self.dp.eval(coord, cell, atype, fparam=self.fparam)
self.results['energy'] = e[0]
self.results['forces'] = f[0]
self.results['stress'] = v[0]
class DP(Calculator):
name = "DP"
implemented_properties = ["energy", "forces", "stress"]
def __init__(self, model, label="DP", type_dict=None, **kwargs):
Calculator.__init__(self, label=label, **kwargs)
self.dp = DeepPot(model)
if type_dict:
self.type_dict = type_dict
else:
self.type_dict = dict(
zip(self.dp.get_type_map(), range(self.dp.get_ntypes())))
def calculate(self,
atoms=None,
properties=["energy", "forces", "stress"],
system_changes=all_changes):
coord = atoms.get_positions().reshape([1, -1])
cell = atoms.get_cell().reshape([1, -1])
symbols = atoms.get_chemical_symbols()
atype = [self.type_dict[k] for k in symbols]
e, f, v = self.dp.eval(coord, cell, atype)
self.results['energy'] = e[0]
self.results['forces'] = f[0]
self.results['stress'] = v[0]
def test_ener (args) :
if args.rand_seed is not None :
np.random.seed(args.rand_seed % (2**32))
data = DataSets (args.system, args.set_prefix, shuffle_test = args.shuffle_test)
test_data = data.get_test ()
numb_test = args.numb_test
natoms = len(test_data["type"][0])
nframes = test_data["box"].shape[0]
numb_test = min(nframes, numb_test)
dp = DeepPot(args.model)
coord = test_data["coord"][:numb_test].reshape([numb_test, -1])
box = test_data["box"][:numb_test]
atype = test_data["type"][0]
if dp.get_dim_fparam() > 0:
fparam = test_data["fparam"][:numb_test]
else :
fparam = None
if dp.get_dim_aparam() > 0:
aparam = test_data["aparam"][:numb_test]
else :
aparam = None
energy, force, virial, ae, av = dp.eval(coord, box, atype, fparam = fparam, aparam = aparam, atomic = True)
energy = energy.reshape([numb_test,1])
force = force.reshape([numb_test,-1])
virial = virial.reshape([numb_test,9])
ae = ae.reshape([numb_test,-1])
av = av.reshape([numb_test,-1])
l2e = (l2err (energy - test_data["energy"][:numb_test].reshape([-1,1])))
l2f = (l2err (force - test_data["force"] [:numb_test]))
l2v = (l2err (virial - test_data["virial"][:numb_test]))
l2ea= l2e/natoms
l2va= l2v/natoms
# print ("# energies: %s" % energy)
print ("# number of test data : %d " % numb_test)
print ("Energy L2err : %e eV" % l2e)
print ("Energy L2err/Natoms : %e eV" % l2ea)
print ("Force L2err : %e eV/A" % l2f)
print ("Virial L2err : %e eV" % l2v)
print ("Virial L2err/Natoms : %e eV" % l2va)
detail_file = args.detail_file
if detail_file is not None :
pe = np.concatenate((np.reshape(test_data["energy"][:numb_test], [-1,1]),
np.reshape(energy, [-1,1])),
axis = 1)
np.savetxt(detail_file+".e.out", pe,
header = 'data_e pred_e')
pf = np.concatenate((np.reshape(test_data["force"] [:numb_test], [-1,3]),
np.reshape(force, [-1,3])),
axis = 1)
np.savetxt(detail_file+".f.out", pf,
header = 'data_fx data_fy data_fz pred_fx pred_fy pred_fz')
pv = np.concatenate((np.reshape(test_data["virial"][:numb_test], [-1,9]),
np.reshape(virial, [-1,9])),
axis = 1)
np.savetxt(detail_file+".v.out", pv,
header = 'data_vxx data_vxy data_vxz data_vyx data_vyy data_vyz data_vzx data_vzy data_vzz pred_vxx pred_vxy pred_vxz pred_vyx pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz')
class DP(Calculator):
name = "DP"
implemented_properties = ["energy", "forces", "virial", "stress"]
def __init__(self, model, label="DP", type_dict=None, **kwargs):
Calculator.__init__(self, label=label, **kwargs)
self.dp = DeepPot(model)
if type_dict:
self.type_dict = type_dict
else:
self.type_dict = dict(
zip(self.dp.get_type_map(), range(self.dp.get_ntypes())))
def calculate(self,
atoms=None,
properties=["energy", "forces", "virial"],
system_changes=all_changes):
coord = atoms.get_positions().reshape([1, -1])
if sum(atoms.get_pbc()) > 0:
cell = atoms.get_cell().reshape([1, -1])
else:
cell = None
symbols = atoms.get_chemical_symbols()
atype = [self.type_dict[k] for k in symbols]
e, f, v = self.dp.eval(coords=coord, cells=cell, atom_types=atype)
self.results['energy'] = e[0][0]
self.results['forces'] = f[0]
self.results['virial'] = v[0].reshape(3, 3)
# convert virial into stress for lattice relaxation
if "stress" in properties:
if sum(atoms.get_pbc()) > 0:
# the usual convention (tensile stress is positive)
# stress = -virial / volume
stress = -0.5 * (v[0].copy() +
v[0].copy().T) / atoms.get_volume()
# Voigt notation
self.results['stress'] = stress.flat[[0, 4, 8, 5, 2, 1]]
else:
raise PropertyNotImplementedError
class TestSet(object):
def __init__(self,
model_file=None,
coord_file=None,
energy_file=None,
force_file=None,
grad_file=None,
box_file=None,
atom_types=None,
length_unit="A",
energy_unit="Eh",
force_unit="Eh/Bohr",
is_pbc=False,
verbose=True):
self.is_pbc = is_pbc
self.verbose = verbose
if isinstance(atom_types, str):
self._atm_type = numpy.loadtxt(atom_types, dtype=int)
else:
self._atm_type = numpy.asarray(atom_types, dtype=int)
self.model_file = model_file
model_type = DeepEval(model_file).model_type
assert model_type == 'ener'
self.dp = DeepPot(model_file)
length_unit, length_unit_converter = get_length_unit_converter(
length_unit)
energy_unit, energy_unit_converter = get_energy_unit_converter(
energy_unit)
force_unit, force_unit_converter = get_force_unit_converter(force_unit)
self._coord_data = numpy.load(coord_file) * length_unit_converter
self._energy_data = numpy.load(energy_file) * energy_unit_converter
if force_file is not None and grad_file is None:
self._force_data = numpy.load(force_file) * force_unit_converter
elif force_file is None and grad_file is not None:
self._force_data = -numpy.load(grad_file) * force_unit_converter
self.nframe = self._coord_data.shape[0]
self.natom = self._coord_data.shape[1]
self._atm_type.reshape(self.natom)
if not is_pbc and (box_file is None):
self._box_data = numpy.zeros([self.nframe, 9])
else:
self._box_data = numpy.load(box_file) * length_unit_converter
assert self._box_data.shape == (self.nframe, 9)
assert self._coord_data.shape == (self.nframe, self.natom, 3)
assert self._force_data.shape == (self.nframe, self.natom, 3)
assert self._energy_data.shape == (self.nframe, )
assert self._atm_type.shape == (self.natom, )
self.dump_info(model_file=model_file,
atom_types=atom_types,
is_pbc=is_pbc,
coord_file=coord_file,
energy_file=energy_file,
force_file=force_file,
grad_file=grad_file,
box_file=box_file,
length_unit=length_unit,
energy_unit=energy_unit,
force_unit=force_unit)
def dump_info(self, **kwargs):
if not self.verbose:
return
dump_info(kwargs)
def build(self, detail_file=None):
numb_test = self.nframe
energy_data = self._energy_data[:numb_test].reshape([-1, 1])
force_data = self._force_data[:numb_test].reshape([numb_test, -1])
coord = self._coord_data.reshape([numb_test, -1])
box = self._box_data.reshape([numb_test, -1])
if not self.is_pbc:
box = None
fparam = None
aparam = None
atype = self._atm_type
atomic = False
ret = self.dp.eval(coord,
box,
atype,
fparam=fparam,
aparam=aparam,
atomic=atomic)
energy = ret[0]
force = ret[1]
virial = ret[2]
energy = energy.reshape([numb_test, 1])
force = force.reshape([numb_test, -1])
virial = virial.reshape([numb_test, 9])
if atomic:
ae = ret[3]
av = ret[4]
ae = ae.reshape([numb_test, -1])
av = av.reshape([numb_test, -1])
l2e = (l2err(energy - energy_data))
l2f = (l2err(force - force_data))
l2ea = l2e / self.natom
print("# number of test data : %d " % numb_test)
print("Energy L2err : %e eV" % l2e)
print("Energy L2err/Natoms : %e eV" % l2ea)
print("Force L2err : %e eV/A" % l2f)
if detail_file is not None:
save_txt_file(detail_file + ".c.out",
coord,
header='data_c',
append=False)
pe = numpy.concatenate(
(numpy.reshape(energy_data, [numb_test, -1]),
numpy.reshape(energy, [numb_test, -1])),
axis=1)
save_txt_file(detail_file + ".e.out",
pe,
header='data_e pred_e',
append=False)
pf = numpy.concatenate((numpy.reshape(force_data, [numb_test, -1]),
numpy.reshape(force, [numb_test, -1])),
axis=1)
save_txt_file(detail_file + ".f.out",
pf,
header='data_f pred_f',
append=False)
coord = test_data["coord"][:numb_test].reshape([numb_test, -1])
box = test_data["box"][:numb_test]
if not data.pbc:
box = None
atype = test_data["type"][0]
if dp.get_dim_fparam() > 0:
fparam = test_data["fparam"][:numb_test]
else :
fparam = None
if dp.get_dim_aparam() > 0:
aparam = test_data["aparam"][:numb_test]
else :
aparam = None
<<<<<<< HEAD
energy, force, virial, ae, av = dp.eval(coord, box, atype, fparam = fparam, aparam = aparam, atomic = True)
=======
detail_file = args.detail_file
if detail_file is not None:
atomic = True
else:
atomic = False
ret = dp.eval(coord, box, atype, fparam = fparam, aparam = aparam, atomic = atomic)
energy = ret[0]
force = ret[1]
virial = ret[2]
>>>>>>> 93aa90768089e265465f93fb3d9f99b494d56b09
energy = energy.reshape([numb_test,1])
force = force.reshape([numb_test,-1])
virial = virial.reshape([numb_test,9])
def test_ener(args):
"""
modify based on args file
"""
if args['rand_seed'] is not None:
np.random.seed(args['rand_seed'] % (2**32))
data = DataSets(args['system'],
args['set_prefix'],
shuffle_test=args['shuffle_test'])
test_data = data.get_test()
numb_test = args['numb_test']
natoms = len(test_data["type"][0])
nframes = test_data["box"].shape[0]
numb_test = min(nframes, numb_test)
dp = DeepPot(args['model'])
coord = test_data["coord"][:numb_test].reshape([numb_test, -1])
box = test_data["box"][:numb_test]
atype = test_data["type"][0]
if dp.get_dim_fparam() > 0:
fparam = test_data["fparam"][:numb_test]
else:
fparam = None
if dp.get_dim_aparam() > 0:
aparam = test_data["aparam"][:numb_test]
else:
aparam = None
detail_file = args['detail_file']
if detail_file is not None:
atomic = True
else:
atomic = False
ret = dp.eval(coord,
box,
atype,
fparam=fparam,
aparam=aparam,
atomic=atomic)
energy = ret[0]
force = ret[1]
virial = ret[2]
energy = energy.reshape([numb_test, 1])
force = force.reshape([numb_test, -1])
virial = virial.reshape([numb_test, 9])
if atomic:
ae = ret[3]
av = ret[4]
ae = ae.reshape([numb_test, -1])
av = av.reshape([numb_test, -1])
l2e = (l2err(energy - test_data["energy"][:numb_test].reshape([-1, 1])))
l2f = (l2err(force - test_data["force"][:numb_test]))
l2v = (l2err(virial - test_data["virial"][:numb_test]))
l2ea = l2e / natoms
l2va = l2v / natoms
# print ("# energies: %s" % energy)
print("# number of test data : %d " % numb_test)
print("Energy L2err : %e eV" % l2e)
print("Energy L2err/Natoms : %e eV" % l2ea)
print("Force L2err : %e eV/A" % l2f)
print("Virial L2err : %e eV" % l2v)
print("Virial L2err/Natoms : %e eV" % l2va)
if detail_file is not None:
pe = np.concatenate((np.reshape(test_data["energy"][:numb_test],
[-1, 1]), np.reshape(energy, [-1, 1])),
axis=1)
np.savetxt(os.path.join(args['system'], detail_file + ".e.out"),
pe,
header='data_e pred_e')
pf = np.concatenate((np.reshape(test_data["force"][:numb_test],
[-1, 3]), np.reshape(force, [-1, 3])),
axis=1)
np.savetxt(os.path.join(args['system'], detail_file + ".f.out"),
pf,
header='data_fx data_fy data_fz pred_fx pred_fy pred_fz')
pv = np.concatenate((np.reshape(test_data["virial"][:numb_test],
[-1, 9]), np.reshape(virial, [-1, 9])),
axis=1)
np.savetxt(
os.path.join(args['system'], detail_file + ".v.out"),
pv,
header=
'data_vxx data_vxy data_vxz data_vyx data_vyy data_vyz data_vzx data_vzy data_vzz pred_vxx pred_vxy pred_vxz pred_vyx pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz'
)
return numb_test, fparam[0][0], natoms, l2e, l2ea, l2f, l2v
def train_ener(inputs):
"""
deepmd-kit has function test_ener which deal with test_data only
`train_ener` are for train data only
"""
if inputs['rand_seed'] is not None:
np.random.seed(inputs['rand_seed'] % (2**32))
data = DataSets(inputs['system'],
inputs['set_prefix'],
shuffle_test=inputs['shuffle_test'])
train_data = get_train_data(data)
numb_test = data.get_sys_numb_batch(
1) ## use 1 batch, # of batches are the numb of train
natoms = len(train_data["type"][0])
nframes = train_data["box"].shape[0]
#print("xxxxx",nframes, numb_test)
numb_test = nframes #, to be investigated, original dp use min, but here should be nframes directly, I think, Jan 18, 21, min(nfames, numb_test)
dp = DeepPot(inputs['model'])
coord = train_data["coord"].reshape([numb_test, -1])
box = train_data["box"]
atype = train_data["type"][0]
if dp.get_dim_fparam() > 0:
fparam = train_data["fparam"]
else:
fparam = None
if dp.get_dim_aparam() > 0:
aparam = train_data["aparam"]
else:
aparam = None
detail_file = inputs['detail_file']
if detail_file is not None:
atomic = True
else:
atomic = False
ret = dp.eval(coord,
box,
atype,
fparam=fparam,
aparam=aparam,
atomic=atomic)
energy = ret[0]
force = ret[1]
virial = ret[2]
energy = energy.reshape([numb_test, 1])
force = force.reshape([numb_test, -1])
virial = virial.reshape([numb_test, 9])
if atomic:
ae = ret[3]
av = ret[4]
ae = ae.reshape([numb_test, -1])
av = av.reshape([numb_test, -1])
l2e = (l2err(energy - train_data["energy"].reshape([-1, 1])))
l2f = (l2err(force - train_data["force"]))
l2v = (l2err(virial - train_data["virial"]))
l2ea = l2e / natoms
l2va = l2v / natoms
# print ("# energies: %s" % energy)
print("# number of train data : %d " % numb_test)
print("Energy L2err : %e eV" % l2e)
print("Energy L2err/Natoms : %e eV" % l2ea)
print("Force L2err : %e eV/A" % l2f)
print("Virial L2err : %e eV" % l2v)
print("Virial L2err/Natoms : %e eV" % l2va)
if detail_file is not None:
pe = np.concatenate((np.reshape(train_data["energy"],
[-1, 1]), np.reshape(energy, [-1, 1])),
axis=1)
np.savetxt(os.path.join(inputs['system'], detail_file + ".e.tr.out"),
pe,
header='data_e pred_e')
pf = np.concatenate((np.reshape(train_data["force"],
[-1, 3]), np.reshape(force, [-1, 3])),
axis=1)
np.savetxt(os.path.join(inputs['system'], detail_file + ".f.tr.out"),
pf,
header='data_fx data_fy data_fz pred_fx pred_fy pred_fz')
pv = np.concatenate((np.reshape(train_data["virial"],
[-1, 9]), np.reshape(virial, [-1, 9])),
axis=1)
np.savetxt(
os.path.join(inputs['system'], detail_file + ".v.tr.out"),
pv,
header=
'data_vxx data_vxy data_vxz data_vyx data_vyy data_vyz data_vzx data_vzy data_vzz pred_vxx pred_vxy pred_vxz pred_vyx pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz'
)
return numb_test, fparam[0][0], natoms, l2e, l2ea, l2f, l2v