refactor(data preprocess): remove the cut off options from info.json and collect the values from input.json

Author: QG-phy

.gitignore

@@ -6,6 +6,9 @@ dptb/tests/**/*.pth
 dptb/tests/**/*.npy
 dptb/tests/**/*.traj
 dptb/tests/**/out*/*
+dptb/tests/**/out*/*
+dptb/tests/**/*lmdb
+dptb/tests/**/*h5
 examples/_*
 *.dat
 *log*

dptb/data/build.py

@@ -109,14 +109,17 @@ def dataset_from_config(config, prefix: str = "dataset") -> AtomicDataset:
 def build_dataset(
         # set_options
         root: str,
+        # dataset_options
+        r_max: float,
+        er_max: float = None,
+        oer_max: float = None,
         type: str = "DefaultDataset",
         prefix: str = None,
         separator:str='.',
         get_Hamiltonian: bool = False,
         get_overlap: bool = False,
         get_DM: bool = False,
         get_eigenvalues: bool = False,
-
         # common_options
         orthogonal: bool = False,
         basis: str = None, 
@@ -224,7 +227,10 @@ def build_dataset(
         # We will sort the info_files here.
         # The order itself is not important, but must be consistant for the same list.
         info_files = {key: info_files[key] for key in sorted(info_files)}
-        
+    
+        for ikey in info_files:
+            info_files[ikey].update({'r_max': r_max, 'er_max': er_max, 'oer_max': oer_max})
+            
         if dataset_type == "DeePHDataset":
             dataset = DeePHE3Dataset(
                 root=root,

dptb/data/dataset/_deeph_dataset.py

@@ -43,13 +43,11 @@ def __init__(
         for file in self.info_files.keys():
             # get the info here
             info = info_files[file]
-            assert "AtomicData_options" in info
-            AtomicData_options = info["AtomicData_options"]
-            assert "r_max" in AtomicData_options
-            assert "pbc" in AtomicData_options
+            assert "r_max" in info
+            assert "pbc" in info
             subdata = os.path.join(self.root, file)
             self.raw_data.append(subdata)
-            self.data_options[subdata] = AtomicData_options
+            self.data_options[subdata] = info
 
         # The AtomicData_options is never used here.
         # Because we always return a list of AtomicData object in `get_data()`.
@@ -68,12 +66,15 @@ def get_data(self):
         for subpath in tqdm(self.raw_data, desc="Loading data"):
             # the type_mapper here is loaded in PyG `dataset` type as `transform` attritube
             # so the OrbitalMapper can be accessed by self.transform here
-            AtomicData_options = self.data_options[subpath]
+            info = self.data_options[subpath]
             atomic_data = AtomicData.from_points(
                 pos = np.loadtxt(os.path.join(subpath, "site_positions.dat")).T,
                 cell = np.loadtxt(os.path.join(subpath, "lat.dat")).T,
                 atomic_numbers = np.loadtxt(os.path.join(subpath, "element.dat")),
-                **AtomicData_options,
+                pbc = info["pbc"],
+                r_max=info["r_max"],
+                er_max=info.get("er_max", None),
+                oer_max=info.get("oer_max", None)
             )
             idp = self.type_mapper
             openmx_to_deeptb(atomic_data, idp, os.path.join(subpath, "./hamiltonians.h5"))

dptb/data/dataset/_default_dataset.py

@@ -40,7 +40,6 @@ class _TrajData(object):
 
     def __init__(self, 
                  root: str, 
-                 AtomicData_options: Dict[str, Any] = {},
                  get_Hamiltonian = False,
                  get_overlap = False,
                  get_DM = False,
@@ -50,13 +49,10 @@ def __init__(self,
 
         assert not get_Hamiltonian * get_DM, "Hamiltonian and Density Matrix can only loaded one at a time, for which will occupy the same attribute in the AtomicData."
         self.root = root
-        self.AtomicData_options = AtomicData_options
         self.info = info
-
         self.data = {}
-        # load cell
-        
-        pbc = AtomicData_options["pbc"]
+        pbc = info["pbc"]
+        # load cell        
         if isinstance(pbc, bool):
             has_cell = pbc
         elif isinstance(pbc, list):
@@ -155,7 +151,6 @@ def __init__(self,
     @classmethod
     def from_ase_traj(cls,
                       root: str, 
-                      AtomicData_options: Dict[str, Any] = {},
                       get_Hamiltonian = False,
                       get_overlap = False,
                       get_DM = False,
@@ -185,7 +180,6 @@ def from_ase_traj(cls,
         np.savetxt(os.path.join(root, "atomic_numbers.dat"), atomic_numbers, fmt='%d')
 
         return cls(root=root,
-                   AtomicData_options=AtomicData_options,
                    get_Hamiltonian=get_Hamiltonian,
                    get_overlap=get_overlap,
                    get_DM=get_DM,
@@ -218,10 +212,11 @@ def toAtomicDataList(self, idp: TypeMapper = None):
                                                                                   dtype=torch.long)
 
             atomic_data = AtomicData.from_points(
+                  r_max = self.info["r_max"],
+                  pbc = self.info["pbc"],
+                  er_max = self.info.get("er_max", None),
+                  oer_max= self.info.get("oer_max", None),
                   **kwargs,
-                  # pbc is stored in AtomicData_options now.
-                  #pbc = self.info["pbc"], 
-                  **self.AtomicData_options
             )
             if "hamiltonian_blocks" in self.data:
                 assert idp is not None, "LCAO Basis must be provided  in `common_option` for loading Hamiltonian."
@@ -300,21 +295,19 @@ def __init__(
         for file in self.info_files.keys():
             # get the info here
             info = info_files[file]
-            assert "AtomicData_options" in info
-            AtomicData_options = info["AtomicData_options"]
-            assert "r_max" in AtomicData_options
-            assert "pbc" in AtomicData_options
+            # assert "AtomicData_options" in info
+            assert "r_max" in info
+            assert "pbc" in info
+            pbc = info["pbc"]
             if info["pos_type"] == "ase":
                 subdata = _TrajData.from_ase_traj(os.path.join(self.root, file), 
-                                AtomicData_options,
                                 get_Hamiltonian, 
                                 get_overlap,
                                 get_DM,
                                 get_eigenvalues,
                                 info=info)
             else:
                 subdata = _TrajData(os.path.join(self.root, file), 
-                                AtomicData_options,
                                 get_Hamiltonian,
                                 get_overlap,
                                 get_DM,

dptb/data/dataset/_hdf5_dataset.py

@@ -38,17 +38,14 @@ class _HDF5_TrajData(object):
 
     def __init__(self, 
                  root: str, 
-                 AtomicData_options: Dict[str, Any] = {},
                  get_Hamiltonian = False,
                  get_overlap = False,
                  get_DM = False,
                  get_eigenvalues = False,
                  info = None):
         assert not get_Hamiltonian * get_DM, "Cannot get both Hamiltonian and DM"
         self.root = root
-        self.AtomicData_options = AtomicData_options
         self.info = info
-
         self.data = {}
 
         assert os.path.exists(os.path.join(root, "structure.pkl")), "structure file not found."
@@ -87,9 +84,11 @@ def toAtomicDataList(self, idp: TypeMapper = None):
                 pos = self.data['structure'][frame]["positions"][:],
                 cell = frame_cell,
                 atomic_numbers = self.data['structure'][frame]["atomic_numbers"][:],
-                # pbc is stored in AtomicData_options now.
-                #pbc = self.info["pbc"], 
-                **self.AtomicData_options)
+                r_max = self.info["r_max"], 
+                er_max = self.info.get("er_max", None),
+                oer_max = self.info.get("oer_max", None),
+                pbc = self.info["pbc"], 
+            )
 
             if "hamiltonian_blocks" in self.data:
                 assert idp is not None, "LCAO Basis must be provided  in `common_option` for loading Hamiltonian."
@@ -171,13 +170,10 @@ def __init__(
         for file in self.info_files.keys():
             # get the info here
             info = info_files[file]
-            assert "AtomicData_options" in info
-            AtomicData_options = info["AtomicData_options"]
-            assert "r_max" in AtomicData_options
-            assert "pbc" in AtomicData_options
+            assert "r_max" in info
+            assert "pbc" in info
             if info["pos_type"] in ["hdf5", 'pickle']:
                 subdata = _HDF5_TrajData(os.path.join(self.root, file), 
-                                AtomicData_options,
                                 get_Hamiltonian, 
                                 get_overlap,
                                 get_DM,

dptb/entrypoints/train.py

@@ -3,7 +3,7 @@
 from dptb.data.build import build_dataset
 from dptb.plugins.monitor import TrainLossMonitor, LearningRateMonitor, Validationer
 from dptb.plugins.train_logger import Logger
-from dptb.utils.argcheck import normalize
+from dptb.utils.argcheck import normalize, collect_cutoffs
 from dptb.plugins.saver import Saver
 from typing import Dict, List, Optional, Any
 from dptb.utils.tools import j_loader, setup_seed, j_must_have
@@ -18,6 +18,7 @@
 import json
 import os
 import time
+import copy
 
 __all__ = ["train"]
 
@@ -147,26 +148,33 @@ def train(
                     jdata["train_options"] = f["config"]["train_options"]
                 if jdata.get("model_options") is None:
                     jdata["model_options"] = f["config"]["model_options"]
+                
+                ## add some warning !
+                for k, v in jdata["model_options"].items():
+                    if k not in f["config"]["model_options"]:
+                        log.warning(f"The model options {k} is not defined in checkpoint, set to {v}.")
+                    else:
+                        deep_dict_difference(k, v, f["config"]["model_options"])
             del f
     else:
         j_must_have(jdata, "model_options")
         j_must_have(jdata, "train_options")
 
-
+    cutoff_options =collect_cutoffs(jdata)
     # setup seed
     setup_seed(seed=jdata["common_options"]["seed"])
 
     # with open(os.path.join(output, "train_config.json"), "w") as fp:
     #     json.dump(jdata, fp, indent=4)
 
     # build dataset
-    train_datasets = build_dataset(**jdata["data_options"]["train"], **jdata["common_options"])
+    train_datasets = build_dataset(**cutoff_options,**jdata["data_options"]["train"], **jdata["common_options"])
     if jdata["data_options"].get("validation"):
-        validation_datasets = build_dataset(**jdata["data_options"]["validation"], **jdata["common_options"])
+        validation_datasets = build_dataset(**cutoff_options, **jdata["data_options"]["validation"], **jdata["common_options"])
     else:
         validation_datasets = None
     if jdata["data_options"].get("reference"):
-        reference_datasets = build_dataset(**jdata["data_options"]["reference"], **jdata["common_options"])
+        reference_datasets = build_dataset(**cutoff_options, **jdata["data_options"]["reference"], **jdata["common_options"])
     else:
         reference_datasets = None
 
@@ -227,3 +235,26 @@ def train(
     log.info(f"wall time: {(end_time - start_time):.3f} s")
 
 
+def deep_dict_difference(base_key, expected_value, model_options):
+    """
+    递归地记录嵌套字典中的选项差异。
+    
+    :param base_key: 基础键名，用于构建警告消息的前缀。
+    :param expected_value: 期望的值，可能是字典或非字典类型。
+    :param model_options: 用于比较的模型选项字典。
+    """
+    target_dict= copy.deepcopy(model_options) # 防止修改原始字典
+    if isinstance(expected_value, dict):
+        for subk, subv in expected_value.items():
+            
+            if  not isinstance(target_dict.get(base_key, {}),dict):
+                log.warning(f"The model option {subk} in {base_key} is not defined in  checkpoint, set to {subv}.")
+            
+            elif subk not in target_dict.get(base_key, {}):
+                log.warning(f"The model option {subk} in {base_key} is not defined in  checkpoint, set to {subv}.")
+            else:
+                target2 = copy.deepcopy(target_dict[base_key])
+                deep_dict_difference(f"{subk}", subv, target2)
+    else:
+        if expected_value != target_dict[base_key]:
+            log.warning(f"The model option {base_key} is set to {expected_value}, but in checkpoint it is {target_dict[base_key]}, make sure it it correct!")

dptb/postprocess/elec_struc_cal.py

@@ -42,6 +42,10 @@ def __init__ (
         self.model.eval()
         self.overlap = hasattr(model, 'overlap')
 
+        if not self.model.transform:
+            log.error('The model.transform is not True, please check the model.')
+            raise RuntimeError('The model.transform is not True, please check the model.')
+        
         if self.overlap:
             self.eigv = Eigenvalues(
                 idp=model.idp,

dptb/tests/data/Sn/soc/dataset/set.0/info.json

@@ -2,12 +2,7 @@
     "nframes": 1,
     "natoms": -1,
     "pos_type": "ase",
-    "AtomicData_options": {
-        "r_max": 6.0,
-        "er_max": 5.0,
-        "oer_max":3.0,
-        "pbc": true
-    },
+    "pbc": true,
     "bandinfo": {
         "band_min": 0,
         "band_max":16,

dptb/tests/data/Sn/soc/input/input_soc.json

@@ -43,7 +43,7 @@
     },
     "model_options": {
         "nnsk": {
-            "onsite": {"method": "strain","rs":6.0, "w": 0.1},
+            "onsite": {"method": "strain","rs":3.0, "w": 0.1},
             "hopping": {"method": "powerlaw", "rs":6.0, "w": 0.1},
             "soc":{"method":"uniform"},
             "push": false,

dptb/tests/data/hBN/dataset/kpath.0/info.json

@@ -2,12 +2,7 @@
     "nframes": 1,
     "natoms": 2,
     "pos_type": "ase",
-    "AtomicData_options": {
-        "r_max": 2.6,
-        "er_max": 2.6,
-        "oer_max":1.6,
-        "pbc": true
-    },
+    "pbc": true,
     "bandinfo": {
         "band_min": 0,
         "band_max": 6,