move E3 statistics initialization into dataset, optmize nested tensor support in hr2hk and eigvals compute

Author: floatingCatty

dptb/data/AtomicData.py

@@ -305,10 +305,6 @@ def _process_dict(kwargs, ignore_fields=[]):
             if num_frames > 1 and v.size(0) != num_frames:
                 raise ValueError(f"Wrong shape for NESTED property {k}")
 
-                
-    
-
-
 
 class AtomicData(Data):
     """A neighbor graph for points in (periodic triclinic) real space.

dptb/data/dataset/_default_dataset.py

@@ -353,7 +353,7 @@ def raw_dir(self):
         # TODO: this is not implemented.
         return self.root
 
-    def E3statistics(self, decay=False):
+    def E3statistics(self, model: torch.nn.Module=None, decay=False):
         assert self.transform is not None
         idp = self.transform
 
@@ -369,6 +369,19 @@ def E3statistics(self, decay=False):
         stats["node"] =  self._E3nodespecies_stat(typed_dataset=typed_dataset)
         stats["edge"] = self._E3edgespecies_stat(typed_dataset=typed_dataset, decay=decay)
 
+        if model is not None:
+            # initilize the model param with statistics
+            scalar_mask = torch.BoolTensor([ir.dim==1 for ir in model.idp.orbpair_irreps])
+            node_shifts = stats["node"]["scalar_ave"]
+            node_scales = stats["node"]["norm_ave"]
+            node_scales[:,scalar_mask] = stats["node"]["scalar_std"]
+
+            edge_shifts = stats["edge"]["scalar_ave"]
+            edge_scales = stats["edge"]["norm_ave"]
+            edge_scales[:,scalar_mask] = stats["edge"]["scalar_std"]
+            model.node_prediction_h.set_scale_shift(scales=node_scales, shifts=node_shifts)
+            model.edge_prediction_h.set_scale_shift(scales=edge_scales, shifts=edge_shifts)
+
         return stats
 
     def _E3edgespecies_stat(self, typed_dataset, decay):

dptb/entrypoints/train.py

@@ -183,7 +183,8 @@ def train(
         # include the init model and from scratch
         # build model will handle the init model cases where the model options provided is not equals to the ones in checkpoint.
         checkpoint = init_model if init_model else None
-        model = build_model(checkpoint=checkpoint, model_options=jdata["model_options"], common_options=jdata["common_options"], statistics=train_datasets.E3statistics())
+        model = build_model(checkpoint=checkpoint, model_options=jdata["model_options"], common_options=jdata["common_options"])
+        train_datasets.E3statistics(model=model)
         trainer = Trainer(
             train_options=jdata["train_options"],
             common_options=jdata["common_options"],

dptb/nn/build.py

@@ -11,8 +11,7 @@
 def build_model(
         checkpoint: str=None,
         model_options: dict={}, 
-        common_options: dict={}, 
-        statistics: dict=None
+        common_options: dict={}
         ):
     """
     The build model method should composed of the following steps:
@@ -141,23 +140,8 @@ def build_model(
     if from_scratch:
         if init_nnenv:
             model = NNENV(**model_options, **common_options)
-
-            # do initialization from statistics if NNENV is e3tb and statistics is provided
-            if model.method == "e3tb" and statistics is not None:
-                scalar_mask = torch.BoolTensor([ir.dim==1 for ir in model.idp.orbpair_irreps])
-                node_shifts = statistics["node"]["scalar_ave"]
-                node_scales = statistics["node"]["norm_ave"]
-                node_scales[:,scalar_mask] = statistics["node"]["scalar_std"]
-
-                edge_shifts = statistics["edge"]["scalar_ave"]
-                edge_scales = statistics["edge"]["norm_ave"]
-                edge_scales[:,scalar_mask] = statistics["edge"]["scalar_std"]
-                model.node_prediction_h.set_scale_shift(scales=node_scales, shifts=node_shifts)
-                model.edge_prediction_h.set_scale_shift(scales=edge_scales, shifts=edge_shifts)
-
         elif init_nnsk:
             model = NNSK(**model_options["nnsk"], **common_options)
-
         elif init_mixed:
             model = MIX(**model_options, **common_options)
         elif init_dftbsk:

dptb/nn/energy.py

@@ -56,13 +56,19 @@ def __init__(
 
 
     def forward(self, data: AtomicDataDict.Type, nk: Optional[int]=None) -> AtomicDataDict.Type:
-        num_k = data[AtomicDataDict.KPOINT_KEY][0].shape[0]
-        kpoints = data[AtomicDataDict.KPOINT_KEY][0] # slice the first dimension, since it is nested tensor by default
+        kpoints = data[AtomicDataDict.KPOINT_KEY]
+        if kpoints.is_nested:
+            nested = True
+            assert kpoints.size(0) == 1
+            kpoints = kpoints[0]
+        else:
+            nested = False
+        num_k = kpoints.shape[0]
         eigvals = []
         if nk is None:
             nk = num_k
         for i in range(int(np.ceil(num_k / nk))):
-            data[AtomicDataDict.KPOINT_KEY] = torch.nested.as_nested_tensor([kpoints[i*nk:(i+1)*nk]])
+            data[AtomicDataDict.KPOINT_KEY] = kpoints[i*nk:(i+1)*nk]
             data = self.h2k(data)
             if self.overlap:
                 data = self.s2k(data)
@@ -74,5 +80,9 @@ def forward(self, data: AtomicDataDict.Type, nk: Optional[int]=None) -> AtomicDa
 
             eigvals.append(torch.linalg.eigvalsh(data[self.h_out_field]))
         data[self.out_field] = torch.nested.as_nested_tensor([torch.cat(eigvals, dim=0)])
+        if nested:
+            data[AtomicDataDict.KPOINT_KEY] = torch.nested.as_nested_tensor([kpoints])
+        else:
+            data[AtomicDataDict.KPOINT_KEY] = kpoints
 
         return data

dptb/nn/hr2hk.py

@@ -56,6 +56,10 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         bondwise_hopping.to(self.device)
         bondwise_hopping.type(self.dtype)
         onsite_block = torch.zeros((len(data[AtomicDataDict.ATOM_TYPE_KEY]), self.idp.full_basis_norb, self.idp.full_basis_norb,), dtype=self.dtype, device=self.device)
+        kpoints = data[AtomicDataDict.KPOINT_KEY]
+        if kpoints.is_nested:
+            assert kpoints.size(0) == 1
+            kpoints = kpoints[0]
 
         soc = data.get(AtomicDataDict.NODE_SOC_SWITCH_KEY, False)
         if isinstance(soc, torch.Tensor):
@@ -111,7 +115,7 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
 
         # R2K procedure can be done for all kpoint at once.
         all_norb = self.idp.atom_norb[data[AtomicDataDict.ATOM_TYPE_KEY]].sum()
-        block = torch.zeros(data[AtomicDataDict.KPOINT_KEY][0].shape[0], all_norb, all_norb, dtype=self.ctype, device=self.device)
+        block = torch.zeros(kpoints.shape[0], all_norb, all_norb, dtype=self.ctype, device=self.device)
         # block = torch.complex(block, torch.zeros_like(block))
         # if data[AtomicDataDict.NODE_SOC_SWITCH_KEY].all():
         #     block_uu = torch.zeros(data[AtomicDataDict.KPOINT_KEY].shape[0], all_norb, all_norb, dtype=self.ctype, device=self.device)
@@ -149,13 +153,13 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
             masked_hblock = hblock[imask][:,jmask]
 
             block[:,iatom_indices,jatom_indices] += masked_hblock.squeeze(0).type_as(block) * \
-                torch.exp(-1j * 2 * torch.pi * (data[AtomicDataDict.KPOINT_KEY][0] @ data[AtomicDataDict.EDGE_CELL_SHIFT_KEY][i])).reshape(-1,1,1)
+                torch.exp(-1j * 2 * torch.pi * (kpoints @ data[AtomicDataDict.EDGE_CELL_SHIFT_KEY][i])).reshape(-1,1,1)
 
         block = block + block.transpose(1,2).conj()
         block = block.contiguous()
 
         if soc:
-            HK_SOC = torch.zeros(data[AtomicDataDict.KPOINT_KEY][0].shape[0], 2*all_norb, 2*all_norb, dtype=self.ctype, device=self.device)
+            HK_SOC = torch.zeros(kpoints.shape[0], 2*all_norb, 2*all_norb, dtype=self.ctype, device=self.device)
             #HK_SOC[:,:all_norb,:all_norb] = block + block_uu
             #HK_SOC[:,:all_norb,all_norb:] = block_ud
             #HK_SOC[:,all_norb:,:all_norb] = block_ud.conj()