Repository: qiskit-community/ffsim (python)
Issue: allow orbital rotation and diag coulomb tuples to be a Numpy array dimension instead
For example, instead of passing (`mat_a`, `mat_b`), allow the matrices to be assembled in a Numpy array. Currently, a Numpy array is assumed to be a single operator. This check should also take the shape into account.
Pull request description: support numpy array input instead of tuple
fixes #195 

Patch:
diff --git a/python/ffsim/gates/diag_coulomb.py b/python/ffsim/gates/diag_coulomb.py
index 32263b8e..c8580502 100644
--- a/python/ffsim/gates/diag_coulomb.py
+++ b/python/ffsim/gates/diag_coulomb.py
@@ -167,7 +167,7 @@ def _get_mat_exp(
     mat_aa: np.ndarray | None
     mat_ab: np.ndarray | None
     mat_bb: np.ndarray | None
-    if isinstance(mat, np.ndarray):
+    if isinstance(mat, np.ndarray) and mat.ndim == 2:
         mat_aa, mat_ab = mat.copy(), mat.copy()
         mat_aa[np.diag_indices(norb)] *= 0.5
         if z_representation:
diff --git a/python/ffsim/gates/orbital_rotation.py b/python/ffsim/gates/orbital_rotation.py
index 7ec48ec7..7152ff77 100644
--- a/python/ffsim/gates/orbital_rotation.py
+++ b/python/ffsim/gates/orbital_rotation.py
@@ -118,7 +118,7 @@ def _get_givens_decomposition(
     tuple[list[GivensRotation], np.ndarray] | None,
     tuple[list[GivensRotation], np.ndarray] | None,
 ]:
-    if isinstance(mat, np.ndarray):
+    if isinstance(mat, np.ndarray) and mat.ndim == 2:
         decomp = givens_decomposition(mat)
         return decomp, decomp
     else:
diff --git a/python/ffsim/qiskit/gates/diag_coulomb.py b/python/ffsim/qiskit/gates/diag_coulomb.py
index bac10b3c..c0203863 100644
--- a/python/ffsim/qiskit/gates/diag_coulomb.py
+++ b/python/ffsim/qiskit/gates/diag_coulomb.py
@@ -33,7 +33,7 @@ def _validate_diag_coulomb_mat(
     rtol: float,
     atol: float,
 ) -> None:
-    if isinstance(mat, np.ndarray):
+    if isinstance(mat, np.ndarray) and mat.ndim == 2:
         if not linalg.is_real_symmetric(mat, rtol=rtol, atol=atol):
             raise ValueError("The input matrix was not real symmetric.")
     else:
@@ -153,7 +153,7 @@ def _diag_coulomb_evo_num_rep_jw(
     mat_aa: np.ndarray | None
     mat_ab: np.ndarray | None
     mat_bb: np.ndarray | None
-    if isinstance(mat, np.ndarray):
+    if isinstance(mat, np.ndarray) and mat.ndim == 2:
         mat_aa, mat_ab, mat_bb = mat, mat, mat
     else:
         mat_aa, mat_ab, mat_bb = mat
@@ -204,7 +204,7 @@ def _diag_coulomb_evo_z_rep_jw(
     mat_aa: np.ndarray | None
     mat_ab: np.ndarray | None
     mat_bb: np.ndarray | None
-    if isinstance(mat, np.ndarray):
+    if isinstance(mat, np.ndarray) and mat.ndim == 2:
         mat_aa, mat_ab, mat_bb = mat, mat, mat
     else:
         mat_aa, mat_ab, mat_bb = mat
diff --git a/python/ffsim/qiskit/gates/orbital_rotation.py b/python/ffsim/qiskit/gates/orbital_rotation.py
index dbbf62fe..aeb2a320 100644
--- a/python/ffsim/qiskit/gates/orbital_rotation.py
+++ b/python/ffsim/qiskit/gates/orbital_rotation.py
@@ -34,7 +34,7 @@ def _validate_orbital_rotation(
     rtol: float,
     atol: float,
 ) -> None:
-    if isinstance(mat, np.ndarray):
+    if isinstance(mat, np.ndarray) and mat.ndim == 2:
         if not linalg.is_unitary(mat, rtol=rtol, atol=atol):
             raise ValueError("The input orbital rotation matrix was not unitary.")
     else:
@@ -104,7 +104,7 @@ class OrbitalRotationJW(Gate):
         if validate:
             _validate_orbital_rotation(orbital_rotation, rtol=rtol, atol=atol)
         self.norb = norb
-        if isinstance(orbital_rotation, np.ndarray):
+        if isinstance(orbital_rotation, np.ndarray) and orbital_rotation.ndim == 2:
             self.orbital_rotation_a = orbital_rotation
             self.orbital_rotation_b = orbital_rotation
         else:


Test patch:
diff --git a/tests/python/gates/diag_coulomb_test.py b/tests/python/gates/diag_coulomb_test.py
index c8f45697..5735db1d 100644
--- a/tests/python/gates/diag_coulomb_test.py
+++ b/tests/python/gates/diag_coulomb_test.py
@@ -190,6 +190,32 @@ def test_apply_diag_coulomb_evolution_num_rep_asymmetric_spin(
             np.testing.assert_allclose(result, expected)
             np.testing.assert_allclose(state, original_state)
 
+            # (None, None, None)
+            result = ffsim.apply_diag_coulomb_evolution(
+                state, (None, None, None), time, norb, nelec
+            )
+            expected = state
+            np.testing.assert_allclose(result, expected)
+            np.testing.assert_allclose(state, original_state)
+
+            # Numpy array input
+            result = ffsim.apply_diag_coulomb_evolution(
+                state, np.stack((mat_aa, mat_ab, mat_bb)), time, norb, nelec
+            )
+            eig = 0
+            for i, j in itertools.product(range(norb), repeat=2):
+                for sigma, tau in itertools.product(range(2), repeat=2):
+                    if i in occupied_orbitals[sigma] and j in occupied_orbitals[tau]:
+                        if (sigma, tau) == (0, 0):
+                            eig += 0.5 * mat_aa[i, j]
+                        elif (sigma, tau) == (1, 1):
+                            eig += 0.5 * mat_bb[i, j]
+                        else:
+                            eig += 0.5 * mat_ab[i, j]
+            expected = np.exp(-1j * eig * time) * state
+            np.testing.assert_allclose(result, expected)
+            np.testing.assert_allclose(state, original_state)
+
 
 @pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(6)))
 def test_apply_diag_coulomb_evolution_z_rep_asymmetric_spin(
diff --git a/tests/python/gates/orbital_rotation_test.py b/tests/python/gates/orbital_rotation_test.py
index a39c2d9f..653ee5cf 100644
--- a/tests/python/gates/orbital_rotation_test.py
+++ b/tests/python/gates/orbital_rotation_test.py
@@ -21,20 +21,15 @@ import scipy.linalg
 import scipy.sparse.linalg
 
 import ffsim
-from ffsim.spin import pair_for_spin
 from ffsim.states import slater_determinant
 
 
-def _orbital_rotation_generator(
-    mat: np.ndarray, spin: ffsim.Spin
-) -> ffsim.FermionOperator:
+def _orbital_rotation_generator(mat: np.ndarray, spin: bool) -> ffsim.FermionOperator:
     norb, _ = mat.shape
     coeffs: dict[tuple[tuple[bool, bool, int], ...], complex] = {}
+    cre, des = (ffsim.cre_b, ffsim.des_b) if spin else (ffsim.cre_a, ffsim.des_a)
     for i, j in itertools.product(range(norb), repeat=2):
-        if spin & ffsim.Spin.ALPHA:
-            coeffs[ffsim.cre_a(i), ffsim.des_a(j)] = mat[i, j]
-        if spin & ffsim.Spin.BETA:
-            coeffs[ffsim.cre_b(i), ffsim.des_b(j)] = mat[i, j]
+        coeffs[cre(i), des(j)] = mat[i, j]
     return ffsim.FermionOperator(coeffs)
 
 
@@ -59,29 +54,51 @@ def test_apply_orbital_rotation_one_body_linop(norb: int, nelec: tuple[int, int]
         np.testing.assert_allclose(result, expected)
 
 
-@pytest.mark.parametrize(
-    "norb, nelec, spin", ffsim.testing.generate_norb_nelec_spin(range(4))
-)
-def test_apply_orbital_rotation_random(
-    norb: int, nelec: tuple[int, int], spin: ffsim.Spin
-):
+@pytest.mark.parametrize("norb, nelec", ffsim.testing.generate_norb_nelec(range(4)))
+def test_apply_orbital_rotation_random(norb: int, nelec: tuple[int, int]):
     """Test applying random orbital rotation yields correct output state."""
     rng = np.random.default_rng()
     dim = ffsim.dim(norb, nelec)
     for _ in range(3):
-        mat = ffsim.random.random_unitary(norb, seed=rng)
+        mat_a = ffsim.random.random_unitary(norb, seed=rng)
+        mat_b = ffsim.random.random_unitary(norb, seed=rng)
         vec = ffsim.random.random_statevector(dim, seed=rng)
 
+        if norb:
+            gen_a = _orbital_rotation_generator(scipy.linalg.logm(mat_a), spin=False)
+            gen_b = _orbital_rotation_generator(scipy.linalg.logm(mat_b), spin=True)
+            op_a = ffsim.linear_operator(gen_a, norb=norb, nelec=nelec)
+            op_b = ffsim.linear_operator(gen_b, norb=norb, nelec=nelec)
+            op_ab = ffsim.linear_operator(gen_a + gen_b, norb=norb, nelec=nelec)
+
+        # (mat_a, mat_b)
+        expected = (
+            scipy.sparse.linalg.expm_multiply(op_ab, vec, traceA=1) if norb else vec
+        )
+        result = ffsim.apply_orbital_rotation(vec, (mat_a, mat_b), norb, nelec)
+        np.testing.assert_allclose(result, expected)
         result = ffsim.apply_orbital_rotation(
-            vec, pair_for_spin(mat, spin), norb, nelec
+            vec, np.stack((mat_a, mat_b)), norb, nelec
         )
-        if norb:
-            generator = _orbital_rotation_generator(scipy.linalg.logm(mat), spin)
-            op = ffsim.linear_operator(generator, norb=norb, nelec=nelec)
-            expected = scipy.sparse.linalg.expm_multiply(op, vec, traceA=1)
-        else:
-            expected = vec
+        np.testing.assert_allclose(result, expected)
+
+        # (mat_a, None)
+        expected = (
+            scipy.sparse.linalg.expm_multiply(op_a, vec, traceA=1) if norb else vec
+        )
+        result = ffsim.apply_orbital_rotation(vec, (mat_a, None), norb, nelec)
+        np.testing.assert_allclose(result, expected)
+
+        # (None, mat_b)
+        expected = (
+            scipy.sparse.linalg.expm_multiply(op_b, vec, traceA=1) if norb else vec
+        )
+        result = ffsim.apply_orbital_rotation(vec, (None, mat_b), norb, nelec)
+        np.testing.assert_allclose(result, expected)
 
+        # (None, None)
+        expected = vec
+        result = ffsim.apply_orbital_rotation(vec, (None, None), norb, nelec)
         np.testing.assert_allclose(result, expected)
 
 
diff --git a/tests/python/qiskit/diag_coulomb_test.py b/tests/python/qiskit/diag_coulomb_test.py
index 44d826e7..3b82830d 100644
--- a/tests/python/qiskit/diag_coulomb_test.py
+++ b/tests/python/qiskit/diag_coulomb_test.py
@@ -114,6 +114,31 @@ def test_random_diag_coulomb_mat(
         )
         np.testing.assert_allclose(result, expected)
 
+        # Numpy array input
+        gate = ffsim.qiskit.DiagCoulombEvolutionJW(
+            norb,
+            np.stack((mat_aa, mat_ab, mat_bb)),
+            time,
+            z_representation=z_representation,
+        )
+        small_vec = ffsim.random.random_statevector(dim, seed=rng)
+        big_vec = ffsim.qiskit.ffsim_vec_to_qiskit_vec(
+            small_vec, norb=norb, nelec=nelec
+        )
+        statevec = Statevector(big_vec).evolve(gate)
+        result = ffsim.qiskit.qiskit_vec_to_ffsim_vec(
+            np.array(statevec), norb=norb, nelec=nelec
+        )
+        expected = ffsim.apply_diag_coulomb_evolution(
+            small_vec,
+            (mat_aa, mat_ab, mat_bb),
+            time,
+            norb=norb,
+            nelec=nelec,
+            z_representation=z_representation,
+        )
+        np.testing.assert_allclose(result, expected)
+
 
 @pytest.mark.parametrize(
     "norb, nelec, z_representation",
diff --git a/tests/python/qiskit/orbital_rotation_test.py b/tests/python/qiskit/orbital_rotation_test.py
index 508f8d6e..97dcf8e8 100644
--- a/tests/python/qiskit/orbital_rotation_test.py
+++ b/tests/python/qiskit/orbital_rotation_test.py
@@ -112,6 +112,21 @@ def test_random_orbital_rotation_asymmetric_spin(norb: int, nelec: tuple[int, in
         )
         np.testing.assert_allclose(result, expected)
 
+        # Numpy array input
+        gate = ffsim.qiskit.OrbitalRotationJW(norb, np.stack([mat_a, mat_b]))
+        small_vec = ffsim.random.random_statevector(dim, seed=rng)
+        big_vec = ffsim.qiskit.ffsim_vec_to_qiskit_vec(
+            small_vec, norb=norb, nelec=nelec
+        )
+        statevec = Statevector(big_vec).evolve(gate)
+        result = ffsim.qiskit.qiskit_vec_to_ffsim_vec(
+            np.array(statevec), norb=norb, nelec=nelec
+        )
+        expected = ffsim.apply_orbital_rotation(
+            small_vec, (mat_a, mat_b), norb=norb, nelec=nelec
+        )
+        np.testing.assert_allclose(result, expected)
+
 
 @pytest.mark.parametrize("norb, nocc", ffsim.testing.generate_norb_nocc(range(5)))
 def test_random_orbital_rotation_spinless(norb: int, nocc: int):