def __init__(self, jdata):
args = ClassArg()\
.add('a', dict, must = True) \
.add('r', dict, must = True)
class_data = args.parse(jdata)
self.param_a = class_data['a']
self.param_r = class_data['r']
self.descrpt_a = DescrptSeA(self.param_a)
self.descrpt_r = DescrptSeR(self.param_r)
assert (self.descrpt_a.get_ntypes() == self.descrpt_r.get_ntypes())
class DescrptSeAR():
def __init__(self, jdata):
args = ClassArg()\
.add('a', dict, must = True) \
.add('r', dict, must = True)
class_data = args.parse(jdata)
self.param_a = class_data['a']
self.param_r = class_data['r']
self.descrpt_a = DescrptSeA(self.param_a)
self.descrpt_r = DescrptSeR(self.param_r)
assert (self.descrpt_a.get_ntypes() == self.descrpt_r.get_ntypes())
self.davg = None
self.dstd = None
def get_rcut(self):
return np.max([self.descrpt_a.get_rcut(), self.descrpt_r.get_rcut()])
def get_ntypes(self):
return self.descrpt_r.get_ntypes()
def get_dim_out(self):
return (self.descrpt_a.get_dim_out() + self.descrpt_r.get_dim_out())
def get_nlist_a(self):
return self.descrpt_a.nlist, self.descrpt_a.rij, self.descrpt_a.sel_a, self.descrpt_a.sel_r
def get_nlist_r(self):
return self.descrpt_r.nlist, self.descrpt_r.rij, self.descrpt_r.sel_a, self.descrpt_r.sel_r
def compute_input_stats(self, data_coord, data_box, data_atype, natoms_vec,
mesh):
self.descrpt_a.compute_input_stats(data_coord, data_box, data_atype,
natoms_vec, mesh)
self.descrpt_r.compute_input_stats(data_coord, data_box, data_atype,
natoms_vec, mesh)
self.davg = [self.descrpt_a.davg, self.descrpt_r.davg]
self.dstd = [self.descrpt_a.dstd, self.descrpt_r.dstd]
def build(self, coord_, atype_, natoms, box, mesh, suffix='', reuse=None):
davg = self.davg
dstd = self.dstd
if davg is None:
davg = [
np.zeros([self.descrpt_a.ntypes, self.descrpt_a.ndescrpt]),
np.zeros([self.descrpt_r.ntypes, self.descrpt_r.ndescrpt])
]
if dstd is None:
dstd = [
np.ones([self.descrpt_a.ntypes, self.descrpt_a.ndescrpt]),
np.ones([self.descrpt_r.ntypes, self.descrpt_r.ndescrpt])
]
# dout
self.dout_a = self.descrpt_a.build(coord_,
atype_,
natoms,
box,
mesh,
suffix=suffix + '_a',
reuse=reuse)
self.dout_r = self.descrpt_r.build(coord_,
atype_,
natoms,
box,
mesh,
suffix=suffix,
reuse=reuse)
self.dout_a = tf.reshape(self.dout_a,
[-1, self.descrpt_a.get_dim_out()])
self.dout_r = tf.reshape(self.dout_r,
[-1, self.descrpt_r.get_dim_out()])
self.dout = tf.concat([self.dout_a, self.dout_r], axis=1)
self.dout = tf.reshape(self.dout, [-1, natoms[0] * self.get_dim_out()])
return self.dout
def prod_force_virial(self, atom_ener, natoms):
f_a, v_a, av_a = self.descrpt_a.prod_force_virial(atom_ener, natoms)
f_r, v_r, av_r = self.descrpt_r.prod_force_virial(atom_ener, natoms)
force = f_a + f_r
virial = v_a + v_r
atom_virial = av_a + av_r
return force, virial, atom_virial
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 _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