update band post process.

QG-phy · QG-phy · commit 9b122822a416 · 2024-08-02T00:03:13.000+08:00
diff --git a/dptb/postprocess/bandstructure/band.py b/dptb/postprocess/bandstructure/band.py
@@ -170,7 +170,7 @@ def __init__(self, model:torch.nn.Module, results_path: str=None, use_gui: bool=
         self.results_path = results_path
         self.use_gui = use_gui
             
-    def get_bands(self, data: Union[AtomicData, ase.Atoms, str], kpath_kwargs: dict, AtomicData_options: dict={}):
+    def get_bands(self, data: Union[AtomicData, ase.Atoms, str], kpath_kwargs: dict, pbc:Union[bool,list]=None, Atomic_options:dict=None):
         kline_type = kpath_kwargs['kline_type']
 
         # get  the ase structure
@@ -208,7 +208,7 @@ def get_bands(self, data: Union[AtomicData, ase.Atoms, str], kpath_kwargs: dict,
             log.error('Error, now, kline_type only support ase_kpath, abacus, or vasp.')
             raise ValueError
         
-        data, eigenvalues = self.get_eigs(data, klist, AtomicData_options)
+        data, eigenvalues = self.get_eigs(data=data, klist=klist, pbc=pbc, Atomic_options=Atomic_options)
         
 
         # get the E_fermi from data
@@ -229,7 +229,7 @@ def get_bands(self, data: Union[AtomicData, ase.Atoms, str], kpath_kwargs: dict,
         #     estimated_E_fermi = None
         if nel_atom is not None:
             data,estimated_E_fermi = self.get_fermi_level(data=data, nel_atom=nel_atom, \
-                        klist = klist, AtomicData_options=AtomicData_options)
+                        klist = klist, pbc=pbc, Atomic_options=Atomic_options)
         else:
             estimated_E_fermi = None
 
diff --git a/dptb/postprocess/elec_struc_cal.py b/dptb/postprocess/elec_struc_cal.py
@@ -9,6 +9,8 @@
 log = logging.getLogger(__name__)
 from dptb.data import AtomicData, AtomicDataDict
 from dptb.nn.energy import Eigenvalues
+from dptb.utils.argcheck import get_cutoffs_from_model_options
+from copy import deepcopy
 
 # This class `ElecStruCal`  is designed to calculate electronic structure properties such as
 # eigenvalues and Fermi energy based on provided input data and model. 
@@ -61,8 +63,9 @@ def __init__ (
                 device=self.device,
                 dtype=model.dtype,
             )
-
-    def get_data(self,data: Union[AtomicData, ase.Atoms, str],AtomicData_options: dict={},device: Union[str, torch.device]=None):
+        r_max, er_max, oer_max  = get_cutoffs_from_model_options(model.model_options)
+        self.cutoffs = {'r_max': r_max, 'er_max': er_max, 'oer_max': oer_max}
+    def get_data(self,data: Union[AtomicData, ase.Atoms, str],pbc:Union[bool,list]=None, device: Union[str, torch.device]=None, Atomic_options:dict=None):
         '''The function `get_data` takes input data in the form of a string, ase.Atoms object, or AtomicData
         object, processes it accordingly, and returns the AtomicData class.
         
@@ -83,15 +86,36 @@ def get_data(self,data: Union[AtomicData, ase.Atoms, str],AtomicData_options: di
             the loaded AtomicData object.
         
         '''
+        atomic_options = deepcopy(self.cutoffs)
+        if pbc is not None:
+            atomic_options.update({'pbc': pbc})
+
+        if Atomic_options is not None:
+            if Atomic_options.get('r_max', None) is not None:
+                if atomic_options['r_max'] != Atomic_options.get('r_max'):
+                    atomic_options['r_max'] = Atomic_options.get('r_max')
+                    log.warning(f'Overwrite the r_max setting in the model with the r_max setting in the Atomic_options: {Atomic_options.get("r_max")}')
+                    log.warning(f'This is very dangerous, please make sure you know what you are doing.')
+            if Atomic_options.get('er_max', None) is not None:
+                if atomic_options['er_max'] != Atomic_options.get('er_max'):
+                    atomic_options['er_max'] = Atomic_options.get('er_max')
+                    log.warning(f'Overwrite the er_max setting in the model with the er_max setting in the Atomic_options: {Atomic_options.get("er_max")}')
+                    log.warning(f'This is very dangerous, please make sure you know what you are doing.')
+            if Atomic_options.get('oer_max', None) is not None:
+                if atomic_options['oer_max'] != Atomic_options.get('oer_max'):
+                    atomic_options['oer_max'] = Atomic_options.get('oer_max')
+                    log.warning(f'Overwrite the oer_max setting in the model with the oer_max setting in the Atomic_options: {Atomic_options.get("oer_max")}')
+                    log.warning(f'This is very dangerous, please make sure you know what you are doing.')
 
         if isinstance(data, str):
             structase = read(data)
-            data = AtomicData.from_ase(structase, **AtomicData_options)
+            data = AtomicData.from_ase(structase, **atomic_options)
         elif isinstance(data, ase.Atoms):
             structase = data
-            data = AtomicData.from_ase(structase, **AtomicData_options)
+            data = AtomicData.from_ase(structase, **atomic_options)
         elif isinstance(data, AtomicData):
             # structase = data.to("cpu").to_ase()
+            log.info('The data is already an instance of AtomicData. Then the data is used directly.')
             data = data
         else:
             raise ValueError('data should be either a string, ase.Atoms, or AtomicData')
@@ -104,7 +128,7 @@ def get_data(self,data: Union[AtomicData, ase.Atoms, str],AtomicData_options: di
         return data
 
 
-    def get_eigs(self, data: Union[AtomicData, ase.Atoms, str], klist: np.ndarray, AtomicData_options: dict={}):
+    def get_eigs(self, data: Union[AtomicData, ase.Atoms, str], klist: np.ndarray, pbc:Union[bool,list]=None, Atomic_options:dict=None):
         '''This function calculates eigenvalues for Hk at specified k-points.
         
         Parameters
@@ -124,7 +148,7 @@ def get_eigs(self, data: Union[AtomicData, ase.Atoms, str], klist: np.ndarray, A
         
         '''
             
-        data  = self.get_data(data=data, AtomicData_options=AtomicData_options, device=self.device)
+        data  = self.get_data(data=data, pbc=pbc, device=self.device,Atomic_options=Atomic_options)
         # set the kpoint of the AtomicData
         data[AtomicDataDict.KPOINT_KEY] = \
             torch.nested.as_nested_tensor([torch.as_tensor(klist, dtype=self.model.dtype, device=self.device)])
@@ -137,7 +161,7 @@ def get_eigs(self, data: Union[AtomicData, ase.Atoms, str], klist: np.ndarray, A
         return data, data[AtomicDataDict.ENERGY_EIGENVALUE_KEY][0].detach().cpu().numpy()
 
     def get_fermi_level(self, data: Union[AtomicData, ase.Atoms, str], nel_atom: dict, \
-                        meshgrid: list = None, klist: np.ndarray=None, AtomicData_options: dict={}):
+                        meshgrid: list = None, klist: np.ndarray=None, pbc:Union[bool,list]=None,Atomic_options:dict=None):
         '''This function calculates the Fermi level based on provided data with iteration method, electron counts per atom, and
         optional parameters like specific k-points and eigenvalues.
         
@@ -188,7 +212,7 @@ def get_fermi_level(self, data: Union[AtomicData, ase.Atoms, str], nel_atom: dic
 
         # eigenvalues would be used if provided, otherwise the eigenvalues would be calculated from the model on the specified k-points
         if not AtomicDataDict.ENERGY_EIGENVALUE_KEY in data:
-            data, eigs = self.get_eigs(data=data, klist=klist, AtomicData_options=AtomicData_options) 
+            data, eigs = self.get_eigs(data=data, klist=klist, pbc=pbc, Atomic_options=Atomic_options) 
             log.info('Getting eigenvalues from the model.')
         else:
             log.info('The eigenvalues are already in data. will use them.')
diff --git a/dptb/tests/test_from_v1json.py b/dptb/tests/test_from_v1json.py
@@ -66,11 +66,8 @@ def test_bands(self):
             device=model.device)
 
         stru_data = f"{rootdir}/json_model/AlAs.vasp"
-        AtomicData_options = {"r_max": 5.2, "pbc": True}
-
         eigenstatus = bcal.get_bands(data=stru_data, 
-               kpath_kwargs=kpath_kwargs, 
-               AtomicData_options=AtomicData_options)
+               kpath_kwargs=kpath_kwargs)
         
         expected_bands =np.array([[-2.48727150e+01, -1.29382324e+01, -1.29382257e+01, -1.29382229e+01, -1.10868120e+01, -8.07862854e+00, -8.07862568e+00, -8.07861805e+00,  9.56408596e+00,  9.56408691e+00,  1.25271873e+01,  1.25271950e+01,  1.25271978e+01,  4.23655891e+01,  4.23656044e+01,  4.32170753e+01,  4.32170792e+01,  4.32170868e+01],
              [-2.41187267e+01, -1.61148472e+01, -1.42793083e+01, -1.42793045e+01, -1.03604565e+01, -8.68612957e+00, -5.90628624e+00, -5.90628576e+00,  2.25617599e+00,  5.51729870e+00,  5.51730347e+00,  5.61441135e+00,  5.90860081e+00,  2.50449829e+01,  2.82622643e+01,  2.82622776e+01,  2.84239502e+01,  3.07470131e+01],
@@ -149,11 +146,10 @@ def test_bands(self):
             device=model.device)
 
         stru_data = f"{rootdir}/json_model/silicon.vasp"
-        AtomicData_options = {"r_max": 2.6, "oer_max":2.5, "pbc": True}
+        AtomicData_options = {"r_max": 2.6, "oer_max":2.5}
 
         eigenstatus = bcal.get_bands(data=stru_data, 
-               kpath_kwargs=kpath_kwargs, 
-               AtomicData_options=AtomicData_options)
+               kpath_kwargs=kpath_kwargs,Atomic_options=AtomicData_options)
         
         expected_bands =np.array([[-20.259584  ,  -8.328452  ,  -8.328452  ,  -8.328451  ,  -5.782879  ,  -5.782879  ,  -5.7828774 ,  -4.800206  ,  -0.8470682 ,  -0.8470663 ,   4.9619126 ,   4.961913  ,   4.9619136 ,   6.4527135 ,   6.452714  ,   6.452715  ,  10.1427765 ,  10.142781  ],
                                   [-19.173727  , -11.876228  , -10.340221  , -10.34022   ,  -6.861969  ,  -4.9920564 ,  -2.1901789 ,  -2.1901765 ,  -0.9258757 ,   0.76235735,   4.2745295 ,   4.2745323 ,   4.990632  ,   5.55916   ,   5.559161  ,   8.533346  ,   8.716906  ,  11.661528  ],
diff --git a/dptb/tests/test_from_v2json.py b/dptb/tests/test_from_v2json.py
@@ -42,11 +42,10 @@ def test_bands(self):
             device=model.device)
 
         stru_data = f"{rootdir}/json_model/AlAs.vasp"
-        AtomicData_options = {"r_max": 5.2, "pbc": True}
+        AtomicData_options = {"r_max": 5.2}
 
         eigenstatus = bcal.get_bands(data=stru_data, 
-               kpath_kwargs=kpath_kwargs, 
-               AtomicData_options=AtomicData_options)
+               kpath_kwargs=kpath_kwargs)
         
         expected_bands =np.array([[-2.48727150e+01, -1.29382324e+01, -1.29382257e+01, -1.29382229e+01, -1.10868120e+01, -8.07862854e+00, -8.07862568e+00, -8.07861805e+00,  9.56408596e+00,  9.56408691e+00,  1.25271873e+01,  1.25271950e+01,  1.25271978e+01,  4.23655891e+01,  4.23656044e+01,  4.32170753e+01,  4.32170792e+01,  4.32170868e+01],
              [-2.41187267e+01, -1.61148472e+01, -1.42793083e+01, -1.42793045e+01, -1.03604565e+01, -8.68612957e+00, -5.90628624e+00, -5.90628576e+00,  2.25617599e+00,  5.51729870e+00,  5.51730347e+00,  5.61441135e+00,  5.90860081e+00,  2.50449829e+01,  2.82622643e+01,  2.82622776e+01,  2.84239502e+01,  3.07470131e+01],
@@ -99,11 +98,11 @@ def test_bands(self):
             device=model.device)
 
         stru_data = f"{rootdir}/json_model/silicon.vasp"
-        AtomicData_options = {"r_max": 2.6, "oer_max":2.5, "pbc": True}
+        AtomicData_options = {"r_max": 2.6, "oer_max":2.5}
 
         eigenstatus = bcal.get_bands(data=stru_data, 
                kpath_kwargs=kpath_kwargs, 
-               AtomicData_options=AtomicData_options)
+               Atomic_options=AtomicData_options)
         
         expected_bands =np.array([[-20.259584  ,  -8.328452  ,  -8.328452  ,  -8.328451  ,  -5.782879  ,  -5.782879  ,  -5.7828774 ,  -4.800206  ,  -0.8470682 ,  -0.8470663 ,   4.9619126 ,   4.961913  ,   4.9619136 ,   6.4527135 ,   6.452714  ,   6.452715  ,  10.1427765 ,  10.142781  ],
                                   [-19.173727  , -11.876228  , -10.340221  , -10.34022   ,  -6.861969  ,  -4.9920564 ,  -2.1901789 ,  -2.1901765 ,  -0.9258757 ,   0.76235735,   4.2745295 ,   4.2745323 ,   4.990632  ,   5.55916   ,   5.559161  ,   8.533346  ,   8.716906  ,  11.661528  ],
diff --git a/dptb/tests/test_get_fermi.py b/dptb/tests/test_get_fermi.py
@@ -13,19 +13,12 @@ def test_get_fermi():
     stru_data = f"{rootdir}/test_get_fermi/PRIMCELL.vasp"
 
     model = build_model(checkpoint=ckpt)
-    AtomicData_options={
-            "r_max": 5.50,
-            "pbc": True
-        }
-
-    AtomicData_options = AtomicData_options
     nel_atom = {"Au":11}
 
     elec_cal = ElecStruCal(model=model,device='cpu')
     _, efermi =elec_cal.get_fermi_level(data=stru_data, 
                     nel_atom = nel_atom,
-                meshgrid=[30,30,30],
-                AtomicData_options=AtomicData_options)
+                meshgrid=[30,30,30])
     
     assert abs(efermi  + 3.25725233554) < 1e-5
 
diff --git a/dptb/tests/test_nrl.py b/dptb/tests/test_nrl.py
@@ -44,16 +44,15 @@ def test_nrl_json_band():
     }
 
     stru_data = f"{rootdir}/json_model/silicon.vasp"
-    AtomicData_options = {"r_max": 5.0, "oer_max":6.6147151362875, "pbc": True}
+    AtomicData_options = {"r_max": 5.0, "oer_max":6.6147151362875}
     kpath_kwargs = jdata["task_options"]
     bcal = Band(model=model, 
                 use_gui=True, 
                 results_path='./', 
                 device=model.device)
 
     eigenstatus = bcal.get_bands(data=stru_data, 
-                   kpath_kwargs=kpath_kwargs, 
-                   AtomicData_options=AtomicData_options)
+                   kpath_kwargs=kpath_kwargs, Atomic_options = AtomicData_options)
     
     expected_eigenvalues = np.array([[-6.1745434 ,  5.282297  ,  5.282303  ,  5.2823052 ,  8.658317  ,  8.6583185 ,  8.658324  ,  9.862869  , 14.152446  , 14.152451  , 15.180438  , 15.180452  , 16.983887  , 16.983889  , 16.983896  , 23.09491   , 23.094921  , 23.094925  ],
                                      [-5.5601606 ,  2.1920488 ,  3.4229636 ,  3.4229672 ,  7.347074  ,  9.382092  , 11.1772175 , 11.177221  , 14.349099  , 14.924912  , 15.062427  , 15.064081  , 16.540335  , 16.54034   , 20.871534  , 20.871536  , 21.472364  , 28.740482  ],
diff --git a/dptb/tests/test_soc.py b/dptb/tests/test_soc.py
@@ -44,11 +44,11 @@ def test_soc_json_band():
                 device=model.device)
     
     stru_data = f"{rootdir}/Sn/soc/dataset/Sn.vasp"
-    AtomicData_options = {"r_max": 6.0, "oer_max":3.0, "pbc": True}
+    AtomicData_options = {"r_max": 6.0, "oer_max":3.0}
     
     eigenstatus = bcal.get_bands(data=stru_data, 
                    kpath_kwargs=kpath_kwargs, 
-                   AtomicData_options=AtomicData_options)
+                   Atomic_options=AtomicData_options)
 
     expected_eigenvalues = np.array([[-18.796585  , -18.796577  ,  -8.796718  ,  -8.796717  ,
          -8.467822  ,  -8.46782   ,  -8.202273  ,  -8.202273  ,
diff --git a/examples/hBN/band_plot.ipynb b/examples/hBN/band_plot.ipynb
