def setUp(self, data):
self.sess = tf.Session()
self.data = data
self.natoms = self.data.get_natoms()
self.ntypes = self.data.get_ntypes()
param_a = {
'sel': [12, 24],
'rcut': 4,
'rcut_smth': 3.5,
'neuron': [5, 10, 20],
'seed': 1,
}
param_r = {
'sel': [20, 40],
'rcut': 6,
'rcut_smth': 6.5,
'neuron': [10, 20, 40],
'seed': 1,
}
param = {'a': param_a, 'r': param_r}
self.descrpt = DescrptSeAR(param)
self.ndescrpt = self.descrpt.get_dim_out()
# davg = np.zeros ([self.ntypes, self.ndescrpt])
# dstd = np.ones ([self.ntypes, self.ndescrpt])
# self.t_avg = tf.constant(davg.astype(np.float64))
# self.t_std = tf.constant(dstd.astype(np.float64))
avg_a = np.zeros([self.ntypes, self.descrpt.descrpt_a.ndescrpt])
std_a = np.ones([self.ntypes, self.descrpt.descrpt_a.ndescrpt])
avg_r = np.zeros([self.ntypes, self.descrpt.descrpt_r.ndescrpt])
std_r = np.ones([self.ntypes, self.descrpt.descrpt_r.ndescrpt])
self.avg = [avg_a, avg_r]
self.std = [std_a, std_r]
self.default_mesh = np.zeros(6, dtype=np.int32)
self.default_mesh[3] = 2
self.default_mesh[4] = 2
self.default_mesh[5] = 2
# make place holder
self.coord = tf.placeholder(global_tf_float_precision,
[None, self.natoms[0] * 3],
name='t_coord')
self.box = tf.placeholder(global_tf_float_precision, [None, 9],
name='t_box')
self.type = tf.placeholder(tf.int32, [None, self.natoms[0]],
name="t_type")
self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms")
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
class Inter():
def setUp(self, data):
self.sess = tf.Session()
self.data = data
self.natoms = self.data.get_natoms()
self.ntypes = self.data.get_ntypes()
param_a = {
'sel': [12, 24],
'rcut': 4,
'rcut_smth': 3.5,
'neuron': [5, 10, 20],
'seed': 1,
}
param_r = {
'sel': [20, 40],
'rcut': 6,
'rcut_smth': 6.5,
'neuron': [10, 20, 40],
'seed': 1,
}
param = {'a': param_a, 'r': param_r}
self.descrpt = DescrptSeAR(param)
self.ndescrpt = self.descrpt.get_dim_out()
# davg = np.zeros ([self.ntypes, self.ndescrpt])
# dstd = np.ones ([self.ntypes, self.ndescrpt])
# self.t_avg = tf.constant(davg.astype(np.float64))
# self.t_std = tf.constant(dstd.astype(np.float64))
avg_a = np.zeros([self.ntypes, self.descrpt.descrpt_a.ndescrpt])
std_a = np.ones([self.ntypes, self.descrpt.descrpt_a.ndescrpt])
avg_r = np.zeros([self.ntypes, self.descrpt.descrpt_r.ndescrpt])
std_r = np.ones([self.ntypes, self.descrpt.descrpt_r.ndescrpt])
self.avg = [avg_a, avg_r]
self.std = [std_a, std_r]
self.default_mesh = np.zeros(6, dtype=np.int32)
self.default_mesh[3] = 2
self.default_mesh[4] = 2
self.default_mesh[5] = 2
# make place holder
self.coord = tf.placeholder(global_tf_float_precision,
[None, self.natoms[0] * 3],
name='t_coord')
self.box = tf.placeholder(global_tf_float_precision, [None, 9],
name='t_box')
self.type = tf.placeholder(tf.int32, [None, self.natoms[0]],
name="t_type")
self.tnatoms = tf.placeholder(tf.int32, [None], name="t_natoms")
def _net(self, inputs, name, reuse=False):
with tf.variable_scope(name, reuse=reuse):
net_w = tf.get_variable('net_w', [self.descrpt.get_dim_out()],
global_tf_float_precision,
tf.constant_initializer(self.net_w_i))
dot_v = tf.matmul(tf.reshape(inputs,
[-1, self.descrpt.get_dim_out()]),
tf.reshape(net_w, [self.descrpt.get_dim_out(), 1]))
return tf.reshape(dot_v, [-1])
def comp_ef(self, dcoord, dbox, dtype, tnatoms, name, reuse=None):
dout = self.descrpt.build(dcoord,
dtype,
tnatoms,
dbox,
self.default_mesh,
suffix=name,
reuse=reuse)
inputs_reshape = tf.reshape(dout, [-1, self.descrpt.get_dim_out()])
atom_ener = self._net(inputs_reshape, name, reuse=reuse)
atom_ener_reshape = tf.reshape(atom_ener, [-1, self.natoms[0]])
energy = tf.reduce_sum(atom_ener_reshape, axis=1)
force, virial, av = self.descrpt.prod_force_virial(
atom_ener_reshape, tnatoms)
return energy, force, virial