[feature] add the support of NRL-TB

dptb/dataprocess/processor.py

@@ -5,7 +5,6 @@
 from dptb.structure.abstract_stracture import AbstractStructure
 
 class Processor(object):
-    # TODO: 现在strain的env 是通过get_env 获得，但是在dptb中的env是有另外的含义。是否已经考虑。
     def __init__(self, structure_list: List[AbstractStructure], kpoint, eigen_list, batchsize: int, wannier_list = None, env_cutoff: float = 3.0, onsitemode=None, 
     onsite_cutoff=None, sorted_bond=None, sorted_onsite=None, sorted_env=None, bandinfo=None, device='cpu', dtype=torch.float32, if_shuffle=True):
         super(Processor, self).__init__()
@@ -69,10 +68,19 @@ def shuffle(self):
             self.__struct_unsampled__ = self.__struct_unsampled__[self.batchsize:]
 
     def get_env(self, cutoff=None, sorted=None):
-        # TODO: the sorted mode should be explained here, in which case, we should use.
         '''It takes the environment of each structure in the workspace and concatenates them into one big
         environment
 
+        Parameters
+        ----------
+        cutoff float : the cutoff radius for the onsite environment.
+        sorted str : the sorted mode for the onsite environment.
+            None: not sorted, return torch tensors.
+            'st': sorted by structure, return a dictionary of tensors. eg. {0: tensor, 1: tensor, ...}
+            'itype-jtype': sorted by itype and jtype, return a dictionary of tensors. eg. {'C-B': tensor, ...}
+
+        for env sorted="st".
+
         Returns
         -------
             A dictionary of the environment for ent type for all the strucutes in  the works sapce.
@@ -125,10 +133,18 @@ def get_env(self, cutoff=None, sorted=None):
         return batch_env # {env_type: (f, itype, i, jtype, j, jtype, Rx, Ry, Rz, s(r), rx, ry, rz)} or [(f, itype, i, jtype, j, jtype, Rx, Ry, Rz, s(r), rx, ry, rz)]
 
     def get_onsitenv(self, cutoff=None, sorted=None):
-        # TODO: the sorted mode should be explained here, in which case, we should use.
         '''It takes the environment of each structure in the workspace and concatenates them into one big
         environment
 
+        Parameters
+        ----------
+        cutoff float : the cutoff radius for the onsite environment.
+        sorted str : the sorted mode for the onsite environment.
+            None: not sorted, return torch tensors.
+            'st': sorted by structure, return a dictionary of tensors. eg. {0: tensor, 1: tensor, ...}
+            'itype-jtype': sorted by itype and jtype, return a dictionary of tensors. eg. {'C-B': tensor, ...}
+
+        for onsiteenv sorted_env="itype-jtype".
         Returns
         -------
             A dictionary of the environment for ent type for all the strucutes in  the works sapce.
@@ -141,7 +157,7 @@ def get_onsitenv(self, cutoff=None, sorted=None):
         if cutoff is None:
             cutoff = self.onsite_cutoff
         else:
-            assert isinstance(cutoff, float)
+            assert isinstance(cutoff, float), "The cutoff should be a float number."
 
         if sorted is None:
             batch_env = []
@@ -183,6 +199,14 @@ def get_onsitenv(self, cutoff=None, sorted=None):
     def get_bond(self, sorted=None):
         '''It takes the bonds of each structure in the workspace and concatenates them into one big dictionary.
 
+        Parameters
+        ----------
+        sorted str : the sorted mode for the onsite environment.
+            None: not sorted, return torch tensors.
+            'st': sorted by structure, return a dictionary of tensors. eg. {0: tensor, 1: tensor, ...}
+
+        For bond sorted="st". 
+
         Returns
         -------
             A Tensor of the bonds lists for bond type for all the strucutes in the works space.
@@ -267,7 +291,7 @@ def __next__(self):
             self.shuffle()
             bond, bond_onsite = self.get_bond(self.sorted_bond)
 
-            if not self.onsitemode == 'strain':
+            if not self.onsitemode in ['strain','NRL']:  # for NRL - TB we also need the onsite env.
                 data = (bond, bond_onsite, self.get_env(sorted=self.sorted_env), None,  self.__struct_workspace__,
                     self.kpoint, self.eigen_list[self.__struct_idx_workspace__].astype(float), self.wannier_list[self.__struct_idx_workspace__])
             else:

dptb/hamiltonian/hamil_eig_sk_crt.py

@@ -56,6 +56,10 @@ def update_hs_list(self, struct, hoppings, onsiteEs, onsiteVs=None, overlaps=Non
             self.onsiteSs = onsiteSs
             self.use_orthogonal_basis = False
 
+        if onsiteSs is not None:
+            log.info(msg='The onsiteSs is not None, But even for non-orthogonal basis, the onsite S matrix part is still identity.')
+            log.info(msg='Therefore the onsiteSs will not be used !!')
+
         if soc_lambdas is None:
             self.soc = False
         else:
@@ -149,7 +153,10 @@ def get_hs_onsite(self, bonds_onsite = None, onsite_envs=None):
 
             sub_hamil_block = th.zeros([self.__struct__.proj_atomtype_norbs[iatype], self.__struct__.proj_atomtype_norbs[jatype]], dtype=self.dtype, device=self.device)
             if not self.use_orthogonal_basis:
-                sub_over_block = th.zeros([self.__struct__.proj_atomtype_norbs[iatype], self.__struct__.proj_atomtype_norbs[jatype]], dtype=self.dtype, device=self.device)
+                # For non - orthogonal basis, the overlap matrix is needed. 
+                # but for the onsite, the overlap matrix is identity.
+                sub_over_block = th.eye(self.__struct__.proj_atomtype_norbs[iatype], dtype=self.dtype, device=self.device)
+
 
             ist = 0
             for ish in self.__struct__.proj_atom_anglr_m[iatype]:     # ['s','p',..]
@@ -159,8 +166,9 @@ def get_hs_onsite(self, bonds_onsite = None, onsite_envs=None):
 
                 indx = self.__struct__.onsite_index_map[iatype][ish] # change onsite index map from {N:{s:}} to {N:{ss:, sp:}}
                 sub_hamil_block[ist:ist+norbi, ist:ist+norbi] = th.eye(norbi, dtype=self.dtype, device=self.device) * self.onsiteEs[ib][indx]
-                if not self.use_orthogonal_basis:
-                    sub_over_block[ist:ist+norbi, ist:ist+norbi] = th.eye(norbi, dtype=self.dtype, device=self.device) * self.onsiteSs[ib][indx]
+                # For non - orthogonal basis, the onsite overlap is identity, we don't need to calculate it.
+                #if not self.use_orthogonal_basis:
+                #    sub_over_block[ist:ist+norbi, ist:ist+norbi] = th.eye(norbi, dtype=self.dtype, device=self.device) * self.onsiteSs[ib][indx]
                 ist = ist + norbi
 
             onsiteH_blocks.append(sub_hamil_block)

dptb/nnops/NN2HRK.py

@@ -86,33 +86,47 @@ def _get_nnsk_HR(self):
         assert self.structure.onsitemode == self.apihost.model_config['onsitemode']
         # TODO: 注意检查 processor 关于 env_cutoff 和 onsite_cutoff.
         predict_process = Processor(structure_list=self.structure, batchsize=1, kpoint=None, eigen_list=None, device=self.device, dtype=self.dtype, 
-                                        env_cutoff=self.apihost.model_config['env_cutoff'], onsitemode=self.apihost.model_config['onsitemode'], onsite_cutoff=self.apihost.model_config['onsite_cutoff'], sorted_onsite="st", sorted_bond="st", sorted_env="st")
+                                        env_cutoff=self.apihost.model_config['env_cutoff'], onsitemode=self.apihost.model_config['onsitemode'], onsite_cutoff=self.apihost.model_config['onsite_cutoff'], sorted_onsite=self.sorted_onsite, sorted_bond=self.sorted_bond, sorted_env=self.sorted_env)
 
         batch_bonds, batch_bond_onsites = predict_process.get_bond(sorted=self.sorted_bond)
-        coeffdict = self.apihost.model(mode='hopping')
+        coeffdict, overlap_coeffdict = self.apihost.model(mode='hopping')
         batch_hoppings = self.apihost.hops_fun.get_skhops(batch_bonds=batch_bonds, coeff_paras=coeffdict, rcut=self.apihost.model_config['skfunction']['sk_cutoff'], w=self.apihost.model_config['skfunction']['sk_decay_w'])
         nn_onsiteE, onsite_coeffdict = self.apihost.model(mode='onsite')
-        batch_onsiteEs = self.apihost.onsite_fun(batch_bonds_onsite=batch_bond_onsites, onsite_db=self.apihost.onsite_db, nn_onsiteE=nn_onsiteE)
+        if self.apihost.overlap:
+            assert overlap_coeffdict is not None, "The overlap_coeffdict should be provided if overlap is True."
+            batch_overlaps = self.apihost.overlap_fun.get_skoverlaps(batch_bonds=batch_bonds, coeff_paras=overlap_coeffdict, rcut=self.apihost.model_config['skfunction']['sk_cutoff'], w=self.apihost.model_config['skfunction']['sk_decay_w'])
+
+
+        if self.apihost.model_config['onsitemode'] in ['strain','NRL']:
+            batch_onsite_envs = predict_process.get_onsitenv(cutoff=self.apihost.model_config['onsite_cutoff'], sorted=self.sorted_onsite)
+        else:
+            batch_onsite_envs = None
+
+        batch_onsiteEs = self.apihost.onsite_fun.get_onsiteEs(batch_bonds_onsite=batch_bond_onsites, onsite_env=batch_onsite_envs, nn_onsite_paras=nn_onsiteE)
 
         if self.apihost.model_config['soc']:
             nn_soc_lambdas, _ = self.apihost.model(mode='soc')
             batch_soc_lambdas = self.apihost.soc_fun(batch_bonds_onsite=batch_bond_onsites, soc_db=self.apihost.soc_db, nn_soc=nn_soc_lambdas)
 
         if self.apihost.model_config['onsitemode'] == 'strain':
-            batch_onsite_envs = predict_process.get_onsitenv(cutoff=self.apihost.model_config['onsite_cutoff'], sorted=self.sorted_onsite)
             batch_onsiteVs = self.apihost.onsitestrain_fun.get_skhops(batch_bonds=batch_onsite_envs, coeff_paras=onsite_coeffdict)
             onsiteEs, hoppings, onsiteVs = batch_onsiteEs[0], batch_hoppings[0], batch_onsiteVs[0]
             onsitenvs = batch_onsite_envs[0][:,1:]
         else:
             onsiteEs, hoppings, onsiteVs = batch_onsiteEs[0], batch_hoppings[0],  None
             onsitenvs = None
 
+        if self.apihost.overlap:
+            overlaps = batch_overlaps[0]
+        else:
+            overlaps = None
+
         if self.apihost.model_config["soc"]:
             soc_lambdas = batch_soc_lambdas[0]
         else:
             soc_lambdas = None
 
-        self.hamileig.update_hs_list(struct=self.structure, hoppings=hoppings, onsiteEs=onsiteEs, onsiteVs=onsiteVs, soc_lambdas=soc_lambdas)
+        self.hamileig.update_hs_list(struct=self.structure, hoppings=hoppings, onsiteEs=onsiteEs, onsiteVs=onsiteVs, overlaps=overlaps, soc_lambdas=soc_lambdas)
         self.hamileig.get_hs_blocks(bonds_onsite=batch_bond_onsites[0][:,1:], bonds_hoppings=batch_bonds[0][:,1:], 
                                     onsite_envs=onsitenvs)
 
@@ -122,37 +136,54 @@ def _get_nnsk_HR(self):
         self.use_orthogonal_basis = self.hamileig.use_orthogonal_basis
         self.allbonds, self.hamil_blocks = self.hamileig.all_bonds, self.hamileig.hamil_blocks
 
-        if not self.hamileig.use_orthogonal_basis:
+        if self.hamileig.use_orthogonal_basis:
             self.overlap_blocks = None
         else:
             self.overlap_blocks = self.hamileig.overlap_blocks
 
     def _get_dptb_HR(self):
         predict_process = Processor(structure_list=self.structure, batchsize=1, kpoint=None, eigen_list=None, device=self.device, dtype=self.dtype, 
-                                        env_cutoff=self.apihost.model_config['env_cutoff'], onsitemode=self.apihost.model_config['onsitemode'], onsite_cutoff=self.apihost.model_config['onsite_cutoff'], sorted_onsite="st", sorted_bond="st", sorted_env="st")
+                                        env_cutoff=self.apihost.model_config['env_cutoff'], onsitemode=self.apihost.model_config['onsitemode'], onsite_cutoff=self.apihost.model_config['onsite_cutoff'], sorted_onsite=self.sorted_onsite, sorted_bond=self.sorted_bond, sorted_env=self.sorted_env)
 
         batch_bonds, batch_bond_onsites = predict_process.get_bond(sorted=self.sorted_bond)
         batch_env = predict_process.get_env(cutoff=self.apihost.model_config['env_cutoff'], sorted=self.sorted_env)
         batch_bond_hoppings, batch_hoppings, batch_bond_onsites, batch_onsiteEs, batch_soc_lambdas = self.apihost.nntb.calc(batch_bonds, batch_env)
 
         if  self.apihost.model_config['use_correction']:
-            coeffdict = self.apihost.sknet(mode='hopping')
+            coeffdict, overlap_coeffdict = self.apihost.sknet(mode='hopping')
             batch_nnsk_hoppings = self.apihost.hops_fun.get_skhops( batch_bond_hoppings, coeffdict, 
                             rcut=self.apihost.model_config["skfunction"]["sk_cutoff"], w=self.apihost.model_config["skfunction"]["sk_decay_w"])
             nnsk_onsiteE, onsite_coeffdict = self.apihost.sknet(mode='onsite')
-            batch_nnsk_onsiteEs = self.apihost.onsite_fun(batch_bonds_onsite=batch_bond_onsites, onsite_db=self.apihost.onsite_db, nn_onsiteE=nnsk_onsiteE)
+
+            if self.apihost.overlap:
+                assert overlap_coeffdict is not None, "The overlap_coeffdict should be provided if overlap is True."
+                batch_nnsk_overlaps = self.apihost.overlap_fun.get_skoverlaps(batch_bonds=batch_bonds, coeff_paras=overlap_coeffdict, 
+                            rcut=self.apihost.model_config['skfunction']['sk_cutoff'], w=self.apihost.model_config['skfunction']['sk_decay_w'])
+
+
+
+            if self.apihost.model_config['onsitemode'] in ['strain','NRL']:
+                batch_onsite_envs = predict_process.get_onsitenv(cutoff=self.apihost.model_config['onsite_cutoff'], sorted=self.sorted_onsite)
+            else:
+                batch_onsite_envs = None
+
+            batch_nnsk_onsiteEs = self.apihost.onsite_fun.get_onsiteEs(batch_bonds_onsite=batch_bond_onsites, onsite_env=batch_onsite_envs, nn_onsite_paras=nnsk_onsiteE)
+
             if self.apihost.model_config["soc"]:
                 nnsk_soc_lambdas, _ = self.apihost.sknet(mode="soc")
                 batch_nnsk_soc_lambdas = self.apihost.soc_fun(batch_bonds_onsite=batch_bond_onsites, soc_db=self.apihost.soc_db, nn_soc=nnsk_soc_lambdas)
 
             if self.apihost.model_config['onsitemode'] == "strain":
-                batch_onsite_envs = predict_process.get_onsitenv(cutoff=self.apihost.model_config['onsite_cutoff'], sorted=self.sorted_onsite)
                 batch_nnsk_onsiteVs = self.apihost.onsitestrain_fun.get_skhops(batch_bonds=batch_onsite_envs, coeff_paras=onsite_coeffdict)
                 onsiteVs = batch_nnsk_onsiteVs[0]
                 onsitenvs = batch_onsite_envs[0][:,1:]
             else:
                 onsiteVs = None
                 onsitenvs = None
+            if self.apihost.overlap:
+                nnsk_overlaps = batch_nnsk_overlaps[0]
+            else:
+                nnsk_overlaps = None
 
             if self.apihost.model_config["soc"] and self.apihost.model_config["dptb"]["soc_env"]:
                 nn_soc_lambdas = batch_soc_lambdas[0]
@@ -165,14 +196,15 @@ def _get_dptb_HR(self):
                     sk_soc_lambdas = None
 
 
-            onsiteEs, hoppings, _, _, soc_lambdas = nnsk_correction(nn_onsiteEs=batch_onsiteEs[0], nn_hoppings=batch_hoppings[0],
+            onsiteEs, hoppings, onsiteSs, overlaps, soc_lambdas = nnsk_correction(nn_onsiteEs=batch_onsiteEs[0], nn_hoppings=batch_hoppings[0],
                                     sk_onsiteEs=batch_nnsk_onsiteEs[0], sk_hoppings=batch_nnsk_hoppings[0],
-                                    sk_onsiteSs=None, sk_overlaps=None, nn_soc_lambdas=nn_soc_lambdas, sk_soc_lambdas=sk_soc_lambdas)
+                                    sk_onsiteSs=None, sk_overlaps=nnsk_overlaps, nn_soc_lambdas=nn_soc_lambdas, sk_soc_lambdas=sk_soc_lambdas)
         else:
-            onsiteEs, hoppings, soc_lambdas, onsiteVs, onsitenvs = batch_onsiteEs[0], batch_hoppings[0], None, None, None
+            assert not self.apihost.overlap, "The overlap should be False if use_correction is False."
+            onsiteEs, hoppings, soc_lambdas, onsiteVs, onsitenvs, onsiteSs, overlaps = batch_onsiteEs[0], batch_hoppings[0], None, None, None, None, None
 
 
-        self.hamileig.update_hs_list(struct=self.structure, hoppings=hoppings, onsiteEs=onsiteEs, onsiteVs=onsiteVs, soc_lambdas=soc_lambdas)
+        self.hamileig.update_hs_list(struct=self.structure, hoppings=hoppings, onsiteEs=onsiteEs, onsiteVs=onsiteVs,overlaps=overlaps, soc_lambdas=soc_lambdas)
         self.hamileig.get_hs_blocks(bonds_onsite=batch_bond_onsites[0][:,1:], bonds_hoppings=batch_bond_hoppings[0][:,1:],
                                     onsite_envs=onsitenvs)
 
@@ -182,7 +214,7 @@ def _get_dptb_HR(self):
         self.use_orthogonal_basis = self.hamileig.use_orthogonal_basis
         self.allbonds, self.hamil_blocks = self.hamileig.all_bonds, self.hamileig.hamil_blocks
 
-        if not self.hamileig.use_orthogonal_basis:
+        if self.hamileig.use_orthogonal_basis:
             self.overlap_blocks = None
         else:
             self.overlap_blocks = self.hamileig.overlap_blocks