def build(self, coord_, atype_, natoms, box_, mesh, davg=None, dstd=None, suffix='', reuse=None): with tf.variable_scope('descrpt_attr' + suffix, reuse=reuse): if davg is None: davg = np.zeros([self.ntypes, self.ndescrpt]) if dstd is None: dstd = np.ones([self.ntypes, self.ndescrpt]) t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), name='rcut', dtype=global_tf_float_precision) t_ntypes = tf.constant(self.ntypes, name='ntypes', dtype=tf.int32) self.t_avg = tf.get_variable( 't_avg', davg.shape, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(davg)) self.t_std = tf.get_variable( 't_std', dstd.shape, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(dstd)) coord = tf.reshape(coord_, [-1, natoms[1] * 3]) box = tf.reshape(box_, [-1, 9]) atype = tf.reshape(atype_, [-1, natoms[1]]) self.descrpt, self.descrpt_deriv, self.rij, self.nlist \ = op_module.descrpt_se_a (coord, atype, natoms, box, mesh, self.t_avg, self.t_std, 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.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, natoms, suffix=suffix, reuse=reuse, trainable=self.trainable) return self.dout
def init_variables(self, graph : tf.Graph, graph_def : tf.GraphDef, model_type : str = "original_model", suffix : str = "", ) -> None: """ Init the embedding net variables with the given frozen model Parameters ---------- graph : tf.Graph The input frozen model graph graph_def : tf.GraphDef The input frozen model graph_def model_type : str the type of the model suffix : str suffix to name scope """ if model_type == 'original_model': self.descrpt.init_variables(graph, graph_def, suffix=suffix) self.fitting.init_variables(graph, graph_def, suffix=suffix) tf.constant("original_model", name = 'model_type', dtype = tf.string) elif model_type == 'compressed_model': self.fitting.init_variables(graph, graph_def, suffix=suffix) tf.constant("compressed_model", name = 'model_type', dtype = tf.string) else: raise RuntimeError("Unknown model type %s" % model_type)
def setUp(self, data, pbc=True, sess=None): self.sess = sess self.data = data self.natoms = self.data.get_natoms() self.ntypes = self.data.get_ntypes() self.sel = [12, 24] self.sel_a = [0, 0] self.rcut_smth = 2.45 self.rcut = 10.0 self.nnei = np.cumsum(self.sel)[-1] self.ndescrpt = self.nnei * 1 davg = np.zeros([self.ntypes, self.ndescrpt]) dstd = np.ones([self.ntypes, self.ndescrpt]) self.t_avg = tf.constant(davg.astype(GLOBAL_NP_FLOAT_PRECISION)) self.t_std = tf.constant(dstd.astype(GLOBAL_NP_FLOAT_PRECISION)) if pbc: self.default_mesh = np.zeros(6, dtype=np.int32) self.default_mesh[3] = 2 self.default_mesh[4] = 2 self.default_mesh[5] = 2 else: self.default_mesh = np.array([], dtype=np.int32) # 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") self.efield = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, self.natoms[0] * 3], name='t_efield')
def _init_from_frz_model(self): # get the model type from the frozen model(self.run_opt.init_frz_model) try: t_model_type = get_tensor_by_name(self.run_opt.init_frz_model, 'model_type') self.model_type = bytes.decode(t_model_type) except GraphWithoutTensorError as e: # throw runtime error if there's no frozen model if not os.path.exists(self.run_opt.init_frz_model): raise RuntimeError( "The input frozen model %s (%s) does not exist! Please check the path of the frozen model. " % (self.run_opt.init_frz_model, os.path.abspath(self.run_opt.init_frz_model)) ) from e # throw runtime error if the frozen_model has no model type information... else: raise RuntimeError( "The input frozen model: %s has no 'model_type' information, " "which is not supported by the 'dp train init-frz-model' interface. " % self.run_opt.init_frz_model ) from e if self.fitting_type != 'ener': raise RuntimeError("The 'dp train init-frz-model' command only supports the 'ener' type fitting net currently!") # self.frz_model will control the self.model to import the descriptor from the given frozen model instead of building from scratch... # initialize fitting net with the given compressed frozen model if self.model_type == 'original_model': self.descrpt.init_variables(self.run_opt.init_frz_model) self.fitting.init_variables(get_fitting_net_variables(self.run_opt.init_frz_model)) tf.constant("original_model", name = 'model_type', dtype = tf.string) elif self.model_type == 'compressed_model': self.frz_model = self.run_opt.init_frz_model self.fitting.init_variables(get_fitting_net_variables(self.frz_model)) tf.constant("compressed_model", name = 'model_type', dtype = tf.string) else: raise RuntimeError("Unknown model type %s" % self.model_type)
def __init__(self, data, comp=0): self.sess = tf.Session() self.data = data self.natoms = self.data.get_natoms() self.ntypes = self.data.get_ntypes() self.sel_a = [12, 24] self.sel_r = [12, 24] self.rcut_a = -1 self.rcut_r = 10.0 self.axis_rule = [0, 1, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0] 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 davg = np.zeros([self.ntypes, self.ndescrpt]) dstd = np.ones([self.ntypes, self.ndescrpt]) self.t_avg = tf.constant(davg.astype(global_np_float_precision)) self.t_std = tf.constant(dstd.astype(global_np_float_precision)) 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 get_nbor_stat(jdata, rcut): max_rcut = get_rcut(jdata) type_map = get_type_map(jdata) if type_map and len(type_map) == 0: type_map = None train_data = get_data(jdata["training"]["training_data"], max_rcut, type_map, None) train_data.get_batch() data_ntypes = train_data.get_ntypes() if type_map is not None: map_ntypes = len(type_map) else: map_ntypes = data_ntypes ntypes = max([map_ntypes, data_ntypes]) neistat = NeighborStat(ntypes, rcut) min_nbor_dist, max_nbor_size = neistat.get_stat(train_data) # moved from traier.py as duplicated # TODO: this is a simple fix but we should have a clear # architecture to call neighbor stat tf.constant(min_nbor_dist, name = 'train_attr/min_nbor_dist', dtype = GLOBAL_ENER_FLOAT_PRECISION) tf.constant(max_nbor_size, name = 'train_attr/max_nbor_size', dtype = tf.int32) return min_nbor_dist, max_nbor_size
def init_variables( self, graph: tf.Graph, graph_def: tf.GraphDef, model_type: str = "original_model", suffix: str = "", ) -> None: """ Init the embedding net variables with the given frozen model Parameters ---------- graph : tf.Graph The input frozen model graph graph_def : tf.GraphDef The input frozen model graph_def model_type : str the type of the model suffix : str suffix to name scope """ # self.frz_model will control the self.model to import the descriptor from the given frozen model instead of building from scratch... # initialize fitting net with the given compressed frozen model if model_type == 'original_model': self.descrpt.init_variables(graph, graph_def, suffix=suffix) self.fitting.init_variables(graph, graph_def, suffix=suffix) tf.constant("original_model", name='model_type', dtype=tf.string) elif model_type == 'compressed_model': self.fitting.init_variables(graph, graph_def, suffix=suffix) tf.constant("compressed_model", name='model_type', dtype=tf.string) else: raise RuntimeError("Unknown model type %s" % model_type) if self.typeebd is not None: self.typeebd.init_variables(graph, graph_def, suffix=suffix)
def test_op_tanh(self): w = tf.constant( [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1, 1.1, 1.2]], dtype='double') x = tf.constant([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [1.0, 1.1, 1.2]], dtype='double') b = tf.constant([[0.1], [0.2], [0.3], [0.4]], dtype='double') xbar = tf.matmul(x, w) + b y = tf.nn.tanh(xbar) dy = op_module.unaggregated_dy_dx_s(y, w, xbar, tf.constant(1)) dy_array = tf.Session().run(dy) answer = np.array( [[ 8.008666403121351973e-02, 1.513925729426658651e-01, 2.134733287761668430e-01, 2.661983049806041501e-01 ], [ 4.010658815015744061e-02, 6.306476628799793926e-02, 7.332167904608145881e-02, 7.494218676568849269e-02 ], [ 1.561705624394135218e-02, 1.994112926507514427e-02, 1.887519955881525671e-02, 1.576442161040989692e-02 ], [ 5.492686739421748753e-03, 5.754985286040992763e-03, 4.493113544969218158e-03, 3.107638130764600777e-03 ]]) places = 18 np.testing.assert_almost_equal(dy_array, answer, places)
def test_op_gelu(self): w = tf.constant( [[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [0.9, 1, 1.1, 1.2]], dtype='double') x = tf.constant([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9], [1.0, 1.1, 1.2]], dtype='double') b = tf.constant([[0.1], [0.2], [0.3], [0.4]], dtype='double') xbar = tf.matmul(x, w) + b y = gelu(xbar) dy = op_module.unaggregated_dy_dx_s(y, w, xbar, tf.constant(2)) dy_array = tf.Session().run(dy) answer = np.array( [[ 8.549286163555620821e-02, 1.782905778685600906e-01, 2.776474599997448833e-01, 3.827650237273348965e-01 ], [ 1.089906023807040714e-01, 2.230820937721638697e-01, 3.381867859682909927e-01, 4.513008399758057232e-01 ], [ 1.124254240556722684e-01, 2.209918074710395253e-01, 3.238894323148118759e-01, 4.220357318198978414e-01 ], [ 1.072173273655498138e-01, 2.082159073100979807e-01, 3.059816075270163083e-01, 4.032981557798429595e-01 ]]) places = 18 np.testing.assert_almost_equal(dy_array, answer, places)
def setUp(self, data, pbc=True): self.sess = tf.Session() self.data = data self.natoms = self.data.get_natoms() self.ntypes = self.data.get_ntypes() self.sel = [12, 24] self.sel_a = [0, 0] self.rcut_smth = 2.45 self.rcut = 10.0 self.nnei = np.cumsum(self.sel)[-1] self.ndescrpt = self.nnei * 1 davg = np.zeros([self.ntypes, self.ndescrpt]) dstd = np.ones([self.ntypes, self.ndescrpt]) self.t_avg = tf.constant(davg.astype(global_np_float_precision)) self.t_std = tf.constant(dstd.astype(global_np_float_precision)) if pbc: self.default_mesh = np.zeros(6, dtype=np.int32) self.default_mesh[3] = 2 self.default_mesh[4] = 2 self.default_mesh[5] = 2 else: self.default_mesh = np.array([], dtype=np.int32) # 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 setUp(self): self.sess = self.test_session().__enter__() self.natoms = [5, 5, 2, 3] self.ntypes = 2 self.sel_a = [12, 24] self.sel_r = [0, 0] self.rcut_a = -1 self.rcut_r_smth = 2.45 self.rcut_r = 10.0 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 davg = np.zeros([self.ntypes, self.ndescrpt]) dstd = np.ones([self.ntypes, self.ndescrpt]) self.t_avg = tf.constant(davg.astype(GLOBAL_NP_FLOAT_PRECISION)) self.t_std = tf.constant(dstd.astype(GLOBAL_NP_FLOAT_PRECISION)) # no pbc self.default_mesh = np.array([], dtype=np.int32) # make place holder self.coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, self.natoms[0] * 3], name='t_coord') self.efield = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, self.natoms[0] * 3], name='t_efield') 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__(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 build (self, coord_ : tf.Tensor, atype_ : tf.Tensor, natoms : tf.Tensor, box_ : tf.Tensor, mesh : tf.Tensor, input_dict : dict, reuse : bool = None, suffix : str = '' ) -> tf.Tensor: """ Build the computational graph for the descriptor Parameters ---------- coord_ The coordinate of atoms atype_ The type of atoms natoms The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms mesh For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed. input_dict Dictionary for additional inputs reuse The weights in the networks should be reused when get the variable. suffix Name suffix to identify this descriptor Returns ------- descriptor The output descriptor """ with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) : t_rcut = tf.constant(self.get_rcut(), name = 'rcut', dtype = GLOBAL_TF_FLOAT_PRECISION) t_ntypes = tf.constant(self.get_ntypes(), name = 'ntypes', dtype = tf.int32) all_dout = [] for idx,ii in enumerate(self.descrpt_list): dout = ii.build(coord_, atype_, natoms, box_, mesh, input_dict, suffix=suffix+f'_{idx}', reuse=reuse) dout = tf.reshape(dout, [-1, ii.get_dim_out()]) all_dout.append(dout) dout = tf.concat(all_dout, axis = 1) dout = tf.reshape(dout, [-1, natoms[0] * self.get_dim_out()]) return dout
def build(self, data=None, stop_batch=0): self.ntypes = self.model.get_ntypes() self.stop_batch = stop_batch if self.numb_fparam > 0: log.info("training with %d frame parameter(s)" % self.numb_fparam) else: log.info("training without frame parameter") if not self.is_compress: # Usually, the type number of the model should be equal to that of the data # However, nt_model > nt_data should be allowed, since users may only want to # train using a dataset that only have some of elements if self.ntypes < data.get_ntypes(): raise ValueError( "The number of types of the training data is %d, but that of the " "model is only %d. The latter must be no less than the former. " "You may need to reset one or both of them. Usually, the former " "is given by `model/type_map` in the training parameter (if set) " "or the maximum number in the training data. The latter is given " "by `model/descriptor/sel` in the training parameter." % (data.get_ntypes(), self.ntypes)) self.type_map = data.get_type_map() self.batch_size = data.get_batch_size() self.model.data_stat(data) # config the init_frz_model command if self.run_opt.init_mode == 'init_from_frz_model': self._init_from_frz_model() # neighbor_stat is moved to train.py as duplicated # TODO: this is a simple fix but we should have a clear # architecture to call neighbor stat else: self.descrpt.enable_compression( self.model_param['compress']["min_nbor_dist"], self.model_param['compress']['model_file'], self.model_param['compress']['table_config'][0], self.model_param['compress']['table_config'][1], self.model_param['compress']['table_config'][2], self.model_param['compress']['table_config'][3]) self.fitting.init_variables( get_fitting_net_variables( self.model_param['compress']['model_file'])) if self.is_compress or self.model_type == 'compressed_model': tf.constant("compressed_model", name='model_type', dtype=tf.string) else: tf.constant("original_model", name='model_type', dtype=tf.string) self._build_lr() self._build_network(data) self._build_training()
def test_embed_atom_type(self): ntypes = 3 natoms = tf.constant([5, 5, 3, 0, 2]) type_embedding = tf.constant([ [1, 2, 3], [3, 2, 1], [7, 7, 7], ]) expected_out = [[1, 2, 3], [1, 2, 3], [1, 2, 3], [7, 7, 7], [7, 7, 7]] atom_embed = embed_atom_type(ntypes, natoms, type_embedding) sess = self.test_session().__enter__() atom_embed = sess.run(atom_embed) np.testing.assert_almost_equal(atom_embed, expected_out, 10)
def __init__(self, jdata, descrpt): # model param self.ntypes = descrpt.get_ntypes() self.dim_descrpt = descrpt.get_dim_out() args = ClassArg()\ .add('numb_fparam', int, default = 0)\ .add('numb_aparam', int, default = 0)\ .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\ .add('resnet_dt', bool, default = True)\ .add('rcond', float, default = 1e-3) \ .add('seed', int) \ .add('atom_ener', list, default = [])\ .add("activation_function", str, default = "tanh")\ .add("precision", str, default = "default")\ .add("trainable", [list, bool], default = True) class_data = args.parse(jdata) self.numb_fparam = class_data['numb_fparam'] self.numb_aparam = class_data['numb_aparam'] self.n_neuron = class_data['neuron'] self.resnet_dt = class_data['resnet_dt'] self.rcond = class_data['rcond'] self.seed = class_data['seed'] self.fitting_activation_fn = get_activation_func( class_data["activation_function"]) self.fitting_precision = get_precision(class_data['precision']) self.trainable = class_data['trainable'] if type(self.trainable) is bool: self.trainable = [self.trainable] * (len(self.n_neuron) + 1) assert (len(self.trainable) == len(self.n_neuron) + 1), 'length of trainable should be that of n_neuron + 1' self.atom_ener = [] for at, ae in enumerate(class_data['atom_ener']): if ae is not None: self.atom_ener.append( tf.constant(ae, global_tf_float_precision, name="atom_%d_ener" % at)) else: self.atom_ener.append(None) self.useBN = False self.bias_atom_e = None # data requirement if self.numb_fparam > 0: add_data_requirement('fparam', self.numb_fparam, atomic=False, must=True, high_prec=False) self.fparam_avg = None self.fparam_std = None self.fparam_inv_std = None if self.numb_aparam > 0: add_data_requirement('aparam', self.numb_aparam, atomic=True, must=True, high_prec=False) self.aparam_avg = None self.aparam_std = None self.aparam_inv_std = None
def comp_ef(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): descrpt, descrpt_deriv, rij, nlist \ = op_module.descrpt_se_r (dcoord, dtype, tnatoms, dbox, tf.constant(self.default_mesh), self.t_avg, self.t_std, rcut = self.rcut, rcut_smth = self.rcut_smth, sel = self.sel) inputs_reshape = tf.reshape(descrpt, [-1, self.ndescrpt]) 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) net_deriv_ = tf.gradients(atom_ener, inputs_reshape) net_deriv = net_deriv_[0] net_deriv_reshape = tf.reshape(net_deriv, [-1, self.natoms[0] * self.ndescrpt]) force = op_module.prod_force_se_r(net_deriv_reshape, descrpt_deriv, nlist, tnatoms) virial, atom_vir = op_module.prod_virial_se_r(net_deriv_reshape, descrpt_deriv, rij, nlist, tnatoms) return energy, force, virial
def test_nopbc_self_built_nlist(self): tem, tem_deriv, trij, tnlist \ = op_module.prod_env_mat_a ( self.tcoord, self.ttype, self.tnatoms, self.tbox, tf.constant(np.zeros(0, dtype = np.int32)), self.t_avg, self.t_std, rcut_a = -1, rcut_r = self.rcut, rcut_r_smth = self.rcut_smth, sel_a = self.sel, sel_r = [0, 0]) self.sess.run(tf.global_variables_initializer()) dem, dem_deriv, drij, dnlist = self.sess.run( [tem, tem_deriv, trij, tnlist], feed_dict={ self.tcoord: self.dcoord, self.ttype: self.dtype, self.tbox: self.dbox, self.tnatoms: self.dnatoms }) self.assertEqual(dem.shape, (self.nframes, self.nloc * self.ndescrpt)) self.assertEqual(dem_deriv.shape, (self.nframes, self.nloc * self.ndescrpt * 3)) self.assertEqual(drij.shape, (self.nframes, self.nloc * self.nnei * 3)) self.assertEqual(dnlist.shape, (self.nframes, self.nloc * self.nnei)) for ff in range(self.nframes): np.testing.assert_almost_equal(dem[ff], self.nopbc_expected_output, 5)
def build(self, coord_, atype_, natoms, box, mesh, input_dict, suffix='', reuse=None): with tf.variable_scope('model_attr' + suffix, reuse=reuse): t_tmap = tf.constant(' '.join(self.type_map), name='tmap', dtype=tf.string) t_st = tf.constant(self.get_sel_type(), name='sel_type', dtype=tf.int32) t_mt = tf.constant(self.model_type, name='model_type', dtype=tf.string) coord = tf.reshape(coord_, [-1, natoms[1] * 3]) atype = tf.reshape(atype_, [-1, natoms[1]]) dout \ = self.descrpt.build(coord_, atype_, natoms, box, mesh, davg = self.davg, dstd = self.dstd, suffix = suffix, reuse = reuse) dout = tf.identity(dout, name='o_descriptor') rot_mat = self.descrpt.get_rot_mat() rot_mat = tf.identity(rot_mat, name='o_rot_mat') polar = self.fitting.build(dout, rot_mat, natoms, reuse=reuse, suffix=suffix) polar = tf.identity(polar, name='o_polar') return {'polar': polar}
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 _make_data(self, xx, idx): with self.sub_graph.as_default(): with self.sub_sess.as_default(): xx = tf.reshape(xx, [xx.size, -1]) for layer in range(self.layer_size): if layer == 0: xbar = tf.matmul( xx, self.matrix["layer_" + str(layer + 1)] [idx]) + self.bias["layer_" + str(layer + 1)][idx] yy = self._layer_0( xx, self.matrix["layer_" + str(layer + 1)][idx], self.bias["layer_" + str(layer + 1)][idx]) dy = op_module.unaggregated_dy_dx_s( yy, self.matrix["layer_" + str(layer + 1)][idx], xbar, tf.constant(self.functype)) dy2 = op_module.unaggregated_dy2_dx_s( yy, dy, self.matrix["layer_" + str(layer + 1)][idx], xbar, tf.constant(self.functype)) else: ybar = tf.matmul( yy, self.matrix["layer_" + str(layer + 1)] [idx]) + self.bias["layer_" + str(layer + 1)][idx] tt, zz = self._layer_1( yy, self.matrix["layer_" + str(layer + 1)][idx], self.bias["layer_" + str(layer + 1)][idx]) dz = op_module.unaggregated_dy_dx( zz - tt, self.matrix["layer_" + str(layer + 1)][idx], dy, ybar, tf.constant(self.functype)) dy2 = op_module.unaggregated_dy2_dx( zz - tt, self.matrix["layer_" + str(layer + 1)][idx], dy, dy2, ybar, tf.constant(self.functype)) dy = dz yy = zz vv = zz.eval() dd = dy.eval() d2 = dy2.eval() return vv, dd, d2
def build_fv_graph(self) -> tf.Tensor: """ Build the computational graph for the force and virial inference. """ with tf.variable_scope('modifier_attr'): t_mdl_name = tf.constant(self.model_name, name='mdl_name', dtype=tf.string) t_modi_type = tf.constant(self.modifier_prefix, name='type', dtype=tf.string) t_mdl_charge_map = tf.constant(' '.join( [str(ii) for ii in self.model_charge_map]), name='mdl_charge_map', dtype=tf.string) t_sys_charge_map = tf.constant(' '.join( [str(ii) for ii in self.sys_charge_map]), name='sys_charge_map', dtype=tf.string) t_ewald_h = tf.constant(self.ewald_h, name='ewald_h', dtype=tf.float64) t_ewald_b = tf.constant(self.ewald_beta, name='ewald_beta', dtype=tf.float64) with self.graph.as_default(): return self._build_fv_graph_inner()
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 test_nopbc_self_built_nlist_deriv(self): hh = 1e-4 tem, tem_deriv, trij, tnlist \ = op_module.prod_env_mat_a ( self.tcoord, self.ttype, self.tnatoms, self.tbox, tf.constant(np.zeros(0, dtype = np.int32)), self.t_avg, self.t_std, rcut_a = -1, rcut_r = self.rcut, rcut_r_smth = self.rcut_smth, sel_a = self.sel, sel_r = [0, 0]) self.sess.run(tf.global_variables_initializer()) self.check_deriv_numerical_deriv(hh, tem, tem_deriv, trij, tnlist)
def comp_ef(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): t_default_mesh = tf.constant(self.default_mesh) descrpt, descrpt_deriv, rij, nlist, axis, rot_mat \ = op_module.descrpt (dcoord, dtype, tnatoms, dbox, t_default_mesh, self.t_avg, self.t_std, 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.axis = axis self.nlist = nlist self.descrpt = descrpt inputs_reshape = tf.reshape(descrpt, [-1, self.ndescrpt]) 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) net_deriv_ = tf.gradients(atom_ener, inputs_reshape) net_deriv = net_deriv_[0] net_deriv_reshape = tf.reshape(net_deriv, [-1, self.natoms[0] * self.ndescrpt]) force = op_module.prod_force(net_deriv_reshape, descrpt_deriv, nlist, axis, tnatoms, n_a_sel=self.nnei_a, n_r_sel=self.nnei_r) virial, atom_vir = op_module.prod_virial(net_deriv_reshape, descrpt_deriv, rij, nlist, axis, tnatoms, n_a_sel=self.nnei_a, n_r_sel=self.nnei_r) return energy, force, virial
def build_efv(self, dcoord, dbox, dtype, tnatoms, name, op, reuse=None): efield = tf.reshape(self.efield, [-1, 3]) efield = self._normalize_3d(efield) efield = tf.reshape(efield, [-1, tnatoms[0] * 3]) if op != op_module.prod_env_mat_a: descrpt = DescrptSeAEfLower( op, **{ 'sel': self.sel_a, 'rcut': 6, 'rcut_smth': 5.5, 'seed': 1, 'uniform_seed': True }) else: descrpt = DescrptSeA( **{ 'sel': self.sel_a, 'rcut': 6, 'rcut_smth': 0.5, 'seed': 1, 'uniform_seed': True }) dout = descrpt.build(dcoord, dtype, tnatoms, dbox, tf.constant(self.default_mesh), {'efield': efield}, suffix=name, reuse=reuse) dout = tf.reshape(dout, [-1, descrpt.get_dim_out()]) atom_ener = tf.reduce_sum(dout, axis=1) atom_ener_reshape = tf.reshape(atom_ener, [-1, self.natoms[0]]) energy = tf.reduce_sum(atom_ener_reshape, axis=1) force, virial, atom_vir \ = descrpt.prod_force_virial (atom_ener, tnatoms) return energy, force, virial, atom_ener, atom_vir
def build_fv_graph(self): with tf.variable_scope('modifier_attr') : t_mdl_name = tf.constant(self.model_name, name = 'mdl_name', dtype = tf.string) t_modi_type = tf.constant(self.modifier_prefix, name = 'type', dtype = tf.string) t_mdl_charge_map = tf.constant(' '.join([str(ii) for ii in self.model_charge_map]), name = 'mdl_charge_map', dtype = tf.string) t_sys_charge_map = tf.constant(' '.join([str(ii) for ii in self.sys_charge_map]), name = 'sys_charge_map', dtype = tf.string) t_ewald_h = tf.constant(self.ewald_h, name = 'ewald_h', dtype = tf.float64) t_ewald_b = tf.constant(self.ewald_beta, name = 'ewald_beta', dtype = tf.float64) with self.graph.as_default(): return self._build_fv_graph_inner()
def build(self, coord_, atype_, natoms, box, mesh, input_dict, suffix='', reuse=None): with tf.variable_scope('model_attr' + suffix, reuse=reuse): t_tmap = tf.constant(' '.join(self.type_map), name='tmap', dtype=tf.string) t_mt = tf.constant(self.model_type, name='model_type', dtype=tf.string) if self.srtab is not None: tab_info, tab_data = self.srtab.get() self.tab_info = tf.get_variable( 't_tab_info', tab_info.shape, dtype=tf.float64, trainable=False, initializer=tf.constant_initializer(tab_info, dtype=tf.float64)) self.tab_data = tf.get_variable( 't_tab_data', tab_data.shape, dtype=tf.float64, trainable=False, initializer=tf.constant_initializer(tab_data, dtype=tf.float64)) coord = tf.reshape(coord_, [-1, natoms[1] * 3]) atype = tf.reshape(atype_, [-1, natoms[1]]) dout \ = self.descrpt.build(coord_, atype_, natoms, box, mesh, davg = self.davg, dstd = self.dstd, suffix = suffix, reuse = reuse) dout = tf.identity(dout, name='o_descriptor') if self.srtab is not None: nlist, rij, sel_a, sel_r = self.descrpt.get_nlist() nnei_a = np.cumsum(sel_a)[-1] nnei_r = np.cumsum(sel_r)[-1] atom_ener = self.fitting.build(dout, input_dict, natoms, bias_atom_e=self.bias_atom_e, reuse=reuse, suffix=suffix) if self.srtab is not None: sw_lambda, sw_deriv \ = op_module.soft_min_switch(atype, rij, nlist, natoms, sel_a = sel_a, sel_r = sel_r, alpha = self.smin_alpha, rmin = self.sw_rmin, rmax = self.sw_rmax) inv_sw_lambda = 1.0 - sw_lambda # NOTICE: # atom energy is not scaled, # force and virial are scaled tab_atom_ener, tab_force, tab_atom_virial \ = op_module.tab_inter(self.tab_info, self.tab_data, atype, rij, nlist, natoms, sw_lambda, sel_a = sel_a, sel_r = sel_r) energy_diff = tab_atom_ener - tf.reshape(atom_ener, [-1, natoms[0]]) tab_atom_ener = tf.reshape(sw_lambda, [-1]) * tf.reshape( tab_atom_ener, [-1]) atom_ener = tf.reshape(inv_sw_lambda, [-1]) * atom_ener energy_raw = tab_atom_ener + atom_ener else: energy_raw = atom_ener energy_raw = tf.reshape(energy_raw, [-1, natoms[0]], name='o_atom_energy' + suffix) energy = tf.reduce_sum(global_cvt_2_ener_float(energy_raw), axis=1, name='o_energy' + suffix) force, virial, atom_virial \ = self.descrpt.prod_force_virial (atom_ener, natoms) if self.srtab is not None: sw_force \ = op_module.soft_min_force(energy_diff, sw_deriv, nlist, natoms, n_a_sel = nnei_a, n_r_sel = nnei_r) force = force + sw_force + tab_force force = tf.reshape(force, [-1, 3 * natoms[1]], name="o_force" + suffix) if self.srtab is not None: sw_virial, sw_atom_virial \ = op_module.soft_min_virial (energy_diff, sw_deriv, rij, nlist, natoms, n_a_sel = nnei_a, n_r_sel = nnei_r) atom_virial = atom_virial + sw_atom_virial + tab_atom_virial virial = virial + sw_virial \ + tf.reduce_sum(tf.reshape(tab_atom_virial, [-1, natoms[1], 9]), axis = 1) virial = tf.reshape(virial, [-1, 9], name="o_virial" + suffix) atom_virial = tf.reshape(atom_virial, [-1, 9 * natoms[1]], name="o_atom_virial" + suffix) model_dict = {} model_dict['energy'] = energy model_dict['force'] = force model_dict['virial'] = virial model_dict['atom_ener'] = energy_raw model_dict['atom_virial'] = atom_virial return model_dict
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 setUp(self): self.sess = self.test_session().__enter__() self.inputs = tf.constant([0., 1., 2.], dtype=tf.float64) self.ndata = 3 self.inputs = tf.reshape(self.inputs, [-1, 1]) self.places = 6