implement dipole fitting by Linfeng

Author: Han Wang

source/train/Fitting.py

@@ -421,3 +421,76 @@ def build (self,
             count += 1
 
         return tf.reshape(outs, [-1])
+
+
+class DipoleFittingSeA () :
+    def __init__ (self, jdata, descrpt) :
+        if not isinstance(descrpt, DescrptSeA) :
+            raise RuntimeError('DipoleFittingSeA only supports DescrptSeA')
+        self.ntypes = descrpt.get_ntypes()
+        self.dim_descrpt = descrpt.get_dim_out()
+        args = ClassArg()\
+               .add('neuron',           list,   default = [120,120,120], alias = 'n_neuron')\
+               .add('resnet_dt',        bool,   default = True)\
+               .add('sel_type',         [list,int],   default = [ii for ii in range(self.ntypes)], alias = 'dipole_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.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
+        self.useBN = False
+
+    def get_sel_type(self):
+        return self.sel_type
+
+    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
+            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 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.reshape(outs, [-1])

source/train/Model.py

@@ -325,6 +325,11 @@ def __init__(self, jdata, descrpt, fitting) :
         TensorModel.__init__(self, jdata, descrpt, fitting, 'wfc')
 
 
+class DipoleModel(TensorModel):
+    def __init__(self, jdata, descrpt, fitting) :
+        TensorModel.__init__(self, jdata, descrpt, fitting, 'dipole')
+
+
 class PolarModel(TensorModel):
     def __init__(self, jdata, descrpt, fitting) :
         TensorModel.__init__(self, jdata, descrpt, fitting, 'polar')

source/train/Trainer.py

@@ -12,12 +12,12 @@
 from deepmd.RunOptions import global_ener_float_precision
 from deepmd.RunOptions import global_cvt_2_tf_float
 from deepmd.RunOptions import global_cvt_2_ener_float
-from deepmd.Fitting import EnerFitting, WFCFitting, PolarFittingLocFrame, PolarFittingSeA
+from deepmd.Fitting import EnerFitting, WFCFitting, PolarFittingLocFrame, PolarFittingSeA, DipoleFittingSeA
 from deepmd.DescrptLocFrame import DescrptLocFrame
 from deepmd.DescrptSeA import DescrptSeA
 from deepmd.DescrptSeR import DescrptSeR
 from deepmd.DescrptSeAR import DescrptSeAR
-from deepmd.Model import Model, WFCModel, PolarModel
+from deepmd.Model import Model, WFCModel, DipoleModel, PolarModel
 from deepmd.Loss import EnerStdLoss, TensorLoss
 from deepmd.LearningRate import LearningRateExp
 
@@ -101,6 +101,11 @@ def _init_param(self, jdata):
                 self.fitting = PolarFittingSeA(fitting_param, self.descrpt)
             else :
                 raise RuntimeError('fitting polar only supports descrptors: loc_frame and se_a')
+        elif fitting_type == 'dipole':
+            if descrpt_type == 'se_a':
+                self.fitting = DipoleFittingSeA(fitting_param, self.descrpt)
+            else :
+                raise RuntimeError('fitting dipole only supports descrptors: se_a')
         else :
             raise RuntimeError('unknow fitting type ' + fitting_type)
 
@@ -112,6 +117,8 @@ def _init_param(self, jdata):
             self.model = WFCModel(model_param, self.descrpt, self.fitting)
         elif fitting_type == 'polar':
             self.model = PolarModel(model_param, self.descrpt, self.fitting)
+        elif fitting_type == 'dipole':
+            self.model = DipoleModel(model_param, self.descrpt, self.fitting)
         else :
             raise RuntimeError('get unknown fitting type when building model')
 
@@ -140,6 +147,12 @@ def _init_param(self, jdata):
                                    tensor_name = 'wfc',
                                    tensor_size = self.model.get_out_size(),
                                    label_name = 'wfc')
+        elif fitting_type == 'dipole':
+            self.loss = TensorLoss(loss_param, 
+                                   model = self.model, 
+                                   tensor_name = 'dipole',
+                                   tensor_size = 3,
+                                   label_name = 'dipole')
         elif fitting_type == 'polar':
             self.loss = TensorLoss(loss_param, 
                                    model = self.model,