Add trace_moment: faster tr(H^k) via identity-multiset enumeration (#80)

[khopade-works]

Closes #80.

Summary

Adds trace_moment(o::Operator, k; scale=0), which computes trace(o^k) by enumerating only the multisets of Pauli strings whose product is proportional to the identity and summing their phases analytically — without ever constructing o^(k/2).

This implements the two pieces requested in #80:

1. Enumerate identity-yielding multisets. A pruned depth-first search over the operator's terms tracks the running XOR of the Pauli strings and keeps only branches that can still return to the identity. Pruning uses (a) the support reachable by the remaining terms and (b) a ceil(active_sites / max_support) lower bound on the terms still needed. Sorting terms by their lowest active site makes the reach prune force low sites to resolve early, keeping the search close to the number of contributing multisets.

2. Analytic per-multiset phase (near closed form). Reordering two Pauli strings only flips the product phase by the sign of their (anti)commutation, so the whole order dependence factors as

   phase_sum(M) = K_ref(M) * Dtilde(M)
   

   where K_ref is the sign of one canonical ordering (O(k)) and Dtilde(M) depends only on the anticommutation graph of the distinct terms. Every term that commutes with all others peels off as a binomial coefficient, leaving a small "anticommuting core" handled by a tiny multiplicity dynamic program. This removes the k! factor of the naive ordered sum.

trace_product is left unchanged; trace_moment is added as a new, non-breaking function.

Performance (issue example: TFIM, N=20, k=1:14)

k	trace_product	trace_moment	speedup
11	2.48s	0.04s	66x
12	2.33s	1.39s	1.7x
13	25.0s	0.25s	99x
14	22.5s	19.3s	1.2x
total k=1:14	52.7s	21.1s	2.5x

• Odd moments are 10–100x faster (the existing path builds H^{(k+1)/2}; the multiset method gets them almost for free).
• Even moments are modestly faster with near-zero memory (the baseline allocates ~6 GiB building H^7).
• The advantage grows with k.

Results agree with trace_product to machine precision (Float64 accumulation only).

Note: the k! saving in the issue is relative to the naive n^k sum; relative to the current tr(H^{k/2} H^{k/2}) implementation the headline wins are odd moments, memory, and scaling, netting 2.5x on the example.

What's included

• src/moments.jl: trace_moment + helpers, with a derivation comment block and docstring.
• src/PauliStrings.jl: export trace_moment.
• docs/src/docstrings.md: added under "Power and moments".
• test/algorithms.jl: trace_moment test set comparing against trace_product (random 1-/2-local operators, the Ising model, complex/non-Hermitian coefficients, an explicit identity term, the scale keyword, k=0, and error handling).
• benchmark/: a TFIM model and trace_moment vs trace_product benchmark entries.

Test plan

• Full existing test suite passes (Pkg.test()).
• New trace_moment test set passes.
• Verified on Julia 1.6 and 1.10 (the CI matrix).
• Phase factorization verified to 0 error against an exact brute-force DP over thousands of random multisets.

limitations

• trace_moment currently targets Operator (the type in the issue example); translation-symmetric OperatorTS still uses the existing path.
• The integer phase bookkeeping is exact up to k ~ 20, beyond which the moment value itself exceeds Float64 range.

[nicolasloizeau]

Nice! do you understand why it only speeds up the odd moments ?
The odd moments of the TFIM are zero, and this is easy to show analytically because the density of states is even, so I'm a bit disappointed the advantage mostly apply to odd moments only.
How about a model that has non-zero odd moments ? e.g. the Heisenberg model :

using PauliStrings

function heisenberg(N)
    H = Operator(N)
    for i in 1:N
        H += "X", i, "X", mod1(i + 1, N)
        H += "Y", i, "Y", mod1(i + 1, N)
        H += "Z", i, "Z", mod1(i + 1, N)
    end
    return H
end


N = 20
H = heisenberg(N)
for k in 1:12
    muk = trace_product(H, k; scale=1)
    println("mu_", k, " = ", muk)
end

Is it possible to generalize the method to trace_product(A, k, B, l) ? For example to compute stuff like :

using PauliStrings

function ising(h, N)
    H = Operator(N)
    for i in 1:N
        H += h, "X", i
    end
    for i in 1:N
        H += "Z", i, "Z", mod1(i + 1, N)
    end
    return H
end

N = 20
H = ising(0.5, N)
O = Operator(N) + ("X", 1)
for k in 1:14
    muk = trace_product(H, k, O, 1; scale=1)
    println("mu_", k, " = ", muk)
end

maybe the advantage it more obvious for problems like these ?

[khopade-works]

Thanks for the thoughtful questions pushed a follow-up commit that generalizes trace_moment to the 4-argument trace(A^k B^l) case.

Why TFIM odd moments look so fast

On TFIM, odd moments are exactly zero (even density of states ⇒ only even powers of the Hamiltonian contribute to the trace). trace_moment still has to enumerate identity-yielding multisets, but the anticommuting-core phase sum vanishes for those odd-k patterns, so the leaf work is essentially free. That is why the TFIM table is dominated by odd-k wins — it is partly a correctness/feature story (cheap zeros), not a uniform speedup over trace_product.

For a model with nonzero odd moments, I ran your Heisenberg example (N=20, k=1:12, scale=1):

k	trace_product (s)	trace_moment (s)	speedup
1	0.003	0.003	1.0×
3	0.12	0.02	6.0×
5	0.55	0.08	6.9×
7	1.8	0.3	6.0×
9	4.5	0.7	6.4×
11	6.7	1.0	6.7×
even k	—	—	~0.8–1.0×

Total k=1:12: 59.4s → 40.0s (1.48×). So nonzero odd moments still benefit, but the Heisenberg operator is denser (60 terms), so even-k moments can be slightly slower than trace_product because we pay enumeration overhead without the “free zero” shortcut.

trace_product(A, k, B, l) generalization

Yes — this is now implemented as trace_moment(A, k, B, l; scale=0). The ordering phase factorizes per Pauli block:

[
\mathrm{tr}(A^k B^l) = \mathrm{scale}\sum_R (-1)^{y(R)} g_A(R), g_B(R)
]

where (g_A, g_B) are the single-operator phase factors from the multiset search. The smaller operator is tabulated by XOR targets; the larger side uses the same pruned DFS.

On your Ising + O = X₁ example (N=20, k=1:14):

k	trace_product (s)	trace_moment (s)	speedup
14	44.2	0.25	175×
total	75.4	2.9	26×

Here the advantage is much clearer: trace_product must build (H^k) (and multiply by (O)), while trace_moment never materializes (H^k).

Happy to tune further for even-k / dense-operator cases (e.g. closed forms for small anticommuting cores) if that would be useful.

[lkdvos]

If the struct is marked as mutable, these can probably just be of type C.

[khopade-works]

Good catch switched to plain extfactor::C and acc::C on the mutable struct.

[lkdvos]

Why do you need to inflate the iszero(u) case? trailing_zeros(u) + 1 should result in a value between 1 and the number of bits, and is its maximal value iff u = 0, so that should already be the largest value I think?

[khopade-works]

You're right. For u = 0, trailing_zeros(u) + 1 is already the maximum (bitwidth + 1), so the identity sorts last without an explicit sentinel. Simplified.

[lkdvos]

Why not use Base.binomial?

[khopade-works]

Dropped the custom helper; now uses Base.binomial (exact for the k range we care about).

[lkdvos]

Since this operation is generally useful, how about adding this to the operations.jl file, and having both commutes(a, b) and anticommutes defined?
I don't think there is a particular need to have this marked as @inline, the compiler is generally smart enough to inline this

[khopade-works]

Moved to operations.jl as commutes / anticommutes (exported). Tests check consistency with commutator / anticommutator.

[lkdvos]

This function might benefit from splitting it up into some smaller functions with a bit more descriptive names, just to keep things somewhat easier to review

[khopade-works]

Split into _peel_commuting_terms!, _anticommuting_core_size!, _core_ordering_phase_sum!.

[lkdvos]

If you start with maxsup = 1, 5 lines above this one, this max is no longer needed.

[khopade-works]

maxsup now initializes to 1; redundant final max removed.

[lkdvos]

This can probably move to the operations.jl file, since it is already used in the prod function as well. The @inline is likely not necessary, since the compiler should be smart enough to do it itself

[khopade-works]

Extracted pauli_prod_phase in operations.jl; both prod overloads and _canonical_sign use it.

[lkdvos]

Do you mind documenting what the "canonical" order is you chose?

[khopade-works]

Documented in the derivation block: operator terms sorted by _minsite; within a multiset, distinct terms in nondecreasing index with consecutive copies (DFS stack order).

[lkdvos]

Could you elaborate a bit more about the choice of this structure for a DFS?

Is there a specific reason to not work with a more stack-based approach? There are several quantities that could be built and "unbuilt" in this way, which would reduce some of the storage requirements.
For example: reach[i + 1] can be computed from reach[i] since xor is invertible, and in general xor is faster than memory access, so it should be beneficial to simply increase/decrease that as you go through the search.
The same is probably through for some of the canonical_sign and coefficients, although I'm less sure about how much that actually matters (except for a very easy early cutoff in case of tiny coefficients).

I think I'm more used to seeing this explicit stack approach, where the vectors are really representing the state of the current search space. I'm not saying this approach is better or worse, mostly just interested in seeing if there was a deliberate design choice here.

[khopade-works]

The search state is an explicit stack idx/mult/depth hold the current multiset. Recursion matches “pick multiplicity for term i, advance to i+1”. reach[i] is a one-time suffix-union table for O(1) support pruning at each node rather than incrementally maintaining it on backtrack (XOR is cheap; this avoids extra memory traffic on deep trees). An iterative frame stack would be equivalent; I went with recursion for clarity. Happy to switch if you prefer.

[lkdvos]

As a more general comment/question:

• Do you have a sense of which part of the algorithms are now dominating? What do profiling results tell you?
• Is there any prospect of multi-threading in this approach? Can you efficiently divide up the search space to make use of a multi-core system?
• Computing the trace of a moment effectively boils down to a search space that can be explored either BFS (which corresponds to simply multiplying the operators, as was previously the case), or DFS (which is the route taken here). Is there any chance of further improving efficiency by making use of a combination of the two? For example, can you make use of H^k for small k to prune the search space earlier?

[khopade-works]

As a more general comment/question:

• Do you have a sense of which part of the algorithms are now dominating? What do profiling results tell you?
• Is there any prospect of multi-threading in this approach? Can you efficiently divide up the search space to make use of a multi-core system?
• Computing the trace of a moment effectively boils down to a search space that can be explored either BFS (which corresponds to simply multiplying the operators, as was previously the case), or DFS (which is the route taken here). Is there any chance of further improving efficiency by making use of a combination of the two? For example, can you make use of H^k for small k to prune the search space earlier?

Good questions I looked at this on the TFIM example (N=20, k=14) and a few smaller sanity checks.

What's dominating: On even moments where we don't get the "free zero" shortcut, _moment_dfs! is clearly the main cost most time is spent walking the search tree and hitting leaves. The per-leaf work (_canonical_sign + _dtilde!) is comparatively cheap for local Hamiltonians because the anticommuting core is usually tiny (often just a handful of terms). On dense models like Heisenberg that's less true, and I think that's part of why even-k moments there end up roughly parity with trace_product we pay enumeration overhead without the odd-k shortcut.

Threading: I think there's a reasonable path here. The most obvious split is in trace_moment(A, k, B, l) where we already loop over target XORs from the smaller operator's table those iterations are independent and could run in parallel with thread-local accumulators. For plain trace_moment(H, k) you could also split at the top-level DFS branches (e.g. "use term 1 zero times" vs "use it once/more"), though the workload balance might be uneven depending on the operator. I haven't implemented this yet for small N the task overhead might not be worth it, but for the large-moment cases where this method is meant to shine it could help.

BFS/DFS hybrid: That's an interesting idea and honestly it's partly what the 4-arg generalization is already doing we fully tabulate the smaller side (a BFS-style buildup of all size-l multisets grouped by XOR target), then DFS only the larger side conditioned on each target. For trace(H^k) alone, precomputing H^m for small m and using its XOR support to prune impossible branches earlier sounds plausible, especially for even k on dense operators where pure DFS is weakest. I haven't explored that yet; it'd be a natural next step if even-k performance on dense Hamiltonians becomes a bottleneck.

Happy to dig deeper into any of these if useful for the PR.

[nicolasloizeau]

Hi, can you please make sure the tests are passing ?

[khopade-works]

@nicolasloizeau Fixed CI is green, Happy to make any more changes if required

[nicolasloizeau]

Nice, would it be also possible to have a translation symmetric version for Operator{<:PauliStringTS} types ?

[khopade-works]

@nicolasloizeau Added trace_moment for Operator{<:PauliStringTS} - uses resum + translation folding over representatives, matches trace_product on TFIM/Heisenberg/2D tests.