def test_ener_shift(self):
dp_random.seed(0)
data = DeepmdDataSystem(['system_0', 'system_1'], 5, 10, 1.0)
data.add('energy', 1, must=True)
ener_shift0 = data.compute_energy_shift(rcond=1)
all_stat = make_stat_input(data, 4, merge_sys=False)
ener_shift1 = EnerFitting._compute_output_stats(all_stat, rcond=1)
np.testing.assert_almost_equal(ener_shift0, ener_shift1)
def test_ntypes(self) :
batch_size = 3
test_size = 2
ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0)
ds.add('test', self.test_ndof, atomic = True, must = True)
ds.add('null', self.test_ndof, atomic = True, must = False)
self.assertEqual(ds.get_ntypes(), 3)
self.assertEqual(ds.get_nbatches(), [2, 4, 3, 2])
self.assertEqual(ds.get_nsystems(), self.nsys)
self.assertEqual(list(ds.get_batch_size()), [batch_size]*4)
def test_ener_shift(self):
dp_random.seed(0)
data = DeepmdDataSystem(['system_0', 'system_1'], 5, 10, 1.0)
data.add('energy', 1, must=True)
ener_shift0 = data.compute_energy_shift(rcond=1)
all_stat = make_stat_input(data, 4, merge_sys=False)
descrpt = DescrptSeA(6.0,
5.8, [46, 92],
neuron=[25, 50, 100],
axis_neuron=16)
fitting = EnerFitting(descrpt, neuron=[240, 240, 240], resnet_dt=True)
ener_shift1 = fitting._compute_output_stats(all_stat, rcond=1)
np.testing.assert_almost_equal(ener_shift0, ener_shift1)
def test_prob_sys_size_1(self) :
batch_size = 1
test_size = 1
ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0)
prob = ds._prob_sys_size_ext("prob_sys_size; 0:2:2; 2:4:8")
self.assertAlmostEqual(np.sum(prob), 1)
self.assertAlmostEqual(np.sum(prob[0:2]), 0.2)
self.assertAlmostEqual(np.sum(prob[2:4]), 0.8)
# number of training set is self.nset-1
# shift is the total number of set size shift...
shift = np.sum(np.arange(self.nset-1))
self.assertAlmostEqual(prob[1]/prob[0], float(self.nframes[1]*(self.nset-1)+shift)/float(self.nframes[0]*(self.nset-1)+shift))
self.assertAlmostEqual(prob[3]/prob[2], float(self.nframes[3]*(self.nset-1)+shift)/float(self.nframes[2]*(self.nset-1)+shift))
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 get_data(jdata: Dict[str, Any], rcut, type_map, modifier):
systems = j_must_have(jdata, "systems")
if isinstance(systems, str):
systems = expand_sys_str(systems)
help_msg = 'Please check your setting for data systems'
# check length of systems
if len(systems) == 0:
msg = 'cannot find valid a data system'
log.fatal(msg)
raise IOError(msg, help_msg)
# rougly check all items in systems are valid
for ii in systems:
ii = DPPath(ii)
if (not ii.is_dir()):
msg = f'dir {ii} is not a valid dir'
log.fatal(msg)
raise IOError(msg, help_msg)
if (not (ii / 'type.raw').is_file()):
msg = f'dir {ii} is not a valid data system dir'
log.fatal(msg)
raise IOError(msg, help_msg)
batch_size = j_must_have(jdata, "batch_size")
sys_probs = jdata.get("sys_probs", None)
auto_prob = jdata.get("auto_prob", "prob_sys_size")
data = DeepmdDataSystem(
systems=systems,
batch_size=batch_size,
test_size=1, # to satisfy the old api
shuffle_test=True, # to satisfy the old api
rcut=rcut,
type_map=type_map,
modifier=modifier,
trn_all_set=True, # sample from all sets
sys_probs=sys_probs,
auto_prob_style=auto_prob
)
data.add_dict(data_requirement)
return data
def neighbor_stat(
*,
system: str,
rcut: float,
type_map: List[str],
**kwargs,
):
"""Calculate neighbor statistics.
Parameters
----------
system : str
system to stat
rcut : float
cutoff radius
type_map : list[str]
type map
Examples
--------
>>> neighbor_stat(system='.', rcut=6., type_map=["C", "H", "O", "N", "P", "S", "Mg", "Na", "HW", "OW", "mNa", "mCl", "mC", "mH", "mMg", "mN", "mO", "mP"])
min_nbor_dist: 0.6599510670195264
max_nbor_size: [23, 26, 19, 16, 2, 2, 1, 1, 72, 37, 5, 0, 31, 29, 1, 21, 20, 5]
"""
all_sys = expand_sys_str(system)
if not len(all_sys):
raise RuntimeError("Did not find valid system")
data = DeepmdDataSystem(
systems=all_sys,
batch_size=1,
test_size=1,
rcut=rcut,
type_map=type_map,
)
data.get_batch()
nei = NeighborStat(data.get_ntypes(), rcut)
min_nbor_dist, max_nbor_size = nei.get_stat(data)
log.info("min_nbor_dist: %f" % min_nbor_dist)
log.info("max_nbor_size: %s" % str(max_nbor_size))
return min_nbor_dist, max_nbor_size
def test_ener_shift_assigned(self):
dp_random.seed(0)
ae0 = dp_random.random()
data = DeepmdDataSystem(['system_0'], 5, 10, 1.0)
data.add('energy', 1, must=True)
all_stat = make_stat_input(data, 4, merge_sys=False)
descrpt = DescrptSeA(6.0,
5.8, [46, 92],
neuron=[25, 50, 100],
axis_neuron=16)
fitting = EnerFitting(descrpt,
neuron=[240, 240, 240],
resnet_dt=True,
atom_ener=[ae0, None, None])
ener_shift1 = fitting._compute_output_stats(all_stat, rcond=1)
# check assigned energy
np.testing.assert_almost_equal(ae0, ener_shift1[0])
# check if total energy are the same
natoms = data.natoms_vec[0][2:]
tot0 = np.dot(data.compute_energy_shift(rcond=1), natoms)
tot1 = np.dot(ener_shift1, natoms)
np.testing.assert_almost_equal(tot0, tot1)
def test_merge_all_stat(self):
dp_random.seed(0)
data0 = DeepmdDataSystem(['system_0', 'system_1'], 5, 10, 1.0)
data0.add('energy', 1, must=True)
dp_random.seed(0)
data1 = DeepmdDataSystem(['system_0', 'system_1'], 5, 10, 1.0)
data1.add('energy', 1, must=True)
dp_random.seed(0)
data2 = DeepmdDataSystem(['system_0', 'system_1'], 5, 10, 1.0)
data2.add('energy', 1, must=True)
dp_random.seed(0)
all_stat_0 = make_stat_input(data0, 10, merge_sys=False)
dp_random.seed(0)
all_stat_1 = make_stat_input(data1, 10, merge_sys=True)
all_stat_2 = merge_sys_stat(all_stat_0)
dp_random.seed(0)
all_stat_3 = _make_all_stat_ref(data2, 10)
####################################
# only check if the energy is concatenated correctly
####################################
dd = 'energy'
# if 'find_' in dd: continue
# if 'natoms_vec' in dd: continue
# if 'default_mesh' in dd: continue
# print(all_stat_2[dd])
# print(dd, all_stat_1[dd])
d1 = np.array(all_stat_1[dd])
d2 = np.array(all_stat_2[dd])
d3 = np.array(all_stat_3[dd])
# print(dd)
# print(d1.shape)
# print(d2.shape)
# self.assertEqual(all_stat_2[dd], all_stat_1[dd])
self._comp_data(d1, d2)
self._comp_data(d1, d3)
def test_get_test(self):
batch_size = 3
test_size = 2
ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0)
ds.add('test', self.test_ndof, atomic = True, must = True)
ds.add('null', self.test_ndof, atomic = True, must = False)
sys_idx = 0
data = ds.get_test(sys_idx=sys_idx)
self.assertEqual(list(data['type'][0]), list(np.sort(self.atom_type[sys_idx])))
self._in_array(np.load('sys_0/set.002/coord.npy'),
ds.get_sys(sys_idx).idx_map,
3,
data['coord'])
self._in_array(np.load('sys_0/set.002/test.npy'),
ds.get_sys(sys_idx).idx_map,
self.test_ndof,
data['test'])
self.assertAlmostEqual(np.linalg.norm(np.zeros([self.nframes[sys_idx]+2,
self.natoms[sys_idx]*self.test_ndof])
-
data['null']
), 0.0)
sys_idx = 2
data = ds.get_test(sys_idx=sys_idx)
self.assertEqual(list(data['type'][0]), list(np.sort(self.atom_type[sys_idx])))
self._in_array(np.load('sys_2/set.002/coord.npy'),
ds.get_sys(sys_idx).idx_map,
3,
data['coord'])
self._in_array(np.load('sys_2/set.002/test.npy'),
ds.get_sys(sys_idx).idx_map,
self.test_ndof,
data['test'])
self.assertAlmostEqual(np.linalg.norm(np.zeros([self.nframes[sys_idx]+2,
self.natoms[sys_idx]*self.test_ndof])
-
data['null']
), 0.0)
def test_batch_size_raise(self):
batch_size = 'foo'
test_size = 2
with self.assertRaises(RuntimeError):
ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0)
def test_batch_size_null(self):
batch_size = 'auto:3'
test_size = 2
ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0)
self.assertEqual(ds.batch_size, [1, 1, 1, 1])