From 592b593d7104824958661ed5a7d1f818c5a5d379 Mon Sep 17 00:00:00 2001
From: "Kevin J. Sung" <kevjsung@umich.edu>
Date: Thu, 21 Sep 2023 18:56:28 -0400
Subject: [PATCH] add fsim gate

---
 python/ffsim/__init__.py                      |  1 +
 python/ffsim/gates/__init__.py                |  1 +
 python/ffsim/gates/basic_gates.py             | 62 ++++++++++++++++++-
 .../{gates_test.py => basic_gates_test.py}    | 32 +++++++++-
 4 files changed, 93 insertions(+), 3 deletions(-)
 rename tests/gates/{gates_test.py => basic_gates_test.py} (91%)

diff --git a/python/ffsim/__init__.py b/python/ffsim/__init__.py
index ceda0a504..2400ec084 100644
--- a/python/ffsim/__init__.py
+++ b/python/ffsim/__init__.py
@@ -32,6 +32,7 @@
 )
 from ffsim.gates import (
     apply_diag_coulomb_evolution,
+    apply_fsim_gate,
     apply_givens_rotation,
     apply_hop_gate,
     apply_num_interaction,
diff --git a/python/ffsim/gates/__init__.py b/python/ffsim/gates/__init__.py
index 068ce6f5a..3caf2a867 100644
--- a/python/ffsim/gates/__init__.py
+++ b/python/ffsim/gates/__init__.py
@@ -11,6 +11,7 @@
 """Fermionic quantum computation gates."""
 
 from ffsim.gates.basic_gates import (
+    apply_fsim_gate,
     apply_givens_rotation,
     apply_hop_gate,
     apply_num_interaction,
diff --git a/python/ffsim/gates/basic_gates.py b/python/ffsim/gates/basic_gates.py
index 299210553..97e1ac484 100644
--- a/python/ffsim/gates/basic_gates.py
+++ b/python/ffsim/gates/basic_gates.py
@@ -323,8 +323,8 @@ def apply_hop_gate(
 ) -> np.ndarray:
     r"""Apply a hop gate.
 
-    A "hop gate" is a Givens rotation gate followed by a number-number interaction
-    gate with angle pi:
+    A "hop gate" is a Givens rotation gate followed by a number-number interaction with
+    angle pi:
 
     .. math::
 
@@ -367,3 +367,61 @@ def apply_hop_gate(
         vec, np.pi, target_orbs, norb=norb, nelec=nelec, copy=False
     )
     return vec
+
+
+def apply_fsim_gate(
+    vec: np.ndarray,
+    theta: float,
+    phi: float,
+    target_orbs: tuple[int, int],
+    norb: int,
+    nelec: tuple[int, int],
+    *,
+    copy: bool = True,
+) -> np.ndarray:
+    r"""Apply an fSim gate.
+
+    An fSim gate consists of a tunneling interaction followed by a number-number
+    interaction (note the negative sign convention for the angles):
+
+    .. math::
+
+        \text{fSim}(\theta, \phi) = \text{NN}(-\phi) \text{T}(-\theta)
+        = \exp\left(-i \phi a^\dagger_i a_i a^\dagger_j a_j\right)
+        \exp\left(-i \theta (a^\dagger_i a_j + a^\dagger_j a_i)\right)
+
+    Under the Jordan-Wigner transform, this gate has the following matrix when applied
+    to neighboring qubits:
+
+    .. math::
+
+        \begin{pmatrix}
+            1 & 0 & 0 & 0 \\
+            0 & \cos(\theta) & -i \sin(\theta) & 0\\
+            0 & -i \sin(\theta) & \cos(\theta) & 0\\
+            0 & 0 & 0 & e^{-i \phi} \\
+        \end{pmatrix}
+
+    Args:
+        vec: The state vector to be transformed.
+        theta: The rotation angle.
+        target_orbs: The orbitals (i, j) to rotate.
+        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_tunneling_interaction(
+        vec, -theta, target_orbs, norb=norb, nelec=nelec, copy=False
+    )
+    vec = apply_num_num_interaction(
+        vec, -phi, target_orbs, norb=norb, nelec=nelec, copy=False
+    )
+    return vec
diff --git a/tests/gates/gates_test.py b/tests/gates/basic_gates_test.py
similarity index 91%
rename from tests/gates/gates_test.py
rename to tests/gates/basic_gates_test.py
index 8b72cd838..498fbd58c 100644
--- a/tests/gates/gates_test.py
+++ b/tests/gates/basic_gates_test.py
@@ -8,7 +8,7 @@
 # copyright notice, and modified files need to carry a notice indicating
 # that they have been altered from the originals.
 
-"""Tests for gates."""
+"""Tests for basic gates."""
 
 from __future__ import annotations
 
@@ -321,3 +321,33 @@ def mat(theta: float) -> np.ndarray:
                     phase_11=phase_11,
                     norb=norb,
                 )
+
+
+def test_apply_fsim_gate_matrix():
+    """Test applying fSim gate matrix."""
+    norb = 4
+    rng = np.random.default_rng()
+
+    def mat(theta: float) -> np.ndarray:
+        c = np.cos(theta)
+        s = np.sin(theta)
+        return np.array([[c, -1j * s], [-1j * s, c]])
+
+    phase_00 = 1
+
+    for _ in range(5):
+        theta = rng.uniform(-10, 10)
+        phi = rng.uniform(-10, 10)
+        phase_11 = np.exp(-1j * phi)
+        for i, j in itertools.combinations(range(norb), 2):
+            for target_orbs in [(i, j), (j, i)]:
+                assert_has_two_orbital_matrix(
+                    lambda vec, norb, nelec: ffsim.apply_fsim_gate(
+                        vec, theta, phi, target_orbs=target_orbs, norb=norb, nelec=nelec
+                    ),
+                    target_orbs=target_orbs,
+                    mat=mat(theta),
+                    phase_00=phase_00,
+                    phase_11=phase_11,
+                    norb=norb,
+                )