Implement trace of FermionOperator

python/ffsim/protocols/trace_protocol.py

@@ -1,4 +1,4 @@
-# (C) Copyright IBM 2023.
+# (C) Copyright IBM 2025.
 #
 # This code is licensed under the Apache License, Version 2.0. You may
 # obtain a copy of this license in the LICENSE.txt file in the root directory
@@ -12,10 +12,13 @@
 
 from __future__ import annotations
 
+import math
 from typing import Any, Protocol
 
 import numpy as np
 
+from ffsim.operators import FermionOperator
+
 
 class SupportsTrace(Protocol):
     """A linear operator whose trace can be computed."""
@@ -32,8 +35,11 @@ def _trace_(self, norb: int, nelec: tuple[int, int]) -> float:
         """
 
 
-def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> float:
+def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> complex:
     """Return the trace of the linear operator."""
+    if isinstance(obj, FermionOperator):
+        return _trace_fermion_operator(obj, norb, nelec)
+
     method = getattr(obj, "_trace_", None)
     if method is not None:
         return method(norb=norb, nelec=nelec)
@@ -45,3 +51,85 @@ def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> float:
         "The object did not have a _trace_ method that returned the trace "
         "or a _diag_ method that returned its diagonal entries."
     )
+
+
+def _term_phase(op: tuple[tuple[bool, bool, int], ...]) -> int:
+    phase = 0
+    while len(op) > 0:
+        action, spin, orb = op[0]
+        conj_idx = op.index((not action, spin, orb))
+        phase += conj_idx - 1
+        op = op[1:conj_idx] + op[conj_idx + 1 :]
+    return (-1) ** phase
+
+
+def _trace_term(
+    op: tuple[tuple[bool, bool, int], ...], norb: int, nelec: tuple[int, int]
+) -> complex:
+    n_alpha, n_beta = nelec
+
+    spin_orbs = set((spin, orb) for _, spin, orb in op)
+    norb_alpha = sum(not spin for spin, _ in spin_orbs)
+    norb_beta = len(spin_orbs) - norb_alpha
+
+    nelec_alpha = 0
+    nelec_beta = 0
+
+    # loop over the support of the operator
+    # assume that each site is either 0 or 1 at the beginning
+    # track the state of the site through the application of the operator
+    # if the state exceed 1 or goes below 0,
+    # the state is not physical and the trace must be 0
+    for this_spin, this_orb in spin_orbs:
+        initial_zero = 0
+        initial_one = 1
+        is_zero = True
+        is_one = True
+        for action, spin, orb in reversed(op):
+            if (spin, orb) != (this_spin, this_orb):
+                continue
+
+            change = action * 2 - 1
+            initial_zero += change
+            initial_one += change
+            if initial_zero < 0 or initial_zero > 1:
+                is_zero = False
+            if initial_one < 0 or initial_one > 1:
+                is_one = False
+
+        # if the operator has support on this_orb,
+        # either the initial state is 0 or 1, but not both
+        assert not is_zero or not is_one
+
+        # return 0 if there is no possible initial state
+        if not is_zero and not is_one:
+            return 0j
+        # the state must return to the initial state, otherwise the trace is zero
+        if (is_zero and initial_zero != 0) or (is_one and initial_one != 1):
+            return 0j
+        # count the number of electrons
+        if is_one:
+            if not this_spin:
+                nelec_alpha += 1
+            else:
+                nelec_beta += 1
+        if nelec_alpha > n_alpha or nelec_beta > n_beta:
+            # the number of electrons exceeds the number of allowed electrons
+            return 0j
+
+    # the trace is nontrival and is a product of
+    # binom(#orbs not in the support of op, #elec on these orbs)
+    # for each spin species
+    # and a phase that depends on the ordering between the actions
+
+    return (
+        _term_phase(op)
+        * math.comb(norb - norb_alpha, n_alpha - nelec_alpha)
+        * math.comb(norb - norb_beta, n_beta - nelec_beta)
+    )
+
+
+def _trace_fermion_operator(
+    ferm: FermionOperator, norb: int, nelec: tuple[int, int]
+) -> complex:
+    return sum(coeff * _trace_term(op, norb, nelec) for op, coeff in ferm.items())

tests/python/operators/fermion_operator_test.py

@@ -643,3 +643,22 @@ def test_mapping_methods():
         ((ffsim.cre_b(1), ffsim.des_b(2)), -0.5j),
         ((ffsim.cre_b(2), ffsim.des_b(1)), 1 - 0.5j),
     }
+
+def test_trace():
+    rng = np.random.default_rng(12345)
+    # compare for random_diagonal_coulomb_hamiltonian
+    norb = 50
+    nelec = (3, 3)
+    ham = ffsim.random.random_diagonal_coulomb_hamiltonian(norb, real=True, seed=rng)
+    t1 = ffsim.trace(ffsim.fermion_operator(ham), norb=norb, nelec=nelec)
+    t2 = ffsim.trace(ham, norb=norb, nelec=nelec)
+    np.testing.assert_allclose(t1, t2)
+
+    # compare for random_molecular_hamiltonian
+    norb = 20
+    nelec = (3, 3)
+    ham = ffsim.random.random_molecular_hamiltonian(norb, seed=rng, dtype=float)
+    t1 = ffsim.trace(ffsim.fermion_operator(ham), norb=norb, nelec=nelec)
+    t2 = ffsim.trace(ham, norb=norb, nelec=nelec)
+
+