初步修改hubbard为free_fermion

qmb/__main__.py

@@ -8,6 +8,7 @@
 from . import openfermion as _  # type: ignore[no-redef]
 from . import fcidump as _  # type: ignore[no-redef]
 from . import hubbard as _  # type: ignore[no-redef]
+from . import free_fermion as _  # type: ignore[no-redef]
 from . import ising as _  # type: ignore[no-redef]
 from . import vmc as _  # type: ignore[no-redef]
 from . import imag as _  # type: ignore[no-redef]

qmb/free_fermion.py

@@ -0,0 +1,113 @@
+"""
+This file offers an interface for defining free fermion models on a two-dimensional lattice.
+"""
+
+import typing
+import logging
+import dataclasses
+import torch
+import tyro
+from .hamiltonian import Hamiltonian
+from .model_dict import model_dict, ModelProto, NetworkConfigProto
+
+
+@dataclasses.dataclass
+class ModelConfig:
+    """
+    The configuration for the free fermion model.
+    """
+
+    # The width of the free fermion lattice
+    m: typing.Annotated[int, tyro.conf.Positional]
+    # The height of the free fermion lattice
+    n: typing.Annotated[int, tyro.conf.Positional]
+
+    # The electron number
+    electron_number: typing.Annotated[int, tyro.conf.arg(aliases=["-e"])]
+
+    def __post_init__(self) -> None:
+        if self.m <= 0 or self.n <= 0:
+            raise ValueError("The dimensions of the free fermion model must be positive integers.")
+
+        if self.electron_number < 0 or self.electron_number > self.m * self.n:
+            raise ValueError(f"The electron number {self.electron_number} is out of bounds for a {self.m}x{self.n} lattice. Each site can host up to one electron.")
+
+
+class Model(ModelProto[ModelConfig]):
+    """
+    This class handles the free fermion model.
+    """
+
+    network_dict: dict[str, type[NetworkConfigProto["Model"]]] = {}
+
+    config_t = ModelConfig
+
+    @classmethod
+    def default_group_name(cls, config: ModelConfig) -> str:
+        return f"FreeFermion_{config.m}x{config.n}_e{config.electron_number}"
+
+    @classmethod
+    def _prepare_hamiltonian(cls, args: ModelConfig) -> dict[tuple[tuple[int, int], ...], complex]:
+
+        def _index(i: int, j: int) -> int:
+            return i + j * args.m
+
+        hamiltonian_dict: dict[tuple[tuple[int, int], ...], complex] = {}
+        for i in range(args.m):
+            for j in range(args.n):
+                # Nearest neighbor hopping
+                if i != 0:
+                    hamiltonian_dict[(_index(i, j), 1), (_index(i - 1, j), 0)] = 1
+                    hamiltonian_dict[(_index(i - 1, j), 1), (_index(i, j), 0)] = 1
+                if j != 0:
+                    hamiltonian_dict[(_index(i, j), 1), (_index(i, j - 1), 0)] = 1
+                    hamiltonian_dict[(_index(i, j - 1), 1), (_index(i, j), 0)] = 1
+
+        return hamiltonian_dict
+
+    def _calculate_ref_energy(self, hamiltonian_dict: dict[tuple[tuple[int, int], ...], complex]) -> float:
+        sites = self.m * self.n
+        hamiltonian = torch.zeros((sites, sites), dtype=torch.complex128)
+        for ((site_1, _), (site_2, _)), value in hamiltonian_dict.items():
+            hamiltonian[site_1, site_2] = torch.tensor(value, dtype=torch.complex128)
+        return torch.linalg.eigh(hamiltonian).eigenvalues[:self.electron_number].sum().item()  # pylint: disable=not-callable
+
+    def __init__(
+        self,
+        args: ModelConfig,
+    ):
+        logging.info("Input arguments successfully parsed")
+
+        assert args.electron_number is not None
+        self.m: int = args.m
+        self.n: int = args.n
+        self.electron_number: int = args.electron_number
+        logging.info("Constructing a free fermion model with dimensions: width = %d, height = %d", self.m, self.n)
+        logging.info("The parameters of the model are: N = %d", args.electron_number)
+
+        logging.info("Initializing Hamiltonian for the lattice")
+        hamiltonian_dict: dict[tuple[tuple[int, int], ...], complex] = self._prepare_hamiltonian(args)
+        logging.info("Hamiltonian initialization complete")
+
+        self.ref_energy: float = self._calculate_ref_energy(hamiltonian_dict)
+        logging.info("The ref energy is %.10f", self.ref_energy)
+
+        logging.info("Converting the Hamiltonian to internal Hamiltonian representation")
+        self.hamiltonian: Hamiltonian = Hamiltonian(hamiltonian_dict, kind="fermi")
+        logging.info("Internal Hamiltonian representation for model has been successfully created")
+
+    def apply_within(self, configs_i: torch.Tensor, psi_i: torch.Tensor, configs_j: torch.Tensor) -> torch.Tensor:
+        return self.hamiltonian.apply_within(configs_i, psi_i, 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 single_relative(self, configs: torch.Tensor) -> torch.Tensor:
+        return self.hamiltonian.single_relative(configs)
+
+    def show_config(self, config: torch.Tensor) -> str:
+        string = "".join(f"{i:08b}"[::-1] for i in config.cpu().numpy())
+        return "[" + ".".join("".join("-" if string[i + j * self.m] == "0" else "x" for i in range(self.m)) for j in range(self.n)) + "]"
+
+
+model_dict["free_fermion"] = Model