Add two-dimensional Fermi-Hubbard model

python/ffsim/__init__.py

@@ -42,6 +42,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

@@ -10,6 +10,8 @@
 
 """Fermi-Hubbard model Hamiltonian."""
 
+from collections import defaultdict
+
 from ffsim._lib import FermionOperator
 from ffsim.operators.fermion_action import cre_a, cre_b, des_a, des_b
 
@@ -85,3 +87,109 @@ 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 x in range(norb_x):
+        for y in range(norb_y):
+            # position in Cartesian coordinates
+            x_right = (x + 1) % norb_x
+            y_up = (y + 1) % norb_y
+
+            # position on C-style chain
+            p = norb_y * x + y
+            p_right = norb_y * x_right + y
+            p_up = norb_y * x + y_up
+
+            if x != norb_x - 1 or periodic:
+                coeffs[(cre_a(p), des_a(p_right))] -= tunneling
+                coeffs[(cre_a(p_right), des_a(p))] -= tunneling
+                coeffs[(cre_b(p), des_b(p_right))] -= tunneling
+                coeffs[(cre_b(p_right), des_b(p))] -= tunneling
+                if nearest_neighbor_interaction:
+                    coeffs[(cre_a(p), des_a(p), cre_a(p_right), des_a(p_right))] = (
+                        nearest_neighbor_interaction
+                    )
+                    coeffs[(cre_a(p), des_a(p), cre_b(p_right), des_b(p_right))] = (
+                        nearest_neighbor_interaction
+                    )
+                    coeffs[(cre_b(p), des_b(p), cre_a(p_right), des_a(p_right))] = (
+                        nearest_neighbor_interaction
+                    )
+                    coeffs[(cre_b(p), des_b(p), cre_b(p_right), des_b(p_right))] = (
+                        nearest_neighbor_interaction
+                    )
+
+            if y != norb_y - 1 or periodic:
+                coeffs[(cre_a(p), des_a(p_up))] -= tunneling
+                coeffs[(cre_a(p_up), des_a(p))] -= tunneling
+                coeffs[(cre_b(p), des_b(p_up))] -= tunneling
+                coeffs[(cre_b(p_up), des_b(p))] -= tunneling
+                if nearest_neighbor_interaction:
+                    coeffs[(cre_a(p), des_a(p), cre_a(p_up), des_a(p_up))] = (
+                        nearest_neighbor_interaction
+                    )
+                    coeffs[(cre_a(p), des_a(p), cre_b(p_up), des_b(p_up))] = (
+                        nearest_neighbor_interaction
+                    )
+                    coeffs[(cre_b(p), des_b(p), cre_a(p_up), des_a(p_up))] = (
+                        nearest_neighbor_interaction
+                    )
+                    coeffs[(cre_b(p), des_b(p), cre_b(p_up), des_b(p_up))] = (
+                        nearest_neighbor_interaction
+                    )
+
+    for p in range(norb_x * norb_y):
+        if interaction:
+            coeffs[(cre_a(p), des_a(p), cre_b(p), des_b(p))] = interaction
+        if chemical_potential:
+            coeffs[(cre_a(p), des_a(p))] = -chemical_potential
+            coeffs[(cre_b(p), des_b(p))] = -chemical_potential
+
+    return FermionOperator(coeffs)

tests/python/operators/__init__.py

@@ -1,4 +1,4 @@
-# (C) Copyright IBM 2023.
+# (C) Copyright IBM 2024.
 #
 # This code is licensed under the Apache License, Version 2.0. You may
 # obtain a copy of this license in the LICENSE.txt file in the root directory