add on-site number-number interaction gate (#86)

python/ffsim/__init__.py

@@ -21,6 +21,7 @@
     apply_num_num_interaction,
     apply_num_op_prod_interaction,
     apply_num_op_sum_evolution,
+    apply_on_site_num_num_interaction,
     apply_orbital_rotation,
     apply_tunneling_interaction,
 )
@@ -98,6 +99,7 @@
     "apply_num_num_interaction",
     "apply_num_op_prod_interaction",
     "apply_num_op_sum_evolution",
+    "apply_on_site_num_num_interaction",
     "apply_orbital_rotation",
     "apply_tunneling_interaction",
     "apply_unitary",

python/ffsim/gates/__init__.py

@@ -17,6 +17,7 @@
     apply_num_interaction,
     apply_num_num_interaction,
     apply_num_op_prod_interaction,
+    apply_on_site_num_num_interaction,
     apply_tunneling_interaction,
 )
 from ffsim.gates.diag_coulomb import apply_diag_coulomb_evolution
@@ -32,6 +33,7 @@
     "apply_num_num_interaction",
     "apply_num_op_prod_interaction",
     "apply_num_op_sum_evolution",
+    "apply_on_site_num_num_interaction",
     "apply_orbital_rotation",
     "apply_tunneling_interaction",
 ]

python/ffsim/gates/basic_gates.py

@@ -273,6 +273,52 @@ def apply_num_num_interaction(
     return vec
 
 
+def apply_on_site_num_num_interaction(
+    vec: np.ndarray,
+    theta: float,
+    target_orb: int,
+    norb: int,
+    nelec: tuple[int, int],
+    *,
+    copy: bool = True,
+):
+    r"""Apply an on-site number-number interaction gate.
+
+    The on-site number-number interaction gate is
+
+    .. math::
+
+        \text{OSNN}(\theta, p) =
+        \exp\left(i \theta a^\dagger_{\alpha, p} a_{\alpha, p}
+        a^\dagger_{\beta, p} a_{\beta, p}\right)
+
+    Args:
+        vec: The state vector to be transformed.
+        theta: The rotation angle.
+        target_orb: The orbital on which to apply the interaction.
+        norb: The number of spatial orbitals.
+        nelec: The number of alpha and beta electrons.
+        copy: Whether to copy the vector before operating on it.
+            - If ``copy=True`` then this function always returns a newly allocated
+            vector and the original vector is left untouched.
+            - If ``copy=False`` then this function may still return a newly allocated
+            vector, but the original vector may have its data overwritten.
+            It is also possible that the original vector is returned,
+            modified in-place.
+    """
+    if copy:
+        vec = vec.copy()
+    vec = apply_num_op_prod_interaction(
+        vec,
+        theta,
+        target_orbs=([target_orb], [target_orb]),
+        norb=norb,
+        nelec=nelec,
+        copy=False,
+    )
+    return vec
+
+
 def apply_num_op_prod_interaction(
     vec: np.ndarray,
     theta: float,

tests/gates/basic_gates_test.py

@@ -226,7 +226,7 @@ def test_apply_num_interaction(norb: int, nelec: tuple[int, int]):
     ],
 )
 def test_apply_num_num_interaction(norb: int, nelec: tuple[int, int]):
-    """Test applying number interaction."""
+    """Test applying number-number interaction."""
     dim = ffsim.dim(norb, nelec)
     rng = np.random.default_rng()
     vec = np.array(ffsim.random.random_statevector(dim, seed=rng))
@@ -267,7 +267,7 @@ def test_apply_num_num_interaction(norb: int, nelec: tuple[int, int]):
     ],
 )
 def test_apply_num_num_interaction_eigenvalues(norb: int, nelec: tuple[int, int]):
-    """Test applying number interaction."""
+    """Test eigenvalues of number-number interaction."""
     rng = np.random.default_rng()
     occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec)
     vec = ffsim.slater_determinant(norb, occupied_orbitals)
@@ -290,6 +290,38 @@ def test_apply_num_num_interaction_eigenvalues(norb: int, nelec: tuple[int, int]
             np.testing.assert_allclose(result, expected)
 
 
+@pytest.mark.parametrize(
+    "norb, nelec",
+    [
+        (2, (1, 1)),
+        (3, (2, 1)),
+    ],
+)
+def test_apply_on_site_num_num_interaction(norb: int, nelec: tuple[int, int]):
+    """Test applying on-site number-number interaction."""
+    dim = ffsim.dim(norb, nelec)
+    rng = np.random.default_rng()
+    vec = np.array(ffsim.random.random_statevector(dim, seed=rng))
+    theta = rng.uniform(-10, 10)
+    for i in range(norb):
+        result = ffsim.apply_on_site_num_num_interaction(
+            vec, theta, i, norb=norb, nelec=nelec
+        )
+        generator = ffsim.FermionOperator(
+            {
+                (
+                    ffsim.cre_a(i),
+                    ffsim.des_a(i),
+                    ffsim.cre_b(i),
+                    ffsim.des_b(i),
+                ): theta
+            }
+        )
+        linop = ffsim.linear_operator(generator, norb=norb, nelec=nelec)
+        expected = scipy.sparse.linalg.expm_multiply(1j * linop, vec, traceA=theta)
+        np.testing.assert_allclose(result, expected)
+
+
 @pytest.mark.parametrize(
     "norb, nelec",
     [