From 2b00da0e952d0168054793acf626ba4d8c358b40 Mon Sep 17 00:00:00 2001 From: "Kevin J. Sung" Date: Thu, 30 Nov 2023 21:20:20 -0500 Subject: [PATCH] add support for complex t2 amplitudes --- .../linalg/double_factorized_decomposition.py | 2 ++ python/ffsim/random/random.py | 20 +++++++++++++++++-- .../double_factorized_decomposition_test.py | 2 +- tests/variational/ucj_test.py | 4 ++-- 4 files changed, 23 insertions(+), 5 deletions(-) diff --git a/python/ffsim/linalg/double_factorized_decomposition.py b/python/ffsim/linalg/double_factorized_decomposition.py index 4dde5bb5b..6f54c3efe 100644 --- a/python/ffsim/linalg/double_factorized_decomposition.py +++ b/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. diff --git a/python/ffsim/random/random.py b/python/ffsim/random/random.py index 102732008..646347833 100644 --- a/python/ffsim/random/random.py +++ b/python/ffsim/random/random.py @@ -199,7 +199,9 @@ 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: @@ -207,17 +209,31 @@ def random_t2_amplitudes(norb: int, nocc: int, *, seed=None) -> np.ndarray: 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 diff --git a/tests/linalg/double_factorized_decomposition_test.py b/tests/linalg/double_factorized_decomposition_test.py index f67150490..5f3ddc050 100644 --- a/tests/linalg/double_factorized_decomposition_test.py +++ b/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) diff --git a/tests/variational/ucj_test.py b/tests/variational/ucj_test.py index 8d41855c0..f279db289 100644 --- a/tests/variational/ucj_test.py +++ b/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)