Merge branch 'main' into dev/rename-preparse-to-default-name

Author: hzhangxyz

qmb/__main__.py

@@ -7,6 +7,7 @@
 import tyro
 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 ising as _  # type: ignore[no-redef]
 from . import vmc as _  # type: ignore[no-redef]
 from . import imag as _  # type: ignore[no-redef]

qmb/_hamiltonian_cpu.cpp

@@ -42,7 +42,7 @@ __device__ bool get_bit(std::uint8_t* data, std::uint8_t index) {
     return ((*data) >> index) & 1;
 }
 
-__device__ bool set_bit(std::uint8_t* data, std::uint8_t index, bool value) {
+__device__ void set_bit(std::uint8_t* data, std::uint8_t index, bool value) {
     if (value) {
         *data |= (1 << index);
     } else {

qmb/_hamiltonian_cuda.cu

@@ -27,6 +27,7 @@ def _load_module(cls, n_qubytes: int = 0, particle_cut: int = 0) -> object:
                 name=name,
                 sources=[
                     f"{folder}/_hamiltonian.cpp",
+                    f"{folder}/_hamiltonian_cpu.cpp",
                     f"{folder}/_hamiltonian_cuda.cu",
                 ],
                 is_python_module=n_qubytes == 0,

qmb/hamiltonian.py

@@ -0,0 +1,234 @@
+"""
+This file offers an interface for defining Hubbard models on a two-dimensional lattice.
+"""
+
+import typing
+import logging
+import dataclasses
+import torch
+import tyro
+from .mlp import WaveFunctionElectronUpDown as MlpWaveFunction
+from .attention import WaveFunctionElectronUpDown as AttentionWaveFunction
+from .hamiltonian import Hamiltonian
+from .model_dict import model_dict, ModelProto, NetworkProto, NetworkConfigProto
+
+
+@dataclasses.dataclass
+class ModelConfig:
+    """
+    The configuration for the Hubbard model.
+    """
+
+    # The width of the hubbard lattice
+    m: typing.Annotated[int, tyro.conf.Positional]
+    # The height of the hubbard lattice
+    n: typing.Annotated[int, tyro.conf.Positional]
+
+    # The coefficient of t
+    t: typing.Annotated[float, tyro.conf.arg(aliases=["-t"])] = 1
+    # The coefficient of U
+    u: typing.Annotated[float, tyro.conf.arg(aliases=["-u"])] = 0
+
+    # The electron number
+    electron_number: typing.Annotated[int | None, tyro.conf.arg(aliases=["-e"])] = None
+
+    # The ref energy of the model
+    ref_energy: typing.Annotated[float, tyro.conf.arg(aliases=["-r"])] = 0
+
+    def __post_init__(self) -> None:
+        if self.electron_number is None:
+            self.electron_number = self.m * self.n
+            logging.info("Electron number is not specified, set to half-filling (one electron per lattice site) %d", self.electron_number)
+
+        if self.m <= 0 or self.n <= 0:
+            raise ValueError("The dimensions of the Hubbard model must be positive integers.")
+
+        if self.electron_number < 0 or self.electron_number > 2 * 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 two electrons (spin up and spin down).")
+
+
+class Model(ModelProto[ModelConfig]):
+    """
+    This class handles the Hubbard model.
+    """
+
+    network_dict: dict[str, type[NetworkConfigProto["Model"]]] = {}
+
+    config_t = ModelConfig
+
+    @classmethod
+    def preparse(cls, input_args: tuple[str, ...]) -> str:
+        args = tyro.cli(ModelConfig, args=input_args)
+        return f"Hubbard_{args.m}x{args.n}_t{args.t}_u{args.u}"
+
+    @classmethod
+    def _prepare_hamiltonian(cls, args: ModelConfig) -> dict[tuple[tuple[int, int], ...], complex]:
+
+        def _index(i: int, j: int, o: int) -> int:
+            return (i + j * args.m) * 2 + o
+
+        hamiltonian_dict: dict[tuple[tuple[int, int], ...], complex] = {}
+        for i in range(args.m):
+            for j in range(args.n):
+                # On-site interaction
+                hamiltonian_dict[(_index(i, j, 0), 1), (_index(i, j, 0), 0), (_index(i, j, 1), 1), (_index(i, j, 1), 0)] = args.u
+
+                # Nearest neighbor hopping
+                if i != 0:
+                    hamiltonian_dict[(_index(i, j, 0), 1), (_index(i - 1, j, 0), 0)] = -args.t
+                    hamiltonian_dict[(_index(i - 1, j, 0), 1), (_index(i, j, 0), 0)] = -args.t
+                    hamiltonian_dict[(_index(i, j, 1), 1), (_index(i - 1, j, 1), 0)] = -args.t
+                    hamiltonian_dict[(_index(i - 1, j, 1), 1), (_index(i, j, 1), 0)] = -args.t
+                if j != 0:
+                    hamiltonian_dict[(_index(i, j, 0), 1), (_index(i, j - 1, 0), 0)] = -args.t
+                    hamiltonian_dict[(_index(i, j - 1, 0), 1), (_index(i, j, 0), 0)] = -args.t
+                    hamiltonian_dict[(_index(i, j, 1), 1), (_index(i, j - 1, 1), 0)] = -args.t
+                    hamiltonian_dict[(_index(i, j - 1, 1), 1), (_index(i, j, 1), 0)] = -args.t
+
+        return hamiltonian_dict
+
+    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 Hubbard model with dimensions: width = %d, height = %d", self.m, self.n)
+        logging.info("The parameters of the model are: t = %.10f, U = %.10f, N = %d", args.t, args.u, 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 = args.ref_energy
+        logging.info("The ref energy is set to %.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(self._show_config_site(string[(i + j * self.m) * 2:(i + j * self.m) * 2 + 2]) for i in range(self.m)) for j in range(self.n)) + "]"
+
+    def _show_config_site(self, string: str) -> str:
+        match string:
+            case "00":
+                return " "
+            case "10":
+                return "↑"
+            case "01":
+                return "↓"
+            case "11":
+                return "↕"
+            case _:
+                raise ValueError(f"Invalid string: {string}")
+
+
+model_dict["hubbard"] = Model
+
+
+@dataclasses.dataclass
+class MlpConfig:
+    """
+    The configuration of the MLP network.
+    """
+
+    # The hidden widths of the network
+    hidden: typing.Annotated[tuple[int, ...], tyro.conf.arg(aliases=["-w"])] = (512,)
+
+    def create(self, model: Model) -> NetworkProto:
+        """
+        Create a MLP network for the model.
+        """
+        logging.info("Hidden layer widths: %a", self.hidden)
+
+        network = MlpWaveFunction(
+            double_sites=model.m * model.n * 2,
+            physical_dim=2,
+            is_complex=True,
+            spin_up=model.electron_number // 2,
+            spin_down=model.electron_number - model.electron_number // 2,
+            hidden_size=self.hidden,
+            ordering=+1,
+        )
+
+        return network
+
+
+Model.network_dict["mlp"] = MlpConfig
+
+
+@dataclasses.dataclass
+class AttentionConfig:
+    """
+    The configuration of the attention network.
+    """
+
+    # Embedding dimension
+    embedding_dim: typing.Annotated[int, tyro.conf.arg(aliases=["-e"])] = 512
+    # Heads number
+    heads_num: typing.Annotated[int, tyro.conf.arg(aliases=["-m"])] = 8
+    # Feedforward dimension
+    feed_forward_dim: typing.Annotated[int, tyro.conf.arg(aliases=["-f"])] = 2048
+    # Shared expert number
+    shared_expert_num: typing.Annotated[int, tyro.conf.arg(aliases=["-s"])] = 1
+    # Routed expert number
+    routed_expert_num: typing.Annotated[int, tyro.conf.arg(aliases=["-r"])] = 0
+    # Selected expert number
+    selected_expert_num: typing.Annotated[int, tyro.conf.arg(aliases=["-c"])] = 0
+    # Network depth
+    depth: typing.Annotated[int, tyro.conf.arg(aliases=["-d"])] = 6
+
+    def create(self, model: Model) -> NetworkProto:
+        """
+        Create an attention network for the model.
+        """
+        logging.info(
+            "Attention network configuration: "
+            "embedding dimension: %d, "
+            "number of heads: %d, "
+            "feed-forward dimension: %d, "
+            "shared expert number: %d, "
+            "routed expert number: %d, "
+            "selected expert number: %d, "
+            "depth: %d",
+            self.embedding_dim,
+            self.heads_num,
+            self.feed_forward_dim,
+            self.shared_expert_num,
+            self.routed_expert_num,
+            self.selected_expert_num,
+            self.depth,
+        )
+
+        network = AttentionWaveFunction(
+            double_sites=model.m * model.n * 2,
+            physical_dim=2,
+            is_complex=True,
+            spin_up=model.electron_number // 2,
+            spin_down=model.electron_number - model.electron_number // 2,
+            embedding_dim=self.embedding_dim,
+            heads_num=self.heads_num,
+            feed_forward_dim=self.feed_forward_dim,
+            shared_num=self.shared_expert_num,
+            routed_num=self.routed_expert_num,
+            selected_num=self.selected_expert_num,
+            depth=self.depth,
+            ordering=+1,
+        )
+
+        return network
+
+
+Model.network_dict["attention"] = AttentionConfig

qmb/hubbard.py

@@ -1,5 +1,5 @@
 """
-This file offers a interface for defining Ising-like models on a two-dimensional lattice.
+This file offers an interface for defining Ising-like models on a two-dimensional lattice.
 """
 
 import typing
@@ -194,12 +194,6 @@ def _z(i: int, j: int) -> tuple[tuple[tuple[tuple[int, int], ...], complex], ...
 
     def __init__(self, args: ModelConfig) -> None:
         logging.info("Input arguments successfully parsed")
-        logging.info("Grid dimensions: width = %d, height = %d", args.m, args.n)
-        logging.info("Element-wise coefficients: X = %.10f, Y = %.10f, Z = %.10f", args.x, args.y, args.z)
-        logging.info("Horizontal bond coefficients: X = %.10f, Y = %.10f, Z = %.10f", args.xh, args.yh, args.zh)
-        logging.info("Vertical bond coefficients: X = %.10f, Y = %.10f, Z = %.10f", args.xv, args.yv, args.zv)
-        logging.info("Diagonal bond coefficients: X = %.10f, Y = %.10f, Z = %.10f", args.xd, args.yd, args.zd)
-        logging.info("Anti-diagonal bond coefficients: X = %.10f, Y = %.10f, Z = %.10f", args.xa, args.ya, args.za)
 
         self.m: int = args.m
         self.n: int = args.n

qmb/ising.py

@@ -16,11 +16,40 @@ class NetworkProto(typing.Protocol):
     def __call__(self, x: torch.Tensor) -> torch.Tensor:
         """
         Calculate the amplitude for the given configurations.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The configurations to calculate the amplitude for.
+            The configurations are a two-dimensional uint8 tensor with first dimension equal to some batch size.
+            The second dimension contains occupation for each qubit which is bitwise encoded.
+
+        Returns
+        -------
+        torch.Tensor
+            The amplitudes of the configurations.
+            The amplitudes are a one-dimensional complex tensor with the only dimension equal to the batch size.
         """
 
     def generate_unique(self, batch_size: int, block_num: int = 1) -> tuple[torch.Tensor, torch.Tensor, None, None]:
         """
         Generate a batch of unique configurations.
+
+        Parameters
+        ----------
+        batch_size : int
+            The number of configurations to generate.
+        block_num : int, default=1
+            The number of batch block to generate. It is used to split the batch into smaller parts to avoid memory issues.
+
+        Returns
+        -------
+        tuple[torch.Tensor, torch.Tensor, None, None]
+            A tuple containing the generated configurations, their amplitudes, and two None values.
+            The configurations are a two-dimensional uint8 tensor with first dimension equal to `batch_size`.
+            The second dimension contains occupation for each qubit which is bitwise encoded.
+            The amplitudes are a one-dimensional complex tensor with the only dimension equal to `batch_size`.
+            The last two None values are reserved for future use.
         """
 
     def load_state_dict(self, data: dict[str, torch.Tensor]) -> typing.Any: