From 9ae946ee765e4460507dda0431d564f18dc36aa7 Mon Sep 17 00:00:00 2001
From: "Kevin J. Sung" <kevjsung@umich.edu>
Date: Wed, 20 Sep 2023 19:49:10 -0400
Subject: [PATCH] test complex one-body tensor for DF Hamiltonian

---
 .../double_factorized_hamiltonian_test.py     |  3 +-
 tests/trotter/qdrift_test.py                  | 38 +++++--------------
 tests/trotter/trotter_test.py                 |  3 +-
 3 files changed, 13 insertions(+), 31 deletions(-)

diff --git a/tests/hamiltonians/double_factorized_hamiltonian_test.py b/tests/hamiltonians/double_factorized_hamiltonian_test.py
index e9797dbdd..f2458a8e8 100644
--- a/tests/hamiltonians/double_factorized_hamiltonian_test.py
+++ b/tests/hamiltonians/double_factorized_hamiltonian_test.py
@@ -27,8 +27,7 @@ def test_double_factorized_hamiltonian(z_representation: bool):
 
     # generate random Hamiltonian
     dim = ffsim.dim(norb, nelec)
-    # TODO test with complex one-body tensor
-    one_body_tensor = np.real(ffsim.random.random_hermitian(norb, seed=2474))
+    one_body_tensor = ffsim.random.random_hermitian(norb, seed=2474)
     two_body_tensor = ffsim.random.random_two_body_tensor_real(norb, seed=7054)
     mol_hamiltonian = ffsim.MolecularHamiltonian(one_body_tensor, two_body_tensor)
     hamiltonian = ffsim.linear_operator(
diff --git a/tests/trotter/qdrift_test.py b/tests/trotter/qdrift_test.py
index c7e492b86..9b6810f01 100644
--- a/tests/trotter/qdrift_test.py
+++ b/tests/trotter/qdrift_test.py
@@ -19,7 +19,6 @@
 from pyscf import ao2mo, gto, mcscf, scf
 
 import ffsim
-from ffsim.contract.hamiltonian import hamiltonian_trace
 from ffsim.trotter.qdrift import (
     one_body_square_decomposition,
     spectral_norm_diag_coulomb,
@@ -264,16 +263,12 @@ def test_simulate_qdrift_double_factorized_h_chain(
     two_body_tensor = ao2mo.restore(1, mc.get_h2cas(), mc.ncas)
     norb, _ = one_body_tensor.shape
     nelec = mol.nelec
-    hamiltonian = ffsim.contract.hamiltonian_linop(
-        one_body_tensor=one_body_tensor,
-        two_body_tensor=two_body_tensor,
-        norb=norb,
-        nelec=nelec,
-    )
+    mol_hamiltonian = ffsim.MolecularHamiltonian(one_body_tensor, two_body_tensor)
+    hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)
 
     # perform double factorization
     df_hamiltonian = ffsim.double_factorized_hamiltonian(
-        ffsim.MolecularHamiltonian(one_body_tensor, two_body_tensor),
+        mol_hamiltonian,
         z_representation=z_representation,
     )
 
@@ -288,12 +283,7 @@ def test_simulate_qdrift_double_factorized_h_chain(
     exact_state = scipy.sparse.linalg.expm_multiply(
         -1j * time * hamiltonian,
         initial_state,
-        traceA=hamiltonian_trace(
-            norb=norb,
-            nelec=nelec,
-            one_body_tensor=one_body_tensor,
-            two_body_tensor=two_body_tensor,
-        ),
+        traceA=ffsim.trace(mol_hamiltonian, norb=norb, nelec=nelec),
     )
 
     # make sure time is not too small
@@ -364,21 +354,18 @@ def test_simulate_qdrift_double_factorized_random(
 ):
     rng = np.random.default_rng(2030)
     # generate random Hamiltonian
-    # TODO test with complex one-body tensor after fixing get_hamiltonian_linop
+    # 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
     )
-    hamiltonian = ffsim.contract.hamiltonian_linop(
-        one_body_tensor=one_body_tensor,
-        two_body_tensor=two_body_tensor,
-        norb=norb,
-        nelec=nelec,
-    )
+    mol_hamiltonian = ffsim.MolecularHamiltonian(one_body_tensor, two_body_tensor)
+    hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)
 
     # perform double factorization
     df_hamiltonian = ffsim.double_factorized_hamiltonian(
-        ffsim.MolecularHamiltonian(one_body_tensor, two_body_tensor),
+        mol_hamiltonian,
         optimize=optimize,
         z_representation=z_representation,
     )
@@ -394,12 +381,7 @@ def test_simulate_qdrift_double_factorized_random(
     exact_state = scipy.sparse.linalg.expm_multiply(
         -1j * time * hamiltonian,
         initial_state,
-        traceA=hamiltonian_trace(
-            norb=norb,
-            nelec=nelec,
-            one_body_tensor=one_body_tensor,
-            two_body_tensor=two_body_tensor,
-        ),
+        traceA=ffsim.trace(mol_hamiltonian, norb=norb, nelec=nelec),
     )
 
     # make sure time is not too small
diff --git a/tests/trotter/trotter_test.py b/tests/trotter/trotter_test.py
index bd66b03f7..81d1a5a5d 100644
--- a/tests/trotter/trotter_test.py
+++ b/tests/trotter/trotter_test.py
@@ -40,7 +40,8 @@ def test_simulate_trotter_double_factorized_random(
 ):
     # generate random Hamiltonian
     dim = ffsim.dim(norb, nelec)
-    # TODO test with complex one-body tensor after fixing get_hamiltonian_linop
+    # 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)
     mol_hamiltonian = ffsim.MolecularHamiltonian(one_body_tensor, two_body_tensor)