add bandgap loss

Author: xuan112358

deepks/model/evaluator.py

@@ -8,7 +8,7 @@
 except ImportError as e:
     sys.path.append(os.path.dirname(os.path.realpath(__file__)) + "/../../")
 from deepks.model.reader import generalized_eigh, eigh_wrapper
-from deepks.model.utils import get_density_matrix, cal_phi_loss, cal_v_delta, cal_vd_masked_loss, get_occ_func, make_loss
+from deepks.model.utils import get_density_matrix, cal_phi_loss, cal_v_delta, cal_vd_masked_loss, cal_bandgap, get_occ_func, make_loss
 
 class Evaluator:
     def __init__(self,
@@ -17,14 +17,15 @@ def __init__(self,
                  v_delta_factor=0., 
                  phi_factor=0., phi_occ=0,
                  band_factor=0.,band_occ=0,
+                 bandgap_factor=0.,bandgap_occ=0,
                  density_m_factor=0.,density_m_occ=0,
                  phi_align_factor=0., phi_align_occ=0,
                  density_factor=0., grad_penalty=0., 
                  energy_lossfn=None, force_lossfn=None, 
                  stress_lossfn=None, orbital_lossfn=None,
                  v_delta_lossfn=None, phi_lossfn=None,
                  phi_align_lossfn=None,
-                 band_lossfn=None, density_m_lossfn=None,
+                 band_lossfn=None, bandgap_lossfn=None, density_m_lossfn=None,
                  energy_per_atom=0,vd_divide_by_nlocal=False,
                  vd_masked_loss=False, 
                  vd_masked_S_threshold=1e-6, vd_masked_H_threshold=1e-6,
@@ -81,6 +82,14 @@ def __init__(self,
         self.band_factor = band_factor
         self.band_lossfn = band_lossfn   
         self.get_band_occ = get_occ_func(band_occ)   
+        # bandgap term
+        if bandgap_lossfn is None:
+            bandgap_lossfn = {}
+        if isinstance(bandgap_lossfn, dict):
+            bandgap_lossfn = make_loss(**bandgap_lossfn)
+        self.bandgap_factor = bandgap_factor
+        self.bandgap_lossfn = bandgap_lossfn   
+        self.get_bandgap_occ = get_occ_func(bandgap_occ)   
         #density matrix term
         if density_m_lossfn is None:
             density_m_lossfn = {}
@@ -136,6 +145,7 @@ def __call__(self, model, sample):
                         or (self.vd_factor > 0 and "lb_vd" in sample)
                         or (self.phi_factor > 0 and "lb_phi" in sample)
                         or (self.band_factor > 0 and "lb_band" in sample)
+                        or (self.bandgap_factor > 0 and "lb_band" in sample)
                         or (self.density_m_factor > 0)
                         or (self.d_factor > 0 and "gldv" in sample)
                         or self.g_penalty > 0)
@@ -178,7 +188,8 @@ def __call__(self, model, sample):
                 loss.append(self.o_factor * self.o_lossfn(o_pred, o_label))
             # optional v_delta/phi/band_energy/density_matrix/phi_alignment calculation
             if (self.vd_factor > 0 and "lb_vd" in sample) or (self.phi_factor > 0 and "lb_phi" in sample) \
-                or (self.band_factor > 0 and "lb_band" in sample) or (self.density_m_factor > 0 and "lb_phi" in sample) \
+                or (self.band_factor > 0 and "lb_band" in sample) or (self.bandgap_factor > 0 and "lb_band" in sample) \
+                or (self.density_m_factor > 0 and "lb_phi" in sample) \
                 or (self.phi_align_factor > 0 and "lb_phi" in sample and "lb_band" in sample):
                 # cal v_delta
                 if "vdp" in sample:
@@ -204,7 +215,7 @@ def __call__(self, model, sample):
                     tot_loss = tot_loss + vd_loss
                     loss.append(vd_loss)
 
-                if (self.phi_factor > 0 and "lb_phi" in sample) or (self.band_factor > 0 and "lb_band" in sample) or (self.density_m_factor > 0 and "lb_phi" in sample):
+                if (self.phi_factor > 0 and "lb_phi" in sample) or (self.band_factor > 0 and "lb_band" in sample) or (self.bandgap_factor > 0 and "lb_band" in sample) or (self.density_m_factor > 0 and "lb_phi" in sample):
                     h_base = sample["h_base"]
                     if "trans_matrix" in sample:
                         trans_matrix=sample["trans_matrix"]
@@ -225,6 +236,15 @@ def __call__(self, model, sample):
                         tot_loss = tot_loss + band_loss
                         # print("occ_band",band_pred[...,:band_occ],band_label[...,:band_occ])
                         loss.append(band_loss)
+                    # optional bandgap calculation
+                    if self.bandgap_factor > 0 and "lb_band" in sample:
+                        band_label=sample["lb_band"]
+                        bandgap_occ=self.get_bandgap_occ(natom)
+                        bandgap_label=cal_bandgap(band_label, bandgap_occ)
+                        bandgap_pred=cal_bandgap(band_pred, bandgap_occ)
+                        bandgap_loss = self.bandgap_factor * self.bandgap_lossfn(bandgap_pred, bandgap_label)
+                        tot_loss = tot_loss + bandgap_loss
+                        loss.append(bandgap_loss)
                     # optional density matrix calculation
                     if self.density_m_factor > 0 and "lb_phi" in sample:
                         # calculate density_m_label every time, kind of waste of time
@@ -283,6 +303,9 @@ def print_head(self,name,data_keys,align_len=20):
         # optional band energy calculation
         if self.band_factor > 0 and "lb_band" in data_keys:
             info+=f"{name}_band".rjust(align_len)
+        # optional bandgap calculation
+        if self.bandgap_factor > 0 and "lb_band" in data_keys:
+            info+=f"{name}_bandgap".rjust(align_len)
         # optional density matrix calculation
         if self.density_m_factor > 0 and "lb_phi" in data_keys:
             info+=f"{name}_dm".rjust(align_len)    

deepks/model/train.py

@@ -16,8 +16,8 @@
 from deepks.model.evaluator import Evaluator, NatomLossList
 
 def train(model, g_reader, n_epoch=1000, test_reader=None, *,
-          energy_factor=1., force_factor=0., stress_factor=0., orbital_factor=0., v_delta_factor=0., phi_factor=0.,phi_occ=0, band_factor=0., band_occ=0, density_m_factor=0., density_m_occ=0, phi_align_factor=0., phi_align_occ=0, density_factor=0.,
-          energy_loss=None, force_loss=None, stress_loss=None, orbital_loss=None, v_delta_loss=None, phi_loss=None, band_loss=None, density_m_loss=None, phi_align_loss=None, grad_penalty=0.,
+          energy_factor=1., force_factor=0., stress_factor=0., orbital_factor=0., v_delta_factor=0., phi_factor=0.,phi_occ=0, band_factor=0., band_occ=0, bandgap_factor=0., bandgap_occ=0, density_m_factor=0., density_m_occ=0, phi_align_factor=0., phi_align_occ=0, density_factor=0.,
+          energy_loss=None, force_loss=None, stress_loss=None, orbital_loss=None, v_delta_loss=None, phi_loss=None, band_loss=None, bandgap_loss=None, density_m_loss=None, phi_align_loss=None, grad_penalty=0.,
           energy_per_atom=0, vd_divide_by_nlocal=False, vd_masked_loss=False, vd_masked_S_threshold=1e-6, vd_masked_H_threshold=1e-6, use_safe_eigh=False,
           start_lr=0.001, decay_steps=100, decay_rate=0.96, stop_lr=None, decay_rate_iter=None,
           weight_decay=0.,  fix_embedding=False,
@@ -53,12 +53,14 @@ def train(model, g_reader, n_epoch=1000, test_reader=None, *,
                           v_delta_factor=v_delta_factor,
                           phi_factor=phi_factor, phi_occ=phi_occ,
                           band_factor=band_factor, band_occ=band_occ,
+                          bandgap_factor=bandgap_factor, bandgap_occ=bandgap_occ,
                           density_m_factor=density_m_factor, density_m_occ=density_m_occ,
                           phi_align_factor=phi_align_factor, phi_align_occ=phi_align_occ,
                           energy_lossfn=energy_loss, force_lossfn=force_loss,
                           stress_lossfn=stress_loss, orbital_lossfn=orbital_loss,
                           v_delta_lossfn=v_delta_loss,phi_lossfn=phi_loss,
-                          band_lossfn=band_loss, density_m_lossfn=density_m_loss,
+                          band_lossfn=band_loss, bandgap_lossfn=bandgap_loss,
+                          density_m_lossfn=density_m_loss,
                           phi_align_lossfn=phi_align_loss,
                           density_factor=density_factor, grad_penalty=grad_penalty, 
                           energy_per_atom=energy_per_atom, vd_divide_by_nlocal=vd_divide_by_nlocal,
@@ -76,13 +78,14 @@ def train(model, g_reader, n_epoch=1000, test_reader=None, *,
                             v_delta_factor=to_one(v_delta_factor),
                             phi_factor=to_one(phi_factor), phi_occ=phi_occ,
                             band_factor=to_one(band_factor), band_occ=band_occ,
+                            bandgap_factor=to_one(bandgap_factor), bandgap_occ=bandgap_occ,
                             density_m_factor=to_one(density_m_factor), density_m_occ=density_m_occ,
                             phi_align_factor=to_one(phi_align_factor), phi_align_occ=phi_align_occ,
                             energy_lossfn=energy_loss, force_lossfn=force_loss,
                             stress_lossfn=stress_loss, orbital_lossfn=orbital_loss,
                             v_delta_lossfn=v_delta_loss,phi_lossfn=phi_loss,
-                            band_lossfn=band_loss, density_m_lossfn=density_m_loss,
-                            phi_align_lossfn=phi_align_loss,
+                            band_lossfn=band_loss, bandgap_lossfn=bandgap_loss,
+                            density_m_lossfn=density_m_loss, phi_align_lossfn=phi_align_loss,
                             density_factor=to_one(density_factor), grad_penalty=grad_penalty,
                             energy_per_atom=energy_per_atom, vd_divide_by_nlocal=vd_divide_by_nlocal,
                             vd_masked_loss=vd_masked_loss, vd_masked_S_threshold=vd_masked_S_threshold, vd_masked_H_threshold=vd_masked_H_threshold,

deepks/model/utils.py

@@ -187,6 +187,9 @@ def cal_vd_masked_loss(H_pred, H_label, S_matrix, S_threshold=1e-6, H_threshold=
 
     return masked_sum / (active_elements + 1e-12)
 
+def cal_bandgap(band, occ):
+    return band[...,occ] - band[...,occ-1]
+
 def get_occ_func(occ):
     # print("type:",type(occ))
     if isinstance(occ, int):