From 8695fee76968a5b32b5779f3b90b9abfdbb6ccfc Mon Sep 17 00:00:00 2001
From: "Kevin J. Sung" <kevjsung@umich.edu>
Date: Mon, 30 Oct 2023 14:56:12 -0400
Subject: [PATCH] support spatial one-body in wick

---
 python/ffsim/states/wick.py | 11 ++++++++--
 tests/states/wick_test.py   | 43 +++++++++++++++++++++----------------
 2 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/python/ffsim/states/wick.py b/python/ffsim/states/wick.py
index 1742eeeda..9867bf0d5 100644
--- a/python/ffsim/states/wick.py
+++ b/python/ffsim/states/wick.py
@@ -15,6 +15,7 @@
 from collections.abc import Sequence
 
 import numpy as np
+import scipy.linalg
 
 
 def expectation_one_body_product(
@@ -48,8 +49,14 @@ def expectation_one_body_product(
     if not n_tensors:
         return 1.0
 
-    norb, _ = one_rdm.shape
-    anti_one_rdm = np.eye(norb) - one_rdm
+    norb, _ = one_body_tensors[0].shape
+    dim, _ = one_rdm.shape
+    if dim == 2 * norb:
+        one_body_tensors = [
+            scipy.linalg.block_diag(mat, mat) for mat in one_body_tensors
+        ]
+
+    anti_one_rdm = np.eye(dim) - one_rdm
 
     alphabet = string.ascii_uppercase + string.ascii_lowercase
     indices = alphabet[: 2 * n_tensors]
diff --git a/tests/states/wick_test.py b/tests/states/wick_test.py
index 752684ebf..9c6b13dbe 100644
--- a/tests/states/wick_test.py
+++ b/tests/states/wick_test.py
@@ -11,6 +11,7 @@
 """Test Wick's theorem utilities."""
 
 import numpy as np
+import pytest
 import scipy.linalg
 import scipy.sparse.linalg
 
@@ -18,12 +19,20 @@
 from ffsim.states.wick import expectation_one_body_power, expectation_one_body_product
 
 
-def test_expectation_product():
+@pytest.mark.parametrize(
+    "norb, occupied_orbitals",
+    [
+        (4, ([0, 1], [0, 1])),
+        (3, ([0], [1, 2])),
+        (2, ([], [0])),
+    ],
+)
+def test_expectation_product(norb: int, occupied_orbitals: tuple[list[int], list[int]]):
     """Test expectation product."""
-    norb = 5
-    nelec = (3, 1)
-    n_alpha, n_beta = nelec
+    occ_a, occ_b = occupied_orbitals
+    nelec = len(occ_a), len(occ_b)
     dim = ffsim.dim(norb, nelec)
+
     rng = np.random.default_rng()
 
     # generate random one-body tensors
@@ -38,28 +47,26 @@ def test_expectation_product():
         linops.append(linop)
 
     # generate a random Slater determinant
-    vecs = ffsim.random.random_unitary(norb, seed=rng)
-    occupied_orbitals_a = vecs[:, :n_alpha]
-    occupied_orbitals_b = vecs[:, :n_beta]
-    one_rdm_a = occupied_orbitals_a.conj() @ occupied_orbitals_a.T
-    one_rdm_b = occupied_orbitals_b.conj() @ occupied_orbitals_b.T
-    one_rdm = scipy.linalg.block_diag(one_rdm_a, one_rdm_b)
+    orbital_rotation = ffsim.random.random_unitary(norb, seed=rng)
+    # TODO test spin-summed
+    one_rdm = ffsim.slater_determinant_rdm(
+        norb,
+        occupied_orbitals,
+        orbital_rotation=orbital_rotation,
+        spin_summed=False,
+    )
 
     # get the full statevector
-    state = ffsim.slater_determinant(norb, (range(n_alpha), range(n_beta)))
-    state = ffsim.apply_orbital_rotation(state, vecs, norb=norb, nelec=nelec)
+    state = ffsim.slater_determinant(
+        norb, occupied_orbitals, orbital_rotation=orbital_rotation
+    )
 
     product_op = scipy.sparse.linalg.LinearOperator(
         shape=(dim, dim), matvec=lambda x: x
     )
-    expanded_one_body_tensors = [
-        scipy.linalg.block_diag(mat, mat) for mat in one_body_tensors
-    ]
     for i in range(n_tensors):
         product_op = product_op @ linops[i]
-        computed = expectation_one_body_product(
-            one_rdm, expanded_one_body_tensors[: i + 1]
-        )
+        computed = expectation_one_body_product(one_rdm, one_body_tensors[: i + 1])
         target = np.vdot(state, product_op @ state)
         np.testing.assert_allclose(computed, target, atol=1e-8)