add support for complex t2 amplitudes

python/ffsim/linalg/double_factorized_decomposition.py

@@ -469,6 +469,8 @@ def double_factorized_t2(
     too small, then the error of the decomposition may exceed the specified
     error threshold.
 
+    Note: Currently, only real-valued t2 amplitudes are supported.
+
     Args:
         t2_amplitudes: The t2 amplitudes tensor.
         tol: Tolerance for error in the decomposition.

python/ffsim/random/random.py

@@ -199,25 +199,41 @@ def random_two_body_tensor(
     return two_body_tensor
 
 
-def random_t2_amplitudes(norb: int, nocc: int, *, seed=None) -> np.ndarray:
+def random_t2_amplitudes(
+    norb: int, nocc: int, *, seed=None, dtype=complex
+) -> np.ndarray:
     """Sample a random t2 amplitudes tensor.
 
     Args:
         norb: The number of orbitals.
         nocc: The number of orbitals that are occupied by an electron.
         seed: A seed to initialize the pseudorandom number generator.
             Should be a valid input to ``np.random.default_rng``.
+        dype: The data type to use for the result.
 
     Returns:
         The sampled t2 amplitudes tensor.
     """
     rng = np.random.default_rng(seed)
     nvrt = norb - nocc
-    t2 = np.zeros((nocc, nocc, nvrt, nvrt))
+    t2 = np.zeros((nocc, nocc, nvrt, nvrt), dtype=dtype)
     pairs = itertools.product(range(nocc), range(nocc, norb))
     for (m, (i, a)), (n, (j, b)) in itertools.product(enumerate(pairs), repeat=2):
         if m <= n and i <= j and a <= b:
             val = rng.standard_normal()
             t2[i, j, a - nocc, b - nocc] = val
             t2[j, i, b - nocc, a - nocc] = val
+    if np.issubdtype(dtype, np.complexfloating):
+        t2_large = np.zeros((norb, norb, norb, norb), dtype=dtype)
+        t2_large[:nocc, :nocc, nocc:, nocc:] = t2
+        orbital_rotation = random_unitary(norb, seed=rng)
+        t2_large = np.einsum(
+            "ijab,iI,jJ,aA,bB->IJAB",
+            t2_large,
+            orbital_rotation.conj(),
+            orbital_rotation.conj(),
+            orbital_rotation,
+            orbital_rotation,
+        )
+        t2 = t2_large[:nocc, :nocc, nocc:, nocc:]
     return t2

tests/linalg/double_factorized_decomposition_test.py

@@ -272,7 +272,7 @@ def test_double_factorized_compressed_constrained():
 @pytest.mark.parametrize("norb, nocc", [(4, 2), (5, 2), (5, 3)])
 def test_double_factorized_t2_amplitudes_random(norb: int, nocc: int):
     """Test double factorization of random t2 amplitudes."""
-    t2 = random_t2_amplitudes(norb, nocc)
+    t2 = random_t2_amplitudes(norb, nocc, dtype=float)
     diag_coulomb_mats, orbital_rotations = double_factorized_t2(t2)
     reconstructed = _reconstruct_t2(diag_coulomb_mats, orbital_rotations, nocc=nocc)
     np.testing.assert_allclose(reconstructed, t2, atol=1e-8)

tests/variational/ucj_test.py

@@ -72,7 +72,7 @@ def test_t_amplitudes_roundtrip():
 
     rng = np.random.default_rng()
 
-    t2 = ffsim.random.random_t2_amplitudes(norb, nocc)
+    t2 = ffsim.random.random_t2_amplitudes(norb, nocc, dtype=float)
     t1 = rng.standard_normal((nocc, norb - nocc))
 
     operator = ffsim.UCJOperator.from_t_amplitudes(t2, t1_amplitudes=t1)
@@ -176,7 +176,7 @@ def test_real_ucj_t_amplitudes_roundtrip():
 
     rng = np.random.default_rng()
 
-    t2 = ffsim.random.random_t2_amplitudes(norb, nocc)
+    t2 = ffsim.random.random_t2_amplitudes(norb, nocc, dtype=float)
     t1 = rng.standard_normal((nocc, norb - nocc))
 
     operator = ffsim.RealUCJOperator.from_t_amplitudes(t2, t1_amplitudes=t1)