def _smth_descriptor(jdata: Dict[str, Any]) -> Dict[str, Any]:
"""Convert data to v1 input for smooth descriptor.
Parameters
----------
jdata : Dict[str, Any]
parsed input json/yaml data
Returns
-------
Dict[str, Any]
dict with descriptor parameters
"""
descriptor = {}
seed = jdata.get("seed", None)
if seed is not None:
descriptor["seed"] = seed
descriptor["type"] = "se_a"
descriptor["sel"] = jdata["sel_a"]
_jcopy(jdata, descriptor, ("rcut", ))
descriptor["rcut_smth"] = jdata.get("rcut_smth", descriptor["rcut"])
descriptor["neuron"] = j_must_have(jdata, "filter_neuron")
descriptor["axis_neuron"] = j_must_have(jdata, "axis_neuron",
["n_axis_neuron"])
descriptor["resnet_dt"] = False
if "resnet_dt" in jdata:
descriptor["resnet_dt"] = jdata["filter_resnet_dt"]
return descriptor
def _learning_rate(jdata):
# learning rate
learning_rate = {}
learning_rate['type'] = 'exp'
learning_rate['decay_steps'] = j_must_have(jdata, 'decay_steps')
learning_rate['decay_rate'] = j_must_have(jdata, 'decay_rate')
learning_rate['start_lr'] = j_must_have(jdata, 'start_lr')
return learning_rate
def compute_efv(jfile):
fp = open(jfile, 'r')
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
rcut = j_must_have(jdata, 'rcut')
data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt)
tot_numb_batches = sum(data.get_nbatches())
lr = LearingRate(jdata, tot_numb_batches)
model = NNPModel(jdata, run_opt=run_opt)
model.build(data, lr)
test_data = data.get_test()
feed_dict_test = {
model.t_prop_c:
test_data["prop_c"],
model.t_energy:
test_data["energy"][:model.numb_test],
model.t_force:
np.reshape(test_data["force"][:model.numb_test, :], [-1]),
model.t_virial:
np.reshape(test_data["virial"][:model.numb_test, :], [-1]),
model.t_atom_ener:
np.reshape(test_data["atom_ener"][:model.numb_test, :], [-1]),
model.t_atom_pref:
np.reshape(test_data["atom_pref"][:model.numb_test, :], [-1]),
model.t_coord:
np.reshape(test_data["coord"][:model.numb_test, :], [-1]),
model.t_box:
test_data["box"][:model.numb_test, :],
model.t_type:
np.reshape(test_data["type"][:model.numb_test, :], [-1]),
model.t_natoms:
test_data["natoms_vec"],
model.t_mesh:
test_data["default_mesh"],
model.t_fparam:
np.reshape(test_data["fparam"][:model.numb_test, :], [-1]),
model.is_training:
False
}
sess = tf.Session()
sess.run(tf.global_variables_initializer())
[e, f, v] = sess.run([model.energy, model.force, model.virial],
feed_dict=feed_dict_test)
return e, f, v
def _fitting_net(jdata: Dict[str, Any]) -> Dict[str, Any]:
"""Convert data to v1 input for fitting net.
Parameters
----------
jdata : Dict[str, Any]
parsed input json/yaml data
Returns
-------
Dict[str, Any]
dict with fitting net parameters
"""
fitting_net = {}
seed = jdata.get("seed", None)
if seed is not None:
fitting_net["seed"] = seed
fitting_net["neuron"] = j_must_have(jdata, "fitting_neuron", ["n_neuron"])
fitting_net["resnet_dt"] = True
if "resnet_dt" in jdata:
fitting_net["resnet_dt"] = jdata["resnet_dt"]
if "fitting_resnet_dt" in jdata:
fitting_net["resnet_dt"] = jdata["fitting_resnet_dt"]
return fitting_net
def _training(jdata):
# training
training = {}
seed = None
if j_have (jdata, 'seed') :
seed = jdata['seed']
training['systems'] = jdata['systems']
training['set_prefix'] = jdata['set_prefix']
training['stop_batch'] = jdata['stop_batch']
training['batch_size'] = jdata['batch_size']
if seed is not None:
training['seed'] = seed
training['disp_file'] = "lcurve.out"
if j_have (jdata, "disp_file") : training['disp_file'] = jdata["disp_file"]
training['disp_freq'] = j_must_have (jdata, 'disp_freq')
training['numb_test'] = j_must_have (jdata, 'numb_test')
training['save_freq'] = j_must_have (jdata, 'save_freq')
training['save_ckpt'] = j_must_have (jdata, 'save_ckpt')
training['disp_training'] = j_must_have (jdata, 'disp_training')
training['time_training'] = j_must_have (jdata, 'time_training')
if j_have (jdata, 'profiling') :
training['profiling'] = jdata['profiling']
if training['profiling'] :
training['profiling_file'] = j_must_have (jdata, 'profiling_file')
return training
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 _training(jdata: Dict[str, Any]) -> Dict[str, Any]:
"""Convert data to v1 input for training.
Parameters
----------
jdata : Dict[str, Any]
parsed input json/yaml data
Returns
-------
Dict[str, Any]
dict with training parameters
"""
training = {}
seed = jdata.get("seed", None)
if seed is not None:
training["seed"] = seed
_jcopy(jdata, training,
("systems", "set_prefix", "stop_batch", "batch_size"))
training["disp_file"] = "lcurve.out"
if "disp_file" in jdata:
training["disp_file"] = jdata["disp_file"]
training["disp_freq"] = j_must_have(jdata, "disp_freq")
training["numb_test"] = j_must_have(jdata, "numb_test")
training["save_freq"] = j_must_have(jdata, "save_freq")
training["save_ckpt"] = j_must_have(jdata, "save_ckpt")
training["disp_training"] = j_must_have(jdata, "disp_training")
training["time_training"] = j_must_have(jdata, "time_training")
if "profiling" in jdata:
training["profiling"] = jdata["profiling"]
if training["profiling"]:
training["profiling_file"] = j_must_have(jdata, "profiling_file")
return training
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 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 _loss(jdata):
# loss
loss = {}
loss['start_pref_e'] = j_must_have (jdata, 'start_pref_e')
loss['limit_pref_e'] = j_must_have (jdata, 'limit_pref_e')
loss['start_pref_f'] = j_must_have (jdata, 'start_pref_f')
loss['limit_pref_f'] = j_must_have (jdata, 'limit_pref_f')
loss['start_pref_v'] = j_must_have (jdata, 'start_pref_v')
loss['limit_pref_v'] = j_must_have (jdata, 'limit_pref_v')
if j_have(jdata, 'start_pref_ae') :
loss['start_pref_ae'] = jdata['start_pref_ae']
if j_have(jdata, 'limit_pref_ae') :
loss['limit_pref_ae'] = jdata['limit_pref_ae']
return loss
def _smth_descriptor(jdata):
descriptor = {}
seed = None
if j_have (jdata, 'seed') :
seed = jdata['seed']
descriptor['type'] = 'se_a'
descriptor['sel'] = jdata['sel_a']
if j_have(jdata, 'rcut_smth') :
descriptor['rcut_smth'] = jdata['rcut_smth']
else :
descriptor['rcut_smth'] = descriptor['rcut']
descriptor['rcut'] = jdata['rcut']
descriptor['neuron'] = j_must_have (jdata, 'filter_neuron')
descriptor['axis_neuron'] = j_must_have_d (jdata, 'axis_neuron', ['n_axis_neuron'])
descriptor['resnet_dt'] = False
if j_have(jdata, 'resnet_dt') :
descriptor['resnet_dt'] = jdata['filter_resnet_dt']
if seed is not None:
descriptor['seed'] = seed
return descriptor
def test_model(self):
jfile = 'water_se_a_aparam.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()
# manually set aparam
test_data['aparam'] = np.load('system/set.000/aparam.npy')
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']],
'aparam': [test_data['aparam']],
}
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')
t_aparam = tf.placeholder(global_tf_float_precision, [None],
name='i_aparam')
is_training = tf.placeholder(tf.bool)
input_dict = {}
input_dict['aparam'] = t_aparam
model_pred\
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
input_dict,
suffix = "se_a_aparam",
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'],
t_aparam: np.reshape(test_data['aparam'][:numb_test, :], [-1]),
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 = [61.35473702079649]
reff = [
7.789591210641927388e-02, 9.411176646369459609e-02,
3.785806413688173194e-03, 1.430830954178063386e-01,
1.146964190520970150e-01, -1.320340288927138173e-02,
-7.308720494747594776e-02, 6.508269338140809657e-02,
5.398739145542804643e-04, 5.863268336973800898e-02,
-1.603409523950408699e-01, -5.083084610994957619e-03,
-2.551569799443983988e-01, 3.087934885732580501e-02,
1.508590526622844222e-02, 4.863249399791078065e-02,
-1.444292753594846324e-01, -1.125098094204559241e-03
]
refv = [
-6.069498397488943819e-01, 1.101778888191114192e-01,
1.981907430646132409e-02, 1.101778888191114608e-01,
-3.315612988100872793e-01, -5.999739184898976799e-03,
1.981907430646132756e-02, -5.999739184898974197e-03,
-1.198656608172396325e-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 is_deepmd_v1_input(jdata):
return "systems" in j_must_have(jdata, "training").keys()
def test_descriptor_one_side_exclude_types(self):
"""When we enable type_one_side, the descriptor should be the same
for different types, when its environments are the same.
Here we generates two data. The only difference is the type:
(1) 0 1 1 1 1 1
(2) 1 1 1 1 1 1
When type_one_side is true, the first atom should have the same descriptor.
Otherwise, it should be different (with random initial variables). We test
both situation.
"""
jfile = 'water_se_a.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 = 1
test_size = 1
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
jdata['model']['descriptor']['type_one_side'] = True
jdata['model']['descriptor']['exclude_types'] = [[0, 0]]
t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c')
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)
# successful
descrpt = Descriptor(**jdata['model']['descriptor'])
dout \
= descrpt.build(
t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
{},
reuse = False,
suffix = "_se_a_1side_exclude_types"
)
# failed
descrpt_failed = Descriptor(
**{
**jdata['model']['descriptor'], "type_one_side": False
})
dout_failed \
= descrpt_failed.build(
t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
{},
reuse = False,
suffix = "_se_a_1side_exclude_types_failed"
)
feed_dict_test1 = {
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
}
feed_dict_test2 = feed_dict_test1.copy()
# original type: 0 0 1 1 1 1
# current: 0 1 1 1 1 1
# current: 1 1 1 1 1 1
new_natoms1 = test_data['natoms_vec'].copy()
new_natoms1[2] = 1
new_natoms1[3] = 5
new_type1 = test_data['type'].copy()
new_type1[:numb_test, 0] = 0
new_type1[:numb_test, 1:6] = 1
new_natoms2 = test_data['natoms_vec'].copy()
new_natoms2[2] = 0
new_natoms2[3] = 6
new_type2 = test_data['type'].copy()
new_type2[:numb_test] = 1
feed_dict_test1[t_type] = np.reshape(new_type1[:numb_test, :], [-1])
feed_dict_test1[t_natoms] = new_natoms1
feed_dict_test2[t_type] = np.reshape(new_type2[:numb_test, :], [-1])
feed_dict_test2[t_natoms] = new_natoms2
print(feed_dict_test1, feed_dict_test2)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
[model_dout1] = sess.run([dout], feed_dict=feed_dict_test1)
[model_dout2] = sess.run([dout], feed_dict=feed_dict_test2)
[model_dout1_failed] = sess.run([dout_failed],
feed_dict=feed_dict_test1)
[model_dout2_failed] = sess.run([dout_failed],
feed_dict=feed_dict_test2)
model_dout1 = model_dout1.reshape([6, -1])
model_dout2 = model_dout2.reshape([6, -1])
model_dout1_failed = model_dout1_failed.reshape([6, -1])
model_dout2_failed = model_dout2_failed.reshape([6, -1])
np.testing.assert_almost_equal(model_dout1[0], model_dout2[0], 10)
with self.assertRaises(AssertionError):
np.testing.assert_almost_equal(model_dout1_failed[0],
model_dout2_failed[0], 10)
def test_model(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 = 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
jdata['model']['descriptor'].pop('type', None)
descrpt = DescrptSeA(**jdata['model']['descriptor'], uniform_seed = True)
jdata['model']['fitting_net']['descrpt'] = descrpt
fitting = EnerFitting(**jdata['model']['fitting_net'], uniform_seed = True)
typeebd_param = jdata['model']['type_embedding']
typeebd = TypeEmbedNet(
neuron = typeebd_param['neuron'],
resnet_dt = typeebd_param['resnet_dt'],
seed = typeebd_param['seed'],
uniform_seed = True)
model = EnerModel(descrpt, fitting, typeebd)
# 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
inputs_dict = {}
model_pred \
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
inputs_dict,
suffix = "se_a_type",
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 = self.test_session().__enter__()
sess.run(tf.global_variables_initializer())
[e, f, v] = sess.run([energy, force, virial],
feed_dict = feed_dict_test)
# print(sess.run(model.type_embedding))
# np.savetxt('tmp.out', sess.run(descrpt.dout, feed_dict = feed_dict_test), fmt='%.10e')
# # print(sess.run(model.atype_embed, feed_dict = feed_dict_test))
# print(sess.run(fitting.inputs, feed_dict = feed_dict_test))
# print(sess.run(fitting.outs, feed_dict = feed_dict_test))
# print(sess.run(fitting.atype_embed, feed_dict = feed_dict_test))
e = e.reshape([-1])
f = f.reshape([-1])
v = v.reshape([-1])
np.savetxt('e.out', e.reshape([1, -1]), delimiter=',')
np.savetxt('f.out', f.reshape([1, -1]), delimiter=',')
np.savetxt('v.out', v.reshape([1, -1]), delimiter=',')
refe = [6.049065170680415804e+01]
reff = [1.021832439441947293e-01,1.122650466359011306e-01,3.927874278714531091e-03,1.407089812207832635e-01,1.312473824343091400e-01,-1.228371057389851181e-02,-1.109672154547165501e-01,6.582735820731049070e-02,1.251568633647655391e-03,7.933758749748777428e-02,-1.831777072317984367e-01,-6.173090134630876760e-03,-2.703597126460742794e-01,4.817856571062521104e-02,1.491963457594796399e-02,5.909711543832503466e-02,-1.743406457563475287e-01,-1.642276779780762769e-03]
refv = [-6.932736357193732823e-01,1.453756052949563837e-01,2.138263139115256783e-02,1.453756052949564392e-01,-3.880901656480436612e-01,-7.782259726407755700e-03,2.138263139115256437e-02,-7.782259726407749628e-03,-1.225285973678705374e-03]
refe = np.reshape(refe, [-1])
reff = np.reshape(reff, [-1])
refv = np.reshape(refv, [-1])
places = 10
np.testing.assert_almost_equal(e, refe, places)
np.testing.assert_almost_equal(f, reff, places)
np.testing.assert_almost_equal(v, refv, places)
def test_model(self):
jfile = 'polar_se_a.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 = 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
jdata['model']['descriptor'].pop('type', None)
jdata['model']['fitting_net'].pop('type', None)
descrpt = DescrptSeA(**jdata['model']['descriptor'], uniform_seed=True)
jdata['model']['fitting_net']['descrpt'] = descrpt
fitting = PolarFittingSeA(**jdata['model']['fitting_net'],
uniform_seed=True)
model = PolarModel(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']],
'fparam': [test_data['fparam']],
}
model._compute_input_stat(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 = "polar_se_a",
reuse = False)
polar = model_pred['polar']
gpolar = model_pred['global_polar']
force = model_pred['force']
virial = model_pred['virial']
atom_virial = model_pred['atom_virial']
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 = self.test_session().__enter__()
sess.run(tf.global_variables_initializer())
[p, gp] = sess.run([polar, gpolar], 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
np.testing.assert_almost_equal(p, refp, places)
gp = gp.reshape([-1])
refgp = np.array(refp).reshape(-1, 9).sum(0)
places = 5
np.testing.assert_almost_equal(gp, refgp, places)
# make sure only one frame is used
feed_dict_single = {
t_prop_c: test_data['prop_c'],
t_coord: np.reshape(test_data['coord'][:1, :], [-1]),
t_box: test_data['box'][:1, :],
t_type: np.reshape(test_data['type'][:1, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
is_training: False
}
[pf, pv, pav] = sess.run([force, virial, atom_virial],
feed_dict=feed_dict_single)
pf, pv = pf.reshape(-1), pv.reshape(-1)
spv = pav.reshape(1, 9, -1, 9).sum(2).reshape(-1)
base_dict = feed_dict_single.copy()
coord0 = base_dict.pop(t_coord)
box0 = base_dict.pop(t_box)
fdf = -finite_difference(
lambda coord: sess.run(gpolar,
feed_dict={
**base_dict, t_coord: coord,
t_box: box0
}).reshape(-1),
test_data['coord'][:numb_test, :].reshape([-1])).reshape(-1)
fdv = -(finite_difference(
lambda box: sess.run(
gpolar,
feed_dict={
**base_dict, t_coord: strerch_box(coord0, box0, box),
t_box: box
}).reshape(-1), test_data['box'][:numb_test, :]).reshape([
-1, 3, 3
]).transpose(0, 2, 1) @ box0.reshape(3, 3)).reshape(-1)
delta = 1e-4
np.testing.assert_allclose(pf, fdf, delta)
np.testing.assert_allclose(pv, fdv, delta)
# make sure atomic virial sum to virial
places = 10
np.testing.assert_almost_equal(pv, spv, places)
def _init_param(self, jdata):
# model config
model_param = j_must_have(jdata, 'model')
descrpt_param = j_must_have(model_param, 'descriptor')
fitting_param = j_must_have(model_param, 'fitting_net')
typeebd_param = model_param.get('type_embedding', None)
self.model_param = model_param
self.descrpt_param = descrpt_param
# descriptor
try:
descrpt_type = descrpt_param['type']
except KeyError:
raise KeyError('the type of descriptor should be set by `type`')
self.descrpt = Descriptor(**descrpt_param)
# fitting net
fitting_type = fitting_param.get('type', 'ener')
self.fitting_type = fitting_type
fitting_param.pop('type', None)
fitting_param['descrpt'] = self.descrpt
if fitting_type == 'ener':
self.fitting = EnerFitting(**fitting_param)
# elif fitting_type == 'wfc':
# self.fitting = WFCFitting(fitting_param, self.descrpt)
elif fitting_type == 'dipole':
if descrpt_type == 'se_e2_a':
self.fitting = DipoleFittingSeA(**fitting_param)
else :
raise RuntimeError('fitting dipole only supports descrptors: se_e2_a')
elif fitting_type == 'polar':
# if descrpt_type == 'loc_frame':
# self.fitting = PolarFittingLocFrame(fitting_param, self.descrpt)
if descrpt_type == 'se_e2_a':
self.fitting = PolarFittingSeA(**fitting_param)
else :
raise RuntimeError('fitting polar only supports descrptors: loc_frame and se_e2_a')
elif fitting_type == 'global_polar':
if descrpt_type == 'se_e2_a':
self.fitting = GlobalPolarFittingSeA(**fitting_param)
else :
raise RuntimeError('fitting global_polar only supports descrptors: loc_frame and se_e2_a')
else :
raise RuntimeError('unknow fitting type ' + fitting_type)
# type embedding
if typeebd_param is not None:
self.typeebd = TypeEmbedNet(
neuron=typeebd_param['neuron'],
resnet_dt=typeebd_param['resnet_dt'],
activation_function=typeebd_param['activation_function'],
precision=typeebd_param['precision'],
trainable=typeebd_param['trainable'],
seed=typeebd_param['seed']
)
else:
self.typeebd = None
# init model
# infer model type by fitting_type
if fitting_type == 'ener':
self.model = EnerModel(
self.descrpt,
self.fitting,
self.typeebd,
model_param.get('type_map'),
model_param.get('data_stat_nbatch', 10),
model_param.get('data_stat_protect', 1e-2),
model_param.get('use_srtab'),
model_param.get('smin_alpha'),
model_param.get('sw_rmin'),
model_param.get('sw_rmax')
)
# elif fitting_type == 'wfc':
# self.model = WFCModel(model_param, self.descrpt, self.fitting)
elif fitting_type == 'dipole':
self.model = DipoleModel(
self.descrpt,
self.fitting,
model_param.get('type_map'),
model_param.get('data_stat_nbatch', 10),
model_param.get('data_stat_protect', 1e-2)
)
elif fitting_type == 'polar':
self.model = PolarModel(
self.descrpt,
self.fitting,
model_param.get('type_map'),
model_param.get('data_stat_nbatch', 10),
model_param.get('data_stat_protect', 1e-2)
)
elif fitting_type == 'global_polar':
self.model = GlobalPolarModel(
self.descrpt,
self.fitting,
model_param.get('type_map'),
model_param.get('data_stat_nbatch', 10),
model_param.get('data_stat_protect', 1e-2)
)
else :
raise RuntimeError('get unknown fitting type when building model')
# learning rate
lr_param = j_must_have(jdata, 'learning_rate')
scale_by_worker = lr_param.get('scale_by_worker', 'linear')
if scale_by_worker == 'linear':
self.scale_lr_coef = float(self.run_opt.world_size)
elif scale_by_worker == 'sqrt':
self.scale_lr_coef = np.sqrt(self.run_opt.world_size).real
else:
self.scale_lr_coef = 1.
lr_type = lr_param.get('type', 'exp')
if lr_type == 'exp':
self.lr = LearningRateExp(lr_param['start_lr'],
lr_param['stop_lr'],
lr_param['decay_steps'])
else :
raise RuntimeError('unknown learning_rate type ' + lr_type)
# loss
# infer loss type by fitting_type
loss_param = jdata.get('loss', None)
loss_type = loss_param.get('type', 'ener')
if fitting_type == 'ener':
loss_param.pop('type', None)
loss_param['starter_learning_rate'] = self.lr.start_lr()
if loss_type == 'ener':
self.loss = EnerStdLoss(**loss_param)
elif loss_type == 'ener_dipole':
self.loss = EnerDipoleLoss(**loss_param)
else:
raise RuntimeError('unknow loss type')
elif fitting_type == 'wfc':
self.loss = TensorLoss(loss_param,
model = self.model,
tensor_name = 'wfc',
tensor_size = self.model.get_out_size(),
label_name = 'wfc')
elif fitting_type == 'dipole':
self.loss = TensorLoss(loss_param,
model = self.model,
tensor_name = 'dipole',
tensor_size = 3,
label_name = 'dipole')
elif fitting_type == 'polar':
self.loss = TensorLoss(loss_param,
model = self.model,
tensor_name = 'polar',
tensor_size = 9,
label_name = 'polarizability')
elif fitting_type == 'global_polar':
self.loss = TensorLoss(loss_param,
model = self.model,
tensor_name = 'global_polar',
tensor_size = 9,
atomic = False,
label_name = 'polarizability')
else :
raise RuntimeError('get unknown fitting type when building loss function')
# training
tr_data = jdata['training']
self.disp_file = tr_data.get('disp_file', 'lcurve.out')
self.disp_freq = tr_data.get('disp_freq', 1000)
self.save_freq = tr_data.get('save_freq', 1000)
self.save_ckpt = tr_data.get('save_ckpt', 'model.ckpt')
self.display_in_training = tr_data.get('disp_training', True)
self.timing_in_training = tr_data.get('time_training', True)
self.profiling = self.run_opt.is_chief and tr_data.get('profiling', False)
self.profiling_file = tr_data.get('profiling_file', 'timeline.json')
self.tensorboard = self.run_opt.is_chief and tr_data.get('tensorboard', False)
self.tensorboard_log_dir = tr_data.get('tensorboard_log_dir', 'log')
self.tensorboard_freq = tr_data.get('tensorboard_freq', 1)
# self.sys_probs = tr_data['sys_probs']
# self.auto_prob_style = tr_data['auto_prob']
self.useBN = False
if fitting_type == 'ener' and self.fitting.get_numb_fparam() > 0 :
self.numb_fparam = self.fitting.get_numb_fparam()
else :
self.numb_fparam = 0
if tr_data.get("validation_data", None) is not None:
self.valid_numb_batch = tr_data["validation_data"].get("numb_btch", 1)
else:
self.valid_numb_batch = 1
# if init the graph with the frozen model
self.frz_model = None
self.model_type = None
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_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 _init_param(self, jdata):
# model config
model_param = j_must_have(jdata, 'model')
descrpt_param = j_must_have(model_param, 'descriptor')
fitting_param = j_must_have(model_param, 'fitting_net')
# descriptor
descrpt_type = j_must_have(descrpt_param, 'type')
if descrpt_type == 'loc_frame':
self.descrpt = DescrptLocFrame(descrpt_param)
elif descrpt_type == 'se_a':
self.descrpt = DescrptSeA(descrpt_param)
elif descrpt_type == 'se_r':
self.descrpt = DescrptSeR(descrpt_param)
elif descrpt_type == 'se_ar':
self.descrpt = DescrptSeAR(descrpt_param)
else:
raise RuntimeError('unknow model type ' + descrpt_type)
# fitting net
try:
fitting_type = fitting_param['type']
except:
fitting_type = 'ener'
if fitting_type == 'ener':
self.fitting = EnerFitting(fitting_param, self.descrpt)
elif fitting_type == 'wfc':
self.fitting = WFCFitting(fitting_param, self.descrpt)
elif fitting_type == 'dipole':
if descrpt_type == 'se_a':
self.fitting = DipoleFittingSeA(fitting_param, self.descrpt)
else:
raise RuntimeError(
'fitting dipole only supports descrptors: se_a')
elif fitting_type == 'polar':
if descrpt_type == 'loc_frame':
self.fitting = PolarFittingLocFrame(fitting_param,
self.descrpt)
elif descrpt_type == 'se_a':
self.fitting = PolarFittingSeA(fitting_param, self.descrpt)
else:
raise RuntimeError(
'fitting polar only supports descrptors: loc_frame and se_a'
)
elif fitting_type == 'global_polar':
if descrpt_type == 'se_a':
self.fitting = GlobalPolarFittingSeA(fitting_param,
self.descrpt)
else:
raise RuntimeError(
'fitting global_polar only supports descrptors: loc_frame and se_a'
)
else:
raise RuntimeError('unknow fitting type ' + fitting_type)
# init model
# infer model type by fitting_type
if fitting_type == Model.model_type:
self.model = Model(model_param, self.descrpt, self.fitting)
elif fitting_type == 'wfc':
self.model = WFCModel(model_param, self.descrpt, self.fitting)
elif fitting_type == 'dipole':
self.model = DipoleModel(model_param, self.descrpt, self.fitting)
elif fitting_type == 'polar':
self.model = PolarModel(model_param, self.descrpt, self.fitting)
elif fitting_type == 'global_polar':
self.model = GlobalPolarModel(model_param, self.descrpt,
self.fitting)
else:
raise RuntimeError('get unknown fitting type when building model')
# learning rate
lr_param = j_must_have(jdata, 'learning_rate')
try:
lr_type = lr_param['type']
except:
lr_type = 'exp'
if lr_type == 'exp':
self.lr = LearningRateExp(lr_param)
else:
raise RuntimeError('unknown learning_rate type ' + lr_type)
# loss
# infer loss type by fitting_type
try:
loss_param = jdata['loss']
loss_type = loss_param.get('type', 'std')
except:
loss_param = None
loss_type = 'std'
if fitting_type == 'ener':
if loss_type == 'std':
self.loss = EnerStdLoss(
loss_param, starter_learning_rate=self.lr.start_lr())
elif loss_type == 'ener_dipole':
self.loss = EnerDipoleLoss(
loss_param, starter_learning_rate=self.lr.start_lr())
else:
raise RuntimeError('unknow loss type')
elif fitting_type == 'wfc':
self.loss = TensorLoss(loss_param,
model=self.model,
tensor_name='wfc',
tensor_size=self.model.get_out_size(),
label_name='wfc')
elif fitting_type == 'dipole':
self.loss = TensorLoss(loss_param,
model=self.model,
tensor_name='dipole',
tensor_size=3,
label_name='dipole')
elif fitting_type == 'polar':
self.loss = TensorLoss(loss_param,
model=self.model,
tensor_name='polar',
tensor_size=9,
label_name='polarizability')
elif fitting_type == 'global_polar':
self.loss = TensorLoss(loss_param,
model=self.model,
tensor_name='global_polar',
tensor_size=9,
atomic=False,
label_name='polarizability')
else:
raise RuntimeError(
'get unknown fitting type when building loss function')
# training
training_param = j_must_have(jdata, 'training')
tr_args = ClassArg()\
.add('numb_test', int, default = 1)\
.add('disp_file', str, default = 'lcurve.out')\
.add('disp_freq', int, default = 100)\
.add('save_freq', int, default = 1000)\
.add('save_ckpt', str, default = 'model.ckpt')\
.add('display_in_training', bool, default = True)\
.add('timing_in_training', bool, default = True)\
.add('profiling', bool, default = False)\
.add('profiling_file',str, default = 'timeline.json')\
.add('sys_probs', list )\
.add('auto_prob_style', str, default = "prob_sys_size")
tr_data = tr_args.parse(training_param)
self.numb_test = tr_data['numb_test']
self.disp_file = tr_data['disp_file']
self.disp_freq = tr_data['disp_freq']
self.save_freq = tr_data['save_freq']
self.save_ckpt = tr_data['save_ckpt']
self.display_in_training = tr_data['display_in_training']
self.timing_in_training = tr_data['timing_in_training']
self.profiling = tr_data['profiling']
self.profiling_file = tr_data['profiling_file']
self.sys_probs = tr_data['sys_probs']
self.auto_prob_style = tr_data['auto_prob_style']
self.useBN = False
if fitting_type == 'ener' and self.fitting.get_numb_fparam() > 0:
self.numb_fparam = self.fitting.get_numb_fparam()
else:
self.numb_fparam = 0
def test_fitting(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 = 1
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
jdata['model']['descriptor']['neuron'] = [5, 5, 5]
jdata['model']['descriptor']['axis_neuron'] = 2
jdata['model']['descriptor'].pop('type', None)
descrpt = DescrptSeA(**jdata['model']['descriptor'], uniform_seed=True)
jdata['model']['fitting_net']['descrpt'] = descrpt
fitting = EnerFitting(**jdata['model']['fitting_net'],
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']]
}
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
dout = np.array([
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
])
type_embedding = np.array([
1.4916816460764615, 0.2720153234707013, -2.4385153754181985,
-1.8454294510880027, 2.874575701113528, 1.1225116575801295,
0.4204818970813372, -2.3784087249787587, -1.5053748251050598,
2.769329403073084
])
dout = dout.reshape([-1, 10])
type_embedding = type_embedding.reshape([ntypes, -1])
atom_ener = fitting.build(
tf.convert_to_tensor(dout),
t_natoms, {'type_embedding': tf.convert_to_tensor(type_embedding)},
reuse=False,
suffix="se_a_type_fit_")
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())
[pred_atom_ener] = sess.run([atom_ener], feed_dict=feed_dict_test)
pred_atom_ener = pred_atom_ener.reshape([-1])
ref_atom_ener = [
10.121733946849165, 10.121712105320634, 10.143275419743475,
10.143299785396826, 10.143311150431957, 10.143295201182019
]
places = 10
np.testing.assert_almost_equal(pred_atom_ener, ref_atom_ener, places)
def _do_work(jdata: Dict[str, Any], run_opt: RunOptions, is_compress: bool = False):
"""Run serial model training.
Parameters
----------
jdata : Dict[str, Any]
arguments read form json/yaml control file
run_opt : RunOptions
object with run configuration
is_compress : Bool
indicates whether in model compress mode
Raises
------
RuntimeError
If unsupported modifier type is selected for model
"""
# make necessary checks
assert "training" in jdata
# init the model
model = DPTrainer(jdata, run_opt=run_opt, is_compress = is_compress)
rcut = model.model.get_rcut()
type_map = model.model.get_type_map()
if len(type_map) == 0:
ipt_type_map = None
else:
ipt_type_map = type_map
# init random seed of data systems
seed = jdata["training"].get("seed", None)
if seed is not None:
# avoid the same batch sequence among workers
seed += run_opt.my_rank
seed = seed % (2 ** 32)
dp_random.seed(seed)
# setup data modifier
modifier = get_modifier(jdata["model"].get("modifier", None))
# decouple the training data from the model compress process
train_data = None
valid_data = None
if not is_compress:
# init data
train_data = get_data(jdata["training"]["training_data"], rcut, ipt_type_map, modifier)
train_data.print_summary("training")
if jdata["training"].get("validation_data", None) is not None:
valid_data = get_data(jdata["training"]["validation_data"], rcut, ipt_type_map, modifier)
valid_data.print_summary("validation")
# get training info
stop_batch = j_must_have(jdata["training"], "numb_steps")
model.build(train_data, stop_batch)
if not is_compress:
# train the model with the provided systems in a cyclic way
start_time = time.time()
model.train(train_data, valid_data)
end_time = time.time()
log.info("finished training")
log.info(f"wall time: {(end_time - start_time):.3f} s")
else:
model.save_compressed()
log.info("finished compressing")
def test_descriptor_one_side(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 = 1
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
jdata['model']['descriptor']['type_one_side'] = True
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_1side")
dout \
= descrpt.build(
t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
{'type_embedding' : type_embedding},
reuse = False,
suffix = "_se_a_type_des_1side"
)
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.0009704469114440277, 0.0007136310372560243,
0.0007136310372560243, 0.000524968274824758,
-0.0019790100690810016, -0.0014556100390424947,
-0.001318691223889266, -0.0009698525512440269,
0.001937780602605409, 0.0014251755182315322, 0.0008158935519461114,
0.0005943870925895051, 0.0005943870925895051,
0.0004340263490412088, -0.0016539827195947239,
-0.0012066241021841376, -0.0011042186455562336,
-0.0008051343572505189, 0.0016229491738044255,
0.0011833923257801077, 0.0006020440527161554,
0.00047526899287409847, 0.00047526899287409847,
0.00037538142786805136, -0.0012811397377036637,
-0.0010116898098710776, -0.0008465095301785942,
-0.0006683577463042215, 0.0012459039620461505,
0.0009836962283627838, 0.00077088529431722, 0.0006105807630364827,
0.0006105807630364827, 0.00048361458700877996,
-0.0016444700616024337, -0.001302510079662288,
-0.0010856603485807576, -0.0008598975276238373,
0.00159730642327918, 0.001265146946434076, 0.0008495806081447204,
0.000671787466824433, 0.000671787466824433, 0.0005312928157964384,
-0.0018105890543181475, -0.001431844407277983,
-0.0011956722392735362, -0.000945544277375045,
0.0017590147511761475, 0.0013910348287283414,
0.0007393644735054756, 0.0005850536182149991,
0.0005850536182149991, 0.0004631887654949332,
-0.0015760302086346792, -0.0012475134925387294,
-0.001041074331192672, -0.0008239586048523492,
0.0015319673563669856, 0.0012124704278707746
]
places = 10
np.testing.assert_almost_equal(model_dout, ref_dout, 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_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 = 'polar_se_a.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 = 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
jdata['model']['descriptor'].pop('type', None)
descrpt = DescrptSeA(**jdata['model']['descriptor'], uniform_seed = True)
jdata['model']['fitting_net'].pop('type', None)
jdata['model']['fitting_net'].pop('fit_diag', None)
jdata['model']['fitting_net']['descrpt'] = descrpt
fitting = DipoleFittingSeA(**jdata['model']['fitting_net'], uniform_seed = True)
model = DipoleModel(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']],
'fparam': [test_data['fparam']],
}
model._compute_input_stat(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 = "dipole_se_a",
reuse = False)
dipole = model_pred['dipole']
gdipole = model_pred['global_dipole']
force = model_pred['force']
virial = model_pred['virial']
atom_virial = model_pred['atom_virial']
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 = self.test_session().__enter__()
sess.run(tf.global_variables_initializer())
[p, gp] = sess.run([dipole, gdipole], feed_dict = feed_dict_test)
p = p.reshape([-1])
refp = [1.616802262298876514e+01,9.809535439521079425e+00,3.572312180768947854e-01,1.336308874095981203e+00,1.057908563208963848e+01,-5.999602350098874881e-01]
places = 10
np.testing.assert_almost_equal(p, refp, places)
gp = gp.reshape([-1])
refgp = np.array(refp).reshape(-1, 3).sum(0)
places = 9
np.testing.assert_almost_equal(gp, refgp, places)
# make sure only one frame is used
feed_dict_single = {t_prop_c: test_data['prop_c'],
t_coord: np.reshape(test_data['coord'] [:1, :], [-1]),
t_box: test_data['box'] [:1, :],
t_type: np.reshape(test_data['type'] [:1, :], [-1]),
t_natoms: test_data['natoms_vec'],
t_mesh: test_data['default_mesh'],
is_training: False}
[pf, pv, pav] = sess.run([force, virial, atom_virial], feed_dict = feed_dict_single)
pf, pv = pf.reshape(-1), pv.reshape(-1)
spv = pav.reshape(1, 3, -1, 9).sum(2).reshape(-1)
base_dict = feed_dict_single.copy()
coord0 = base_dict.pop(t_coord)
box0 = base_dict.pop(t_box)
fdf = - finite_difference(
lambda coord: sess.run(gdipole,
feed_dict={**base_dict,
t_coord:coord,
t_box:box0}).reshape(-1),
test_data['coord'][:numb_test, :].reshape([-1])).reshape(-1)
fdv = - (finite_difference(
lambda box: sess.run(gdipole,
feed_dict={**base_dict,
t_coord:strerch_box(coord0, box0, box),
t_box:box}).reshape(-1),
test_data['box'][:numb_test, :]).reshape([-1,3,3]).transpose(0,2,1)
@ box0.reshape(3,3)).reshape(-1)
delta = 1e-5
np.testing.assert_allclose(pf, fdf, delta)
np.testing.assert_allclose(pv, fdv, delta)
# make sure atomic virial sum to virial
places = 10
np.testing.assert_almost_equal(pv, spv, 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_model(self):
jfile = 'water_se_a_srtab.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 = 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
jdata['model']['descriptor'].pop('type', None)
descrpt = DescrptSeA(**jdata['model']['descriptor'], uniform_seed = True)
jdata['model']['fitting_net']['descrpt'] = descrpt
fitting = EnerFitting(**jdata['model']['fitting_net'], uniform_seed = True)
# descrpt = DescrptSeA(jdata['model']['descriptor'])
# fitting = EnerFitting(jdata['model']['fitting_net'], descrpt)
model = EnerModel(
descrpt,
fitting,
None,
jdata['model'].get('type_map'),
jdata['model'].get('data_stat_nbatch'),
jdata['model'].get('data_stat_protect'),
jdata['model'].get('use_srtab'),
jdata['model'].get('smin_alpha'),
jdata['model'].get('sw_rmin'),
jdata['model'].get('sw_rmax')
)
# 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_a_srtab",
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 = self.test_session().__enter__()
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.141610882066236599e+02]
reff = [-1.493121233165248043e+02,-1.831419491743885715e+02,-8.439542992300344437e+00,-1.811987095947552859e+02,-1.476380826187439084e+02,1.264271856742560018e+01,1.544377958934875323e+02,-7.816520233903435866e+00,1.287925245463442225e+00,-4.000393268449002449e+00,1.910748885843098890e+02,7.134789955349889468e+00,1.826908441979261113e+02,3.677156386479059513e+00,-1.122312112141401741e+01,-2.617413911684622008e+00,1.438445070562470391e+02,-1.402769654524568033e+00]
refv = [3.585047655925112622e+02,-7.569252978336677984e+00,-1.068382043878426124e+01,-7.569252978336677096e+00,3.618439481685132932e+02,5.448668500896081568e+00,-1.068382043878426302e+01,5.448668500896082456e+00,1.050393462151727686e+00]
refe = np.reshape(refe, [-1])
reff = np.reshape(reff, [-1])
refv = np.reshape(refv, [-1])
places = 10
np.testing.assert_almost_equal(e, refe, places)
np.testing.assert_almost_equal(f, reff, places)
np.testing.assert_almost_equal(v, refv, places)
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_model(self):
jfile = 'water_se_a_fparam.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']])
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')
t_fparam = tf.placeholder(global_tf_float_precision, [None],
name='i_fparam')
is_training = tf.placeholder(tf.bool)
input_dict = {}
input_dict['fparam'] = t_fparam
model_pred\
= model.build (t_coord,
t_type,
t_natoms,
t_box,
t_mesh,
input_dict,
suffix = "se_a_fparam",
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'],
t_fparam: np.reshape(test_data['fparam'][:numb_test, :], [-1]),
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.135136929183754972e+01]
reff = [
7.761477777656561328e-02, 9.383013575207051205e-02,
3.776776376267230399e-03, 1.428268971463224069e-01,
1.143858253900619654e-01, -1.318441687719179231e-02,
-7.271897092708884403e-02, 6.494907553857684479e-02,
5.355599592111062821e-04, 5.840910251709752199e-02,
-1.599042555763417750e-01, -5.067165555590445389e-03,
-2.546246315216804113e-01, 3.073296814647456451e-02,
1.505994759166155023e-02, 4.849282500878367153e-02,
-1.439937492508420736e-01, -1.120701494357654411e-03
]
refv = [
-6.054303146013112480e-01, 1.097859194719944115e-01,
1.977605183964963390e-02, 1.097859194719943976e-01,
-3.306167096812382966e-01, -5.978855662865613894e-03,
1.977605183964964083e-02, -5.978855662865616497e-03,
-1.196331922996723236e-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)
