From 3ef035b458ce05c999b65f82f8d16c20103b82ba Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Wed, 8 Jan 2025 15:26:40 -0700
Subject: [PATCH 01/16] Fixed issue 348: implement trace of FermionOperator

---
 python/ffsim/protocols/trace_protocol.py | 88 ++++++++++++++++++++++++
 tests/python/operators/trace_test.py     | 23 +++++++
 2 files changed, 111 insertions(+)
 create mode 100644 tests/python/operators/trace_test.py

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index 95d421dab..caf0aa3a9 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/python/ffsim/protocols/trace_protocol.py
@@ -14,8 +14,12 @@
 
 from typing import Any, Protocol
 
+from ffsim.operators import FermionOperator
+
 import numpy as np
 
+from math import comb
+
 
 class SupportsTrace(Protocol):
     """A linear operator whose trace can be computed."""
@@ -34,6 +38,9 @@ def _trace_(self, norb: int, nelec: tuple[int, int]) -> float:
 
 def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> float:
     """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 +52,84 @@ 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 _trace_fermion_operator(
+    A: FermionOperator, norb: int, nelec: tuple[int, int]
+):
+    n_elec_alpha, n_elec_beta = nelec
+    if len(A)>1:
+        trace = 0
+        for op in A:
+            B = FermionOperator({op:A[op]})
+            trace+= _trace_fermion_operator(B, norb, nelec)
+    else:
+        trace = 0
+        op = list(A.keys())[0]
+        coeff = A[op]
+        alpha_indices = []
+        beta_indices = []
+        for a in op:
+            if not a[1]:
+                alpha_indices.append(a[2])
+            else:
+                beta_indices.append(a[2])
+        alpha_indices = list(set(alpha_indices))
+        beta_indices = list(set(beta_indices))
+        combined_indices = [(i,False) for i in alpha_indices]  +[(i,True) for i in beta_indices]
+        
+        alpha_mapping = dict([(j,i) for i, j in enumerate(alpha_indices)])
+        beta_mapping = dict([(j,i) for i, j in enumerate(beta_indices)])
+        
+        n_alpha = len(alpha_indices)
+        n_beta = len(beta_indices)
+
+        no_configuration = False
+        eta_alpha = 0
+        eta_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 i, spin in combined_indices:
+            initial_zero = 0
+            initial_one = 1
+            is_zero = True
+            is_one = True
+            for a in reversed(op):  
+                if a[1] == spin and a[2] == i:
+                    change = a[0]*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 not is_zero and not is_one:
+                    print((is_zero,is_one))
+                    no_configuration = True
+                    break
+            if (is_zero and initial_zero != 0) or (is_one and initial_one != 1):
+                no_configuration = True
+            assert not is_zero or not is_one # only one case is possible if the operator has support on site i
+            # count the number of electrons
+            if is_one:
+                if spin == False:
+                    eta_alpha += 1
+                else:
+                    eta_beta += 1
+            # the number of electrons exceed the number of allowed electrons
+            if eta_alpha > n_elec_alpha or eta_beta > n_elec_beta: 
+                no_configuration = True
+            if no_configuration:
+                trace = 0
+                break
+                
+        if not no_configuration:
+            # the trace is nontrival and is a product of
+            # coeff, and
+            # binom(number of orbitals not in the support of A, number of electrons allowed on these orbitals) for both spin species
+            trace = coeff*comb(norb - len(alpha_indices),n_elec_alpha - eta_alpha)*comb(norb - len(beta_indices),n_elec_beta - eta_beta)
+
+            
+    return trace
diff --git a/tests/python/operators/trace_test.py b/tests/python/operators/trace_test.py
new file mode 100644
index 000000000..4acf1a933
--- /dev/null
+++ b/tests/python/operators/trace_test.py
@@ -0,0 +1,23 @@
+import numpy as np
+import time
+import ffsim
+
+def test_trace(norb = 50, nelec = (3,3)):
+    n = norb
+    
+    h = np.matrix(np.random.rand(n,n))
+    h = h + h.H
+    V = np.array([np.matrix(np.random.rand(n,n)),np.matrix(np.random.rand(n,n))])
+    H = ffsim.DiagonalCoulombHamiltonian(h,V)
+    
+    tic = time.time()
+    t1 = ffsim.trace(ffsim.fermion_operator(H),norb=norb, nelec=nelec)
+    toc = time.time()
+    print(f"Trace = {t1}. Took {toc-tic}s by converting DiagonalCoulombHamiltonian to FermionOperator")
+    
+    tic = time.time()
+    t2 = ffsim.trace(H, norb=n, nelec=nelec)
+    toc = time.time()
+    print(f"Trace = {t2}. Took {toc-tic}s using DiagonalCoulombHamiltonian._trace()")
+
+    assert np.abs(t1-t2)<1e-3, "The two methods do not match!"
\ No newline at end of file

From 11d906f8acdb256e8d10f17e3f0084586cd4ce48 Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Thu, 9 Jan 2025 09:35:35 -0700
Subject: [PATCH 02/16] addressing KS's comments

---
 tests/python/operators/trace_test.py | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/tests/python/operators/trace_test.py b/tests/python/operators/trace_test.py
index 4acf1a933..88d8f2d96 100644
--- a/tests/python/operators/trace_test.py
+++ b/tests/python/operators/trace_test.py
@@ -1,13 +1,26 @@
+# (C) Copyright IBM 2023.
+#
+# 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
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+
 import numpy as np
 import time
 import ffsim
 
-def test_trace(norb = 50, nelec = (3,3)):
-    n = norb
-    
-    h = np.matrix(np.random.rand(n,n))
+def test_trace():
+    norb = 50
+    nelec = (3,3)
+
+    rng = np.random.default_rng(12345)
+    h = np.matrix(rng.random((norb,norb)))
     h = h + h.H
-    V = np.array([np.matrix(np.random.rand(n,n)),np.matrix(np.random.rand(n,n))])
+    V = np.array([np.matrix(rng.random((norb,norb))),np.matrix(rng.random((norb,norb)))])
     H = ffsim.DiagonalCoulombHamiltonian(h,V)
     
     tic = time.time()
@@ -16,7 +29,7 @@ def test_trace(norb = 50, nelec = (3,3)):
     print(f"Trace = {t1}. Took {toc-tic}s by converting DiagonalCoulombHamiltonian to FermionOperator")
     
     tic = time.time()
-    t2 = ffsim.trace(H, norb=n, nelec=nelec)
+    t2 = ffsim.trace(H, norb=norb, nelec=nelec)
     toc = time.time()
     print(f"Trace = {t2}. Took {toc-tic}s using DiagonalCoulombHamiltonian._trace()")
 

From 41db00ec1db7ecccc33d699a3ea9970904703189 Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Thu, 9 Jan 2025 09:53:39 -0700
Subject: [PATCH 03/16] addressing KS's comment

---
 python/ffsim/protocols/trace_protocol.py | 146 +++++++++++------------
 tests/python/operators/trace_test.py     |   9 +-
 2 files changed, 74 insertions(+), 81 deletions(-)

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index caf0aa3a9..1b1b61a4c 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/python/ffsim/protocols/trace_protocol.py
@@ -18,7 +18,7 @@
 
 import numpy as np
 
-from math import comb
+import math
 
 
 class SupportsTrace(Protocol):
@@ -53,83 +53,83 @@ def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> float:
         "or a _diag_ method that returned its diagonal entries."
     )
 
-def _trace_fermion_operator(
+def _trace_fermion_operator_single(
     A: FermionOperator, norb: int, nelec: tuple[int, int]
 ):
     n_elec_alpha, n_elec_beta = nelec
-    if len(A)>1:
-        trace = 0
-        for op in A:
-            B = FermionOperator({op:A[op]})
-            trace+= _trace_fermion_operator(B, norb, nelec)
-    else:
-        trace = 0
-        op = list(A.keys())[0]
-        coeff = A[op]
-        alpha_indices = []
-        beta_indices = []
-        for a in op:
-            if not a[1]:
-                alpha_indices.append(a[2])
-            else:
-                beta_indices.append(a[2])
-        alpha_indices = list(set(alpha_indices))
-        beta_indices = list(set(beta_indices))
-        combined_indices = [(i,False) for i in alpha_indices]  +[(i,True) for i in beta_indices]
-        
-        alpha_mapping = dict([(j,i) for i, j in enumerate(alpha_indices)])
-        beta_mapping = dict([(j,i) for i, j in enumerate(beta_indices)])
-        
-        n_alpha = len(alpha_indices)
-        n_beta = len(beta_indices)
-
-        no_configuration = False
-        eta_alpha = 0
-        eta_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 i, spin in combined_indices:
-            initial_zero = 0
-            initial_one = 1
-            is_zero = True
-            is_one = True
-            for a in reversed(op):  
-                if a[1] == spin and a[2] == i:
-                    change = a[0]*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 not is_zero and not is_one:
-                    print((is_zero,is_one))
-                    no_configuration = True
-                    break
-            if (is_zero and initial_zero != 0) or (is_one and initial_one != 1):
-                no_configuration = True
-            assert not is_zero or not is_one # only one case is possible if the operator has support on site i
-            # count the number of electrons
-            if is_one:
-                if spin == False:
-                    eta_alpha += 1
-                else:
-                    eta_beta += 1
-            # the number of electrons exceed the number of allowed electrons
-            if eta_alpha > n_elec_alpha or eta_beta > n_elec_beta: 
+    op = list(A.keys())[0]
+    coeff = A[op]
+    alpha_indices = []
+    beta_indices = []
+    for _, spin, orb in op:
+        if not spin:
+            alpha_indices.append(orb)
+        else:
+            beta_indices.append(orb)
+    alpha_indices = list(set(alpha_indices))
+    beta_indices = list(set(beta_indices))
+    combined_indices = [(i,False) for i in alpha_indices]  +[(i,True) for i in beta_indices]
+    
+    
+    n_alpha = len(alpha_indices)
+    n_beta = len(beta_indices)
+
+    no_configuration = False
+    eta_alpha = 0
+    eta_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 i, spin in combined_indices:
+        initial_zero = 0
+        initial_one = 1
+        is_zero = True
+        is_one = True
+        for action, aspin, orb in reversed(op):  
+            if aspin == spin and orb == i:
+                change = a[0]*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 not is_zero and not is_one:
                 no_configuration = True
-            if no_configuration:
-                trace = 0
                 break
-                
-        if not no_configuration:
-            # the trace is nontrival and is a product of
-            # coeff, and
-            # binom(number of orbitals not in the support of A, number of electrons allowed on these orbitals) for both spin species
-            trace = coeff*comb(norb - len(alpha_indices),n_elec_alpha - eta_alpha)*comb(norb - len(beta_indices),n_elec_beta - eta_beta)
+        if (is_zero and initial_zero != 0) or (is_one and initial_one != 1):
+            no_configuration = True
+        assert not is_zero or not is_one # only one case is possible if the operator has support on site i
+        # count the number of electrons
+        if is_one:
+            if spin == False:
+                eta_alpha += 1
+            else:
+                eta_beta += 1
+        # the number of electrons exceed the number of allowed electrons
+        if eta_alpha > n_elec_alpha or eta_beta > n_elec_beta: 
+            no_configuration = True
+        if no_configuration:
+            trace = 0
+            break
+            
+    if not no_configuration:
+        # the trace is nontrival and is a product of
+        # coeff, and
+        # binom(number of orbitals not in the support of A, number of electrons allowed on these orbitals) for both spin species
+        trace = coeff*math.comb(norb - len(alpha_indices),n_elec_alpha - eta_alpha)*math.comb(norb - len(beta_indices),n_elec_beta - eta_beta)
+        
+    return trace
 
+def _trace_fermion_operator(
+    A: FermionOperator, norb: int, nelec: tuple[int, int]
+):
+    n_elec_alpha, n_elec_beta = nelec
+    for op in A:
+        B = FermionOperator({op:A[op]})
+        trace+= _trace_fermion_operator_single(B, norb, nelec)
+    
             
     return trace
diff --git a/tests/python/operators/trace_test.py b/tests/python/operators/trace_test.py
index 88d8f2d96..ee6c4557e 100644
--- a/tests/python/operators/trace_test.py
+++ b/tests/python/operators/trace_test.py
@@ -10,7 +10,6 @@
 
 
 import numpy as np
-import time
 import ffsim
 
 def test_trace():
@@ -23,14 +22,8 @@ def test_trace():
     V = np.array([np.matrix(rng.random((norb,norb))),np.matrix(rng.random((norb,norb)))])
     H = ffsim.DiagonalCoulombHamiltonian(h,V)
     
-    tic = time.time()
     t1 = ffsim.trace(ffsim.fermion_operator(H),norb=norb, nelec=nelec)
-    toc = time.time()
-    print(f"Trace = {t1}. Took {toc-tic}s by converting DiagonalCoulombHamiltonian to FermionOperator")
     
-    tic = time.time()
     t2 = ffsim.trace(H, norb=norb, nelec=nelec)
-    toc = time.time()
-    print(f"Trace = {t2}. Took {toc-tic}s using DiagonalCoulombHamiltonian._trace()")
 
-    assert np.abs(t1-t2)<1e-3, "The two methods do not match!"
\ No newline at end of file
+    np.testing.assert_allclose(t1, t2)
\ No newline at end of file

From eda4f9147208fda383bd5e3211c2d78394d78959 Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Thu, 9 Jan 2025 10:06:07 -0700
Subject: [PATCH 04/16] addressing KS's comment

---
 python/ffsim/protocols/trace_protocol.py | 3 ++-
 tests/python/operators/trace_test.py     | 5 +----
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index 1b1b61a4c..ee7294e90 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/python/ffsim/protocols/trace_protocol.py
@@ -89,7 +89,7 @@ def _trace_fermion_operator_single(
         is_one = True
         for action, aspin, orb in reversed(op):  
             if aspin == spin and orb == i:
-                change = a[0]*2-1
+                change = action*2-1
                 initial_zero += change
                 initial_one += change
                 if initial_zero < 0 or initial_zero > 1:
@@ -127,6 +127,7 @@ def _trace_fermion_operator(
     A: FermionOperator, norb: int, nelec: tuple[int, int]
 ):
     n_elec_alpha, n_elec_beta = nelec
+    trace = 0
     for op in A:
         B = FermionOperator({op:A[op]})
         trace+= _trace_fermion_operator_single(B, norb, nelec)
diff --git a/tests/python/operators/trace_test.py b/tests/python/operators/trace_test.py
index ee6c4557e..ee0d826ca 100644
--- a/tests/python/operators/trace_test.py
+++ b/tests/python/operators/trace_test.py
@@ -17,10 +17,7 @@ def test_trace():
     nelec = (3,3)
 
     rng = np.random.default_rng(12345)
-    h = np.matrix(rng.random((norb,norb)))
-    h = h + h.H
-    V = np.array([np.matrix(rng.random((norb,norb))),np.matrix(rng.random((norb,norb)))])
-    H = ffsim.DiagonalCoulombHamiltonian(h,V)
+    H = ffsim.random.random_diagonal_coulomb_hamiltonian(norb, real=True, seed=rng)
     
     t1 = ffsim.trace(ffsim.fermion_operator(H),norb=norb, nelec=nelec)
     

From b532693d7f35691240464f09cf0f5fc030cfc398 Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Thu, 9 Jan 2025 10:27:19 -0700
Subject: [PATCH 05/16] formatting

---
 python/ffsim/protocols/trace_protocol.py | 61 ++++++++++++------------
 tests/python/operators/trace_test.py     | 16 ++++---
 2 files changed, 40 insertions(+), 37 deletions(-)

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index ee7294e90..07ece090a 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/python/ffsim/protocols/trace_protocol.py
@@ -12,13 +12,12 @@
 
 from __future__ import annotations
 
+import math
 from typing import Any, Protocol
 
-from ffsim.operators import FermionOperator
-
 import numpy as np
 
-import math
+from ffsim.operators import FermionOperator
 
 
 class SupportsTrace(Protocol):
@@ -40,7 +39,7 @@ def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> float:
     """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)
@@ -53,12 +52,13 @@ def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> float:
         "or a _diag_ method that returned its diagonal entries."
     )
 
+
 def _trace_fermion_operator_single(
-    A: FermionOperator, norb: int, nelec: tuple[int, int]
+    ferm: FermionOperator, norb: int, nelec: tuple[int, int]
 ):
     n_elec_alpha, n_elec_beta = nelec
-    op = list(A.keys())[0]
-    coeff = A[op]
+    op = list(ferm.keys())[0]
+    coeff = ferm[op]
     alpha_indices = []
     beta_indices = []
     for _, spin, orb in op:
@@ -68,11 +68,9 @@ def _trace_fermion_operator_single(
             beta_indices.append(orb)
     alpha_indices = list(set(alpha_indices))
     beta_indices = list(set(beta_indices))
-    combined_indices = [(i,False) for i in alpha_indices]  +[(i,True) for i in beta_indices]
-    
-    
-    n_alpha = len(alpha_indices)
-    n_beta = len(beta_indices)
+    combined_indices = [(i, False) for i in alpha_indices] + [
+        (i, True) for i in beta_indices
+    ]
 
     no_configuration = False
     eta_alpha = 0
@@ -81,15 +79,16 @@ def _trace_fermion_operator_single(
     # 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
+    # if the state exceed 1 or goes below 0,
+    # the state is not physical and the trace must be 0
     for i, spin in combined_indices:
         initial_zero = 0
         initial_one = 1
         is_zero = True
         is_one = True
-        for action, aspin, orb in reversed(op):  
+        for action, aspin, orb in reversed(op):
             if aspin == spin and orb == i:
-                change = action*2-1
+                change = action * 2 - 1
                 initial_zero += change
                 initial_one += change
                 if initial_zero < 0 or initial_zero > 1:
@@ -101,36 +100,38 @@ def _trace_fermion_operator_single(
                 break
         if (is_zero and initial_zero != 0) or (is_one and initial_one != 1):
             no_configuration = True
-        assert not is_zero or not is_one # only one case is possible if the operator has support on site i
+        # only one case is possible if the operator has support on site i
+        assert not is_zero or not is_one
         # count the number of electrons
         if is_one:
-            if spin == False:
+            if not spin:
                 eta_alpha += 1
             else:
                 eta_beta += 1
         # the number of electrons exceed the number of allowed electrons
-        if eta_alpha > n_elec_alpha or eta_beta > n_elec_beta: 
+        if eta_alpha > n_elec_alpha or eta_beta > n_elec_beta:
             no_configuration = True
         if no_configuration:
             trace = 0
             break
-            
+
     if not no_configuration:
         # the trace is nontrival and is a product of
         # coeff, and
-        # binom(number of orbitals not in the support of A, number of electrons allowed on these orbitals) for both spin species
-        trace = coeff*math.comb(norb - len(alpha_indices),n_elec_alpha - eta_alpha)*math.comb(norb - len(beta_indices),n_elec_beta - eta_beta)
-        
+        # binom(#orbs not in the support of A, #elec on these orbs)
+        # for each spin species
+        trace = coeff
+        trace = trace * math.comb(norb - len(alpha_indices), n_elec_alpha - eta_alpha)
+        trace = trace * math.comb(norb - len(beta_indices), n_elec_beta - eta_beta)
+
     return trace
 
-def _trace_fermion_operator(
-    A: FermionOperator, norb: int, nelec: tuple[int, int]
-):
+
+def _trace_fermion_operator(ferm: FermionOperator, norb: int, nelec: tuple[int, int]):
     n_elec_alpha, n_elec_beta = nelec
     trace = 0
-    for op in A:
-        B = FermionOperator({op:A[op]})
-        trace+= _trace_fermion_operator_single(B, norb, nelec)
-    
-            
+    for op in ferm:
+        single_term = FermionOperator({op: ferm[op]})
+        trace += _trace_fermion_operator_single(single_term, norb, nelec)
+
     return trace
diff --git a/tests/python/operators/trace_test.py b/tests/python/operators/trace_test.py
index ee0d826ca..b4d9be4a6 100644
--- a/tests/python/operators/trace_test.py
+++ b/tests/python/operators/trace_test.py
@@ -10,17 +10,19 @@
 
 
 import numpy as np
+
 import ffsim
 
+
 def test_trace():
     norb = 50
-    nelec = (3,3)
+    nelec = (3, 3)
 
     rng = np.random.default_rng(12345)
-    H = ffsim.random.random_diagonal_coulomb_hamiltonian(norb, real=True, seed=rng)
-    
-    t1 = ffsim.trace(ffsim.fermion_operator(H),norb=norb, nelec=nelec)
-    
-    t2 = ffsim.trace(H, norb=norb, nelec=nelec)
+    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)
\ No newline at end of file
+    np.testing.assert_allclose(t1, t2)

From 147646ec4be21c1452b12b580911c9a5553ce65c Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Thu, 9 Jan 2025 10:45:49 -0700
Subject: [PATCH 06/16] fix assignment type error

---
 python/ffsim/protocols/trace_protocol.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index 07ece090a..82d5a61de 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/python/ffsim/protocols/trace_protocol.py
@@ -112,7 +112,7 @@ def _trace_fermion_operator_single(
         if eta_alpha > n_elec_alpha or eta_beta > n_elec_beta:
             no_configuration = True
         if no_configuration:
-            trace = 0
+            trace = 0j
             break
 
     if not no_configuration:
@@ -120,8 +120,8 @@ def _trace_fermion_operator_single(
         # coeff, and
         # binom(#orbs not in the support of A, #elec on these orbs)
         # for each spin species
-        trace = coeff
-        trace = trace * math.comb(norb - len(alpha_indices), n_elec_alpha - eta_alpha)
+
+        trace = coeff * math.comb(norb - len(alpha_indices), n_elec_alpha - eta_alpha)
         trace = trace * math.comb(norb - len(beta_indices), n_elec_beta - eta_beta)
 
     return trace
@@ -129,7 +129,7 @@ def _trace_fermion_operator_single(
 
 def _trace_fermion_operator(ferm: FermionOperator, norb: int, nelec: tuple[int, int]):
     n_elec_alpha, n_elec_beta = nelec
-    trace = 0
+    trace = 0j
     for op in ferm:
         single_term = FermionOperator({op: ferm[op]})
         trace += _trace_fermion_operator_single(single_term, norb, nelec)

From 4815a8fb86a486826cc9ecda345cb10d128842fa Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Thu, 9 Jan 2025 13:38:00 -0700
Subject: [PATCH 07/16] address KS's comments round 2

---
 python/ffsim/protocols/trace_protocol.py | 75 ++++++++++--------------
 1 file changed, 30 insertions(+), 45 deletions(-)

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index 82d5a61de..7c08b9980 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/python/ffsim/protocols/trace_protocol.py
@@ -53,26 +53,15 @@ def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> float:
     )
 
 
-def _trace_fermion_operator_single(
-    ferm: FermionOperator, norb: int, nelec: tuple[int, int]
+def _trace_term(
+    op: tuple[tuple[bool, bool, int], ...], norb: int, nelec: tuple[int, int]
 ):
     n_elec_alpha, n_elec_beta = nelec
-    op = list(ferm.keys())[0]
-    coeff = ferm[op]
-    alpha_indices = []
-    beta_indices = []
-    for _, spin, orb in op:
-        if not spin:
-            alpha_indices.append(orb)
-        else:
-            beta_indices.append(orb)
-    alpha_indices = list(set(alpha_indices))
-    beta_indices = list(set(beta_indices))
-    combined_indices = [(i, False) for i in alpha_indices] + [
-        (i, True) for i in beta_indices
-    ]
-
-    no_configuration = False
+
+    combined_indices = set([(orb, spin) for _, spin, orb in op])
+    norb_alpha = sum([not spin for _, spin in combined_indices])
+    norb_beta = sum([spin for _, spin in combined_indices])
+
     eta_alpha = 0
     eta_beta = 0
 
@@ -87,19 +76,21 @@ def _trace_fermion_operator_single(
         is_zero = True
         is_one = True
         for action, aspin, orb in reversed(op):
-            if aspin == spin and orb == i:
-                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 not is_zero and not is_one:
-                no_configuration = True
-                break
+            if (aspin, orb) != (spin, i):
+                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 not is_zero and not is_one:  # no possible initial state
+                return 0j
+        # if neither the state 0 nor 1 returns to the initial state, the trace must be 0
         if (is_zero and initial_zero != 0) or (is_one and initial_one != 1):
-            no_configuration = True
+            return 0j
         # only one case is possible if the operator has support on site i
         assert not is_zero or not is_one
         # count the number of electrons
@@ -110,28 +101,22 @@ def _trace_fermion_operator_single(
                 eta_beta += 1
         # the number of electrons exceed the number of allowed electrons
         if eta_alpha > n_elec_alpha or eta_beta > n_elec_beta:
-            no_configuration = True
-        if no_configuration:
-            trace = 0j
-            break
+            return 0j
 
-    if not no_configuration:
-        # the trace is nontrival and is a product of
-        # coeff, and
-        # binom(#orbs not in the support of A, #elec on these orbs)
-        # for each spin species
+    # the trace is nontrival and is a product of
+    # coeff, and
+    # binom(#orbs not in the support of A, #elec on these orbs)
+    # for each spin species
 
-        trace = coeff * math.comb(norb - len(alpha_indices), n_elec_alpha - eta_alpha)
-        trace = trace * math.comb(norb - len(beta_indices), n_elec_beta - eta_beta)
+    trace = math.comb(norb - norb_alpha, n_elec_alpha - eta_alpha) * math.comb(
+        norb - norb_beta, n_elec_beta - eta_beta
+    )
 
     return trace
 
 
 def _trace_fermion_operator(ferm: FermionOperator, norb: int, nelec: tuple[int, int]):
     n_elec_alpha, n_elec_beta = nelec
-    trace = 0j
-    for op in ferm:
-        single_term = FermionOperator({op: ferm[op]})
-        trace += _trace_fermion_operator_single(single_term, norb, nelec)
+    trace = sum([coeff * _trace_term(op, norb, nelec) for op, coeff in ferm.items()])
 
     return trace

From 3e517809c39369664931c98b071f990416b59ce8 Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Thu, 9 Jan 2025 15:42:33 -0700
Subject: [PATCH 08/16] addressing KS's comments round 3

---
 python/ffsim/protocols/trace_protocol.py | 52 +++++++++++-------------
 tests/python/operators/trace_test.py     |  4 --
 2 files changed, 24 insertions(+), 32 deletions(-)

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index 7c08b9980..2a6418ace 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/python/ffsim/protocols/trace_protocol.py
@@ -56,27 +56,27 @@ def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> float:
 def _trace_term(
     op: tuple[tuple[bool, bool, int], ...], norb: int, nelec: tuple[int, int]
 ):
-    n_elec_alpha, n_elec_beta = nelec
+    n_alpha, n_beta = nelec
 
-    combined_indices = set([(orb, spin) for _, spin, orb in op])
-    norb_alpha = sum([not spin for _, spin in combined_indices])
-    norb_beta = sum([spin for _, spin in combined_indices])
+    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
 
-    eta_alpha = 0
-    eta_beta = 0
+    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 i, spin in combined_indices:
+    for this_spin, this_orb in spin_orbs:
         initial_zero = 0
         initial_one = 1
         is_zero = True
         is_one = True
-        for action, aspin, orb in reversed(op):
-            if (aspin, orb) != (spin, i):
+        for action, spin, orb in reversed(op):
+            if (spin, orb) != (this_spin, this_orb):
                 continue
 
             change = action * 2 - 1
@@ -86,37 +86,33 @@ def _trace_term(
                 is_zero = False
             if initial_one < 0 or initial_one > 1:
                 is_one = False
-            if not is_zero and not is_one:  # no possible initial state
+            # return 0 immediately if there is no possible initial state
+            if not is_zero and not is_one:  
                 return 0j
-        # if neither the state 0 nor 1 returns to the initial state, the trace must be 0
+                
+        # if the operator has support on site i, either the initial state is 0 or 1, but not both
+        assert not is_zero or not is_one
+        # 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
-        # only one case is possible if the operator has support on site i
-        assert not is_zero or not is_one
         # count the number of electrons
         if is_one:
-            if not spin:
-                eta_alpha += 1
+            if not this_spin:
+                nelec_alpha += 1
             else:
-                eta_beta += 1
-        # the number of electrons exceed the number of allowed electrons
-        if eta_alpha > n_elec_alpha or eta_beta > n_elec_beta:
+                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
-    # coeff, and
-    # binom(#orbs not in the support of A, #elec on these orbs)
+    # binom(#orbs not in the support of op, #elec on these orbs)
     # for each spin species
 
-    trace = math.comb(norb - norb_alpha, n_elec_alpha - eta_alpha) * math.comb(
-        norb - norb_beta, n_elec_beta - eta_beta
+    return math.comb(norb - norb_alpha, n_alpha - nelec_alpha) * math.comb(
+        norb - norb_beta, n_beta - nelec_beta
     )
 
-    return trace
-
 
 def _trace_fermion_operator(ferm: FermionOperator, norb: int, nelec: tuple[int, int]):
-    n_elec_alpha, n_elec_beta = nelec
-    trace = sum([coeff * _trace_term(op, norb, nelec) for op, coeff in ferm.items()])
-
-    return trace
+    return sum([coeff * _trace_term(op, norb, nelec) for op, coeff in ferm.items()])
diff --git a/tests/python/operators/trace_test.py b/tests/python/operators/trace_test.py
index b4d9be4a6..32eae0a2c 100644
--- a/tests/python/operators/trace_test.py
+++ b/tests/python/operators/trace_test.py
@@ -17,12 +17,8 @@
 def test_trace():
     norb = 50
     nelec = (3, 3)
-
     rng = np.random.default_rng(12345)
     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)

From 01a8add53f6785ec5c5061d99a63eb7794d231ee Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Thu, 9 Jan 2025 15:44:59 -0700
Subject: [PATCH 09/16] format

---
 python/ffsim/protocols/trace_protocol.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index 2a6418ace..507edb712 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/python/ffsim/protocols/trace_protocol.py
@@ -87,10 +87,11 @@ def _trace_term(
             if initial_one < 0 or initial_one > 1:
                 is_one = False
             # return 0 immediately if there is no possible initial state
-            if not is_zero and not is_one:  
+            if not is_zero and not is_one:
                 return 0j
-                
-        # if the operator has support on site i, either the initial state is 0 or 1, but not both
+
+        # 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
         # 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):

From 6b21d84479751e1654d9d9e6cf62090989693992 Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Thu, 9 Jan 2025 19:54:16 -0700
Subject: [PATCH 10/16] addressing KS's comments round 4

---
 python/ffsim/protocols/trace_protocol.py      | 19 ++++++++-------
 .../python/operators/fermion_operator_test.py |  9 +++++++
 tests/python/operators/trace_test.py          | 24 -------------------
 3 files changed, 20 insertions(+), 32 deletions(-)
 delete mode 100644 tests/python/operators/trace_test.py

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index 507edb712..f53be287e 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/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
@@ -35,7 +35,7 @@ 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)
@@ -55,7 +55,7 @@ def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> float:
 
 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)
@@ -86,13 +86,14 @@ def _trace_term(
                 is_zero = False
             if initial_one < 0 or initial_one > 1:
                 is_one = False
-            # return 0 immediately if there is no possible initial state
-            if not is_zero and not is_one:
-                return 0j
 
         # 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
@@ -115,5 +116,7 @@ def _trace_term(
     )
 
 
-def _trace_fermion_operator(ferm: FermionOperator, norb: int, nelec: tuple[int, int]):
-    return sum([coeff * _trace_term(op, norb, nelec) for op, coeff in ferm.items()])
+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())
diff --git a/tests/python/operators/fermion_operator_test.py b/tests/python/operators/fermion_operator_test.py
index 0e36c3939..1bec6fdf3 100644
--- a/tests/python/operators/fermion_operator_test.py
+++ b/tests/python/operators/fermion_operator_test.py
@@ -643,3 +643,12 @@ 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():
+    norb = 50
+    nelec = (3, 3)
+    rng = np.random.default_rng(12345)
+    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)
diff --git a/tests/python/operators/trace_test.py b/tests/python/operators/trace_test.py
deleted file mode 100644
index 32eae0a2c..000000000
--- a/tests/python/operators/trace_test.py
+++ /dev/null
@@ -1,24 +0,0 @@
-# (C) Copyright IBM 2023.
-#
-# 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
-# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
-#
-# Any modifications or derivative works of this code must retain this
-# copyright notice, and modified files need to carry a notice indicating
-# that they have been altered from the originals.
-
-
-import numpy as np
-
-import ffsim
-
-
-def test_trace():
-    norb = 50
-    nelec = (3, 3)
-    rng = np.random.default_rng(12345)
-    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)

From f1dafb954f630a0d1676ae16d2a26be3078bf342 Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Fri, 10 Jan 2025 10:34:25 -0700
Subject: [PATCH 11/16] fix the phase bug for trace

---
 python/ffsim/protocols/trace_protocol.py        | 17 +++++++++++++++--
 tests/python/operators/fermion_operator_test.py | 12 +++++++++++-
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index f53be287e..4b1a62ccd 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/python/ffsim/protocols/trace_protocol.py
@@ -53,6 +53,16 @@ def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> complex:
     )
 
 
+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:
@@ -110,9 +120,12 @@ def _trace_term(
     # 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 math.comb(norb - norb_alpha, n_alpha - nelec_alpha) * math.comb(
-        norb - norb_beta, n_beta - nelec_beta
+    return (
+        _term_phase(op)
+        * math.comb(norb - norb_alpha, n_alpha - nelec_alpha)
+        * math.comb(norb - norb_beta, n_beta - nelec_beta)
     )
 
 
diff --git a/tests/python/operators/fermion_operator_test.py b/tests/python/operators/fermion_operator_test.py
index 1bec6fdf3..8f4550516 100644
--- a/tests/python/operators/fermion_operator_test.py
+++ b/tests/python/operators/fermion_operator_test.py
@@ -645,10 +645,20 @@ def test_mapping_methods():
     }
 
 def test_trace():
+    rng = np.random.default_rng(12345)
+    # compare for random_diagonal_coulomb_hamiltonian
     norb = 50
     nelec = (3, 3)
-    rng = np.random.default_rng(12345)
     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)
+
+    

From 8e8e2c50242b163229c98af1d38a77f5e81a3fc2 Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Fri, 10 Jan 2025 12:59:38 -0700
Subject: [PATCH 12/16] added _trace_ to FermionOperator

---
 src/fermion_operator.rs | 94 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 94 insertions(+)

diff --git a/src/fermion_operator.rs b/src/fermion_operator.rs
index e4f47906d..54fa5f23d 100644
--- a/src/fermion_operator.rs
+++ b/src/fermion_operator.rs
@@ -423,6 +423,100 @@ impl FermionOperator {
         }
         true
     }
+
+    fn _trace_(&self, norb: usize, nelec: (usize, usize)) -> Complex64 {
+        self.coeffs
+            .iter()
+            .map(|(op, &coeff)| coeff * term_trace(op, norb, nelec) as f64)
+            .sum()
+    }
+}
+
+fn term_phase(op: &[(bool, bool, i32)]) -> i32 {
+    let mut phase = 0;
+    let mut op = op.to_vec(); // Work with a mutable copy
+    while !op.is_empty() {
+        let (action, spin, orb) = op[0];
+        let conj_idx = op
+            .iter()
+            .position(|&(a, s, o)| (a, s, o) == (!action, spin, orb))
+            .unwrap();
+        phase += conj_idx - 1;
+        op = [&op[1..conj_idx], &op[conj_idx + 1..]].concat();
+    }
+    (-1i32).pow(phase as u32)
+}
+
+fn term_trace(op: &[(bool, bool, i32)], norb: usize, nelec: (usize, usize)) -> i32 {
+    let (n_alpha, n_beta) = nelec;
+
+    let spin_orbs = op
+        .iter()
+        .map(|&(_, spin, orb)| (spin, orb))
+        .collect::<std::collections::HashSet<_>>();
+    let norb_alpha = spin_orbs.iter().filter(|&&(spin, _)| !spin).count();
+    let norb_beta = spin_orbs.len() - norb_alpha;
+
+    let mut nelec_alpha = 0;
+    let mut nelec_beta = 0;
+
+    for &(this_spin, this_orb) in &spin_orbs {
+        let mut initial_zero = 0;
+        let mut initial_one = 1;
+        let mut is_zero = true;
+        let mut is_one = true;
+
+        for &(action, spin, orb) in op.iter().rev() {
+            if (spin, orb) != (this_spin, this_orb) {
+                continue;
+            }
+
+            let change = if action { 1 } else { -1 };
+            initial_zero += change;
+            initial_one += change;
+
+            if !(0..=1).contains(&initial_zero) {
+                is_zero = false;
+            }
+            if !(0..=1).contains(&initial_one) {
+                is_one = false;
+            }
+        }
+
+        assert!(!is_zero || !is_one);
+
+        if !is_zero && !is_one {
+            return 0;
+        }
+
+        if (is_zero && initial_zero != 0) || (is_one && initial_one != 1) {
+            return 0;
+        }
+
+        if is_one {
+            if !this_spin {
+                nelec_alpha += 1;
+            } else {
+                nelec_beta += 1;
+            }
+        }
+
+        if nelec_alpha > n_alpha || nelec_beta > n_beta {
+            return 0;
+        }
+    }
+
+    let binomial = |n, k| {
+        if k > n {
+            0
+        } else {
+            (1..=n).rev().take(k).product::<usize>() / (1..=k).product::<usize>()
+        }
+    };
+
+    term_phase(op)
+        * binomial(norb - norb_alpha, n_alpha - nelec_alpha) as i32
+        * binomial(norb - norb_beta, n_beta - nelec_beta) as i32
 }
 
 fn _normal_ordered_term(term: &[(bool, bool, i32)], coeff: &Complex64) -> FermionOperator {

From 2b9ba3c1ef6b51c6a7d7d0d22b57444284055dce Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Fri, 10 Jan 2025 13:21:20 -0700
Subject: [PATCH 13/16] revert trace_protocol back to before, added comments
 for rust

---
 python/ffsim/protocols/trace_protocol.py | 92 +-----------------------
 src/fermion_operator.rs                  | 12 +++-
 2 files changed, 13 insertions(+), 91 deletions(-)

diff --git a/python/ffsim/protocols/trace_protocol.py b/python/ffsim/protocols/trace_protocol.py
index 4b1a62ccd..95d421dab 100644
--- a/python/ffsim/protocols/trace_protocol.py
+++ b/python/ffsim/protocols/trace_protocol.py
@@ -1,4 +1,4 @@
-# (C) Copyright IBM 2025.
+# (C) Copyright IBM 2023.
 #
 # 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,13 +12,10 @@
 
 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."""
@@ -35,11 +32,8 @@ def _trace_(self, norb: int, nelec: tuple[int, int]) -> float:
         """
 
 
-def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> complex:
+def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> float:
     """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)
@@ -51,85 +45,3 @@ def trace(obj: Any, norb: int, nelec: tuple[int, int]) -> complex:
         "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())
diff --git a/src/fermion_operator.rs b/src/fermion_operator.rs
index 54fa5f23d..3bc9adab5 100644
--- a/src/fermion_operator.rs
+++ b/src/fermion_operator.rs
@@ -460,6 +460,11 @@ fn term_trace(op: &[(bool, bool, i32)], norb: usize, nelec: (usize, usize)) -> i
     let mut nelec_alpha = 0;
     let mut 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 {
         let mut initial_zero = 0;
         let mut initial_one = 1;
@@ -482,17 +487,21 @@ fn term_trace(op: &[(bool, bool, i32)], norb: usize, nelec: (usize, usize)) -> i
                 is_one = false;
             }
         }
-
+        // if the operator has support on this_orb,
+        // either the initial state is 0 or 1, but not both
         assert!(!is_zero || !is_one);
 
+        // return 0 if there is no possible initial state
         if !is_zero && !is_one {
             return 0;
         }
 
+        // the state must return to the initial state, otherwise the trace is zero
         if (is_zero && initial_zero != 0) || (is_one && initial_one != 1) {
             return 0;
         }
 
+        // count the number of electrons in the support of op
         if is_one {
             if !this_spin {
                 nelec_alpha += 1;
@@ -502,6 +511,7 @@ fn term_trace(op: &[(bool, bool, i32)], norb: usize, nelec: (usize, usize)) -> i
         }
 
         if nelec_alpha > n_alpha || nelec_beta > n_beta {
+            // the number of electrons exceeds the number of allowed electrons
             return 0;
         }
     }

From cd9ce59f2e0cbf473fdf0f019e7eb37fd9d47709 Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Fri, 10 Jan 2025 14:38:10 -0700
Subject: [PATCH 14/16] import binomial from num-integer

---
 Cargo.lock              | 1 +
 Cargo.toml              | 1 +
 src/fermion_operator.rs | 9 +--------
 3 files changed, 3 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e4c71b49c..d7637bcb7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -194,6 +194,7 @@ dependencies = [
  "blas",
  "blas-src",
  "ndarray",
+ "num-integer",
  "numpy",
  "openblas-src",
  "pyo3",
diff --git a/Cargo.toml b/Cargo.toml
index 2b4b56d74..237a64b41 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,6 +18,7 @@ pyo3 = { version = "0.21", features = [
     "num-complex",
     "abi3-py39",
 ] }
+num-integer = "0.1"
 
 [target.'cfg(target_os = "linux")'.dependencies]
 blas-src = { version = "0.10", features = ["openblas"] }
diff --git a/src/fermion_operator.rs b/src/fermion_operator.rs
index 3bc9adab5..263cd2b19 100644
--- a/src/fermion_operator.rs
+++ b/src/fermion_operator.rs
@@ -8,6 +8,7 @@
 // copyright notice, and modified files need to carry a notice indicating
 // that they have been altered from the originals.
 
+use num_integer::binomial;
 use numpy::Complex64;
 use pyo3::class::basic::CompareOp;
 use pyo3::exceptions::PyKeyError;
@@ -516,14 +517,6 @@ fn term_trace(op: &[(bool, bool, i32)], norb: usize, nelec: (usize, usize)) -> i
         }
     }
 
-    let binomial = |n, k| {
-        if k > n {
-            0
-        } else {
-            (1..=n).rev().take(k).product::<usize>() / (1..=k).product::<usize>()
-        }
-    };
-
     term_phase(op)
         * binomial(norb - norb_alpha, n_alpha - nelec_alpha) as i32
         * binomial(norb - norb_beta, n_beta - nelec_beta) as i32

From caed565179dc5ac88405c0e2507ceb519014844c Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Mon, 13 Jan 2025 10:18:07 -0700
Subject: [PATCH 15/16] reformat fermion_operator_test.py

---
 tests/python/operators/fermion_operator_test.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/python/operators/fermion_operator_test.py b/tests/python/operators/fermion_operator_test.py
index 8f4550516..952b9db5d 100644
--- a/tests/python/operators/fermion_operator_test.py
+++ b/tests/python/operators/fermion_operator_test.py
@@ -644,6 +644,7 @@ def test_mapping_methods():
         ((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
@@ -660,5 +661,3 @@ def test_trace():
     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)
-
-    

From 753c788490a968930c0b19e925c81798807c561a Mon Sep 17 00:00:00 2001
From: Minh Tran <minhtran@ibm.com>
Date: Tue, 21 Jan 2025 11:05:04 -0700
Subject: [PATCH 16/16] fix check

---
 tests/python/operators/fermion_operator_test.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/python/operators/fermion_operator_test.py b/tests/python/operators/fermion_operator_test.py
index 952b9db5d..9585d5f91 100644
--- a/tests/python/operators/fermion_operator_test.py
+++ b/tests/python/operators/fermion_operator_test.py
@@ -647,6 +647,8 @@ def test_mapping_methods():
 
 def test_trace():
     rng = np.random.default_rng(12345)
+    ham: ffsim.DiagonalCoulombHamiltonian | ffsim.MolecularHamiltonian
+
     # compare for random_diagonal_coulomb_hamiltonian
     norb = 50
     nelec = (3, 3)