Add two-dimensional Fermi-Hubbard model

python/ffsim/__init__.py

@@ -41,6 +41,7 @@
     des_a,
     des_b,
     fermi_hubbard_1d,
+    fermi_hubbard_2d,
     number_operator,
 )
 from ffsim.protocols import (
@@ -126,6 +127,7 @@
     "expectation_one_body_power",
     "expectation_one_body_product",
     "fermi_hubbard_1d",
+    "fermi_hubbard_2d",
     "fermion_operator",
     "hartree_fock_state",
     "indices_to_strings",

python/ffsim/operators/__init__.py

@@ -11,7 +11,7 @@
 """Operators."""
 
 from ffsim.operators.common_operators import number_operator
-from ffsim.operators.fermi_hubbard import fermi_hubbard_1d
+from ffsim.operators.fermi_hubbard import fermi_hubbard_1d, fermi_hubbard_2d
 from ffsim.operators.fermion_action import (
     FermionAction,
     cre,
@@ -33,5 +33,6 @@
     "des_a",
     "des_b",
     "fermi_hubbard_1d",
+    "fermi_hubbard_2d",
     "number_operator",
 ]

python/ffsim/operators/fermi_hubbard.py

@@ -87,3 +87,107 @@ def fermi_hubbard_1d(
             coeffs[(cre_b(p), des_b(p))] = -chemical_potential
 
     return FermionOperator(coeffs)
+
+
+def fermi_hubbard_2d(
+    norb_x: int,
+    norb_y: int,
+    tunneling: float,
+    interaction: float,
+    *,
+    chemical_potential: float = 0,
+    nearest_neighbor_interaction: float = 0,
+    periodic: bool = False,
+) -> FermionOperator:
+    r"""Two-dimensional Fermi-Hubbard model Hamiltonian.
+
+    The Hamiltonian for the two-dimensional Fermi-Hubbard model on a square lattice with
+    :math:`N = N_x \times N_y` spatial orbitals is given by
+
+    .. math::
+
+        H = -t \sum_{\sigma} \sum_{\braket{pq}}
+        (a^\dagger_{\sigma, p} a_{\sigma, q} + \text{h.c.})
+        + U \sum_{p=1}^{N} n_{\alpha, p} n_{\beta, p}
+        - \mu \sum_{p=1}^{N} (n_{\alpha, p} + n_{\beta, p})
+        + V \sum_{\sigma, \sigma'} \sum_{\braket{pq}} n_{\sigma, p} n_{\sigma', q}
+
+    where :math:`\braket{\dots}` denotes nearest-neighbor pairs and
+    :math:`n_{\sigma, p} = a_{\sigma, p}^\dagger a_{\sigma, p}` is the number operator
+    on orbital :math:`p` with spin :math:`\sigma`.
+
+    References:
+        - `The Hubbard Model`_
+
+    Args:
+        norb_x: The number of spatial orbitals in the x-direction :math:`N_x`.
+        norb_y: The number of spatial orbitals in the y-direction :math:`N_y`.
+        tunneling: The tunneling amplitude :math:`t`.
+        interaction: The onsite interaction strength :math:`U`.
+        chemical_potential: The chemical potential :math:`\mu`.
+        nearest_neighbor_interaction: The nearest-neighbor interaction strength
+            :math:`V`.
+        periodic: Whether to use periodic boundary conditions.
+
+    Returns:
+        The two-dimensional Fermi-Hubbard model Hamiltonian.
+
+    .. _The Hubbard Model: https://doi.org/10.1146/annurev-conmatphys-031620-102024
+    """
+    coeffs: dict[tuple[tuple[bool, bool, int], ...], complex] = defaultdict(float)
+
+    for orb in range(norb_x * norb_y):
+        # map from row-major ordering to Cartesian coordinates
+        y, x = divmod(orb, norb_x)
+
+        # get the coordinates/orbitals to the right and up
+        x_right = (x + 1) % norb_x
+        y_up = (y + 1) % norb_y
+        orb_right = x_right + norb_x * y
+        orb_up = x + norb_x * y_up
+
+        if x != norb_x - 1 or periodic:
+            coeffs[(cre_a(orb), des_a(orb_right))] -= tunneling
+            coeffs[(cre_a(orb_right), des_a(orb))] -= tunneling
+            coeffs[(cre_b(orb), des_b(orb_right))] -= tunneling
+            coeffs[(cre_b(orb_right), des_b(orb))] -= tunneling
+            if nearest_neighbor_interaction:
+                coeffs[(cre_a(orb), des_a(orb), cre_a(orb_right), des_a(orb_right))] = (
+                    nearest_neighbor_interaction
+                )
+                coeffs[(cre_a(orb), des_a(orb), cre_b(orb_right), des_b(orb_right))] = (
+                    nearest_neighbor_interaction
+                )
+                coeffs[(cre_b(orb), des_b(orb), cre_a(orb_right), des_a(orb_right))] = (
+                    nearest_neighbor_interaction
+                )
+                coeffs[(cre_b(orb), des_b(orb), cre_b(orb_right), des_b(orb_right))] = (
+                    nearest_neighbor_interaction
+                )
+
+        if y != norb_y - 1 or periodic:
+            coeffs[(cre_a(orb), des_a(orb_up))] -= tunneling
+            coeffs[(cre_a(orb_up), des_a(orb))] -= tunneling
+            coeffs[(cre_b(orb), des_b(orb_up))] -= tunneling
+            coeffs[(cre_b(orb_up), des_b(orb))] -= tunneling
+            if nearest_neighbor_interaction:
+                coeffs[(cre_a(orb), des_a(orb), cre_a(orb_up), des_a(orb_up))] = (
+                    nearest_neighbor_interaction
+                )
+                coeffs[(cre_a(orb), des_a(orb), cre_b(orb_up), des_b(orb_up))] = (
+                    nearest_neighbor_interaction
+                )
+                coeffs[(cre_b(orb), des_b(orb), cre_a(orb_up), des_a(orb_up))] = (
+                    nearest_neighbor_interaction
+                )
+                coeffs[(cre_b(orb), des_b(orb), cre_b(orb_up), des_b(orb_up))] = (
+                    nearest_neighbor_interaction
+                )
+
+        if interaction:
+            coeffs[(cre_a(orb), des_a(orb), cre_b(orb), des_b(orb))] = interaction
+        if chemical_potential:
+            coeffs[(cre_a(orb), des_a(orb))] = -chemical_potential
+            coeffs[(cre_b(orb), des_b(orb))] = -chemical_potential
+
+    return FermionOperator(coeffs)