def build(self, input_d, rot_mat, natoms, reuse=None, suffix=''): start_index = 0 inputs = tf.cast( tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]]) count = 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]) rot_mat_i = tf.slice(rot_mat, [0, start_index * self.dim_rot_mat], [-1, natoms[2 + type_i] * self.dim_rot_mat]) rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) start_index += natoms[2 + type_i] if not type_i in self.sel_type: continue layer = inputs_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) 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) # (nframes x natoms) x naxis final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn=None, name='final_layer_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed, precision=self.fitting_precision) # (nframes x natoms) x 1 * naxis final_layer = tf.reshape(final_layer, [ tf.shape(inputs)[0] * natoms[2 + type_i], 1, self.dim_rot_mat_1 ]) # (nframes x natoms) x 1 x 3(coord) final_layer = tf.matmul(final_layer, rot_mat_i) # nframes x natoms x 3 final_layer = tf.reshape( final_layer, [tf.shape(inputs)[0], natoms[2 + type_i], 3]) # concat the results if count == 0: outs = final_layer else: outs = tf.concat([outs, final_layer], axis=1) count += 1 return tf.cast(tf.reshape(outs, [-1]), global_tf_float_precision)
def build(self, input_d, rot_mat, natoms, reuse=None, suffix=''): start_index = 0 inputs = tf.cast( tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]]) count = 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]) rot_mat_i = tf.slice(rot_mat, [0, start_index * self.dim_rot_mat], [-1, natoms[2 + type_i] * self.dim_rot_mat]) rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) start_index += natoms[2 + type_i] if not type_i in self.sel_type: continue layer = inputs_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) 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) if self.fit_diag: bavg = np.zeros(self.dim_rot_mat_1) # bavg[0] = self.avgeig[0] # bavg[1] = self.avgeig[1] # bavg[2] = self.avgeig[2] # (nframes x natoms) x naxis final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn=None, name='final_layer_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed, bavg=bavg, precision=self.fitting_precision) # (nframes x natoms) x naxis final_layer = tf.reshape(final_layer, [ tf.shape(inputs)[0] * natoms[2 + type_i], self.dim_rot_mat_1 ]) # (nframes x natoms) x naxis x naxis final_layer = tf.matrix_diag(final_layer) else: bavg = np.zeros(self.dim_rot_mat_1 * self.dim_rot_mat_1) # bavg[0*self.dim_rot_mat_1+0] = self.avgeig[0] # bavg[1*self.dim_rot_mat_1+1] = self.avgeig[1] # bavg[2*self.dim_rot_mat_1+2] = self.avgeig[2] # (nframes x natoms) x (naxis x naxis) final_layer = one_layer( layer, self.dim_rot_mat_1 * self.dim_rot_mat_1, activation_fn=None, name='final_layer_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed, bavg=bavg, precision=self.fitting_precision) # (nframes x natoms) x naxis x naxis final_layer = tf.reshape(final_layer, [ tf.shape(inputs)[0] * natoms[2 + type_i], self.dim_rot_mat_1, self.dim_rot_mat_1 ]) # (nframes x natoms) x naxis x naxis final_layer = final_layer + tf.transpose(final_layer, perm=[0, 2, 1]) # (nframes x natoms) x naxis x 3(coord) final_layer = tf.matmul(final_layer, rot_mat_i) # (nframes x natoms) x 3(coord) x 3(coord) final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a=True) # nframes x natoms x 3 x 3 final_layer = tf.reshape( final_layer, [tf.shape(inputs)[0], natoms[2 + type_i], 3, 3]) # shift and scale sel_type_idx = self.sel_type.index(type_i) final_layer = final_layer * self.scale[sel_type_idx] final_layer = final_layer + self.diag_shift[sel_type_idx] * tf.eye( 3, batch_shape=[tf.shape(inputs)[0], natoms[2 + type_i]], dtype=global_tf_float_precision) # concat the results if count == 0: outs = final_layer else: outs = tf.concat([outs, final_layer], axis=1) count += 1 return tf.cast(tf.reshape(outs, [-1]), global_tf_float_precision)
def build(self, inputs, input_dict, natoms, bias_atom_e=None, reuse=None, suffix=''): with tf.variable_scope('fitting_attr' + suffix, reuse=reuse): t_dfparam = tf.constant(self.numb_fparam, name='dfparam', dtype=tf.int32) start_index = 0 inputs = tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]]) shape = inputs.get_shape().as_list() if bias_atom_e is not None: assert (len(bias_atom_e) == self.ntypes) 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]) 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] layer = inputs_i if self.numb_fparam > 0: fparam = input_dict['fparam'] ext_fparam = tf.reshape(fparam, [-1, self.numb_fparam]) ext_fparam = tf.tile(ext_fparam, [1, natoms[2 + type_i]]) ext_fparam = tf.reshape(ext_fparam, [-1, self.numb_fparam]) layer = tf.concat([layer, ext_fparam], 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=reuse, seed=self.seed, use_timestep=self.resnet_dt) else: layer = one_layer(layer, self.n_neuron[ii], name='layer_' + str(ii) + '_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed) 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) 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.reshape(outs, [-1])
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, input_d, rot_mat, natoms, reuse=None, suffix=''): start_index = 0 inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]) rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]]) shape = inputs.get_shape().as_list() count = 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]) rot_mat_i = tf.slice(rot_mat, [0, start_index * self.dim_rot_mat], [-1, natoms[2 + type_i] * self.dim_rot_mat]) rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) start_index += natoms[2 + type_i] if not type_i in self.sel_type: continue layer = inputs_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) else: layer = one_layer(layer, self.n_neuron[ii], name='layer_' + str(ii) + '_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed) # (nframes x natoms) x (naxis x naxis) final_layer = one_layer(layer, self.dim_rot_mat_1 * self.dim_rot_mat_1, activation_fn=None, name='final_layer_type_' + str(type_i) + suffix, reuse=reuse, seed=self.seed) # (nframes x natoms) x naxis x naxis final_layer = tf.reshape(final_layer, [ tf.shape(inputs)[0] * natoms[2 + type_i], self.dim_rot_mat_1, self.dim_rot_mat_1 ]) # (nframes x natoms) x naxis x naxis final_layer = final_layer + tf.transpose(final_layer, perm=[0, 2, 1]) # (nframes x natoms) x naxis x 3(coord) final_layer = tf.matmul(final_layer, rot_mat_i) # (nframes x natoms) x 3(coord) x 3(coord) final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a=True) # nframes x natoms x 3 x 3 final_layer = tf.reshape( final_layer, [tf.shape(inputs)[0], natoms[2 + type_i], 3, 3]) # concat the results if count == 0: outs = final_layer else: outs = tf.concat([outs, final_layer], axis=1) count += 1 return tf.reshape(outs, [-1])