polarizability: add option to only fit the diag part

Author: Han Wang

examples/water/train/polar_se_a.json

@@ -17,7 +17,8 @@
 	},
 	"fitting_net": {
 	    "type":		"polar",
-	    "pol_type":		[0],
+	    "sel_type":		[0],
+	    "fit_diag":		true,
 	    "neuron":		[100, 100, 100],
 	    "resnet_dt":	true,
 	    "seed":		1,
@@ -28,7 +29,7 @@
 
     "learning_rate" :{
 	"type":		"exp",
-	"start_lr":	0.001,
+	"start_lr":	0.01,
 	"decay_steps":	5000,
 	"decay_rate":	0.95,
 	"_comment":	"that's all"

source/train/Fitting.py

@@ -352,12 +352,14 @@ def __init__ (self, jdata, descrpt) :
         args = ClassArg()\
                .add('neuron',           list,   default = [120,120,120], alias = 'n_neuron')\
                .add('resnet_dt',        bool,   default = True)\
+               .add('fit_diag',         bool,   default = True)\
                .add('sel_type',         [list,int],   default = [ii for ii in range(self.ntypes)], alias = 'pol_type')\
                .add('seed',             int)
         class_data = args.parse(jdata)
         self.n_neuron = class_data['neuron']
         self.resnet_dt = class_data['resnet_dt']
         self.sel_type = class_data['sel_type']
+        self.fit_diag = class_data['fit_diag']
         self.seed = class_data['seed']
         self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
         self.dim_rot_mat = self.dim_rot_mat_1 * 3
@@ -400,12 +402,20 @@ def build (self,
                     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])
+            if self.fit_diag :
+                # (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)
+                # (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 :                
+                # (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)