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