add functions and tests for random two-body tensors

python/ffsim/random/__init__.py

@@ -18,7 +18,7 @@
     random_special_orthogonal,
     random_statevector,
     random_t2_amplitudes,
-    random_two_body_tensor_real,
+    random_two_body_tensor,
     random_unitary,
 )
 
@@ -30,6 +30,6 @@
     "random_special_orthogonal",
     "random_statevector",
     "random_t2_amplitudes",
-    "random_two_body_tensor_real",
+    "random_two_body_tensor",
     "random_unitary",
 ]

python/ffsim/random/random.py

@@ -10,6 +10,8 @@
 
 from __future__ import annotations
 
+import itertools
+
 import numpy as np
 
 
@@ -160,10 +162,10 @@ def random_antihermitian(dim: int, *, seed=None, dtype=complex) -> np.ndarray:
     return mat - mat.T.conj()
 
 
-def random_two_body_tensor_real(
-    dim: int, *, rank: int | None = None, seed=None, dtype=float
+def random_two_body_tensor(
+    dim: int, *, rank: int | None = None, seed=None, dtype=complex
 ) -> np.ndarray:
-    """Sample a random two-body tensor with real-valued orbitals.
+    """Sample a random two-body tensor.
 
     Args:
         dim: The dimension of the tensor. The shape of the returned tensor will be
@@ -182,7 +184,19 @@ def random_two_body_tensor_real(
         rank = dim * (dim + 1) // 2
     cholesky_vecs = rng.standard_normal((rank, dim, dim)).astype(dtype, copy=False)
     cholesky_vecs += cholesky_vecs.transpose((0, 2, 1))
-    return np.einsum("ipr,iqs->prqs", cholesky_vecs, cholesky_vecs)
+    two_body_tensor = np.einsum("ipr,iqs->prqs", cholesky_vecs, cholesky_vecs)
+    if np.issubdtype(dtype, np.complexfloating):
+        orbital_rotation = random_unitary(dim, seed=rng)
+        two_body_tensor = np.einsum(
+            "abcd,aA,bB,cC,dD->ABCD",
+            two_body_tensor,
+            orbital_rotation,
+            orbital_rotation.conj(),
+            orbital_rotation,
+            orbital_rotation.conj(),
+            optimize=True,
+        )
+    return two_body_tensor
 
 
 def random_t2_amplitudes(norb: int, nocc: int, *, seed=None) -> np.ndarray:
@@ -199,7 +213,11 @@ def random_t2_amplitudes(norb: int, nocc: int, *, seed=None) -> np.ndarray:
     """
     rng = np.random.default_rng(seed)
     nvrt = norb - nocc
-    t2 = rng.standard_normal((nocc, nocc, nvrt, nvrt))
-    t2 -= np.transpose(t2, (0, 1, 3, 2))
-    t2 -= np.transpose(t2, (1, 0, 2, 3))
+    t2 = np.zeros((nocc, nocc, nvrt, nvrt))
+    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
     return t2

tests/hamiltonians/double_factorized_hamiltonian_test.py

@@ -28,7 +28,7 @@ def test_linear_operator(z_representation: bool):
 
     dim = ffsim.dim(norb, nelec)
     one_body_tensor = ffsim.random.random_hermitian(norb, seed=rng)
-    two_body_tensor = ffsim.random.random_two_body_tensor_real(norb, seed=rng)
+    two_body_tensor = ffsim.random.random_two_body_tensor(norb, seed=rng, dtype=float)
     constant = rng.standard_normal()
     mol_hamiltonian = ffsim.MolecularHamiltonian(
         one_body_tensor, two_body_tensor, constant
@@ -55,7 +55,7 @@ def test_z_representation_round_trip():
     norb = 4
 
     one_body_tensor = ffsim.random.random_hermitian(norb, seed=2474)
-    two_body_tensor = ffsim.random.random_two_body_tensor_real(norb, seed=7054)
+    two_body_tensor = ffsim.random.random_two_body_tensor(norb, seed=7054, dtype=float)
 
     df_hamiltonian = ffsim.DoubleFactorizedHamiltonian.from_molecular_hamiltonian(
         ffsim.MolecularHamiltonian(one_body_tensor, two_body_tensor)

tests/hamiltonians/molecular_hamiltonian_test.py

@@ -90,7 +90,8 @@ def test_fermion_operator(norb: int, nelec: tuple[int, int]):
     rng = np.random.default_rng()
 
     one_body_tensor = ffsim.random.random_hermitian(norb, seed=rng)
-    two_body_tensor = ffsim.random.random_two_body_tensor_real(norb, seed=rng)
+    # TODO remove dtype=float after adding support for complex
+    two_body_tensor = ffsim.random.random_two_body_tensor(norb, seed=rng, dtype=float)
     constant = rng.standard_normal()
     mol_hamiltonian = ffsim.MolecularHamiltonian(
         one_body_tensor, two_body_tensor, constant=constant

tests/hamiltonians/single_factorized_hamiltonian_test.py

@@ -35,7 +35,7 @@ def test_linear_operator(norb: int, nelec: tuple[int, int], cholesky: bool):
 
     dim = ffsim.dim(norb, nelec)
     one_body_tensor = ffsim.random.random_hermitian(norb, seed=rng)
-    two_body_tensor = ffsim.random.random_two_body_tensor_real(norb, seed=rng)
+    two_body_tensor = ffsim.random.random_two_body_tensor(norb, seed=rng, dtype=float)
     constant = rng.standard_normal()
     mol_hamiltonian = ffsim.MolecularHamiltonian(
         one_body_tensor, two_body_tensor, constant

tests/linalg/double_factorized_decomposition_test.py

@@ -16,6 +16,7 @@
 import pytest
 from pyscf import ao2mo, cc, gto, mcscf, scf
 
+import ffsim
 from ffsim.linalg import (
     double_factorized,
     modified_cholesky,
@@ -24,11 +25,7 @@
     double_factorized_t2,
     optimal_diag_coulomb_mats,
 )
-from ffsim.random import (
-    random_t2_amplitudes,
-    random_two_body_tensor_real,
-    random_unitary,
-)
+from ffsim.random import random_t2_amplitudes, random_unitary
 
 
 def _reconstruct_t2(
@@ -70,7 +67,7 @@ def test_modified_cholesky(dim: int):
 )
 def test_double_factorized_random(dim: int, cholesky: bool):
     """Test double-factorized decomposition on a random tensor."""
-    two_body_tensor = random_two_body_tensor_real(dim, seed=9825)
+    two_body_tensor = ffsim.random.random_two_body_tensor(dim, seed=9825, dtype=float)
     diag_coulomb_mats, orbital_rotations = double_factorized(
         two_body_tensor, cholesky=cholesky
     )
@@ -149,7 +146,7 @@ def test_double_factorized_tol_max_vecs(cholesky: bool):
 def test_optimal_diag_coulomb_mats_exact():
     """Test optimal diag Coulomb matrices on exact decomposition."""
     dim = 5
-    two_body_tensor = random_two_body_tensor_real(dim, seed=8386)
+    two_body_tensor = ffsim.random.random_two_body_tensor(dim, seed=8386, dtype=float)
 
     _, orbital_rotations = double_factorized(two_body_tensor)
     diag_coulomb_mats_optimal = optimal_diag_coulomb_mats(
@@ -169,7 +166,7 @@ def test_optimal_diag_coulomb_mats_exact():
 def test_optimal_diag_coulomb_mats_approximate():
     """Test optimal diag Coulomb matrices on approximate decomposition."""
     dim = 5
-    two_body_tensor = random_two_body_tensor_real(dim, seed=3718)
+    two_body_tensor = ffsim.random.random_two_body_tensor(dim, seed=3718, dtype=float)
 
     diag_coulomb_mats, orbital_rotations = double_factorized(
         two_body_tensor, max_vecs=3
@@ -202,7 +199,7 @@ def test_double_factorized_compressed():
     """Test compressed double factorization"""
     # TODO test on simple molecule like ethylene
     dim = 2
-    two_body_tensor = random_two_body_tensor_real(dim, seed=8364)
+    two_body_tensor = ffsim.random.random_two_body_tensor(dim, seed=8364, dtype=float)
 
     diag_coulomb_mats, orbital_rotations = double_factorized(
         two_body_tensor, max_vecs=2
@@ -235,7 +232,7 @@ def test_double_factorized_compressed_constrained():
     """Test constrained compressed double factorization"""
     # TODO test on simple molecule like ethylene
     dim = 3
-    two_body_tensor = random_two_body_tensor_real(dim, seed=2927)
+    two_body_tensor = ffsim.random.random_two_body_tensor(dim, seed=2927, dtype=float)
 
     diag_coulomb_mats, orbital_rotations = double_factorized(
         two_body_tensor, max_vecs=2

tests/random_test.py

@@ -15,37 +15,88 @@
 import itertools
 
 import numpy as np
+from pyscf import cc, gto, scf
 
 import ffsim
 
 
-def test_random_two_body_tensor_symmetry():
-    """Test random two-body tensor symmetry."""
-    n_orbitals = 5
-    two_body_tensor = ffsim.random.random_two_body_tensor_real(n_orbitals)
-    for p, q, r, s in itertools.product(range(n_orbitals), repeat=4):
-        val = two_body_tensor[p, q, r, s]
-        np.testing.assert_allclose(two_body_tensor[r, s, p, q], val)
-        np.testing.assert_allclose(two_body_tensor[q, p, s, r], val.conjugate())
-        np.testing.assert_allclose(two_body_tensor[s, r, q, p], val.conjugate())
-        np.testing.assert_allclose(two_body_tensor[q, p, r, s], val)
-        np.testing.assert_allclose(two_body_tensor[s, r, p, q], val)
-        np.testing.assert_allclose(two_body_tensor[p, q, s, r], val)
-        np.testing.assert_allclose(two_body_tensor[r, s, q, p], val)
+def assert_t2_has_correct_symmetry(t2: np.ndarray):
+    nocc, _, nvrt, _ = t2.shape
+    norb = nocc + nvrt
+    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:
+            np.testing.assert_allclose(
+                t2[i, j, a - nocc, b - nocc], t2[j, i, b - nocc, a - nocc]
+            )
+
+
+def test_assert_t2_has_correct_symmetry():
+    """Test that t2 amplitudes from a real molecule passes our symmetry test."""
+    # Build a stretched ethene molecule
+    bond_distance = 2.678
+    a = 0.5 * bond_distance
+    b = a + 0.5626
+    c = 0.9289
+    mol = gto.Mole()
+    mol.build(
+        atom=[
+            ["C", (0, 0, a)],
+            ["C", (0, 0, -a)],
+            ["H", (0, c, b)],
+            ["H", (0, -c, b)],
+            ["H", (0, c, -b)],
+            ["H", (0, -c, -b)],
+        ],
+        basis="sto-6g",
+        symmetry="d2h",
+    )
+    hartree_fock = scf.RHF(mol)
+    hartree_fock.kernel()
+    # Define active space
+    active_space = range(mol.nelectron // 2 - 2, mol.nelectron // 2 + 2)
+    # Get CCSD t2 amplitudes for initializing the ansatz
+    ccsd = cc.CCSD(
+        hartree_fock,
+        frozen=[i for i in range(mol.nao_nr()) if i not in active_space],
+    )
+    _, _, t2 = ccsd.kernel()
+    assert_t2_has_correct_symmetry(t2)
 
 
 def test_random_t2_amplitudes_symmetry():
     """Test random t2 amplitudes symmetry."""
-    # TODO t2 amplitudes from pySCF don't satisfy these symmetries
     norb = 5
     nocc = 3
     t2 = ffsim.random.random_t2_amplitudes(norb, nocc)
-    for i, j in itertools.product(range(nocc), repeat=2):
-        for a, b in itertools.product(range(norb - nocc), repeat=2):
-            val = t2[i, j, a, b]
-            np.testing.assert_allclose(t2[i, j, b, a], -val)
-            np.testing.assert_allclose(t2[j, i, a, b], -val)
-            np.testing.assert_allclose(t2[j, i, b, a], val)
+    assert_t2_has_correct_symmetry(t2)
+
+
+def test_random_two_body_tensor_symmetry_real():
+    """Test random real two-body tensor symmetry."""
+    n_orbitals = 5
+    two_body_tensor = ffsim.random.random_two_body_tensor(n_orbitals, dtype=float)
+    assert np.issubdtype(two_body_tensor.dtype, np.floating)
+    for i, j, k, ell in itertools.product(range(n_orbitals), repeat=4):
+        val = two_body_tensor[i, j, k, ell]
+        np.testing.assert_allclose(two_body_tensor[k, ell, i, j], val)
+        np.testing.assert_allclose(two_body_tensor[j, i, ell, k], val.conjugate())
+        np.testing.assert_allclose(two_body_tensor[ell, k, j, i], val.conjugate())
+        np.testing.assert_allclose(two_body_tensor[j, i, k, ell], val)
+        np.testing.assert_allclose(two_body_tensor[ell, k, i, j], val)
+        np.testing.assert_allclose(two_body_tensor[i, j, ell, k], val)
+        np.testing.assert_allclose(two_body_tensor[k, ell, j, i], val)
+
+
+def test_random_two_body_tensor_symmetry():
+    """Test random two-body tensor symmetry."""
+    n_orbitals = 5
+    two_body_tensor = ffsim.random.random_two_body_tensor(n_orbitals)
+    for i, j, k, ell in itertools.product(range(n_orbitals), repeat=4):
+        val = two_body_tensor[i, j, k, ell]
+        np.testing.assert_allclose(two_body_tensor[k, ell, i, j], val)
+        np.testing.assert_allclose(two_body_tensor[j, i, ell, k], val.conjugate())
+        np.testing.assert_allclose(two_body_tensor[ell, k, j, i], val.conjugate())
 
 
 def test_random_unitary():

tests/trotter/qdrift_test.py

@@ -354,8 +354,8 @@ def test_simulate_qdrift_double_factorized_random(
     # TODO test with complex one-body tensor fails due to the following issue
     # https://github.com/qiskit-community/ffsim/issues/14
     one_body_tensor = ffsim.random.random_real_symmetric_matrix(norb, seed=rng)
-    two_body_tensor = ffsim.random.random_two_body_tensor_real(
-        norb, rank=norb, seed=rng
+    two_body_tensor = ffsim.random.random_two_body_tensor(
+        norb, rank=norb, seed=rng, dtype=float
     )
     mol_hamiltonian = ffsim.MolecularHamiltonian(one_body_tensor, two_body_tensor)
     hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)

tests/trotter/trotter_test.py

@@ -43,7 +43,7 @@ def test_simulate_trotter_double_factorized_random(
     # TODO test with complex one-body tensor fails due to the following issue
     # https://github.com/qiskit-community/ffsim/issues/14
     one_body_tensor = np.real(ffsim.random.random_hermitian(norb, seed=2474))
-    two_body_tensor = ffsim.random.random_two_body_tensor_real(norb, seed=7054)
+    two_body_tensor = ffsim.random.random_two_body_tensor(norb, seed=7054, dtype=float)
     mol_hamiltonian = ffsim.MolecularHamiltonian(one_body_tensor, two_body_tensor)
     hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)