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 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 """ nei_type = np.array([]) for ii in range(self.ntypes): nei_type = np.append(nei_type, ii * np.ones(self.sel_a[ii])) self.nei_type = tf.get_variable( 't_nei_type', [self.nnei], dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(nei_type)) self.dout = DescrptSeA.build(self, coord_, atype_, natoms, box_, mesh, input_dict, suffix=suffix, reuse=reuse) tf.summary.histogram('embedding_net_output', self.dout) return self.dout
def _net(self, inputs, name, reuse=False): with tf.variable_scope(name, reuse=reuse): net_w = tf.get_variable('net_w', [self.ndescrpt], GLOBAL_TF_FLOAT_PRECISION, tf.constant_initializer(self.net_w_i)) dot_v = tf.matmul(tf.reshape(inputs, [-1, self.ndescrpt]), tf.reshape(net_w, [self.ndescrpt, 1])) return tf.reshape(dot_v, [-1])
def _net(self, inputs, name, reuse=False): with tf.variable_scope(name, reuse=reuse): net_w = tf.get_variable('net_w', [self.ndescrpt], global_tf_float_precision, tf.constant_initializer(self.net_w_i)) dot_v = tf.matmul(tf.reshape(inputs, [-1, self.ndescrpt]), tf.reshape(net_w, [self.ndescrpt, 1])) return tf.reshape(dot_v, [-1])
def comp_v_dw(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): energy, force, virial = self.comp_ef(dcoord, dbox, dtype, tnatoms, name, reuse) with tf.variable_scope(name, reuse=True): net_w = tf.get_variable('net_w', [self.ndescrpt], GLOBAL_TF_FLOAT_PRECISION, tf.constant_initializer(self.net_w_i)) v_mag = tf.reduce_sum(virial) v_mag_dw = tf.gradients(v_mag, net_w) assert (len(v_mag_dw) == 1), "length of dw is wrong" return v_mag, v_mag_dw[0]
def comp_f_dw(self, dcoord, dbox, dtype, tnatoms, name, reuse=None): energy, force, virial = self.comp_ef(dcoord, dbox, dtype, tnatoms, name, reuse) with tf.variable_scope(name, reuse=True): net_w = tf.get_variable('net_w', [self.ndescrpt], global_tf_float_precision, tf.constant_initializer(self.net_w_i)) f_mag = tf.reduce_sum(tf.nn.tanh(force)) f_mag_dw = tf.gradients(f_mag, net_w) assert (len(f_mag_dw) == 1), "length of dw is wrong" return f_mag, f_mag_dw[0]
def __init__(self, data): # tabulated Inter.__init__(self, data) _make_tab(data.get_ntypes()) self.srtab = TabInter('tab.xvg') self.smin_alpha = 0.3 self.sw_rmin = 1 self.sw_rmax = 3.45 tab_info, tab_data = self.srtab.get() with tf.variable_scope('tab', reuse=tf.AUTO_REUSE): self.tab_info = tf.get_variable( 't_tab_info', tab_info.shape, dtype=tf.float64, trainable=False, initializer=tf.constant_initializer(tab_info)) self.tab_data = tf.get_variable( 't_tab_data', tab_data.shape, dtype=tf.float64, trainable=False, initializer=tf.constant_initializer(tab_data))
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 build( self, inputs: tf.Tensor, natoms: tf.Tensor, input_dict: dict = None, reuse: bool = None, suffix: str = '', ) -> tf.Tensor: """ Build the computational graph for fitting net Parameters ---------- inputs The input descriptor input_dict Additional dict for inputs. if numb_fparam > 0, should have input_dict['fparam'] if numb_aparam > 0, should have input_dict['aparam'] 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 reuse The weights in the networks should be reused when get the variable. suffix Name suffix to identify this descriptor Returns ------- ener The system energy """ if input_dict is None: input_dict = {} bias_atom_e = self.bias_atom_e if self.numb_fparam > 0 and (self.fparam_avg is None or self.fparam_inv_std is None): raise RuntimeError( 'No data stat result. one should do data statisitic, before build' ) if self.numb_aparam > 0 and (self.aparam_avg is None or self.aparam_inv_std is None): raise RuntimeError( 'No data stat result. one should do data statisitic, before build' ) with tf.variable_scope('fitting_attr' + suffix, reuse=reuse): t_dfparam = tf.constant(self.numb_fparam, name='dfparam', dtype=tf.int32) t_daparam = tf.constant(self.numb_aparam, name='daparam', dtype=tf.int32) if self.numb_fparam > 0: t_fparam_avg = tf.get_variable( 't_fparam_avg', self.numb_fparam, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(self.fparam_avg)) t_fparam_istd = tf.get_variable( 't_fparam_istd', self.numb_fparam, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(self.fparam_inv_std)) if self.numb_aparam > 0: t_aparam_avg = tf.get_variable( 't_aparam_avg', self.numb_aparam, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(self.aparam_avg)) t_aparam_istd = tf.get_variable( 't_aparam_istd', self.numb_aparam, dtype=GLOBAL_TF_FLOAT_PRECISION, trainable=False, initializer=tf.constant_initializer(self.aparam_inv_std)) inputs = tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]]) if len(self.atom_ener): # only for atom_ener nframes = input_dict.get('nframes') if nframes is not None: # like inputs, but we don't want to add a dependency on inputs inputs_zero = tf.zeros((nframes, self.dim_descrpt * natoms[0]), dtype=self.fitting_precision) else: inputs_zero = tf.zeros_like(inputs, dtype=self.fitting_precision) if bias_atom_e is not None: assert (len(bias_atom_e) == self.ntypes) fparam = None aparam = None if self.numb_fparam > 0: fparam = input_dict['fparam'] fparam = tf.reshape(fparam, [-1, self.numb_fparam]) fparam = (fparam - t_fparam_avg) * t_fparam_istd if self.numb_aparam > 0: aparam = input_dict['aparam'] aparam = tf.reshape(aparam, [-1, self.numb_aparam]) aparam = (aparam - t_aparam_avg) * t_aparam_istd aparam = tf.reshape(aparam, [-1, self.numb_aparam * natoms[0]]) type_embedding = input_dict.get('type_embedding', None) if type_embedding is not None: atype_embed = embed_atom_type(self.ntypes, natoms, type_embedding) atype_embed = tf.tile(atype_embed, [tf.shape(inputs)[0], 1]) else: atype_embed = None if atype_embed is None: start_index = 0 outs_list = [] for type_i in range(self.ntypes): if bias_atom_e is None: type_bias_ae = 0.0 else: type_bias_ae = bias_atom_e[type_i] final_layer = self._build_lower(start_index, natoms[2 + type_i], inputs, fparam, aparam, bias_atom_e=type_bias_ae, suffix='_type_' + str(type_i) + suffix, reuse=reuse) # concat the results if type_i < len( self.atom_ener) and self.atom_ener[type_i] is not None: zero_layer = self._build_lower(start_index, natoms[2 + type_i], inputs_zero, fparam, aparam, bias_atom_e=type_bias_ae, suffix='_type_' + str(type_i) + suffix, reuse=True) final_layer += self.atom_ener[type_i] - zero_layer final_layer = tf.reshape( final_layer, [tf.shape(inputs)[0], natoms[2 + type_i]]) outs_list.append(final_layer) start_index += natoms[2 + type_i] # concat the results # concat once may be faster than multiple concat outs = tf.concat(outs_list, axis=1) # with type embedding else: if len(self.atom_ener) > 0: raise RuntimeError( "setting atom_ener is not supported by type embedding") atype_embed = tf.cast(atype_embed, self.fitting_precision) type_shape = atype_embed.get_shape().as_list() inputs = tf.concat( [tf.reshape(inputs, [-1, self.dim_descrpt]), atype_embed], axis=1) self.dim_descrpt = self.dim_descrpt + type_shape[1] inputs = tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]]) final_layer = self._build_lower(0, natoms[0], inputs, fparam, aparam, bias_atom_e=0.0, suffix=suffix, reuse=reuse) outs = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[0]]) # add atom energy bias; TF will broadcast to all batches # tf.repeat is avaiable in TF>=2.1 or TF 1.15 _TF_VERSION = Version(TF_VERSION) if (Version('1.15') <= _TF_VERSION < Version('2') or _TF_VERSION >= Version('2.1')) and self.bias_atom_e is not None: outs += tf.repeat( tf.Variable(self.bias_atom_e, dtype=self.fitting_precision, trainable=False, name="bias_atom_ei"), natoms[2:]) if self.tot_ener_zero: force_tot_ener = 0.0 outs = tf.reshape(outs, [-1, natoms[0]]) outs_mean = tf.reshape(tf.reduce_mean(outs, axis=1), [-1, 1]) outs_mean = outs_mean - tf.ones_like( outs_mean, dtype=GLOBAL_TF_FLOAT_PRECISION) * ( force_tot_ener / global_cvt_2_tf_float(natoms[0])) outs = outs - outs_mean outs = tf.reshape(outs, [-1]) tf.summary.histogram('fitting_net_output', outs) return tf.reshape(outs, [-1])
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 """ davg = self.davg dstd = self.dstd 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(self.rcut, name='rcut', dtype=GLOBAL_TF_FLOAT_PRECISION) t_ntypes = tf.constant(self.ntypes, name='ntypes', dtype=tf.int32) t_ndescrpt = tf.constant(self.ndescrpt, name='ndescrpt', dtype=tf.int32) t_sel = tf.constant(self.sel_a, name='sel', 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.prod_env_mat_r(coord, atype, natoms, box, mesh, self.t_avg, self.t_std, rcut = self.rcut, rcut_smth = self.rcut_smth, sel = self.sel_r) self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) self._identity_tensors(suffix=suffix) # only used when tensorboard was set as true tf.summary.histogram('descrpt', self.descrpt) tf.summary.histogram('rij', self.rij) tf.summary.histogram('nlist', self.nlist) self.dout = self._pass_filter(self.descrpt_reshape, natoms, suffix=suffix, reuse=reuse, trainable=self.trainable) tf.summary.histogram('embedding_net_output', self.dout) return self.dout
def one_layer(inputs, outputs_size, activation_fn=tf.nn.tanh, precision=GLOBAL_TF_FLOAT_PRECISION, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, use_timestep=False, trainable=True, useBN=False, uniform_seed=False, initial_variables=None, mixed_prec=None, final_layer=False): # For good accuracy, the last layer of the fitting network uses a higher precision neuron network. if mixed_prec is not None and final_layer: inputs = tf.cast(inputs, get_precision(mixed_prec['output_prec'])) with tf.variable_scope(name, reuse=reuse): shape = inputs.get_shape().as_list() w_initializer = tf.random_normal_initializer( stddev=stddev / np.sqrt(shape[1] + outputs_size), seed=seed if (seed is None or uniform_seed) else seed + 0) b_initializer = tf.random_normal_initializer( stddev=stddev, mean=bavg, seed=seed if (seed is None or uniform_seed) else seed + 1) if initial_variables is not None: w_initializer = tf.constant_initializer( initial_variables[name + '/matrix']) b_initializer = tf.constant_initializer(initial_variables[name + '/bias']) w = tf.get_variable('matrix', [shape[1], outputs_size], precision, w_initializer, trainable=trainable) variable_summaries(w, 'matrix') b = tf.get_variable('bias', [outputs_size], precision, b_initializer, trainable=trainable) variable_summaries(b, 'bias') if mixed_prec is not None and not final_layer: inputs = tf.cast(inputs, get_precision(mixed_prec['compute_prec'])) w = tf.cast(w, get_precision(mixed_prec['compute_prec'])) b = tf.cast(b, get_precision(mixed_prec['compute_prec'])) hidden = tf.nn.bias_add(tf.matmul(inputs, w), b) if activation_fn != None and use_timestep: idt_initializer = tf.random_normal_initializer( stddev=0.001, mean=0.1, seed=seed if (seed is None or uniform_seed) else seed + 2) if initial_variables is not None: idt_initializer = tf.constant_initializer( initial_variables[name + '/idt']) idt = tf.get_variable('idt', [outputs_size], precision, idt_initializer, trainable=trainable) variable_summaries(idt, 'idt') if activation_fn != None: if useBN: None # hidden_bn = self._batch_norm(hidden, name=name+'_normalization', reuse=reuse) # return activation_fn(hidden_bn) else: if use_timestep: if mixed_prec is not None and not final_layer: idt = tf.cast( idt, get_precision(mixed_prec['compute_prec'])) hidden = tf.reshape(activation_fn(hidden), [-1, outputs_size]) * idt else: hidden = tf.reshape(activation_fn(hidden), [-1, outputs_size]) if mixed_prec is not None: hidden = tf.cast(hidden, get_precision(mixed_prec['output_prec'])) return hidden
def embedding_net(xx, network_size, precision, activation_fn=tf.nn.tanh, resnet_dt=False, name_suffix='', stddev=1.0, bavg=0.0, seed=None, trainable=True, uniform_seed=False, initial_variables=None, mixed_prec=None): r"""The embedding network. The embedding network function :math:`\mathcal{N}` is constructed by is the composition of multiple layers :math:`\mathcal{L}^{(i)}`: .. math:: \mathcal{N} = \mathcal{L}^{(n)} \circ \mathcal{L}^{(n-1)} \circ \cdots \circ \mathcal{L}^{(1)} A layer :math:`\mathcal{L}` is given by one of the following forms, depending on the number of nodes: [1]_ .. math:: \mathbf{y}=\mathcal{L}(\mathbf{x};\mathbf{w},\mathbf{b})= \begin{cases} \boldsymbol{\phi}(\mathbf{x}^T\mathbf{w}+\mathbf{b}) + \mathbf{x}, & N_2=N_1 \\ \boldsymbol{\phi}(\mathbf{x}^T\mathbf{w}+\mathbf{b}) + (\mathbf{x}, \mathbf{x}), & N_2 = 2N_1\\ \boldsymbol{\phi}(\mathbf{x}^T\mathbf{w}+\mathbf{b}), & \text{otherwise} \\ \end{cases} where :math:`\mathbf{x} \in \mathbb{R}^{N_1}`$` is the input vector and :math:`\mathbf{y} \in \mathbb{R}^{N_2}` is the output vector. :math:`\mathbf{w} \in \mathbb{R}^{N_1 \times N_2}` and :math:`\mathbf{b} \in \mathbb{R}^{N_2}`$` are weights and biases, respectively, both of which are trainable if `trainable` is `True`. :math:`\boldsymbol{\phi}` is the activation function. Parameters ---------- xx : Tensor Input tensor :math:`\mathbf{x}` of shape [-1,1] network_size: list of int Size of the embedding network. For example [16,32,64] precision: Precision of network weights. For example, tf.float64 activation_fn: Activation function :math:`\boldsymbol{\phi}` resnet_dt: boolean Using time-step in the ResNet construction name_suffix: str The name suffix append to each variable. stddev: float Standard deviation of initializing network parameters bavg: float Mean of network intial bias seed: int Random seed for initializing network parameters trainable: boolean If the network is trainable uniform_seed : boolean Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed initial_variables : dict The input dict which stores the embedding net variables mixed_prec The input dict which stores the mixed precision setting for the embedding net References ---------- .. [1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identitymappings in deep residual networks. InComputer Vision – ECCV 2016,pages 630–645. Springer International Publishing, 2016. """ input_shape = xx.get_shape().as_list() outputs_size = [input_shape[1]] + network_size for ii in range(1, len(outputs_size)): w_initializer = tf.random_normal_initializer( stddev=stddev / np.sqrt(outputs_size[ii] + outputs_size[ii - 1]), seed=seed if (seed is None or uniform_seed) else seed + ii * 3 + 0) b_initializer = tf.random_normal_initializer( stddev=stddev, mean=bavg, seed=seed if (seed is None or uniform_seed) else seed + 3 * ii + 1) if initial_variables is not None: scope = tf.get_variable_scope().name w_initializer = tf.constant_initializer( initial_variables[scope + '/matrix_' + str(ii) + name_suffix]) b_initializer = tf.constant_initializer( initial_variables[scope + '/bias_' + str(ii) + name_suffix]) w = tf.get_variable('matrix_' + str(ii) + name_suffix, [outputs_size[ii - 1], outputs_size[ii]], precision, w_initializer, trainable=trainable) variable_summaries(w, 'matrix_' + str(ii) + name_suffix) b = tf.get_variable('bias_' + str(ii) + name_suffix, [outputs_size[ii]], precision, b_initializer, trainable=trainable) variable_summaries(b, 'bias_' + str(ii) + name_suffix) if mixed_prec is not None: xx = tf.cast(xx, get_precision(mixed_prec['compute_prec'])) w = tf.cast(w, get_precision(mixed_prec['compute_prec'])) b = tf.cast(b, get_precision(mixed_prec['compute_prec'])) hidden = tf.reshape(activation_fn(tf.nn.bias_add(tf.matmul(xx, w), b)), [-1, outputs_size[ii]]) if resnet_dt: idt_initializer = tf.random_normal_initializer( stddev=0.001, mean=1.0, seed=seed if (seed is None or uniform_seed) else seed + 3 * ii + 2) if initial_variables is not None: scope = tf.get_variable_scope().name idt_initializer = tf.constant_initializer( initial_variables[scope + '/idt_' + str(ii) + name_suffix]) idt = tf.get_variable('idt_' + str(ii) + name_suffix, [1, outputs_size[ii]], precision, idt_initializer, trainable=trainable) variable_summaries(idt, 'idt_' + str(ii) + name_suffix) if mixed_prec is not None: idt = tf.cast(idt, get_precision(mixed_prec['compute_prec'])) if outputs_size[ii] == outputs_size[ii - 1]: if resnet_dt: xx += hidden * idt else: xx += hidden elif outputs_size[ii] == outputs_size[ii - 1] * 2: if resnet_dt: xx = tf.concat([xx, xx], 1) + hidden * idt else: xx = tf.concat([xx, xx], 1) + hidden else: xx = hidden if mixed_prec is not None: xx = tf.cast(xx, get_precision(mixed_prec['output_prec'])) return xx
def build(self, inputs, input_dict, natoms, reuse=None, suffix=''): bias_atom_e = self.bias_atom_e if self.numb_fparam > 0 and (self.fparam_avg is None or self.fparam_inv_std is None): raise RuntimeError( 'No data stat result. one should do data statisitic, before build' ) if self.numb_aparam > 0 and (self.aparam_avg is None or self.aparam_inv_std is None): raise RuntimeError( 'No data stat result. one should do data statisitic, before build' ) with tf.variable_scope('fitting_attr' + suffix, reuse=reuse): t_dfparam = tf.constant(self.numb_fparam, name='dfparam', dtype=tf.int32) t_daparam = tf.constant(self.numb_aparam, name='daparam', dtype=tf.int32) if self.numb_fparam > 0: t_fparam_avg = tf.get_variable( 't_fparam_avg', self.numb_fparam, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(self.fparam_avg)) t_fparam_istd = tf.get_variable( 't_fparam_istd', self.numb_fparam, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(self.fparam_inv_std)) if self.numb_aparam > 0: t_aparam_avg = tf.get_variable( 't_aparam_avg', self.numb_aparam, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(self.aparam_avg)) t_aparam_istd = tf.get_variable( 't_aparam_istd', self.numb_aparam, dtype=global_tf_float_precision, trainable=False, initializer=tf.constant_initializer(self.aparam_inv_std)) start_index = 0 inputs = tf.cast( tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) if bias_atom_e is not None: assert (len(bias_atom_e) == self.ntypes) if self.numb_fparam > 0: fparam = input_dict['fparam'] fparam = tf.reshape(fparam, [-1, self.numb_fparam]) fparam = (fparam - t_fparam_avg) * t_fparam_istd if self.numb_aparam > 0: aparam = input_dict['aparam'] aparam = tf.reshape(aparam, [-1, self.numb_aparam]) aparam = (aparam - t_aparam_avg) * t_aparam_istd aparam = tf.reshape(aparam, [-1, self.numb_aparam * natoms[0]]) for type_i in range(self.ntypes): # cut-out inputs inputs_i = tf.slice(inputs, [0, start_index * self.dim_descrpt], [-1, natoms[2 + type_i] * self.dim_descrpt]) inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt]) layer = inputs_i if self.numb_fparam > 0: ext_fparam = tf.tile(fparam, [1, natoms[2 + type_i]]) ext_fparam = tf.reshape(ext_fparam, [-1, self.numb_fparam]) layer = tf.concat([layer, ext_fparam], axis=1) if self.numb_aparam > 0: ext_aparam = tf.slice( aparam, [0, start_index * self.numb_aparam], [-1, natoms[2 + type_i] * self.numb_aparam]) ext_aparam = tf.reshape(ext_aparam, [-1, self.numb_aparam]) layer = tf.concat([layer, ext_aparam], axis=1) start_index += natoms[2 + type_i] if bias_atom_e is None: type_bias_ae = 0.0 else: type_bias_ae = bias_atom_e[type_i] for ii in range(0, len(self.n_neuron)): if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii - 1]: layer += one_layer( layer, self.n_neuron[ii], name='layer_' + str(ii) + '_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed, use_timestep=self.resnet_dt, activation_fn=self.fitting_activation_fn, precision=self.fitting_precision, trainable=self.trainable[ii]) else: layer = one_layer(layer, self.n_neuron[ii], name='layer_' + str(ii) + '_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed, activation_fn=self.fitting_activation_fn, precision=self.fitting_precision, trainable=self.trainable[ii]) final_layer = one_layer(layer, 1, activation_fn=None, bavg=type_bias_ae, name='final_layer_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed, precision=self.fitting_precision, trainable=self.trainable[-1]) if type_i < len( self.atom_ener) and self.atom_ener[type_i] is not None: inputs_zero = tf.zeros_like(inputs_i, dtype=global_tf_float_precision) layer = inputs_zero if self.numb_fparam > 0: layer = tf.concat([layer, ext_fparam], axis=1) if self.numb_aparam > 0: layer = tf.concat([layer, ext_aparam], axis=1) for ii in range(0, len(self.n_neuron)): if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii - 1]: layer += one_layer( layer, self.n_neuron[ii], name='layer_' + str(ii) + '_type_' + str(type_i) + suffix, reuse=True, seed=self.seed, use_timestep=self.resnet_dt, activation_fn=self.fitting_activation_fn, precision=self.fitting_precision, trainable=self.trainable[ii]) else: layer = one_layer( layer, self.n_neuron[ii], name='layer_' + str(ii) + '_type_' + str(type_i) + suffix, reuse=True, seed=self.seed, activation_fn=self.fitting_activation_fn, precision=self.fitting_precision, trainable=self.trainable[ii]) zero_layer = one_layer(layer, 1, activation_fn=None, bavg=type_bias_ae, name='final_layer_type_' + str(type_i) + suffix, reuse=True, seed=self.seed, precision=self.fitting_precision, trainable=self.trainable[-1]) final_layer += self.atom_ener[type_i] - zero_layer final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2 + type_i]]) # concat the results if type_i == 0: outs = final_layer else: outs = tf.concat([outs, final_layer], axis=1) return tf.cast(tf.reshape(outs, [-1]), global_tf_float_precision)
def build(self, coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, 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) t_ver = tf.constant(MODEL_VERSION, name='model_version', 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]]) # type embedding if any if self.typeebd is not None: type_embedding = self.typeebd.build( self.ntypes, reuse=reuse, suffix=suffix, ) input_dict['type_embedding'] = type_embedding if frz_model == None: dout \ = self.descrpt.build(coord_, atype_, natoms, box, mesh, input_dict, suffix = suffix, reuse = reuse) dout = tf.identity(dout, name='o_descriptor') else: tf.constant(self.rcut, name='descrpt_attr/rcut', dtype=GLOBAL_TF_FLOAT_PRECISION) tf.constant(self.ntypes, name='descrpt_attr/ntypes', dtype=tf.int32) feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh) return_elements = [ *self.descrpt.get_tensor_names(), 'o_descriptor:0' ] imported_tensors \ = self._import_graph_def_from_frz_model(frz_model, feed_dict, return_elements) dout = imported_tensors[-1] self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1]) 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, natoms, input_dict, 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.pair_tab(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 model_dict['coord'] = coord model_dict['atype'] = atype return model_dict
def build(self, coord_, atype_, natoms, box_, mesh, input_dict, suffix='', reuse=None): efield = input_dict['efield'] davg = self.davg dstd = self.dstd 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) t_ndescrpt = tf.constant(self.ndescrpt, name='ndescrpt', dtype=tf.int32) t_sel = tf.constant(self.sel_a, name='sel', 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]]) efield = tf.reshape(efield, [-1, 3]) efield = self._normalize_3d(efield) efield = tf.reshape(efield, [-1, natoms[0] * 3]) self.descrpt, self.descrpt_deriv, self.rij, self.nlist \ = self.op (coord, atype, natoms, box, mesh, efield, 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.descrpt_reshape = tf.identity(self.descrpt_reshape, name='o_rmat') self.descrpt_deriv = tf.identity(self.descrpt_deriv, name='o_rmat_deriv') self.rij = tf.identity(self.rij, name='o_rij') self.nlist = tf.identity(self.nlist, name='o_nlist') # only used when tensorboard was set as true tf.summary.histogram('descrpt', self.descrpt) tf.summary.histogram('rij', self.rij) tf.summary.histogram('nlist', self.nlist) self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, atype, natoms, input_dict, suffix=suffix, reuse=reuse, trainable=self.trainable) tf.summary.histogram('embedding_net_output', self.dout) return self.dout
def one_layer(inputs, outputs_size, activation_fn=tf.nn.tanh, precision = GLOBAL_TF_FLOAT_PRECISION, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, use_timestep = False, trainable = True, useBN = False, uniform_seed = False, initial_variables = None): with tf.variable_scope(name, reuse=reuse): shape = inputs.get_shape().as_list() w_initializer = tf.random_normal_initializer( stddev=stddev / np.sqrt(shape[1] + outputs_size), seed=seed if (seed is None or uniform_seed) else seed + 0) b_initializer = tf.random_normal_initializer( stddev=stddev, mean=bavg, seed=seed if (seed is None or uniform_seed) else seed + 1) if initial_variables is not None: w_initializer = tf.constant_initializer(initial_variables[name + '/matrix']) b_initializer = tf.constant_initializer(initial_variables[name + '/bias']) w = tf.get_variable('matrix', [shape[1], outputs_size], precision, w_initializer, trainable = trainable) variable_summaries(w, 'matrix') b = tf.get_variable('bias', [outputs_size], precision, b_initializer, trainable = trainable) variable_summaries(b, 'bias') hidden = tf.matmul(inputs, w) + b if activation_fn != None and use_timestep : idt_initializer = tf.random_normal_initializer( stddev=0.001, mean=0.1, seed=seed if (seed is None or uniform_seed) else seed + 2) if initial_variables is not None: idt_initializer = tf.constant_initializer(initial_variables[name + '/idt']) idt = tf.get_variable('idt', [outputs_size], precision, idt_initializer, trainable = trainable) variable_summaries(idt, 'idt') if activation_fn != None: if useBN: None # hidden_bn = self._batch_norm(hidden, name=name+'_normalization', reuse=reuse) # return activation_fn(hidden_bn) else: if use_timestep : return tf.reshape(activation_fn(hidden), [-1, outputs_size]) * idt else : return tf.reshape(activation_fn(hidden), [-1, outputs_size]) else: if useBN: None # return self._batch_norm(hidden, name=name+'_normalization', reuse=reuse) else: return hidden