add fitting global polarizability

Author: Han Wang

source/train/Fitting.py

@@ -423,6 +423,34 @@ def build (self,
         return tf.reshape(outs, [-1])
 
 
+class GlobalPolarFittingSeA () :
+    def __init__ (self, jdata, descrpt) :
+        if not isinstance(descrpt, DescrptSeA) :
+            raise RuntimeError('GlobalPolarFittingSeA only supports DescrptSeA')
+        self.ntypes = descrpt.get_ntypes()
+        self.dim_descrpt = descrpt.get_dim_out()
+        self.polar_fitting = PolarFittingSeA(jdata, descrpt)
+
+    def get_sel_type(self):
+        return self.polar_fitting.get_sel_type()
+
+    def get_out_size(self):
+        return self.polar_fitting.get_out_size()
+
+    def build (self,
+               input_d,
+               rot_mat,
+               natoms,
+               reuse = None,
+               suffix = '') :
+        inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]])
+        outs = self.polar_fitting.build(input_d, rot_mat, natoms, reuse, suffix)
+        # nframes x natoms x 9
+        outs = tf.reshape(outs, [tf.shape(inputs)[0], -1, 9])
+        outs = tf.reduce_sum(outs, axis = 1)
+        return tf.reshape(outs, [-1])
+
+
 class DipoleFittingSeA () :
     def __init__ (self, jdata, descrpt) :
         if not isinstance(descrpt, DescrptSeA) :

source/train/Loss.py

@@ -189,10 +189,11 @@ def __init__ (self, jdata, **kwarg) :
         self.tensor_name = kwarg['tensor_name']
         self.tensor_size = kwarg['tensor_size']
         self.label_name = kwarg['label_name']
+        self.atomic = kwarg.get('atomic', True)
         # data required
         add_data_requirement(self.label_name, 
                              self.tensor_size, 
-                             atomic=True,  
+                             atomic=self.atomic,  
                              must=True, 
                              high_prec=False, 
                              type_sel = type_sel)
@@ -206,6 +207,9 @@ def build (self,
         polar_hat = label_dict[self.label_name]
         polar = model_dict[self.tensor_name]
         l2_loss = tf.reduce_mean( tf.square(polar - polar_hat), name='l2_'+suffix)
+        if not self.atomic :
+            atom_norm  = 1./ global_cvt_2_tf_float(natoms[0]) 
+            l2_loss = l2_loss * atom_norm
         self.l2_l = l2_loss
         more_loss = {}
 

source/train/Model.py

@@ -333,3 +333,10 @@ def __init__(self, jdata, descrpt, fitting) :
 class PolarModel(TensorModel):
     def __init__(self, jdata, descrpt, fitting) :
         TensorModel.__init__(self, jdata, descrpt, fitting, 'polar')
+
+
+class GlobalPolarModel(TensorModel):
+    def __init__(self, jdata, descrpt, fitting) :
+        TensorModel.__init__(self, jdata, descrpt, fitting, 'global_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, DipoleFittingSeA
+from deepmd.Fitting import EnerFitting, WFCFitting, PolarFittingLocFrame, PolarFittingSeA, GlobalPolarFittingSeA, 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, DipoleModel, PolarModel
+from deepmd.Model import Model, WFCModel, DipoleModel, PolarModel, GlobalPolarModel
 from deepmd.Loss import EnerStdLoss, TensorLoss
 from deepmd.LearningRate import LearningRateExp
 
@@ -94,18 +94,23 @@ def _init_param(self, jdata):
             self.fitting = EnerFitting(fitting_param, self.descrpt)
         elif fitting_type == 'wfc':            
             self.fitting = WFCFitting(fitting_param, self.descrpt)
+        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')
         elif fitting_type == 'polar':
             if descrpt_type == 'loc_frame':
                 self.fitting = PolarFittingLocFrame(fitting_param, self.descrpt)
             elif descrpt_type == 'se_a':
                 self.fitting = PolarFittingSeA(fitting_param, self.descrpt)
             else :
                 raise RuntimeError('fitting polar only supports descrptors: loc_frame and se_a')
-        elif fitting_type == 'dipole':
+        elif fitting_type == 'global_polar':
             if descrpt_type == 'se_a':
-                self.fitting = DipoleFittingSeA(fitting_param, self.descrpt)
+                self.fitting = GlobalPolarFittingSeA(fitting_param, self.descrpt)
             else :
-                raise RuntimeError('fitting dipole only supports descrptors: se_a')
+                raise RuntimeError('fitting global_polar only supports descrptors: loc_frame and se_a')
         else :
             raise RuntimeError('unknow fitting type ' + fitting_type)
 
@@ -115,10 +120,12 @@ def _init_param(self, jdata):
             self.model = Model(model_param, self.descrpt, self.fitting)
         elif fitting_type == 'wfc':
             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)
+        elif fitting_type == 'polar':
+            self.model = PolarModel(model_param, self.descrpt, self.fitting)
+        elif fitting_type == 'global_polar':
+            self.model = GlobalPolarModel(model_param, self.descrpt, self.fitting)
         else :
             raise RuntimeError('get unknown fitting type when building model')
 
@@ -159,6 +166,13 @@ def _init_param(self, jdata):
                                    tensor_name = 'polar',
                                    tensor_size = 9,
                                    label_name = 'polarizability')
+        elif fitting_type == 'global_polar':
+            self.loss = TensorLoss(loss_param, 
+                                   model = self.model, 
+                                   tensor_name = 'global_polar',
+                                   tensor_size = 9,
+                                   atomic = False,
+                                   label_name = 'polarizability')
         else :
             raise RuntimeError('get unknown fitting type when building loss function')