feat(loss): add new loss options to control number bands to exclude in in dft bands.

dptb/nnops/loss.py

@@ -84,6 +84,8 @@ def __init__(
             diff_on: bool=False,
             eout_weight: float=0.01,
             diff_weight: float=0.01,
+            diff_valence: dict=None,
+            spin_deg: int = 2,
             dtype: Union[str, torch.dtype] = torch.float32, 
             device: Union[str, torch.device] = torch.device("cpu"),
             **kwargs,
@@ -94,6 +96,8 @@ def __init__(
         self.diff_on = diff_on
         self.eout_weight = eout_weight
         self.diff_weight = diff_weight
+        self.diff_valence = diff_valence  
+        self.spin_deg = spin_deg  
 
         if basis is not None:
             self.idp = OrbitalMapper(basis, method="e3tb", device=self.device)
@@ -157,6 +161,15 @@ def forward(
             eig_pred = data[AtomicDataDict.ENERGY_EIGENVALUE_KEY][0] # (n_kpt, n_band)
             eig_label = ref_data[AtomicDataDict.ENERGY_EIGENVALUE_KEY][0] # (n_kpt, n_band_dft/n_band)
 
+            if self.diff_valence is not None and isinstance(self.diff_valence, dict):
+                nbands_exclude = sum([self.diff_valence[self.idp.type_to_chemical_symbol[int(ii)]] for ii in ref_data['atom_types']])
+                assert nbands_exclude % self.spin_deg == 0
+                nbands_exclude = nbands_exclude // self.spin_deg
+            else:
+                nbands_exclude = 0
+
+            eig_label = eig_label[:,nbands_exclude:]
+
             norbs = eig_pred.shape[-1]
             nbanddft = eig_label.shape[-1]
             num_kp = eig_label.shape[-2]

dptb/utils/argcheck.py

@@ -714,6 +714,8 @@ def loss_options():
         Argument("diff_on", bool, optional=True, default=False, doc="Whether to use random differences in loss function. Default: False"),
         Argument("eout_weight", float, optional=True, default=0.01, doc="The weight of eigenvalue out of range. Default: 0.01"),
         Argument("diff_weight", float, optional=True, default=0.01, doc="The weight of eigenvalue difference. Default: 0.01"),
+        Argument("diff_valence", [dict,None], optional=True, default=None, doc="set the difference of the number of valence electrons in DFT and TB. eg {'A':6,'B':7}, Default: None, which means no difference"),
+        Argument("spin_deg", int, optional=True, default=2, doc="The spin degeneracy of band structure. Default: 2"),
     ]
 
     skints = [

dptb/utils/config_sk.py

@@ -30,6 +30,8 @@
         "loss_options": {
             "train": {
                 "method": "eigvals",
+                "diff_valence":None,
+                "spin_deg":2,
                 "diff_on": False,
                 "eout_weight": 0.01,
                 "diff_weight": 0.01