Pr/44 (#19)

* add comments QG

* add comment QG

dptb/data/transforms.py

@@ -165,6 +165,26 @@ def has_chemical_symbols(self) -> bool:
     def format(
         data: list, type_names: List[str], element_formatter: str = ".6f"
     ) -> str:
+        """
+            Formats a list of data elements along with their type names.
+
+        Parameters:
+            data (list): The data elements to be formatted. This should be a list of numbers.
+            type_names (List[str]): The type names corresponding to the data elements. This should be a list of strings.
+            element_formatter (str, optional): The format in which the data elements should be displayed. Defaults to ".6f".
+
+        Returns:
+            str: A string representation of the data elements along with their type names.
+
+        Raises:
+            ValueError: If `data` is not None, not 0-dimensional, or not 1-dimensional with length equal to the length of `type_names`.
+        
+        Example:
+            >>> data = [1.123456789, 2.987654321]
+            >>> type_names = ['Type1', 'Type2']
+            >>> print(TypeMapper.format(data, type_names))
+                [Type1: 1.123457, Type2: 2.987654]    
+        """ 
         data = torch.as_tensor(data) if data is not None else None
         if data is None:
             return f"[{', '.join(type_names)}: None]"
@@ -434,6 +454,8 @@ def __init__(
             self.reduced_matrix_element = int(((orbtype_count["s"] + 9 * orbtype_count["p"] + 25 * orbtype_count["d"] + 49 * orbtype_count["f"]) + \
                                                     self.full_basis_norb ** 2)/2) # reduce onsite elements by blocks. we cannot reduce it by element since the rme will pass into CG basis to form the whole block
         else:
+            # two factor: this outside one is the number of min(l,l')+1, ie. the number of sk integrals for each orbital pair.
+            # the inside one the type of bond considering the interaction between different orbitals. s-p -> p-s. there are 2 types of bond. and 1 type of s-s.
             self.reduced_matrix_element = (
                 1 * orbtype_count["s"] * orbtype_count["s"] + \
                 2 * orbtype_count["s"] * orbtype_count["p"] + \

dptb/nn/base.py

@@ -248,6 +248,9 @@ def __init__(
         self.out_field = out_field
         self.layers = torch.nn.ModuleList([])
         for kk in range(len(config)-1):
+            # the first layer will take the in_field as key to take `data[in_field]` and output the out_field, data[out_field] = layer(data[in_field])
+            # the rest of the layers will take the out_field as key to take `data[out_field]` and output the out_field, data[out_field] = layer(data[out_field])
+            # That why we need to set the in_field and out_field for 1st layer and the rest of the layers.
             if kk == 0:
                 self.layers.append(
                     AtomicResBlock(

dptb/nn/deeptb.py

@@ -16,7 +16,30 @@
 """
 
 def get_neuron_config(nl):
+    """Extracts the configuration of a neural network from a list of layer sizes.
+
+    Args:
+        nl: A list of integers representing the number of neurons in each layer.
+            If the list has an even number of elements, the last element is assumed
+            to be the output layer size.
+
+    Returns:
+        A list of dictionaries, where each dictionary describes the configuration of
+        a layer in the neural network. Each dictionary has the following keys:
+        - in_features: The number of input neurons for the layer.
+        - hidden_features: The number of hidden neurons for the layer (if applicable).
+        - out_features: The number of output neurons for the layer.
+    
+        e.g.
+        [1, 2, 3, 4, 5, 6] -> [{'in_features': 1, 'hidden_features': 2, 'out_features': 3}, 
+                               {'in_features': 3, 'hidden_features': 4, 'out_features': 5}, 
+                               {'in_features': 5, 'out_features': 6}]
+        [1, 2, 3, 4, 5]    -> [{'in_features': 1, 'hidden_features': 2, 'out_features': 3}, 
+                               {'in_features': 3, 'hidden_features': 4, 'out_features': 5}]
+    """
+
     n = len(nl)
+    assert n > 1, "The neuron config should have at least 2 layers."
     if n % 2 == 0:
         d_out = nl[-1]
         nl = nl[:-1]

dptb/nn/hamiltonian.py

@@ -296,9 +296,14 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
                 rme[:,None, None, :, :], dim=-2) # shape (N, 2l1+1, 2l2+1, n_pair)
 
             # rotation
+            # when get the angle, the xyz vector should be transformed to yzx.
             angle = xyz_to_angles(data[AtomicDataDict.EDGE_VECTORS_KEY][:,[1,2,0]]) # (tensor(N), tensor(N))
+            # The roataion matrix is SO3 rotation, therefore Irreps(l,1), is used here.
             rot_mat_L = Irrep(int(l1), 1).D_from_angles(angle[0].cpu(), angle[1].cpu(), torch.tensor(0., dtype=self.dtype)).to(self.device) # tensor(N, 2l1+1, 2l1+1)
             rot_mat_R = Irrep(int(l2), 1).D_from_angles(angle[0].cpu(), angle[1].cpu(), torch.tensor(0., dtype=self.dtype)).to(self.device) # tensor(N, 2l2+1, 2l2+1)
+
+            # Here The string to control einsum is important, the order of the index should be the same as the order of the tensor
+            # H_z = torch.einsum("nlm, nmoq, nko -> nqlk", rot_mat_L, H_z, rot_mat_R) # shape (N, n_pair, 2l1+1, 2l2+1)
             HR = torch.einsum("nlm, nmoq, nko -> nqlk", rot_mat_L, H_z, rot_mat_R).reshape(n_edge, -1) # shape (N, n_pair * 2l2+1 * 2l2+1)
 
             if l1 < l2:

dptb/nn/nnsk.py

@@ -151,6 +151,9 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         for k in self.idp_sk.orbpair_maps.keys():
             iorb, jorb = k.split("-")
             if iorb == jorb:
+                # This is to keep the symmetry of the hopping parameters for the same orbital pairs
+                # As-Bs = Bs-As; we need to do this because for different orbital pairs, we only have one set of parameters, 
+                # eg. we only have As-Bp and Bs-Ap, but not Ap-Bs and Bp-As; and we will use Ap-Bs = Bs-Ap and Bp-As = As-Bp to calculate the hopping integral
                 self.hopping_param.data[:,self.idp_sk.orbpair_maps[k],:] = 0.5 * (params[:,self.idp_sk.orbpair_maps[k],:] + reflect_params[:,self.idp_sk.orbpair_maps[k],:])
         if hasattr(self, "overlap"):
             params = self.overlap_param.data
@@ -169,6 +172,10 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         edge_number = self.idp_sk.untransform_bond(edge_index).T
         edge_index = self.idp_sk.transform_bond(*edge_number)
 
+        # the edge number is the atomic number of the two atoms in the bond.
+        # The bond length list is actually the nucli radius (unit of angstrom) at the atomic number.
+        # now this bond length list is only available for the first 83 elements.
+
         r0 = 0.5*bond_length_list.type(self.dtype).to(self.device)[edge_number-1].sum(0)
 
         data[AtomicDataDict.EDGE_FEATURES_KEY] = self.hopping_fn.get_skhij(
@@ -183,6 +190,9 @@ def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
             for orbpair_key, slices in self.idp_sk.orbpair_maps.items():
                 if orbpair_key.split("-")[0] == orbpair_key.split("-")[1]:
                     equal_orbpair[slices] = 1.0
+            # this paraconst is to make sure the overlap between the same orbital pairs of the save atom is 1.0 
+            # this is taken from the formula of NRL-TB. 
+            # the overlap tag now is only designed to be used in the NRL-TB case. In the future, we may need to change this.
             paraconst = edge_number[0].eq(edge_number[1]).float().view(-1, 1) * equal_orbpair.unsqueeze(0)
 
             data[AtomicDataDict.EDGE_OVERLAP_KEY] = self.overlap_fn.get_sksij(

dptb/nn/rescale.py

@@ -371,6 +371,8 @@ def __init__(
         start = 0
         start_scalar = 0
         for mul, ir in irreps_in:
+            # only the scalar irreps can be shifted
+            # all the irreps can be scaled
             if str(ir) == "0e":
                 self.num_scalar += mul
                 self.shift_index += list(range(start_scalar, start_scalar + mul))