apply UCJ operator directly

python/ffsim/tenpy/__init__.py

@@ -11,8 +11,10 @@
 """Code that uses TeNPy, e.g. for emulating quantum circuits."""
 
 from ffsim.tenpy.circuits.gates import (
-    gate1,
-    gate2,
+    apply_diag_coulomb_evolution,
+    apply_gate1,
+    apply_gate2,
+    apply_orbital_rotation,
     givens_rotation,
     num_interaction,
     num_num_interaction,
@@ -25,8 +27,10 @@
 from ffsim.tenpy.util import product_state_as_mps
 
 __all__ = [
-    "gate1",
-    "gate2",
+    "apply_diag_coulomb_evolution",
+    "apply_gate1",
+    "apply_gate2",
+    "apply_orbital_rotation",
     "givens_rotation",
     "lucj_circuit_as_mps",
     "MolecularChain",

python/ffsim/tenpy/circuits/gates.py

@@ -6,6 +6,7 @@
 from tenpy.networks.mps import MPS
 from tenpy.networks.site import SpinHalfFermionSite
 
+from ffsim.linalg import givens_decomposition
 from ffsim.spin import Spin
 
 # ignore lowercase argument and variable checks to maintain TeNPy naming conventions
@@ -247,7 +248,7 @@ def num_num_interaction(theta: float, spin: Spin) -> np.ndarray:
     return NNgate_sym
 
 
-def gate1(U1: np.ndarray, site: int, psi: MPS) -> None:
+def apply_gate1(U1: np.ndarray, site: int, psi: MPS) -> None:
     r"""Apply a single-site gate to a
     `TeNPy MPS <https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mps.MPS.html#tenpy.networks.mps.MPS>`__
     wavefunction.
@@ -257,7 +258,8 @@ def gate1(U1: np.ndarray, site: int, psi: MPS) -> None:
         site: The gate will be applied to `site` on the
             `TeNPy MPS <https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mps.MPS.html#tenpy.networks.mps.MPS>`__
             wavefunction.
-        psi: The wavefunction MPS.
+        psi: The `TeNPy MPS <https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mps.MPS.html#tenpy.networks.mps.MPS>`__
+            wavefunction.
 
     Returns:
         None
@@ -268,13 +270,13 @@ def gate1(U1: np.ndarray, site: int, psi: MPS) -> None:
     psi.apply_local_op(site, U1_npc)
 
 
-def gate2(
+def apply_gate2(
     U2: np.ndarray,
     site: int,
     psi: MPS,
     eng: TEBDEngine,
     chi_list: list,
-    norm_tol: float,
+    norm_tol: float = 1e-5,
 ) -> None:
     r"""Apply a two-site gate to a `TeNPy MPS <https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mps.MPS.html#tenpy.networks.mps.MPS>`__
     wavefunction.
@@ -308,3 +310,110 @@ def gate2(
     # recanonicalize psi if below error threshold
     if np.linalg.norm(psi.norm_test()) > norm_tol:
         psi.canonical_form_finite()
+
+
+def apply_orbital_rotation(
+    mat: np.ndarray,
+    psi: MPS,
+    eng: TEBDEngine,
+    chi_list: list,
+    norm_tol: float = 1e-5,
+) -> None:
+    r"""Apply an orbital rotation gate to a
+    `TeNPy MPS <https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mps.MPS.html#tenpy.networks.mps.MPS>`__
+    wavefunction.
+
+    The orbital rotation gate is defined in
+    `apply_orbital_rotation <https://qiskit-community.github.io/ffsim/api/ffsim.html#ffsim.apply_orbital_rotation>`__.
+
+    Args:
+        mat: The orbital rotation matrix of dimension `(norb, norb)`.
+        psi: The `TeNPy MPS <https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mps.MPS.html#tenpy.networks.mps.MPS>`__
+            wavefunction.
+        eng: The
+            `TeNPy TEBDEngine <https://tenpy.readthedocs.io/en/latest/reference/tenpy.algorithms.tebd.TEBDEngine.html#tenpy.algorithms.tebd.TEBDEngine>`__.
+        chi_list: The list to which to append the MPS bond dimensions as the circuit is
+            evaluated.
+        norm_tol: The norm error above which we recanonicalize the wavefunction, as
+            defined in the
+            `TeNPy documentation <https://tenpy.readthedocs.io/en/latest/reference/tenpy.algorithms.dmrg.DMRGEngine.html#cfg-option-DMRGEngine.norm_tol>`__.
+
+    Returns:
+        None
+    """
+
+    # Givens decomposition
+    givens_list, diag_mat = givens_decomposition(mat)
+
+    # apply the Givens rotation gates
+    for gate in givens_list:
+        theta = np.arccos(gate.c)
+        s = np.conj(gate.s)
+        phi = np.real(1j * np.log(-s / np.sin(theta))) if theta else 0
+        conj = True if gate.j < gate.i else False
+        apply_gate2(
+            givens_rotation(theta, Spin.ALPHA_AND_BETA, conj, phi=phi),
+            max(gate.i, gate.j),
+            psi,
+            eng,
+            chi_list,
+            norm_tol,
+        )
+
+    # apply the number interaction gates
+    for i, z in enumerate(diag_mat):
+        theta = float(np.angle(z))
+        apply_gate1(
+            np.exp(1j * theta) * num_interaction(-theta, Spin.ALPHA_AND_BETA), i, psi
+        )
+
+
+def apply_diag_coulomb_evolution(
+    mat: np.ndarray, psi: MPS, eng: TEBDEngine, chi_list: list, norm_tol: float = 1e-5
+) -> None:
+    r"""Apply a diagonal Coulomb evolution gate to a
+    `TeNPy MPS <https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mps.MPS.html#tenpy.networks.mps.MPS>`__
+    wavefunction.
+
+    The diagonal Coulomb evolution gate is defined in
+    `apply_diag_coulomb_evolution <https://qiskit-community.github.io/ffsim/api/ffsim.html#ffsim.apply_diag_coulomb_evolution>`__.
+
+    Args:
+        mat: The diagonal Coulomb matrices of dimension `(2, norb, norb)`.
+        psi: The `TeNPy MPS <https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mps.MPS.html#tenpy.networks.mps.MPS>`__
+            wavefunction.
+        eng: The
+            `TeNPy TEBDEngine <https://tenpy.readthedocs.io/en/latest/reference/tenpy.algorithms.tebd.TEBDEngine.html#tenpy.algorithms.tebd.TEBDEngine>`__.
+        chi_list: The list to which to append the MPS bond dimensions as the circuit is
+            evaluated.
+        norm_tol: The norm error above which we recanonicalize the wavefunction, as
+            defined in the
+            `TeNPy documentation <https://tenpy.readthedocs.io/en/latest/reference/tenpy.algorithms.dmrg.DMRGEngine.html#cfg-option-DMRGEngine.norm_tol>`__.
+
+    Returns:
+        None
+    """
+
+    # extract norb
+    assert np.shape(mat)[1] == np.shape(mat)[2]
+    norb = np.shape(mat)[1]
+
+    # unpack alpha-alpha and alpha-beta matrices
+    mat_aa, mat_ab = mat
+
+    # apply alpha-alpha gates
+    for i in range(norb):
+        for j in range(norb):
+            if j > i and mat_aa[i, j]:
+                apply_gate2(
+                    num_num_interaction(-mat_aa[i, j], Spin.ALPHA_AND_BETA),
+                    j,
+                    psi,
+                    eng,
+                    chi_list,
+                    norm_tol,
+                )
+
+    # apply alpha-beta gates
+    for i in range(norb):
+        apply_gate1(on_site_interaction(-mat_ab[i, i]), i, psi)

python/ffsim/tenpy/circuits/lucj_circuit.py

@@ -1,27 +1,20 @@
 from __future__ import annotations
 
 import numpy as np
-from qiskit.circuit import QuantumCircuit, QuantumRegister
 from tenpy.algorithms.tebd import TEBDEngine
 from tenpy.networks.mps import MPS
 
-import ffsim
-from ffsim.spin import Spin
 from ffsim.tenpy.circuits.gates import (
-    gate1,
-    gate2,
-    givens_rotation,
-    num_interaction,
-    num_num_interaction,
-    on_site_interaction,
+    apply_diag_coulomb_evolution,
+    apply_orbital_rotation,
 )
 from ffsim.tenpy.util import product_state_as_mps
 from ffsim.variational.ucj_spin_balanced import UCJOpSpinBalanced
 
 
 def lucj_circuit_as_mps(
     norb: int,
-    nelec: tuple,
+    nelec: int | tuple[int, int],
     ucj_op: UCJOpSpinBalanced,
     options: dict,
     norm_tol: float = 1e-5,
@@ -52,82 +45,20 @@ def lucj_circuit_as_mps(
     # prepare initial Hartree-Fock state
     psi = product_state_as_mps(norb, nelec, 0)
 
-    # construct the qiskit circuit
-    qubits = QuantumRegister(2 * norb)
-    circuit = QuantumCircuit(qubits)
-    circuit.append(ffsim.qiskit.UCJOpSpinBalancedJW(ucj_op), qubits)
-
     # define the TEBD engine
     eng = TEBDEngine(psi, None, options)
 
-    # execute the tenpy circuit
-    for ins in circuit.decompose(reps=2):
-        if ins.operation.name == "p":
-            qubit = ins.qubits[0]
-            idx = qubit._index
-            spin_flag = Spin.ALPHA if idx < norb else Spin.BETA
-            lmbda = ins.operation.params[0]
-            gate1(
-                np.exp(1j * lmbda) * num_interaction(-lmbda, spin_flag), idx % norb, psi
-            )
-        elif ins.operation.name == "xx_plus_yy":
-            qubit0 = ins.qubits[0]
-            qubit1 = ins.qubits[1]
-            idx0, idx1 = qubit0._index, qubit1._index
-            if idx0 < norb and idx1 < norb:
-                spin_flag = Spin.ALPHA
-            elif idx0 >= norb and idx1 >= norb:
-                spin_flag = Spin.BETA
-            else:
-                raise ValueError("XXPlusYY gate not allowed across spin sectors")
-            theta_val = ins.operation.params[0]
-            beta_val = ins.operation.params[1]
-            # directionality important when beta!=0
-            conj_flag = True if idx0 > idx1 else False
-            gate2(
-                givens_rotation(
-                    theta_val / 2, spin_flag, conj_flag, phi=beta_val - np.pi / 2
-                ),
-                max(idx0 % norb, idx1 % norb),
-                psi,
-                eng,
-                chi_list,
-                norm_tol,
-            )
-        elif ins.operation.name == "cp":
-            qubit0 = ins.qubits[0]
-            qubit1 = ins.qubits[1]
-            idx0, idx1 = qubit0._index, qubit1._index
-            lmbda = ins.operation.params[0]
-            # onsite (different spins)
-            if np.abs(idx0 - idx1) == norb:
-                gate1(on_site_interaction(-lmbda), min(idx0, idx1), psi)
-            # NN (up spins)
-            elif np.abs(idx0 - idx1) == 1 and idx0 < norb and idx1 < norb:
-                gate2(
-                    num_num_interaction(-lmbda, Spin.ALPHA),
-                    max(idx0, idx1),
-                    psi,
-                    eng,
-                    chi_list,
-                    norm_tol,
-                )
-            # NN (down spins)
-            elif np.abs(idx0 - idx1) == 1 and idx0 >= norb and idx1 >= norb:
-                gate2(
-                    num_num_interaction(-lmbda, Spin.BETA),
-                    max(idx0 % norb, idx1 % norb),
-                    psi,
-                    eng,
-                    chi_list,
-                    norm_tol,
-                )
-            else:
-                raise ValueError(
-                    "CPhase only implemented onsite (different spins) "
-                    "and NN (same spins)"
-                )
-        else:
-            raise ValueError(f"gate {ins.operation.name} not implemented.")
+    # construct the LUCJ MPS
+    n_reps = np.shape(ucj_op.orbital_rotations)[0]
+    for i in range(n_reps):
+        apply_orbital_rotation(
+            np.conj(ucj_op.orbital_rotations[i]).T, psi, eng, chi_list, norm_tol
+        )
+        apply_diag_coulomb_evolution(
+            ucj_op.diag_coulomb_mats[i], psi, eng, chi_list, norm_tol
+        )
+        apply_orbital_rotation(
+            ucj_op.orbital_rotations[i], psi, eng, chi_list, norm_tol
+        )
 
     return psi, chi_list

tests/python/tenpy/lucj_circuit_test.py

@@ -75,7 +75,7 @@ def test_lucj_circuit_as_mps(norb: int, nelec: tuple[int, int], connectivity: st
         interaction_pairs=_interaction_pairs_spin_balanced_(
             connectivity=connectivity, norb=norb
         ),
-        with_final_orbital_rotation=True,
+        with_final_orbital_rotation=False,
     )
     params = rng.uniform(-10, 10, size=n_params)
     lucj_op = ffsim.UCJOpSpinBalanced.from_parameters(
@@ -85,7 +85,7 @@ def test_lucj_circuit_as_mps(norb: int, nelec: tuple[int, int], connectivity: st
         interaction_pairs=_interaction_pairs_spin_balanced_(
             connectivity=connectivity, norb=norb
         ),
-        with_final_orbital_rotation=True,
+        with_final_orbital_rotation=False,
     )
 
     # generate the corresponding LUCJ circuit