feat: add vd_masked_loss_width

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, cal_bandgap, get_occ_func, make_loss
+from deepks.model.utils import get_density_matrix, cal_phi_loss, cal_v_delta, cal_vd_masked_loss_hs, cal_vd_masked_loss_width, cal_bandgap, get_occ_func, make_loss
 
 class Evaluator:
     def __init__(self,
@@ -27,8 +27,9 @@ def __init__(self,
                  phi_align_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_loss=0, 
                  vd_masked_S_threshold=1e-6, vd_masked_H_threshold=1e-6,
+                 vd_masked_width=1,
                  use_safe_eigh=False):
         # energy term
         if energy_lossfn is None:
@@ -119,6 +120,8 @@ def __init__(self,
         # threshold for vd_masked_loss
         self.vd_masked_S_threshold=vd_masked_S_threshold
         self.vd_masked_H_threshold=vd_masked_H_threshold
+        # width for vd_masked_loss_width
+        self.vd_masked_width=vd_masked_width
 
     def __call__(self, model, sample):
         _dref = next(model.parameters()).device
@@ -205,8 +208,11 @@ def __call__(self, model, sample):
                 # optional v_delta calculation
                 if self.vd_factor > 0 and "lb_vd" in sample:
                     vd_label = sample["lb_vd"]
-                    if self.vd_masked_loss and "overlap" in sample:
-                        vd_loss = self.vd_factor * cal_vd_masked_loss(vd_pred, vd_label, sample["overlap"], self.vd_masked_S_threshold, self.vd_masked_H_threshold)
+                    if self.vd_masked_loss :
+                        if self.vd_masked_loss == 1 and "overlap" in sample:
+                            vd_loss = self.vd_factor * cal_vd_masked_loss_hs(vd_pred, vd_label, sample["overlap"], self.vd_masked_S_threshold, self.vd_masked_H_threshold)
+                        elif self.vd_masked_loss == 2:
+                            vd_loss = self.vd_factor * cal_vd_masked_loss_width(vd_pred, vd_label, self.vd_masked_width)
                     else:
                         vd_loss = self.vd_factor * self.vd_lossfn(vd_pred, vd_label)
                         # original: mean method,divide by nlocal**2. vd_divide_by_nlocal:divide by nlocal

deepks/model/train.py

@@ -18,7 +18,7 @@
 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, 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,
+          energy_per_atom=0, vd_divide_by_nlocal=False, vd_masked_loss=0, vd_masked_S_threshold=1e-6, vd_masked_H_threshold=1e-6, vd_masked_width=1, 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,
           display_epoch=100, display_detail_test=0, display_natom_loss=False, ckpt_file="model.pth",
@@ -64,7 +64,8 @@ def train(model, g_reader, n_epoch=1000, test_reader=None, *,
                           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,
-                          vd_masked_loss=vd_masked_loss, vd_masked_S_threshold=vd_masked_S_threshold, vd_masked_H_threshold=vd_masked_H_threshold,
+                          vd_masked_loss=vd_masked_loss, vd_masked_S_threshold=vd_masked_S_threshold,
+                          vd_masked_H_threshold=vd_masked_H_threshold, vd_masked_width=vd_masked_width,
                           use_safe_eigh=use_safe_eigh)
     if not display_detail_test:
         # make test evaluator that only returns l2loss of energy
@@ -88,7 +89,8 @@ def train(model, g_reader, n_epoch=1000, test_reader=None, *,
                             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,
+                            vd_masked_loss=vd_masked_loss, vd_masked_S_threshold=vd_masked_S_threshold,
+                            vd_masked_H_threshold=vd_masked_H_threshold,vd_masked_width=vd_masked_width,
                             use_safe_eigh=use_safe_eigh)
 
     print("# epoch      trn_err   tst_err        lr  trn_time  tst_time",end='')

deepks/model/utils.py

@@ -157,7 +157,7 @@ def cal_phi_loss(phi_pred,phi_label,phi_occ):
     #print("loss.shape:",loss.shape)
     return loss
 
-def cal_vd_masked_loss(H_pred, H_label, S_matrix, S_threshold=1e-6, H_threshold=1e-6):
+def cal_vd_masked_loss_hs(H_pred, H_label, S_matrix, S_threshold=1e-6, H_threshold=1e-6):
     """
     Computes the Mean Squared Error of Hamiltonian elements filtered by 
     the Overlap matrix and Hamiltonian matrix magnitude. Supports 4D tensors (nframe, nks, nlocal, nlocal).
@@ -175,7 +175,7 @@ def cal_vd_masked_loss(H_pred, H_label, S_matrix, S_threshold=1e-6, H_threshold=
     with torch.no_grad():
         # Generate mask where |S| > threshold or |H| > threshold. 
         # Diagonals (S_ii=1) are naturally included if threshold < 1.
-        mask = ( (torch.abs(S_matrix) > S_threshold) | (torch.abs(H_label) > H_threshold) ).to(H_pred.dtype)
+        mask = ( (torch.abs(S_matrix) > S_threshold) | (torch.abs(H_label) > H_threshold) ).to(device=H_pred.device, dtype=H_pred.dtype)
 
     # Compute element-wise squared difference
     diff_sq = (H_pred - H_label) ** 2
@@ -187,6 +187,34 @@ 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_vd_masked_loss_width(H_pred, H_label, width=1):
+    """
+    Computes the Mean Squared Error of Hamiltonian elements filtered by width .
+    Args:
+        H_pred (torch.Tensor): Predicted Hamiltonian.
+        H_label (torch.Tensor): Label Hamiltonian.
+        width (int): Width of the mask
+    Returns:
+        torch.Tensor: Scalar loss value.
+    """
+    with torch.no_grad():
+        nlocal = H_pred.size(-1)
+        i = torch.arange(nlocal, device=H_pred.device).view(nlocal, 1)  # (nlocal, 1)
+        j = torch.arange(nlocal, device=H_pred.device).view(1, nlocal)  # (1, nlocal)
+        
+        # calculate the shortest distance on the circle
+        diff = torch.abs(i - j)
+        dist = torch.minimum(diff, nlocal - diff)
+        
+        # mask is 1 if distance < width, 0 otherwise
+        mask = (dist < width).to(dtype=H_pred.dtype)
+    
+    diff_sq = (H_pred - H_label) ** 2
+    masked_sum = torch.sum(diff_sq * mask)
+    active_elements = torch.sum(mask)
+
+    return masked_sum / (active_elements + 1e-12)
+
 def cal_bandgap(band, occ):
     return band[...,occ] - band[...,occ-1]