diff --git a/python/ffsim/qiskit/slater_determinant.py b/python/ffsim/qiskit/slater_determinant.py index 697143d37..9b87fdae0 100644 --- a/python/ffsim/qiskit/slater_determinant.py +++ b/python/ffsim/qiskit/slater_determinant.py @@ -78,10 +78,10 @@ class PrepareSlaterDeterminantJW(Gate): where :math:`\mathcal{U}` is an :doc:`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 diff --git a/python/ffsim/states/states.py b/python/ffsim/states/states.py index ca801765a..f00c37679 100644 --- a/python/ffsim/states/states.py +++ b/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 ` 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) diff --git a/tests/python/qiskit/slater_determinant_test.py b/tests/python/qiskit/slater_determinant_test.py index e5d0a47a5..afad20748 100644 --- a/tests/python/qiskit/slater_determinant_test.py +++ 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.""" diff --git a/tests/python/states/states_test.py b/tests/python/states/states_test.py index a3c139cd5..b65ab8fac 100644 --- a/tests/python/states/states_test.py +++ b/tests/python/states/states_test.py @@ -12,23 +12,25 @@ from __future__ import annotations +import itertools + import numpy as np import pyscf import pytest -import scipy.linalg import ffsim -def test_slater_determinant(): - """Test Slater determinant.""" - norb = 5 - nelec = ffsim.testing.random_nelec(norb) - occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec) +@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(5))) +def test_slater_determinant_same_rotation(norb: int, nelec: tuple[int, int]): + """Test Slater determinant with same rotation for both spins.""" + rng = np.random.default_rng() + + occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng) occ_a, occ_b = occupied_orbitals - one_body_tensor = ffsim.random.random_hermitian(norb) - eigs, orbital_rotation = scipy.linalg.eigh(one_body_tensor) + one_body_tensor = ffsim.random.random_hermitian(norb, seed=rng) + eigs, orbital_rotation = np.linalg.eigh(one_body_tensor) eig = sum(eigs[occ_a]) + sum(eigs[occ_b]) state = ffsim.slater_determinant( norb, occupied_orbitals, orbital_rotation=orbital_rotation @@ -38,6 +40,36 @@ def test_slater_determinant(): np.testing.assert_allclose(hamiltonian @ state, eig * state) +@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(5))) +def test_slater_determinant_diff_rotation(norb: int, nelec: tuple[int, int]): + """Test Slater determinant with different rotations for each spin.""" + rng = np.random.default_rng() + + occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng) + occ_a, occ_b = occupied_orbitals + + orbital_rotation_a = ffsim.random.random_unitary(norb, seed=rng) + orbital_rotation_b = ffsim.random.random_unitary(norb, seed=rng) + + state = ffsim.slater_determinant( + norb, + occupied_orbitals, + orbital_rotation=(orbital_rotation_a, orbital_rotation_b), + ) + state_a = ffsim.slater_determinant( + norb, + (occ_a, []), + orbital_rotation=orbital_rotation_a, + ) + state_b = ffsim.slater_determinant( + norb, + ([], occ_b), + orbital_rotation=orbital_rotation_b, + ) + + np.testing.assert_allclose(state, np.kron(state_a, state_b)) + + def test_hartree_fock_state(): """Test Hartree-Fock state.""" mol = pyscf.gto.Mole() @@ -60,24 +92,24 @@ def test_hartree_fock_state(): @pytest.mark.parametrize( - "norb, occupied_orbitals, spin_summed", + "norb, nelec, spin_summed", [ - (4, ([0, 1], [0, 1]), True), - (4, ([0, 1], [0, 1]), False), - (3, ([0], [1, 2]), True), - (3, ([0], [1, 2]), False), - (2, ([], [0]), True), - (2, ([], [0]), False), + (norb, nelec, spin_summed) + for (norb, nelec), spin_summed in itertools.product( + ffsim.testing.generate_norb_nelec(range(5)), [False, True] + ) ], ) -def test_slater_determinant_one_rdm( - norb: int, occupied_orbitals: tuple[list[int], list[int]], spin_summed: bool +def test_slater_determinant_one_rdm_same_rotation( + norb: int, nelec: tuple[int, int], spin_summed: bool ): """Test Slater determinant 1-RDM.""" + rng = np.random.default_rng() + + occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng) occ_a, occ_b = occupied_orbitals nelec = len(occ_a), len(occ_b) - rng = np.random.default_rng() orbital_rotation = ffsim.random.random_unitary(norb, seed=rng) vec = ffsim.slater_determinant( @@ -94,6 +126,44 @@ def test_slater_determinant_one_rdm( np.testing.assert_allclose(rdm, expected, atol=1e-12) +@pytest.mark.parametrize( + "norb, nelec, spin_summed", + [ + (norb, nelec, spin_summed) + for (norb, nelec), spin_summed in itertools.product( + ffsim.testing.generate_norb_nelec(range(5)), [False, True] + ) + ], +) +def test_slater_determinant_one_rdm_diff_rotation( + norb: int, nelec: tuple[int, int], spin_summed: bool +): + """Test Slater determinant 1-RDM.""" + rng = np.random.default_rng() + + occupied_orbitals = ffsim.testing.random_occupied_orbitals(norb, nelec, seed=rng) + occ_a, occ_b = occupied_orbitals + nelec = len(occ_a), len(occ_b) + + orbital_rotation_a = ffsim.random.random_unitary(norb, seed=rng) + orbital_rotation_b = ffsim.random.random_unitary(norb, seed=rng) + + vec = ffsim.slater_determinant( + norb, + occupied_orbitals, + orbital_rotation=(orbital_rotation_a, orbital_rotation_b), + ) + rdm = ffsim.slater_determinant_rdm( + norb, + occupied_orbitals, + orbital_rotation=(orbital_rotation_a, orbital_rotation_b), + spin_summed=spin_summed, + ) + expected = ffsim.rdm(vec, norb, nelec, spin_summed=spin_summed) + + np.testing.assert_allclose(rdm, expected, atol=1e-12) + + def test_indices_to_strings(): """Test converting statevector indices to strings.""" norb = 3