Add interface for diagonal term.

Author: hzhangxyz

qmb/fcidump.py

@@ -192,6 +192,9 @@ def apply_within(self, configs_i: torch.Tensor, psi_i: torch.Tensor, configs_j:
     def find_relative(self, configs_i: torch.Tensor, psi_i: torch.Tensor, count_selected: int, configs_exclude: torch.Tensor | None = None) -> torch.Tensor:
         return self.hamiltonian.find_relative(configs_i, psi_i, count_selected, configs_exclude)
 
+    def diagonal_term(self, configs: torch.Tensor) -> torch.Tensor:
+        return self.hamiltonian.diagonal_term(configs)
+
     def single_relative(self, configs: torch.Tensor) -> torch.Tensor:
         return self.hamiltonian.single_relative(configs)
 

qmb/hamiltonian.py

@@ -165,6 +165,26 @@ def find_relative(
         configs_j = _find_relative(configs_i, torch.view_as_real(psi_i), count_selected, self.site, self.kind, self.coef, configs_exclude)
         return configs_j
 
+    def diagonal_term(self, configs: torch.Tensor) -> torch.Tensor:
+        """
+        Get the diagonal term of the Hamiltonian for the given configurations.
+
+        Parameters
+        ----------
+        configs : torch.Tensor
+            A uint8 tensor of shape [batch_size, n_qubytes] representing the input configurations.
+
+        Returns
+        -------
+        torch.Tensor
+            A complex64 tensor of shape [batch_size] representing the diagonal term of the Hamiltonian for the given configurations.
+        """
+        self._prepare_data(configs.device)
+        _diagonal_term: typing.Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
+        _diagonal_term = getattr(self._load_module(configs.device.type, configs.size(1), self.particle_cut), "diagonal_term")
+        psi_result = torch.view_as_complex(_diagonal_term(configs, self.site, self.kind, self.coef))
+        return psi_result
+
     def single_relative(self, configs: torch.Tensor) -> torch.Tensor:
         """
         Find a single relative configuration for each configurations.

qmb/hubbard.py

@@ -112,6 +112,9 @@ def apply_within(self, configs_i: torch.Tensor, psi_i: torch.Tensor, configs_j:
     def find_relative(self, configs_i: torch.Tensor, psi_i: torch.Tensor, count_selected: int, configs_exclude: torch.Tensor | None = None) -> torch.Tensor:
         return self.hamiltonian.find_relative(configs_i, psi_i, count_selected, configs_exclude)
 
+    def diagonal_term(self, configs: torch.Tensor) -> torch.Tensor:
+        return self.hamiltonian.diagonal_term(configs)
+
     def single_relative(self, configs: torch.Tensor) -> torch.Tensor:
         return self.hamiltonian.single_relative(configs)
 

qmb/ising.py

@@ -221,6 +221,9 @@ def apply_within(self, configs_i: torch.Tensor, psi_i: torch.Tensor, configs_j:
     def find_relative(self, configs_i: torch.Tensor, psi_i: torch.Tensor, count_selected: int, configs_exclude: torch.Tensor | None = None) -> torch.Tensor:
         return self.hamiltonian.find_relative(configs_i, psi_i, count_selected, configs_exclude)
 
+    def diagonal_term(self, configs: torch.Tensor) -> torch.Tensor:
+        return self.hamiltonian.diagonal_term(configs)
+
     def single_relative(self, configs: torch.Tensor) -> torch.Tensor:
         return self.hamiltonian.single_relative(configs)
 

qmb/model_dict.py

@@ -173,6 +173,21 @@ def find_relative(self, configs_i: torch.Tensor, psi_i: torch.Tensor, count_sele
             The relative configurations.
         """
 
+    def diagonal_term(self, configs: torch.Tensor) -> torch.Tensor:
+        """
+        Calculate the diagonal term for the given configurations.
+
+        Parameters
+        ----------
+        configs : torch.Tensor
+            The configurations to calculate the diagonal term for.
+
+        Returns
+        -------
+        torch.Tensor
+            The diagonal term of the configurations.
+        """
+
     def single_relative(self, configs: torch.Tensor) -> torch.Tensor:
         """
         Find a single relative configuration for each configurations.

qmb/openfermion.py

@@ -86,6 +86,9 @@ def apply_within(self, configs_i: torch.Tensor, psi_i: torch.Tensor, configs_j:
     def find_relative(self, configs_i: torch.Tensor, psi_i: torch.Tensor, count_selected: int, configs_exclude: torch.Tensor | None = None) -> torch.Tensor:
         return self.hamiltonian.find_relative(configs_i, psi_i, count_selected, configs_exclude)
 
+    def diagonal_term(self, configs: torch.Tensor) -> torch.Tensor:
+        return self.hamiltonian.diagonal_term(configs)
+
     def single_relative(self, configs: torch.Tensor) -> torch.Tensor:
         return self.hamiltonian.single_relative(configs)