add energy dipole loss

Author: Han Wang

source/train/Loss.py

@@ -178,6 +178,100 @@ def print_on_training(self,
         return print_str        
 
 
+class EnerDipoleLoss () :
+    def __init__ (self, jdata, **kwarg) :
+        self.starter_learning_rate = kwarg['starter_learning_rate']
+        args = ClassArg()\
+            .add('start_pref_e',        float,  must = True, default = 0.1) \
+            .add('limit_pref_e',        float,  must = True, default = 1.00)\
+            .add('start_pref_ed',       float,  must = True, default = 1.00)\
+            .add('limit_pref_ed',       float,  must = True, default = 1.00)
+        class_data = args.parse(jdata)
+        self.start_pref_e = class_data['start_pref_e']
+        self.limit_pref_e = class_data['limit_pref_e']
+        self.start_pref_ed = class_data['start_pref_ed']
+        self.limit_pref_ed = class_data['limit_pref_ed']
+        # data required
+        add_data_requirement('energy', 1, atomic=False, must=True, high_prec=True)
+        add_data_requirement('energy_dipole', 3, atomic=False, must=True, high_prec=False)
+
+    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 print_header(self) :
+        prop_fmt = '   %9s %9s'
+        print_str = ''
+        print_str += prop_fmt % ('l2_tst', 'l2_trn')
+        print_str += prop_fmt % ('l2_e_tst', 'l2_e_trn')
+        print_str += prop_fmt % ('l2_ed_tst', 'l2_ed_trn')
+        return print_str
+
+
+    def print_on_training(self, 
+                          sess, 
+                          natoms,
+                          feed_dict_test,
+                          feed_dict_batch) :
+        error_test, error_e_test, error_ed_test\
+            = sess.run([self.l2_l, \
+                        self.l2_more['l2_ener_loss'], \
+                        self.l2_more['l2_ener_dipole_loss']],
+                       feed_dict=feed_dict_test)
+        error_train, error_e_train, error_ed_train\
+            = sess.run([self.l2_l, \
+                        self.l2_more['l2_ener_loss'], \
+                        self.l2_more['l2_ener_dipole_loss']],
+                       feed_dict=feed_dict_batch)
+        print_str = ""
+        prop_fmt = "   %9.2e %9.2e"
+        print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train))
+        print_str += prop_fmt % (np.sqrt(error_e_test) / natoms[0], np.sqrt(error_e_train) / natoms[0])
+        print_str += prop_fmt % (np.sqrt(error_ed_test), np.sqrt(error_ed_train))
+        return print_str        
+
 
 class TensorLoss () :
     def __init__ (self, jdata, **kwarg) :

source/train/Model.py

@@ -270,6 +270,8 @@ def build (self,
         model_dict['virial'] = virial
         model_dict['atom_ener'] = energy_raw
         model_dict['atom_virial'] = atom_virial
+        model_dict['coord'] = coord
+        model_dict['atype'] = atype
 
         return model_dict
 

source/train/Trainer.py

@@ -18,7 +18,7 @@
 from deepmd.DescrptSeR import DescrptSeR
 from deepmd.DescrptSeAR import DescrptSeAR
 from deepmd.Model import Model, WFCModel, DipoleModel, PolarModel, GlobalPolarModel
-from deepmd.Loss import EnerStdLoss, TensorLoss
+from deepmd.Loss import EnerStdLoss, EnerDipoleLoss, TensorLoss
 from deepmd.LearningRate import LearningRateExp
 
 from tensorflow.python.framework import ops
@@ -146,8 +146,15 @@ def _init_param(self, jdata):
             loss_param = jdata['loss']
         except:
             loss_param = None
+        loss_type = loss_param.get('type', 'std')
+
         if fitting_type == 'ener':
-            self.loss = EnerStdLoss(loss_param, starter_learning_rate = self.lr.start_lr())
+            if loss_type == 'std':
+                self.loss = EnerStdLoss(loss_param, starter_learning_rate = self.lr.start_lr())
+            elif loss_type == 'ener_dipole':
+                self.loss = EnerDipoleLoss(loss_param, starter_learning_rate = self.lr.start_lr())
+            else:
+                raise RuntimeError('unknow loss type')
         elif fitting_type == 'wfc':
             self.loss = TensorLoss(loss_param, 
                                    model = self.model, 
@@ -262,7 +269,6 @@ def _build_network(self, data):
         self.place_holders['natoms_vec']        = tf.placeholder(tf.int32,   [self.ntypes+2], name='t_natoms')
         self.place_holders['default_mesh']      = tf.placeholder(tf.int32,   [None], name='t_mesh')
         self.place_holders['is_training']       = tf.placeholder(tf.bool)
-
         self.model_pred\
             = self.model.build (self.place_holders['coord'], 
                                 self.place_holders['type'],