def _load_graph(self,
frozen_graph_filename,
prefix='load',
default_tf_graph=False):
# We load the protobuf file from the disk and parse it to retrieve the
# unserialized graph_def
with tf.gfile.GFile(frozen_graph_filename, "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
if default_tf_graph:
tf.import_graph_def(graph_def,
input_map=None,
return_elements=None,
name=prefix,
producer_op_list=None)
graph = tf.get_default_graph()
else:
# Then, we can use again a convenient built-in function to import a graph_def into the
# current default Graph
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def,
input_map=None,
return_elements=None,
name=prefix,
producer_op_list=None)
# for ii in graph.as_graph_def().node:
# print(ii.name)
return graph
def load_graph(graphName):
graph_def = tf.GraphDef()
with open(graphName, "rb") as f:
graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def, name="")
return graph
def __init__(self, hh, beta):
"""
Constructor
Parameters
----------
hh
Grid spacing of the reciprocal part of Ewald sum. Unit: A
beta
Splitting parameter of the Ewald sum. Unit: A^{-1}
"""
self.hh = hh
self.beta = beta
with tf.Graph().as_default() as graph:
# place holders
self.t_nloc = tf.placeholder(tf.int32, [1], name="t_nloc")
self.t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None],
name='t_coord')
self.t_charge = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None],
name='t_charge')
self.t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None],
name='t_box')
# output
self.t_energy, self.t_force, self.t_virial \
= op_module.ewald_recp(self.t_coord, self.t_charge, self.t_nloc, self.t_box,
ewald_h = self.hh,
ewald_beta = self.beta)
self.sess = tf.Session(graph=graph, config=default_tf_session_config)
def _compute_dstats_sys_smth (self,
data_coord,
data_box,
data_atype,
natoms_vec,
mesh) :
avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(global_np_float_precision)
std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(global_np_float_precision)
sub_graph = tf.Graph()
with sub_graph.as_default():
descrpt, descrpt_deriv, rij, nlist \
= op_module.descrpt_se_a (tf.constant(data_coord),
tf.constant(data_atype),
tf.constant(natoms_vec, dtype = tf.int32),
tf.constant(data_box),
tf.constant(mesh),
tf.constant(avg_zero),
tf.constant(std_ones),
rcut_a = self.rcut_a,
rcut_r = self.rcut_r,
rcut_r_smth = self.rcut_r_smth,
sel_a = self.sel_a,
sel_r = self.sel_r)
# self.sess.run(tf.global_variables_initializer())
# sub_sess = tf.Session(graph = sub_graph,
# config=tf.ConfigProto(intra_op_parallelism_threads=self.run_opt.num_intra_threads,
# inter_op_parallelism_threads=self.run_opt.num_inter_threads
# ))
sub_sess = tf.Session(graph = sub_graph)
dd_all = sub_sess.run(descrpt)
sub_sess.close()
natoms = natoms_vec
dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]])
start_index = 0
sysr = []
sysa = []
sysn = []
sysr2 = []
sysa2 = []
for type_i in range(self.ntypes):
end_index = start_index + self.ndescrpt * natoms[2+type_i]
dd = dd_all[:, start_index:end_index]
dd = np.reshape(dd, [-1, self.ndescrpt])
start_index = end_index
# compute
dd = np.reshape (dd, [-1, 4])
ddr = dd[:,:1]
dda = dd[:,1:]
sumr = np.sum(ddr)
suma = np.sum(dda) / 3.
sumn = dd.shape[0]
sumr2 = np.sum(np.multiply(ddr, ddr))
suma2 = np.sum(np.multiply(dda, dda)) / 3.
sysr.append(sumr)
sysa.append(suma)
sysn.append(sumn)
sysr2.append(sumr2)
sysa2.append(suma2)
return sysr, sysr2, sysa, sysa2, sysn
def _load_graph(frozen_graph_filename: "Path",
prefix: str = "load",
default_tf_graph: bool = False):
# We load the protobuf file from the disk and parse it to retrieve the
# unserialized graph_def
with tf.gfile.GFile(str(frozen_graph_filename), "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
if default_tf_graph:
tf.import_graph_def(graph_def,
input_map=None,
return_elements=None,
name=prefix,
producer_op_list=None)
graph = tf.get_default_graph()
else:
# Then, we can use again a convenient built-in function to import
# a graph_def into the current default Graph
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def,
input_map=None,
return_elements=None,
name=prefix,
producer_op_list=None)
return graph
def __init__(self, ntypes: int, rcut: float) -> None:
"""
Constructor
"""
self.rcut = rcut
self.ntypes = ntypes
self.place_holders = {}
sub_graph = tf.Graph()
with sub_graph.as_default():
for ii in ['coord', 'box']:
self.place_holders[ii] = tf.placeholder(
GLOBAL_NP_FLOAT_PRECISION, [None, None], name='t_' + ii)
self.place_holders['type'] = tf.placeholder(tf.int32, [None, None],
name='t_type')
self.place_holders['natoms_vec'] = tf.placeholder(
tf.int32, [self.ntypes + 2], name='t_natoms')
self.place_holders['default_mesh'] = tf.placeholder(tf.int32,
[None],
name='t_mesh')
self._max_nbor_size, self._min_nbor_dist \
= op_module.neighbor_stat(self.place_holders['coord'],
self.place_holders['type'],
self.place_holders['natoms_vec'],
self.place_holders['box'],
self.place_holders['default_mesh'],
rcut = self.rcut)
self.sub_sess = tf.Session(graph=sub_graph,
config=default_tf_session_config)
def __init__(self,
rcut: float,
sel_a : List[int],
sel_r : List[int],
axis_rule : List[int]
) -> None:
"""
Constructor
"""
# args = ClassArg()\
# .add('sel_a', list, must = True) \
# .add('sel_r', list, must = True) \
# .add('rcut', float, default = 6.0) \
# .add('axis_rule',list, must = True)
# class_data = args.parse(jdata)
self.sel_a = sel_a
self.sel_r = sel_r
self.axis_rule = axis_rule
self.rcut_r = rcut
# ntypes and rcut_a === -1
self.ntypes = len(self.sel_a)
assert(self.ntypes == len(self.sel_r))
self.rcut_a = -1
# numb of neighbors and numb of descrptors
self.nnei_a = np.cumsum(self.sel_a)[-1]
self.nnei_r = np.cumsum(self.sel_r)[-1]
self.nnei = self.nnei_a + self.nnei_r
self.ndescrpt_a = self.nnei_a * 4
self.ndescrpt_r = self.nnei_r * 1
self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
self.davg = None
self.dstd = None
self.place_holders = {}
avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
sub_graph = tf.Graph()
with sub_graph.as_default():
name_pfx = 'd_lf_'
for ii in ['coord', 'box']:
self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii)
self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type')
self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms')
self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh')
self.stat_descrpt, descrpt_deriv, rij, nlist, axis, rot_mat \
= op_module.descrpt (self.place_holders['coord'],
self.place_holders['type'],
self.place_holders['natoms_vec'],
self.place_holders['box'],
self.place_holders['default_mesh'],
tf.constant(avg_zero),
tf.constant(std_ones),
rcut_a = self.rcut_a,
rcut_r = self.rcut_r,
sel_a = self.sel_a,
sel_r = self.sel_r,
axis_rule = self.axis_rule)
self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config)
def _compute_dstats_sys_nonsmth (self,
data_coord,
data_box,
data_atype,
natoms_vec,
mesh,
reuse = None) :
avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(global_np_float_precision)
std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(global_np_float_precision)
sub_graph = tf.Graph()
with sub_graph.as_default():
descrpt, descrpt_deriv, rij, nlist, axis, rot_mat \
= op_module.descrpt (tf.constant(data_coord),
tf.constant(data_atype),
tf.constant(natoms_vec, dtype = tf.int32),
tf.constant(data_box),
tf.constant(mesh),
tf.constant(avg_zero),
tf.constant(std_ones),
rcut_a = self.rcut_a,
rcut_r = self.rcut_r,
sel_a = self.sel_a,
sel_r = self.sel_r,
axis_rule = self.axis_rule)
# self.sess.run(tf.global_variables_initializer())
# sub_sess = tf.Session(graph = sub_graph,
# config=tf.ConfigProto(intra_op_parallelism_threads=self.run_opt.num_intra_threads,
# inter_op_parallelism_threads=self.run_opt.num_inter_threads
# ))
sub_sess = tf.Session(graph = sub_graph)
dd_all = sub_sess.run(descrpt)
sub_sess.close()
natoms = natoms_vec
dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]])
start_index = 0
sysv = []
sysn = []
sysv2 = []
for type_i in range(self.ntypes):
end_index = start_index + self.ndescrpt * natoms[2+type_i]
dd = dd_all[:, start_index:end_index]
dd = np.reshape(dd, [-1, self.ndescrpt])
start_index = end_index
# compute
sumv = np.sum(dd, axis = 0)
sumn = dd.shape[0]
sumv2 = np.sum(np.multiply(dd,dd), axis = 0)
sysv.append(sumv)
sysn.append(sumn)
sysv2.append(sumv2)
return sysv, sysv2, sysn
def __init__(self, hh, beta):
self.hh = hh
self.beta = beta
with tf.Graph().as_default() as graph:
# place holders
self.t_nloc = tf.placeholder(tf.int32, [1], name="t_nloc")
self.t_coord = tf.placeholder(global_tf_float_precision, [None],
name='t_coord')
self.t_charge = tf.placeholder(global_tf_float_precision, [None],
name='t_charge')
self.t_box = tf.placeholder(global_tf_float_precision, [None],
name='t_box')
# output
self.t_energy, self.t_force, self.t_virial \
= op_module.ewald_recp(self.t_coord, self.t_charge, self.t_nloc, self.t_box,
ewald_h = self.hh,
ewald_beta = self.beta)
self.sess = tf.Session(graph=graph, config=default_tf_session_config)
def load_graph(graph_name: str) -> tf.Graph:
"""Load graph from passed in path.
Parameters
----------
graph_name : str
path to frozen graph on disk
Returns
-------
tf.Graph
tf graph object
"""
graph_def = tf.GraphDef()
with open(graph_name, "rb") as f:
graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def, name="")
return graph
def load_graph_def(model_file: str) -> Tuple[tf.Graph, tf.GraphDef]:
"""
Load graph as well as the graph_def from the frozen model(model_file)
Parameters
----------
model_file : str
The input frozen model path
Returns
-------
tf.Graph
The graph loaded from the frozen model
tf.GraphDef
The graph_def loaded from the frozen model
"""
graph_def = tf.GraphDef()
with open(model_file, "rb") as f:
graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def, name="")
return graph, graph_def
def __init__(self,
rcut: float,
rcut_smth: float,
sel: List[str],
neuron: List[int] = [24, 48, 96],
resnet_dt: bool = False,
trainable: bool = True,
seed: int = None,
type_one_side: bool = True,
exclude_types: List[List[int]] = [],
set_davg_zero: bool = False,
activation_function: str = 'tanh',
precision: str = 'default',
uniform_seed: bool = False) -> None:
"""
Constructor
"""
# args = ClassArg()\
# .add('sel', list, must = True) \
# .add('rcut', float, default = 6.0) \
# .add('rcut_smth',float, default = 0.5) \
# .add('neuron', list, default = [10, 20, 40]) \
# .add('resnet_dt',bool, default = False) \
# .add('trainable',bool, default = True) \
# .add('seed', int) \
# .add('type_one_side', bool, default = False) \
# .add('exclude_types', list, default = []) \
# .add('set_davg_zero', bool, default = False) \
# .add("activation_function", str, default = "tanh") \
# .add("precision", str, default = "default")
# class_data = args.parse(jdata)
self.sel_r = sel
self.rcut = rcut
self.rcut_smth = rcut_smth
self.filter_neuron = neuron
self.filter_resnet_dt = resnet_dt
self.seed = seed
self.uniform_seed = uniform_seed
self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron)
self.trainable = trainable
self.filter_activation_fn = get_activation_func(activation_function)
self.filter_precision = get_precision(precision)
exclude_types = exclude_types
self.exclude_types = set()
for tt in exclude_types:
assert (len(tt) == 2)
self.exclude_types.add((tt[0], tt[1]))
self.exclude_types.add((tt[1], tt[0]))
self.set_davg_zero = set_davg_zero
self.type_one_side = type_one_side
# descrpt config
self.sel_a = [0 for ii in range(len(self.sel_r))]
self.ntypes = len(self.sel_r)
# numb of neighbors and numb of descrptors
self.nnei_a = np.cumsum(self.sel_a)[-1]
self.nnei_r = np.cumsum(self.sel_r)[-1]
self.nnei = self.nnei_a + self.nnei_r
self.ndescrpt_a = self.nnei_a * 4
self.ndescrpt_r = self.nnei_r * 1
self.ndescrpt = self.nnei_r
self.useBN = False
self.davg = None
self.dstd = None
self.embedding_net_variables = None
self.place_holders = {}
avg_zero = np.zeros([self.ntypes,
self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
std_ones = np.ones([self.ntypes,
self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
sub_graph = tf.Graph()
with sub_graph.as_default():
name_pfx = 'd_ser_'
for ii in ['coord', 'box']:
self.place_holders[ii] = tf.placeholder(
GLOBAL_NP_FLOAT_PRECISION, [None, None],
name=name_pfx + 't_' + ii)
self.place_holders['type'] = tf.placeholder(tf.int32, [None, None],
name=name_pfx +
't_type')
self.place_holders['natoms_vec'] = tf.placeholder(
tf.int32, [self.ntypes + 2], name=name_pfx + 't_natoms')
self.place_holders['default_mesh'] = tf.placeholder(
tf.int32, [None], name=name_pfx + 't_mesh')
self.stat_descrpt, descrpt_deriv, rij, nlist \
= op_module.prod_env_mat_r(self.place_holders['coord'],
self.place_holders['type'],
self.place_holders['natoms_vec'],
self.place_holders['box'],
self.place_holders['default_mesh'],
tf.constant(avg_zero),
tf.constant(std_ones),
rcut = self.rcut,
rcut_smth = self.rcut_smth,
sel = self.sel_r)
self.sub_sess = tf.Session(graph=sub_graph,
config=default_tf_session_config)
def _load_sub_graph(self):
sub_graph_def = tf.GraphDef()
with tf.Graph().as_default() as sub_graph:
tf.import_graph_def(sub_graph_def, name="")
return sub_graph, sub_graph_def
def __init__(
self,
op,
rcut: float,
rcut_smth: float,
sel: List[str],
neuron: List[int] = [24, 48, 96],
axis_neuron: int = 8,
resnet_dt: bool = False,
trainable: bool = True,
seed: int = None,
type_one_side: bool = True,
exclude_types: List[List[int]] = [],
set_davg_zero: bool = False,
activation_function: str = 'tanh',
precision: str = 'default',
uniform_seed: bool = False,
) -> None:
DescrptSeA.__init__(self, rcut, rcut_smth, sel, neuron, axis_neuron,
resnet_dt, trainable, seed, type_one_side,
exclude_types, set_davg_zero, activation_function,
precision, uniform_seed)
# DescrptSeA.__init__(self, **jdata)
# args = ClassArg()\
# .add('sel', list, must = True) \
# .add('rcut', float, default = 6.0) \
# .add('rcut_smth',float, default = 5.5) \
# .add('neuron', list, default = [10, 20, 40]) \
# .add('axis_neuron', int, default = 4, alias = 'n_axis_neuron') \
# .add('resnet_dt',bool, default = False) \
# .add('trainable',bool, default = True) \
# .add('seed', int)
# class_data = args.parse(jdata)
# self.sel_a = class_data['sel']
# self.rcut_r = class_data['rcut']
# self.rcut_r_smth = class_data['rcut_smth']
# self.filter_neuron = class_data['neuron']
# self.n_axis_neuron = class_data['axis_neuron']
# self.filter_resnet_dt = class_data['resnet_dt']
# self.seed = class_data['seed']
# self.trainable = class_data['trainable']
self.sel_a = sel
self.rcut_r = rcut
self.rcut_r_smth = rcut_smth
self.filter_neuron = neuron
self.n_axis_neuron = axis_neuron
self.filter_resnet_dt = resnet_dt
self.seed = seed
self.trainable = trainable
self.op = op
# descrpt config
self.sel_r = [0 for ii in range(len(self.sel_a))]
self.ntypes = len(self.sel_a)
assert (self.ntypes == len(self.sel_r))
self.rcut_a = -1
# numb of neighbors and numb of descrptors
self.nnei_a = np.cumsum(self.sel_a)[-1]
self.nnei_r = np.cumsum(self.sel_r)[-1]
self.nnei = self.nnei_a + self.nnei_r
self.ndescrpt_a = self.nnei_a * 4
self.ndescrpt_r = self.nnei_r * 1
self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
self.useBN = False
self.dstd = None
self.davg = None
add_data_requirement('efield',
3,
atomic=True,
must=True,
high_prec=False)
self.place_holders = {}
avg_zero = np.zeros([self.ntypes,
self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
std_ones = np.ones([self.ntypes,
self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
sub_graph = tf.Graph()
with sub_graph.as_default():
name_pfx = 'd_sea_ef_'
for ii in ['coord', 'box']:
self.place_holders[ii] = tf.placeholder(
GLOBAL_NP_FLOAT_PRECISION, [None, None],
name=name_pfx + 't_' + ii)
self.place_holders['type'] = tf.placeholder(tf.int32, [None, None],
name=name_pfx +
't_type')
self.place_holders['natoms_vec'] = tf.placeholder(
tf.int32, [self.ntypes + 2], name=name_pfx + 't_natoms')
self.place_holders['default_mesh'] = tf.placeholder(
tf.int32, [None], name=name_pfx + 't_mesh')
self.place_holders['efield'] = tf.placeholder(
GLOBAL_NP_FLOAT_PRECISION, [None, None],
name=name_pfx + 't_efield')
self.stat_descrpt, descrpt_deriv, rij, nlist \
= self.op(self.place_holders['coord'],
self.place_holders['type'],
self.place_holders['natoms_vec'],
self.place_holders['box'],
self.place_holders['default_mesh'],
self.place_holders['efield'],
tf.constant(avg_zero),
tf.constant(std_ones),
rcut_a = self.rcut_a,
rcut_r = self.rcut_r,
rcut_r_smth = self.rcut_r_smth,
sel_a = self.sel_a,
sel_r = self.sel_r)
self.sub_sess = tf.Session(graph=sub_graph,
config=default_tf_session_config)
def __init__(self,
rcut: float,
rcut_smth: float,
sel: List[str],
neuron: List[int] = [24, 48, 96],
resnet_dt: bool = False,
trainable: bool = True,
seed: int = None,
set_davg_zero: bool = False,
activation_function: str = 'tanh',
precision: str = 'default',
uniform_seed: bool = False) -> None:
"""
Constructor
"""
self.sel_a = sel
self.rcut_r = rcut
self.rcut_r_smth = rcut_smth
self.filter_neuron = neuron
self.filter_resnet_dt = resnet_dt
self.seed = seed
self.uniform_seed = uniform_seed
self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron)
self.trainable = trainable
self.filter_activation_fn = get_activation_func(activation_function)
self.filter_precision = get_precision(precision)
# self.exclude_types = set()
# for tt in exclude_types:
# assert(len(tt) == 2)
# self.exclude_types.add((tt[0], tt[1]))
# self.exclude_types.add((tt[1], tt[0]))
self.set_davg_zero = set_davg_zero
# descrpt config
self.sel_r = [0 for ii in range(len(self.sel_a))]
self.ntypes = len(self.sel_a)
assert (self.ntypes == len(self.sel_r))
self.rcut_a = -1
# numb of neighbors and numb of descrptors
self.nnei_a = np.cumsum(self.sel_a)[-1]
self.nnei_r = np.cumsum(self.sel_r)[-1]
self.nnei = self.nnei_a + self.nnei_r
self.ndescrpt_a = self.nnei_a * 4
self.ndescrpt_r = self.nnei_r * 1
self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
self.useBN = False
self.dstd = None
self.davg = None
self.place_holders = {}
avg_zero = np.zeros([self.ntypes,
self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
std_ones = np.ones([self.ntypes,
self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
sub_graph = tf.Graph()
with sub_graph.as_default():
name_pfx = 'd_sea_'
for ii in ['coord', 'box']:
self.place_holders[ii] = tf.placeholder(
GLOBAL_NP_FLOAT_PRECISION, [None, None],
name=name_pfx + 't_' + ii)
self.place_holders['type'] = tf.placeholder(tf.int32, [None, None],
name=name_pfx +
't_type')
self.place_holders['natoms_vec'] = tf.placeholder(
tf.int32, [self.ntypes + 2], name=name_pfx + 't_natoms')
self.place_holders['default_mesh'] = tf.placeholder(
tf.int32, [None], name=name_pfx + 't_mesh')
self.stat_descrpt, descrpt_deriv, rij, nlist \
= op_module.prod_env_mat_a(self.place_holders['coord'],
self.place_holders['type'],
self.place_holders['natoms_vec'],
self.place_holders['box'],
self.place_holders['default_mesh'],
tf.constant(avg_zero),
tf.constant(std_ones),
rcut_a = self.rcut_a,
rcut_r = self.rcut_r,
rcut_r_smth = self.rcut_r_smth,
sel_a = self.sel_a,
sel_r = self.sel_r)
self.sub_sess = tf.Session(graph=sub_graph,
config=default_tf_session_config)
def __init__(self, jdata):
args = ClassArg()\
.add('sel', list, must = True) \
.add('rcut', float, default = 6.0) \
.add('rcut_smth',float, default = 5.5) \
.add('neuron', list, default = [10, 20, 40]) \
.add('axis_neuron', int, default = 4, alias = 'n_axis_neuron') \
.add('resnet_dt',bool, default = False) \
.add('trainable',bool, default = True) \
.add('seed', int) \
.add('exclude_types', list, default = []) \
.add('set_davg_zero', bool, default = False) \
.add('activation_function', str, default = 'tanh') \
.add('precision', str, default = "default")
class_data = args.parse(jdata)
self.sel_a = class_data['sel']
self.rcut_r = class_data['rcut']
self.rcut_r_smth = class_data['rcut_smth']
self.filter_neuron = class_data['neuron']
self.n_axis_neuron = class_data['axis_neuron']
self.filter_resnet_dt = class_data['resnet_dt']
self.seed = class_data['seed']
self.trainable = class_data['trainable']
self.filter_activation_fn = get_activation_func(
class_data['activation_function'])
self.filter_precision = get_precision(class_data['precision'])
exclude_types = class_data['exclude_types']
self.exclude_types = set()
for tt in exclude_types:
assert (len(tt) == 2)
self.exclude_types.add((tt[0], tt[1]))
self.exclude_types.add((tt[1], tt[0]))
self.set_davg_zero = class_data['set_davg_zero']
# descrpt config
self.sel_r = [0 for ii in range(len(self.sel_a))]
self.ntypes = len(self.sel_a)
assert (self.ntypes == len(self.sel_r))
self.rcut_a = -1
# numb of neighbors and numb of descrptors
self.nnei_a = np.cumsum(self.sel_a)[-1]
self.nnei_r = np.cumsum(self.sel_r)[-1]
self.nnei = self.nnei_a + self.nnei_r
self.ndescrpt_a = self.nnei_a * 4
self.ndescrpt_r = self.nnei_r * 1
self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
self.useBN = False
self.dstd = None
self.davg = None
self.place_holders = {}
avg_zero = np.zeros([self.ntypes,
self.ndescrpt]).astype(global_np_float_precision)
std_ones = np.ones([self.ntypes,
self.ndescrpt]).astype(global_np_float_precision)
sub_graph = tf.Graph()
with sub_graph.as_default():
name_pfx = 'd_sea_'
for ii in ['coord', 'box']:
self.place_holders[ii] = tf.placeholder(
global_np_float_precision, [None, None],
name=name_pfx + 't_' + ii)
self.place_holders['type'] = tf.placeholder(tf.int32, [None, None],
name=name_pfx +
't_type')
self.place_holders['natoms_vec'] = tf.placeholder(
tf.int32, [self.ntypes + 2], name=name_pfx + 't_natoms')
self.place_holders['default_mesh'] = tf.placeholder(
tf.int32, [None], name=name_pfx + 't_mesh')
self.stat_descrpt, descrpt_deriv, rij, nlist \
= op_module.descrpt_se_a(self.place_holders['coord'],
self.place_holders['type'],
self.place_holders['natoms_vec'],
self.place_holders['box'],
self.place_holders['default_mesh'],
tf.constant(avg_zero),
tf.constant(std_ones),
rcut_a = self.rcut_a,
rcut_r = self.rcut_r,
rcut_r_smth = self.rcut_r_smth,
sel_a = self.sel_a,
sel_r = self.sel_r)
self.sub_sess = tf.Session(graph=sub_graph,
config=default_tf_session_config)
def __init__(self,
rcut: float,
rcut_smth: float,
sel: List[str],
neuron: List[int] = [24, 48, 96],
axis_neuron: int = 8,
resnet_dt: bool = False,
trainable: bool = True,
seed: int = None,
type_one_side: bool = True,
exclude_types: List[List[int]] = [],
set_davg_zero: bool = False,
activation_function: str = 'tanh',
precision: str = 'default',
uniform_seed: bool = False) -> None:
"""
Constructor
"""
if rcut < rcut_smth:
raise RuntimeError(
"rcut_smth (%f) should be no more than rcut (%f)!" %
(rcut_smth, rcut))
self.sel_a = sel
self.rcut_r = rcut
self.rcut_r_smth = rcut_smth
self.filter_neuron = neuron
self.n_axis_neuron = axis_neuron
self.filter_resnet_dt = resnet_dt
self.seed = seed
self.uniform_seed = uniform_seed
self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron)
self.trainable = trainable
self.compress_activation_fn = get_activation_func(activation_function)
self.filter_activation_fn = get_activation_func(activation_function)
self.filter_precision = get_precision(precision)
self.exclude_types = set()
for tt in exclude_types:
assert (len(tt) == 2)
self.exclude_types.add((tt[0], tt[1]))
self.exclude_types.add((tt[1], tt[0]))
self.set_davg_zero = set_davg_zero
self.type_one_side = type_one_side
# descrpt config
self.sel_r = [0 for ii in range(len(self.sel_a))]
self.ntypes = len(self.sel_a)
assert (self.ntypes == len(self.sel_r))
self.rcut_a = -1
# numb of neighbors and numb of descrptors
self.nnei_a = np.cumsum(self.sel_a)[-1]
self.nnei_r = np.cumsum(self.sel_r)[-1]
self.nnei = self.nnei_a + self.nnei_r
self.ndescrpt_a = self.nnei_a * 4
self.ndescrpt_r = self.nnei_r * 1
self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r
self.useBN = False
self.dstd = None
self.davg = None
self.compress = False
self.embedding_net_variables = None
self.mixed_prec = None
self.place_holders = {}
nei_type = np.array([])
for ii in range(self.ntypes):
nei_type = np.append(nei_type,
ii * np.ones(self.sel_a[ii])) # like a mask
self.nei_type = tf.constant(nei_type, dtype=tf.int32)
avg_zero = np.zeros([self.ntypes,
self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
std_ones = np.ones([self.ntypes,
self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION)
sub_graph = tf.Graph()
with sub_graph.as_default():
name_pfx = 'd_sea_'
for ii in ['coord', 'box']:
self.place_holders[ii] = tf.placeholder(
GLOBAL_NP_FLOAT_PRECISION, [None, None],
name=name_pfx + 't_' + ii)
self.place_holders['type'] = tf.placeholder(tf.int32, [None, None],
name=name_pfx +
't_type')
self.place_holders['natoms_vec'] = tf.placeholder(
tf.int32, [self.ntypes + 2], name=name_pfx + 't_natoms')
self.place_holders['default_mesh'] = tf.placeholder(
tf.int32, [None], name=name_pfx + 't_mesh')
self.stat_descrpt, descrpt_deriv, rij, nlist \
= op_module.prod_env_mat_a(self.place_holders['coord'],
self.place_holders['type'],
self.place_holders['natoms_vec'],
self.place_holders['box'],
self.place_holders['default_mesh'],
tf.constant(avg_zero),
tf.constant(std_ones),
rcut_a = self.rcut_a,
rcut_r = self.rcut_r,
rcut_r_smth = self.rcut_r_smth,
sel_a = self.sel_a,
sel_r = self.sel_r)
self.sub_sess = tf.Session(graph=sub_graph,
config=default_tf_session_config)
self.original_sel = None
