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, 9 * natoms[0]]) count = 0 outs_list = [] 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* 9], [-1, natoms[2+type_i]* 9] ) rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 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 9 final_layer = one_layer(layer, 9, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, final_layer = True) # (nframes x natoms) x 3 x 3 final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 3, 3]) # (nframes x natoms) x 3 x 3 final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1]) # (nframes x natoms) x 3 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 outs_list.append(final_layer) count += 1 outs = tf.concat(outs_list, axis = 1) tf.summary.histogram('fitting_net_output', outs) return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
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 _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 _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] if self.neuron[0] == 1: yy = self._layer_0( xx, self.matrix["layer_" + str(layer + 1)][idx], self.bias["layer_" + str(layer + 1)][idx]) + xx dy = op_module.unaggregated_dy_dx_s( yy, self.matrix["layer_" + str(layer + 1)][idx], xbar, tf.constant(self.functype)) + tf.ones([1, 1], yy.dtype) dy2 = op_module.unaggregated_dy2_dx_s( yy, dy, self.matrix["layer_" + str(layer + 1)][idx], xbar, tf.constant(self.functype)) elif self.neuron[0] == 2: tt, yy = self._layer_1( xx, self.matrix["layer_" + str(layer + 1)][idx], self.bias["layer_" + str(layer + 1)][idx]) dy = op_module.unaggregated_dy_dx_s( yy - tt, self.matrix["layer_" + str(layer + 1)][idx], xbar, tf.constant(self.functype)) + tf.ones([1, 2], yy.dtype) dy2 = op_module.unaggregated_dy2_dx_s( yy - tt, dy, self.matrix["layer_" + str(layer + 1)][idx], xbar, tf.constant(self.functype)) else: 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 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): with tf.variable_scope(name, reuse=reuse): shape = inputs.get_shape().as_list() w = tf.get_variable( 'matrix', [shape[1], outputs_size], precision, tf.random_normal_initializer(stddev=stddev / np.sqrt(shape[1] + outputs_size), seed=seed), trainable=trainable) b = tf.get_variable('bias', [outputs_size], precision, tf.random_normal_initializer(stddev=stddev, mean=bavg, seed=seed), trainable=trainable) hidden = tf.matmul(inputs, w) + b if activation_fn != None and use_timestep: idt = tf.get_variable('idt', [outputs_size], precision, tf.random_normal_initializer(stddev=0.001, mean=0.1, seed=seed), trainable=trainable) 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
def build(self, learning_rate, natoms, model_dict, label_dict, suffix): coord = model_dict['coord'] energy = model_dict['energy'] atom_ener = model_dict['atom_ener'] nframes = tf.shape(atom_ener)[0] natoms = tf.shape(atom_ener)[1] # build energy dipole atom_ener0 = atom_ener - tf.reshape( tf.tile( tf.reshape(energy / global_cvt_2_ener_float(natoms), [-1, 1]), [1, natoms]), [nframes, natoms]) coord = tf.reshape(coord, [nframes, natoms, 3]) atom_ener0 = tf.reshape(atom_ener0, [nframes, 1, natoms]) ener_dipole = tf.matmul(atom_ener0, coord) ener_dipole = tf.reshape(ener_dipole, [nframes, 3]) energy_hat = label_dict['energy'] ener_dipole_hat = label_dict['energy_dipole'] find_energy = label_dict['find_energy'] find_ener_dipole = label_dict['find_energy_dipole'] l2_ener_loss = tf.reduce_mean(tf.square(energy - energy_hat), name='l2_' + suffix) ener_dipole_reshape = tf.reshape(ener_dipole, [-1]) ener_dipole_hat_reshape = tf.reshape(ener_dipole_hat, [-1]) l2_ener_dipole_loss = tf.reduce_mean( tf.square(ener_dipole_reshape - ener_dipole_hat_reshape), name='l2_' + suffix) # atom_norm_ener = 1./ global_cvt_2_ener_float(natoms[0]) atom_norm_ener = 1. / global_cvt_2_ener_float(natoms) pref_e = global_cvt_2_ener_float( find_energy * (self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * learning_rate / self.starter_learning_rate)) pref_ed = global_cvt_2_tf_float( find_ener_dipole * (self.limit_pref_ed + (self.start_pref_ed - self.limit_pref_ed) * learning_rate / self.starter_learning_rate)) l2_loss = 0 more_loss = {} l2_loss += atom_norm_ener * (pref_e * l2_ener_loss) l2_loss += global_cvt_2_ener_float(pref_ed * l2_ener_dipole_loss) more_loss['l2_ener_loss'] = l2_ener_loss more_loss['l2_ener_dipole_loss'] = l2_ener_dipole_loss self.l2_l = l2_loss self.l2_more = more_loss return l2_loss, more_loss
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, 9 * 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* 9], [-1, natoms[2+type_i]* 9] ) rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 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, uniform_seed = self.uniform_seed) 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, uniform_seed = self.uniform_seed) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x (nwfc x 3) final_layer = one_layer(layer, self.wfc_numb * 3, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (nframes x natoms) x nwfc(wc) x 3(coord_local) final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.wfc_numb, 3]) # (nframes x natoms) x nwfc(wc) x 3(coord) final_layer = tf.matmul(final_layer, rot_mat_i) # nframes x natoms x nwfc(wc) x 3(coord_local) final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], self.wfc_numb, 3]) # concat the results if count == 0: outs = final_layer else: outs = tf.concat([outs, final_layer], axis = 1) count += 1 tf.summary.histogram('fitting_net_output', outs) return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
def _layer_1(self, x, w, b): t = tf.concat([x, x], axis=1) return t, self.activation_fn(tf.matmul(x, w) + b) + t
def _layer_0(self, x, w, b): return self.activation_fn(tf.matmul(x, w) + b)
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 _filter(self, inputs, type_input, natoms, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, trainable=True): # natom x (nei x 4) shape = inputs.get_shape().as_list() outputs_size = [1] + self.filter_neuron outputs_size_2 = self.n_axis_neuron with tf.variable_scope(name, reuse=reuse): start_index = 0 xyz_scatter_total = [] for type_i in range(self.ntypes): # cut-out inputs # with natom x (nei_type_i x 4) inputs_i = tf.slice(inputs, [0, start_index * 4], [-1, self.sel_a[type_i] * 4]) start_index += self.sel_a[type_i] shape_i = inputs_i.get_shape().as_list() # with (natom x nei_type_i) x 4 inputs_reshape = tf.reshape(inputs_i, [-1, 4]) xyz_scatter = tf.reshape( tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1]) if (type_input, type_i) not in self.exclude_types: for ii in range(1, len(outputs_size)): w = tf.get_variable( 'matrix_' + str(ii) + '_' + str(type_i), [outputs_size[ii - 1], outputs_size[ii]], self.filter_precision, tf.random_normal_initializer( stddev=stddev / np.sqrt(outputs_size[ii] + outputs_size[ii - 1]), seed=seed), trainable=trainable) b = tf.get_variable( 'bias_' + str(ii) + '_' + str(type_i), [1, outputs_size[ii]], self.filter_precision, tf.random_normal_initializer(stddev=stddev, mean=bavg, seed=seed), trainable=trainable) if self.filter_resnet_dt: idt = tf.get_variable( 'idt_' + str(ii) + '_' + str(type_i), [1, outputs_size[ii]], self.filter_precision, tf.random_normal_initializer(stddev=0.001, mean=1.0, seed=seed), trainable=trainable) if outputs_size[ii] == outputs_size[ii - 1]: if self.filter_resnet_dt: xyz_scatter += activation_fn( tf.matmul(xyz_scatter, w) + b) * idt else: xyz_scatter += activation_fn( tf.matmul(xyz_scatter, w) + b) elif outputs_size[ii] == outputs_size[ii - 1] * 2: if self.filter_resnet_dt: xyz_scatter = tf.concat( [xyz_scatter, xyz_scatter], 1) + activation_fn( tf.matmul(xyz_scatter, w) + b) * idt else: xyz_scatter = tf.concat( [xyz_scatter, xyz_scatter], 1) + activation_fn( tf.matmul(xyz_scatter, w) + b) else: xyz_scatter = activation_fn( tf.matmul(xyz_scatter, w) + b) else: w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=global_tf_float_precision) xyz_scatter = tf.matmul(xyz_scatter, w) # natom x nei_type_i x out_size xyz_scatter = tf.reshape( xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1])) xyz_scatter_total.append(xyz_scatter) # natom x nei x outputs_size xyz_scatter = tf.concat(xyz_scatter_total, axis=1) # natom x nei x 4 inputs_reshape = tf.reshape(inputs, [-1, shape[1] // 4, 4]) # natom x 4 x outputs_size xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a=True) xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape[1]) # natom x 4 x outputs_size_2 xyz_scatter_2 = tf.slice(xyz_scatter_1, [0, 0, 0], [-1, -1, outputs_size_2]) # # natom x 3 x outputs_size_2 # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1]) # natom x 3 x outputs_size_1 qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1]) # natom x outputs_size_2 x 3 qmat = tf.transpose(qmat, perm=[0, 2, 1]) # natom x outputs_size x outputs_size_2 result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True) # natom x (outputs_size x outputs_size_2) result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]]) return result, qmat
def _filter_r(self, inputs, type_input, natoms, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, trainable=True): # natom x nei outputs_size = [1] + self.filter_neuron with tf.variable_scope(name, reuse=reuse): start_index = 0 xyz_scatter_total = [] for type_i in range(self.ntypes): # cut-out inputs # with natom x nei_type_i inputs_i = tf.slice(inputs, [0, start_index], [-1, self.sel_r[type_i]]) start_index += self.sel_r[type_i] shape_i = inputs_i.get_shape().as_list() # with (natom x nei_type_i) x 1 xyz_scatter = tf.reshape(inputs_i, [-1, 1]) if (type_input, type_i) not in self.exclude_types: for ii in range(1, len(outputs_size)): w = tf.get_variable( 'matrix_' + str(ii) + '_' + str(type_i), [outputs_size[ii - 1], outputs_size[ii]], self.filter_precision, tf.random_normal_initializer( stddev=stddev / np.sqrt(outputs_size[ii] + outputs_size[ii - 1]), seed=seed), trainable=trainable) b = tf.get_variable( 'bias_' + str(ii) + '_' + str(type_i), [1, outputs_size[ii]], self.filter_precision, tf.random_normal_initializer(stddev=stddev, mean=bavg, seed=seed), trainable=trainable) if self.filter_resnet_dt: idt = tf.get_variable( 'idt_' + str(ii) + '_' + str(type_i), [1, outputs_size[ii]], self.filter_precision, tf.random_normal_initializer(stddev=0.001, mean=1.0, seed=seed), trainable=trainable) if outputs_size[ii] == outputs_size[ii - 1]: if self.filter_resnet_dt: xyz_scatter += activation_fn( tf.matmul(xyz_scatter, w) + b) * idt else: xyz_scatter += activation_fn( tf.matmul(xyz_scatter, w) + b) elif outputs_size[ii] == outputs_size[ii - 1] * 2: if self.filter_resnet_dt: xyz_scatter = tf.concat( [xyz_scatter, xyz_scatter], 1) + activation_fn( tf.matmul(xyz_scatter, w) + b) * idt else: xyz_scatter = tf.concat( [xyz_scatter, xyz_scatter], 1) + activation_fn( tf.matmul(xyz_scatter, w) + b) else: xyz_scatter = activation_fn( tf.matmul(xyz_scatter, w) + b) else: w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=global_tf_float_precision) xyz_scatter = tf.matmul(xyz_scatter, w) # natom x nei_type_i x out_size xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1], outputs_size[-1])) xyz_scatter_total.append(xyz_scatter) # natom x nei x outputs_size xyz_scatter = tf.concat(xyz_scatter_total, axis=1) # natom x outputs_size # res_rescale = 1. / 5. result = tf.reduce_mean(xyz_scatter, axis=1) * res_rescale return result
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, input_d : tf.Tensor, rot_mat : tf.Tensor, natoms : tf.Tensor, reuse : bool = None, suffix : str = '') : """ Build the computational graph for fitting net Parameters ---------- input_d The input descriptor rot_mat The rotation matrix from the descriptor. 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 ------- atomic_polar The atomic polarizability """ 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]]) count = 0 outs_list = [] 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, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec) 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, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift 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, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (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, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (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.constant_matrix[sel_type_idx] * tf.eye(3, batch_shape=[tf.shape(inputs)[0], natoms[2+type_i]], dtype = GLOBAL_TF_FLOAT_PRECISION) # concat the results outs_list.append(final_layer) count += 1 outs = tf.concat(outs_list, axis = 1) tf.summary.histogram('fitting_net_output', outs) return tf.reshape(outs, [-1])
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
def _ebd_filter(self, inputs, atype, natoms, input_dict, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, trainable=True): outputs_size = self.filter_neuron[-1] outputs_size_2 = self.n_axis_neuron # nf x natom x (nei x 4) nframes = tf.shape(inputs)[0] shape = tf.reshape(inputs, [-1, self.ndescrpt]).get_shape().as_list() # nf x natom x nei x outputs_size mat_g = self._embedding_net(inputs, natoms, self.filter_neuron, activation_fn=activation_fn, stddev=stddev, bavg=bavg, name=name, reuse=reuse, seed=seed, trainable=trainable) # nf x natom x nei x outputs_size mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size]) # (nf x natom) x nei x outputs_size if self.type_one_side: if self.numb_aparam > 0: aparam = input_dict['aparam'] xyz_scatter \ = self._type_embedding_net_one_side_aparam(mat_g, atype, natoms, aparam, name = name, reuse = reuse, seed = seed, trainable = trainable) else: xyz_scatter \ = self._type_embedding_net_one_side(mat_g, atype, natoms, name = name, reuse = reuse, seed = seed, trainable = trainable) else: xyz_scatter \ = self._type_embedding_net_two_sides(mat_g, atype, natoms, name = name, reuse = reuse, seed = seed, trainable = trainable) # natom x nei x 4 inputs_reshape = tf.reshape(inputs, [-1, shape[1] // 4, 4]) # natom x 4 x outputs_size xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a=True) xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape[1]) # natom x 4 x outputs_size_2 xyz_scatter_2 = tf.slice(xyz_scatter_1, [0, 0, 0], [-1, -1, outputs_size_2]) # # natom x 3 x outputs_size_2 # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1]) # natom x 3 x outputs_size_1 qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1]) # natom x outputs_size_2 x 3 qmat = tf.transpose(qmat, perm=[0, 2, 1]) # natom x outputs_size x outputs_size_2 result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True) # natom x (outputs_size x outputs_size_2) result = tf.reshape(result, [-1, outputs_size_2 * outputs_size]) return result, qmat
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 build(self, input_d: tf.Tensor, rot_mat: tf.Tensor, natoms: tf.Tensor, reuse: bool = None, suffix: str = '') -> tf.Tensor: """ Build the computational graph for fitting net Parameters ---------- input_d The input descriptor rot_mat The rotation matrix from the descriptor. 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 ------- dipole The atomic dipole. """ 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, uniform_seed=self.uniform_seed) 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, uniform_seed=self.uniform_seed) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (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, uniform_seed=self.uniform_seed) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift # (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 tf.summary.histogram('fitting_net_output', outs) return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION)
def _build_fv_graph_inner(self): self.t_ef = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_ef') nf = 10 nfxnas = 64 * nf nfxna = 192 * nf nf = -1 nfxnas = -1 nfxna = -1 self.t_box_reshape = tf.reshape(self.t_box, [-1, 9]) t_nframes = tf.shape(self.t_box_reshape)[0] # (nframes x natoms_sel) x 1 x 3 self.t_ef_reshape = tf.reshape(self.t_ef, [nfxnas, 1, 3]) # (nframes x natoms) x ndescrpt self.descrpt = self.graph.get_tensor_by_name( os.path.join(self.modifier_prefix, 'o_rmat:0')) self.descrpt_deriv = self.graph.get_tensor_by_name( os.path.join(self.modifier_prefix, 'o_rmat_deriv:0')) self.nlist = self.graph.get_tensor_by_name( os.path.join(self.modifier_prefix, 'o_nlist:0')) self.rij = self.graph.get_tensor_by_name( os.path.join(self.modifier_prefix, 'o_rij:0')) # self.descrpt_reshape = tf.reshape(self.descrpt, [nf, 192 * self.ndescrpt]) # self.descrpt_deriv = tf.reshape(self.descrpt_deriv, [nf, 192 * self.ndescrpt * 3]) # nframes x (natoms_sel x 3) self.t_tensor_reshpe = tf.reshape(self.t_tensor, [t_nframes, -1]) # nframes x (natoms x 3) self.t_tensor_reshpe = self._enrich(self.t_tensor_reshpe, dof=3) # (nframes x natoms) x 3 self.t_tensor_reshpe = tf.reshape(self.t_tensor_reshpe, [nfxna, 3]) # (nframes x natoms) x 1 self.t_dipole_x = tf.slice(self.t_tensor_reshpe, [0, 0], [nfxna, 1]) self.t_dipole_y = tf.slice(self.t_tensor_reshpe, [0, 1], [nfxna, 1]) self.t_dipole_z = tf.slice(self.t_tensor_reshpe, [0, 2], [nfxna, 1]) self.t_dipole_z = tf.reshape(self.t_dipole_z, [nfxna, 1]) # (nframes x natoms) x ndescrpt [self.t_dipole_x_d] = tf.gradients(self.t_dipole_x, self.descrpt) [self.t_dipole_y_d] = tf.gradients(self.t_dipole_y, self.descrpt) [self.t_dipole_z_d] = tf.gradients(self.t_dipole_z, self.descrpt) # nframes x (natoms x ndescrpt) self.t_dipole_x_d = tf.reshape(self.t_dipole_x_d, [-1, self.t_natoms[0] * self.ndescrpt]) self.t_dipole_y_d = tf.reshape(self.t_dipole_y_d, [-1, self.t_natoms[0] * self.ndescrpt]) self.t_dipole_z_d = tf.reshape(self.t_dipole_z_d, [-1, self.t_natoms[0] * self.ndescrpt]) # nframes x (natoms_sel x ndescrpt) self.t_dipole_x_d = self._slice_descrpt_deriv(self.t_dipole_x_d) self.t_dipole_y_d = self._slice_descrpt_deriv(self.t_dipole_y_d) self.t_dipole_z_d = self._slice_descrpt_deriv(self.t_dipole_z_d) # (nframes x natoms_sel) x ndescrpt self.t_dipole_x_d = tf.reshape(self.t_dipole_x_d, [nfxnas, self.ndescrpt]) self.t_dipole_y_d = tf.reshape(self.t_dipole_y_d, [nfxnas, self.ndescrpt]) self.t_dipole_z_d = tf.reshape(self.t_dipole_z_d, [nfxnas, self.ndescrpt]) # (nframes x natoms_sel) x 3 x ndescrpt self.t_dipole_d = tf.concat( [self.t_dipole_x_d, self.t_dipole_y_d, self.t_dipole_z_d], axis=1) self.t_dipole_d = tf.reshape(self.t_dipole_d, [nfxnas, 3 * self.ndescrpt]) # (nframes x natoms_sel) x 3 x ndescrpt self.t_dipole_d = tf.reshape(self.t_dipole_d, [-1, 3, self.ndescrpt]) # (nframes x natoms_sel) x 1 x ndescrpt self.t_ef_d = tf.matmul(self.t_ef_reshape, self.t_dipole_d) # nframes x (natoms_sel x ndescrpt) self.t_ef_d = tf.reshape(self.t_ef_d, [t_nframes, -1]) # nframes x (natoms x ndescrpt) self.t_ef_d = self._enrich(self.t_ef_d, dof=self.ndescrpt) self.t_ef_d = tf.reshape(self.t_ef_d, [nf, self.t_natoms[0] * self.ndescrpt]) # t_ef_d is force (with -1), prod_forc takes deriv, so we need the opposite self.t_ef_d_oppo = -self.t_ef_d force = op_module.prod_force_se_a(self.t_ef_d_oppo, self.descrpt_deriv, self.nlist, self.t_natoms, n_a_sel=self.nnei_a, n_r_sel=self.nnei_r) virial, atom_virial \ = op_module.prod_virial_se_a (self.t_ef_d_oppo, self.descrpt_deriv, self.rij, self.nlist, self.t_natoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) force = tf.identity(force, name='o_dm_force') virial = tf.identity(virial, name='o_dm_virial') atom_virial = tf.identity(atom_virial, name='o_dm_av') return force, virial, atom_virial
def _filter_lower( self, type_i, type_input, start_index, incrs_index, inputs, nframes, natoms, type_embedding=None, is_exclude=False, activation_fn=None, bavg=0.0, stddev=1.0, trainable=True, suffix='', ): """ input env matrix, returns R.G """ outputs_size = [1] + self.filter_neuron # cut-out inputs # with natom x (nei_type_i x 4) inputs_i = tf.slice(inputs, [0, start_index * 4], [-1, incrs_index * 4]) shape_i = inputs_i.get_shape().as_list() natom = tf.shape(inputs_i)[0] # with (natom x nei_type_i) x 4 inputs_reshape = tf.reshape(inputs_i, [-1, 4]) # with (natom x nei_type_i) x 1 xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1]) if type_embedding is not None: xyz_scatter = self._concat_type_embedding(xyz_scatter, nframes, natoms, type_embedding) if self.compress: raise RuntimeError( 'compression of type embedded descriptor is not supported at the moment' ) # natom x 4 x outputs_size if self.compress and (not is_exclude): info = [ self.lower, self.upper, self.upper * self.table_config[0], self.table_config[1], self.table_config[2], self.table_config[3] ] if self.type_one_side: net = 'filter_-1_net_' + str(type_i) else: net = 'filter_' + str(type_input) + '_net_' + str(type_i) return op_module.tabulate_fusion_se_a( tf.cast(self.table.data[net], self.filter_precision), info, xyz_scatter, tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), last_layer_size=outputs_size[-1]) else: if (not is_exclude): # with (natom x nei_type_i) x out_size xyz_scatter = embedding_net( xyz_scatter, self.filter_neuron, self.filter_precision, activation_fn=activation_fn, resnet_dt=self.filter_resnet_dt, name_suffix=suffix, stddev=stddev, bavg=bavg, seed=self.seed, trainable=trainable, uniform_seed=self.uniform_seed, initial_variables=self.embedding_net_variables, mixed_prec=self.mixed_prec) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift else: # we can safely return the final xyz_scatter filled with zero directly return tf.cast(tf.fill((natom, 4, outputs_size[-1]), 0.), self.filter_precision) # natom x nei_type_i x out_size xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1])) # When using tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]) below # [588 24] -> [588 6 4] correct # but if sel is zero # [588 0] -> [147 0 4] incorrect; the correct one is [588 0 4] # So we need to explicitly assign the shape to tf.shape(inputs_i)[0] instead of -1 # natom x 4 x outputs_size return tf.matmul(tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), xyz_scatter, transpose_a=True)
def _filter(self, inputs, type_input, natoms, type_embedding=None, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, trainable=True): nframes = tf.shape(tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]))[0] # natom x (nei x 4) shape = inputs.get_shape().as_list() outputs_size = [1] + self.filter_neuron outputs_size_2 = self.n_axis_neuron all_excluded = all([(type_input, type_i) in self.exclude_types for type_i in range(self.ntypes)]) if all_excluded: # all types are excluded so result and qmat should be zeros # we can safaly return a zero matrix... # See also https://stackoverflow.com/a/34725458/9567349 # result: natom x outputs_size x outputs_size_2 # qmat: natom x outputs_size x 3 natom = tf.shape(inputs)[0] result = tf.cast( tf.fill((natom, outputs_size_2, outputs_size[-1]), 0.), GLOBAL_TF_FLOAT_PRECISION) qmat = tf.cast(tf.fill((natom, outputs_size[-1], 3), 0.), GLOBAL_TF_FLOAT_PRECISION) return result, qmat with tf.variable_scope(name, reuse=reuse): start_index = 0 type_i = 0 # natom x 4 x outputs_size if type_embedding is None: rets = [] for type_i in range(self.ntypes): ret = self._filter_lower(type_i, type_input, start_index, self.sel_a[type_i], inputs, nframes, natoms, type_embedding=type_embedding, is_exclude=(type_input, type_i) in self.exclude_types, activation_fn=activation_fn, stddev=stddev, bavg=bavg, trainable=trainable, suffix="_" + str(type_i)) if (type_input, type_i) not in self.exclude_types: # add zero is meaningless; skip rets.append(ret) start_index += self.sel_a[type_i] # faster to use accumulate_n than multiple add xyz_scatter_1 = tf.accumulate_n(rets) else: xyz_scatter_1 = self._filter_lower( type_i, type_input, start_index, np.cumsum(self.sel_a)[-1], inputs, nframes, natoms, type_embedding=type_embedding, is_exclude=False, activation_fn=activation_fn, stddev=stddev, bavg=bavg, trainable=trainable) # natom x nei x outputs_size # xyz_scatter = tf.concat(xyz_scatter_total, axis=1) # natom x nei x 4 # inputs_reshape = tf.reshape(inputs, [-1, shape[1]//4, 4]) # natom x 4 x outputs_size # xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a = True) if self.original_sel is None: # shape[1] = nnei * 4 nnei = shape[1] / 4 else: nnei = tf.cast( tf.Variable(np.sum(self.original_sel), dtype=tf.int32, trainable=False, name="nnei"), self.filter_precision) xyz_scatter_1 = xyz_scatter_1 / nnei # natom x 4 x outputs_size_2 xyz_scatter_2 = tf.slice(xyz_scatter_1, [0, 0, 0], [-1, -1, outputs_size_2]) # # natom x 3 x outputs_size_2 # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1]) # natom x 3 x outputs_size_1 qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1]) # natom x outputs_size_1 x 3 qmat = tf.transpose(qmat, perm=[0, 2, 1]) # natom x outputs_size x outputs_size_2 result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True) # natom x (outputs_size x outputs_size_2) result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]]) return result, qmat
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 _filter_type_ext(self, inputs, natoms, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, trainable=True): # natom x (nei x 4) outputs_size = [1] + self.filter_neuron outputs_size_2 = self.n_axis_neuron with tf.variable_scope(name, reuse=reuse): start_index = 0 result_all = [] xyz_scatter_1_all = [] xyz_scatter_2_all = [] for type_i in range(self.ntypes): # cut-out inputs # with natom x (nei_type_i x 4) inputs_i = tf.slice(inputs, [0, start_index * 4], [-1, self.sel_a[type_i] * 4]) start_index += self.sel_a[type_i] shape_i = inputs_i.get_shape().as_list() # with (natom x nei_type_i) x 4 inputs_reshape = tf.reshape(inputs_i, [-1, 4]) xyz_scatter = tf.reshape( tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1]) for ii in range(1, len(outputs_size)): w = tf.get_variable( 'matrix_' + str(ii) + '_' + str(type_i), [outputs_size[ii - 1], outputs_size[ii]], self.filter_precision, tf.random_normal_initializer( stddev=stddev / np.sqrt(outputs_size[ii] + outputs_size[ii - 1]), seed=seed), trainable=trainable) b = tf.get_variable('bias_' + str(ii) + '_' + str(type_i), [1, outputs_size[ii]], self.filter_precision, tf.random_normal_initializer( stddev=stddev, mean=bavg, seed=seed), trainable=trainable) if self.filter_resnet_dt: idt = tf.get_variable( 'idt_' + str(ii) + '_' + str(type_i), [1, outputs_size[ii]], self.filter_precision, tf.random_normal_initializer(stddev=0.001, mean=1.0, seed=seed), trainable=trainable) if outputs_size[ii] == outputs_size[ii - 1]: if self.filter_resnet_dt: xyz_scatter += activation_fn( tf.matmul(xyz_scatter, w) + b) * idt else: xyz_scatter += activation_fn( tf.matmul(xyz_scatter, w) + b) elif outputs_size[ii] == outputs_size[ii - 1] * 2: if self.filter_resnet_dt: xyz_scatter = tf.concat( [xyz_scatter, xyz_scatter], 1) + activation_fn( tf.matmul(xyz_scatter, w) + b) * idt else: xyz_scatter = tf.concat( [xyz_scatter, xyz_scatter], 1) + activation_fn( tf.matmul(xyz_scatter, w) + b) else: xyz_scatter = activation_fn( tf.matmul(xyz_scatter, w) + b) # natom x nei_type_i x out_size xyz_scatter = tf.reshape( xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1])) # natom x nei_type_i x 4 inputs_i_reshape = tf.reshape(inputs_i, [-1, shape_i[1] // 4, 4]) # natom x 4 x outputs_size xyz_scatter_1 = tf.matmul(inputs_i_reshape, xyz_scatter, transpose_a=True) xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape_i[1]) # natom x 4 x outputs_size_2 xyz_scatter_2 = tf.slice(xyz_scatter_1, [0, 0, 0], [-1, -1, outputs_size_2]) xyz_scatter_1_all.append(xyz_scatter_1) xyz_scatter_2_all.append(xyz_scatter_2) # for type_i in range(self.ntypes): # for type_j in range(type_i, self.ntypes): # # natom x outputs_size x outputs_size_2 # result = tf.matmul(xyz_scatter_1_all[type_i], xyz_scatter_2_all[type_j], transpose_a = True) # # natom x (outputs_size x outputs_size_2) # result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]]) # result_all.append(tf.identity(result)) xyz_scatter_2_coll = tf.concat(xyz_scatter_2_all, axis=2) for type_i in range(self.ntypes): # natom x outputs_size x (outputs_size_2 x ntypes) result = tf.matmul(xyz_scatter_1_all[type_i], xyz_scatter_2_coll, transpose_a=True) # natom x (outputs_size x outputs_size_2 x ntypes) result = tf.reshape( result, [-1, outputs_size_2 * self.ntypes * outputs_size[-1]]) result_all.append(tf.identity(result)) # natom x (ntypes x outputs_size x outputs_size_2 x ntypes) result_all = tf.concat(result_all, axis=1) return result_all
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])