accept tuple of orbital rotations in slater determinant (#155)

python/ffsim/qiskit/slater_determinant.py

@@ -78,10 +78,10 @@ class PrepareSlaterDeterminantJW(Gate):
 
     where :math:`\mathcal{U}` is an
     :doc:`orbital rotation </explanations/orbital-rotation>` and
-    :math:`\lvert x \rangle` is a computational basis state. The reason this gate
-    exists (when :class:`OrbitalRotationJW` already exists) is that the preparation
-    of a Slater determinant has a more optimized circuit than a generic orbital
-    rotation.
+    :math:`\lvert x \rangle` is an electronic configuration (computational basis state).
+    The reason this gate exists (when :class:`OrbitalRotationJW` already exists) is that
+    the preparation of a Slater determinant has a more optimized circuit than a generic
+    orbital rotation.
 
     This gate is meant to be applied to the all zeros state. Its behavior when applied
     to any other state is not guaranteed. The global phase of the prepared state may
@@ -100,7 +100,7 @@ def __init__(
         self,
         norb: int,
         occupied_orbitals: tuple[Sequence[int], Sequence[int]],
-        orbital_rotation: np.ndarray | None = None,
+        orbital_rotation: np.ndarray | tuple[np.ndarray, np.ndarray] | None = None,
         label: str | None = None,
         validate: bool = True,
         rtol: float = 1e-5,
@@ -110,12 +110,14 @@ def __init__(
 
         Args:
             norb: The number of spatial orbitals.
-            occupied_orbitals: A pair of lists of integers. The first list specifies
-                the occupied alpha orbitals and the second list specifies the occupied
-                beta orbitals.
-            orbital_rotation: An optional orbital rotation to apply to the
-                electron configuration. In other words, this is a unitary matrix that
-                describes the orbitals of the Slater determinant.
+            occupied_orbitals: The occupied orbitals in the electonic configuration.
+                This is a pair of lists of integers, where the first list specifies the
+                spin alpha orbitals and the second list specifies the spin beta
+                orbitals.
+            orbital_rotation: The orbital rotation. You can pass either a single Numpy
+                array specifying the orbital rotation to apply to both spin sectors, or
+                you can pass a pair of Numpy arrays to specify independent orbital
+                rotations for spin alpha and spin beta.
             label: The label of the gate.
             validate: Whether to validate the inputs.
             rtol: Relative numerical tolerance for input validation.
@@ -125,12 +127,23 @@ def __init__(
             ValueError: orbital_coeffs must be a 2-dimensional array.
             ValueError: orbital_coeffs must have orthonormal columns.
         """
-        if (
-            validate
-            and orbital_rotation is not None
-            and not linalg.is_unitary(orbital_rotation, rtol=rtol, atol=atol)
-        ):
-            raise ValueError("The input orbital rotation matrix was not unitary.")
+        if validate and orbital_rotation is not None:
+            if isinstance(orbital_rotation, np.ndarray):
+                if not linalg.is_unitary(orbital_rotation, rtol=rtol, atol=atol):
+                    raise ValueError(
+                        "The input orbital rotation matrix was not unitary."
+                    )
+            else:
+                orbital_rotation_a, orbital_rotation_b = orbital_rotation
+                if not linalg.is_unitary(orbital_rotation_a, rtol=rtol, atol=atol):
+                    raise ValueError(
+                        "The orbital rotation matrix for spin alpha was not unitary."
+                    )
+                if not linalg.is_unitary(orbital_rotation_b, rtol=rtol, atol=atol):
+                    raise ValueError(
+                        "The orbital rotation matrix for spin beta was not unitary."
+                    )
+
         self.norb = norb
         self.occupied_orbitals = occupied_orbitals
         self.orbital_rotation = orbital_rotation
@@ -150,12 +163,17 @@ def _define(self):
             for instruction in _prepare_configuration_jw(beta_qubits, occ_b):
                 circuit.append(instruction)
         else:
+            if isinstance(self.orbital_rotation, np.ndarray):
+                orbital_rotation_a = self.orbital_rotation
+                orbital_rotation_b = self.orbital_rotation
+            else:
+                orbital_rotation_a, orbital_rotation_b = self.orbital_rotation
             for instruction in _prepare_slater_determinant_jw(
-                alpha_qubits, self.orbital_rotation.T[list(occ_a)]
+                alpha_qubits, orbital_rotation_a.T[list(occ_a)]
             ):
                 circuit.append(instruction)
             for instruction in _prepare_slater_determinant_jw(
-                beta_qubits, self.orbital_rotation.T[list(occ_b)]
+                beta_qubits, orbital_rotation_b.T[list(occ_b)]
             ):
                 circuit.append(instruction)
         self.definition = circuit

python/ffsim/states/states.py

@@ -21,6 +21,7 @@
 from pyscf.fci.spin_op import contract_ss
 
 from ffsim.gates.orbital_rotation import apply_orbital_rotation
+from ffsim.spin import Spin
 
 
 def dims(norb: int, nelec: tuple[int, int]) -> tuple[int, int]:
@@ -72,18 +73,30 @@ def one_hot(shape: int | tuple[int, ...], index, *, dtype=complex):
 def slater_determinant(
     norb: int,
     occupied_orbitals: tuple[Sequence[int], Sequence[int]],
-    orbital_rotation: np.ndarray | None = None,
+    orbital_rotation: np.ndarray | tuple[np.ndarray, np.ndarray] | None = None,
 ) -> np.ndarray:
-    """Return a Slater determinant.
+    r"""Return a Slater determinant.
+
+    A Slater determinant is a state of the form
+
+    .. math::
+
+        \mathcal{U} \lvert x \rangle,
+
+    where :math:`\mathcal{U}` is an
+    :doc:`orbital rotation </explanations/orbital-rotation>` and
+    :math:`\lvert x \rangle` is an electronic configuration.
 
     Args:
         norb: The number of spatial orbitals.
-        occupied_orbitals: A pair of lists of integers. The first list specifies
-            the occupied alpha orbitals and the second list specifies the occupied
-            beta orbitals.
-        orbital_rotation: An optional orbital rotation to apply to the
-            electron configuration. In other words, this is a unitary matrix that
-            describes the orbitals of the Slater determinant.
+        occupied_orbitals: The occupied orbitals in the electonic configuration.
+            This is a pair of lists of integers, where the first list specifies the
+            spin alpha orbitals and the second list specifies the spin beta
+            orbitals.
+        orbital_rotation: The orbital rotation. You can pass either a single Numpy
+            array specifying the orbital rotation to apply to both spin sectors, or
+            you can pass a pair of Numpy arrays to specify independent orbital
+            rotations for spin alpha and spin beta.
 
     Returns:
         The Slater determinant as a statevector.
@@ -106,7 +119,28 @@ def slater_determinant(
     beta_index = cistring.str2addr(norb, n_beta, beta_string)
     vec = one_hot((dim1, dim2), (alpha_index, beta_index), dtype=complex).reshape(-1)
     if orbital_rotation is not None:
-        vec = apply_orbital_rotation(vec, orbital_rotation, norb=norb, nelec=nelec)
+        if isinstance(orbital_rotation, np.ndarray):
+            vec = apply_orbital_rotation(
+                vec, orbital_rotation, norb=norb, nelec=nelec, copy=False
+            )
+        else:
+            orbital_rotation_a, orbital_rotation_b = orbital_rotation
+            vec = apply_orbital_rotation(
+                vec,
+                orbital_rotation_a,
+                norb=norb,
+                nelec=nelec,
+                spin=Spin.ALPHA,
+                copy=False,
+            )
+            vec = apply_orbital_rotation(
+                vec,
+                orbital_rotation_b,
+                norb=norb,
+                nelec=nelec,
+                spin=Spin.BETA,
+                copy=False,
+            )
     return vec
 
 
@@ -129,28 +163,33 @@ def hartree_fock_state(norb: int, nelec: tuple[int, int]) -> np.ndarray:
 def slater_determinant_rdm(
     norb: int,
     occupied_orbitals: tuple[Sequence[int], Sequence[int]],
-    orbital_rotation: np.ndarray | None = None,
+    orbital_rotation: np.ndarray | tuple[np.ndarray, np.ndarray] | None = None,
     rank: int = 1,
     spin_summed: bool = True,
 ) -> np.ndarray:
-    """Return the reduced density matrix of a Slater determinant.
+    """Return the reduced density matrix of a `Slater determinant`_.
 
     Note:
         Currently, only rank 1 is supported.
 
     Args:
         norb: The number of spatial orbitals.
-        occupied_orbitals: A tuple of two sequences of integers. The first
-            sequence contains the indices of the occupied alpha orbitals, and
-            the second sequence similarly for the beta orbitals.
-        orbital_rotation: An optional orbital rotation to apply to the
-            electron configuration. In other words, this is a unitary matrix that
-            describes the orbitals of the Slater determinant.
-        rank: The rank of the reduced density matrix.
+        occupied_orbitals: The occupied orbitals in the electonic configuration.
+            This is a pair of lists of integers, where the first list specifies the
+            spin alpha orbitals and the second list specifies the spin beta
+            orbitals.
+        orbital_rotation: The orbital rotation. You can pass either a single Numpy
+            array specifying the orbital rotation to apply to both spin sectors, or
+            you can pass a pair of Numpy arrays to specify independent orbital
+            rotations for spin alpha and spin beta.
+        rank: The rank of the reduced density matrix. I.e., rank 1 corresponds to the
+            one-particle RDM, rank 2 corresponds to the 2-particle RDM, etc.
         spin_summed: Whether to sum over the spin index.
 
     Returns:
         The reduced density matrix of the Slater determinant.
+
+    .. _Slater determinant: ffsim.html#ffsim.slater_determinant
     """
     if rank == 1:
         rdm_a = np.zeros((norb, norb), dtype=complex)
@@ -162,8 +201,13 @@ def slater_determinant_rdm(
         if len(beta_orbitals):
             rdm_b[(beta_orbitals, beta_orbitals)] = 1
         if orbital_rotation is not None:
-            rdm_a = orbital_rotation.conj() @ rdm_a @ orbital_rotation.T
-            rdm_b = orbital_rotation.conj() @ rdm_b @ orbital_rotation.T
+            if isinstance(orbital_rotation, np.ndarray):
+                orbital_rotation_a = orbital_rotation
+                orbital_rotation_b = orbital_rotation
+            else:
+                orbital_rotation_a, orbital_rotation_b = orbital_rotation
+            rdm_a = orbital_rotation_a.conj() @ rdm_a @ orbital_rotation_a.T
+            rdm_b = orbital_rotation_b.conj() @ rdm_b @ orbital_rotation_b.T
         if spin_summed:
             return rdm_a + rdm_b
         return scipy.linalg.block_diag(rdm_a, rdm_b)

tests/python/qiskit/slater_determinant_test.py

@@ -35,7 +35,7 @@ def test_prepare_hartree_fock_jw(norb: int, nelec: tuple[int, int]):
 
 
 @pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(5)))
-def test_random_slater_determinant(norb: int, nelec: tuple[int, int]):
+def test_random_slater_determinant_same_rotation(norb: int, nelec: tuple[int, int]):
     """Test random Slater determinant circuit gives correct output state."""
     rng = np.random.default_rng()
     for _ in range(3):
@@ -57,6 +57,35 @@ def test_random_slater_determinant(norb: int, nelec: tuple[int, int]):
         ffsim.testing.assert_allclose_up_to_global_phase(result, expected)
 
 
+@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(5)))
+def test_random_slater_determinant_diff_rotation(norb: int, nelec: tuple[int, int]):
+    """Test random Slater determinant circuit with independent orbital rotations."""
+    rng = np.random.default_rng()
+    for _ in range(3):
+        occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec)
+        orbital_rotation_a = ffsim.random.random_unitary(norb, seed=rng)
+        orbital_rotation_b = ffsim.random.random_unitary(norb, seed=rng)
+
+        gate = ffsim.qiskit.PrepareSlaterDeterminantJW(
+            norb,
+            occupied_orbitals,
+            orbital_rotation=(orbital_rotation_a, orbital_rotation_b),
+        )
+
+        statevec = Statevector.from_int(0, 2 ** (2 * norb)).evolve(gate)
+        result = ffsim.qiskit.qiskit_vec_to_ffsim_vec(
+            np.array(statevec), norb=norb, nelec=nelec
+        )
+
+        expected = ffsim.slater_determinant(
+            norb,
+            occupied_orbitals,
+            orbital_rotation=(orbital_rotation_a, orbital_rotation_b),
+        )
+
+        ffsim.testing.assert_allclose_up_to_global_phase(result, expected)
+
+
 @pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(5)))
 def test_slater_determinant_no_rotation(norb: int, nelec: tuple[int, int]):
     """Test Slater determinant circuit with no orbital rotation."""