Skip to content

Avoid building per-term QubitOperators in jordan_wigner interaction transform#1427

Open
blazingphoenix7 wants to merge 1 commit into
quantumlib:mainfrom
blazingphoenix7:speed-up-jordan-wigner-interaction-op
Open

Avoid building per-term QubitOperators in jordan_wigner interaction transform#1427
blazingphoenix7 wants to merge 1 commit into
quantumlib:mainfrom
blazingphoenix7:speed-up-jordan-wigner-interaction-op

Conversation

@blazingphoenix7

Copy link
Copy Markdown
Contributor

Third of the operator hot-path findings from #1407; the earlier two are #1415 and #1417.

jordan_wigner on an InteractionOperator maps the one- and two-body integrals
straight to Pauli terms instead of going through a FermionOperator first, and
#587 points at exactly this routine as the fast path to build on ("there is a
special routine that maps the InteractionOperators directly to qubits ... and
that is faster"). #121 is the other side of it: it flagged this same routine as
"surprisingly slow" back in 2017 and got closed without a fix.

The transform walks every one-body pair and every pair of pairs, so the two-body
part is O(N^4). For each Pauli string it built a fresh single-term
QubitOperator and summed that into the running total. Each of those
constructions parses the term, validates every (index, action) factor, runs the
Pauli _simplify, and in the three-unique-index case multiplies by a Z
operator, which deepcopies. None of that is needed to add a term to a sum: a
QubitOperator's terms is just a dict from a sorted tuple of (index, action)
factors to a coefficient, and the strings this routine emits are already sorted
with their factors already combined, so there is nothing for the constructor to
normalize.

This accumulates the pauli_string -> coefficient pairs into a plain dict with a
small _add_term helper that does what __iadd__ / __isub__ do to a single
term (add onto the running value, drop it once it falls under EQ_TOLERANCE), and
builds one QubitOperator at the end by assigning the finished dict to .terms.
Accumulating into a dict instead of building and summing single-term operators
avoids that per-term construction and validation, and it measures about 2x
faster, with identical output.

Public functions

jordan_wigner_one_body and jordan_wigner_two_body are public and re-exported,
and bksf_test and jordan_wigner_test call them directly, so they keep their
signatures and their output. I pulled the term-building logic into two private
helpers (_one_body_terms, _two_body_terms) that return the dict, and the
public functions and the interaction-op driver both go through them, so there is
one copy of the logic and the public results are unchanged. The Z times
hopping-term product in the three-unique case is a single-qubit relabeling (the
shared qubit carries at most a parity Z, and Z squared is the identity), so it
is a direct toggle of that index in the string now rather than an operator
multiply and deepcopy.

Correctness

Before editing the source I captured the current jordan_wigner output on 104
random real InteractionOperators and saved it. The mix is physical
(symmetric, molecular-style) operators, dense generic real tensors, sparse and
very sparse tensors, one-body-only, constant-only, and the exact-zero operator,
at N = 1, 2, 3, 4, 6, 8, 10, 12. The new code reproduces that captured output
exactly on every case: identical term-key sets, max absolute coefficient
difference 0, and (new - old).induced_norm(1) equal to 0. It also matches an
independent dict-based reference I wrote separately, to the same 0. Seeds are
fixed, so the check is deterministic.

jordan_wigner_test.py and bksf_test.py pass (the two public helpers are
checked there against the FermionOperator path over a full index grid with
complex coefficients), and transforms/ passes as a whole. black and mypy are
clean on the file. Every changed line is hit by the existing suite except the
pre-existing n_qubits < count_qubits(iop) guard, which the old code did not
cover either and which I left untouched.

The routine is still documented as real-only and this does not change that. It
computes the same coefficients from the same integrals; only the container they
land in changed.

Benchmarks

Random real InteractionOperators, cumulative jordan_wigner time. Medians over 5
repeats within each of 6 alternated upstream/branch rounds, fixed seeds, on an
Intel Core Ultra 5 225 laptop (Windows, Python 3.12, numpy 2.5).

N     upstream    ours       speedup
8     16 ms       7.6 ms     2.1x
12    101 ms      48 ms      2.1x
16    359 ms      177 ms     2.0x
20    963 ms      476 ms     2.0x
24    2222 ms     1093 ms    2.0x

…ransform

The interaction-operator path built a fresh single-term QubitOperator for every
Pauli string in the O(N^4) two-body loop and summed them in, so each term paid
the constructor's parse, validation, simplify and copy. This accumulates the
Pauli strings into a plain dict and builds the QubitOperator once at the end. The
public jordan_wigner_one_body and jordan_wigner_two_body share the same term
builders, so their output is unchanged, verified bit-exact against the previous
implementation.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request refactors the Jordan-Wigner transformation in jordan_wigner.py to accumulate Pauli terms in a dictionary instead of repeatedly instantiating and summing QubitOperator objects, which improves performance. The reviewer suggests further optimizing this by passing the terms dictionary directly into _one_body_terms and _two_body_terms for in-place accumulation, thereby eliminating the need for the _merge_terms helper function.

Comment on lines +145 to +148
def _merge_terms(terms, contribution):
"""Add every term of a contribution dict into terms."""
for key, value in contribution.items():
_add_term(terms, key, value)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Since we are passing the terms dictionary directly to _one_body_terms and _two_body_terms to accumulate terms in-place, the _merge_terms helper function is no longer needed and can be safely removed.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, but I'd keep _merge_terms and the two-step here. The aim is output that stays bit-for-bit identical to the current transform, and the two levels of dropping are what make that hold. Upstream builds each contribution as its own operator via jordan_wigner_one_body/two_body (whose += drops sub-tolerance terms on its own, against zero), then folds that into the running total with a second += that drops again. _one_body_terms/_two_body_terms plus _merge_terms mirror exactly that. Accumulating straight into the shared dict removes the first drop, so a sub-EQ_TOLERANCE contribution gets tested against the live total instead of against zero. Two places it shows up: an all-zero operator comes out as {} instead of {(): 0.0}, and a sub-tolerance term folded onto an existing value shifts it by up to EQ_TOLERANCE. Both stay within isclose, but they'd break the exact match, so I'll keep the helper.

Comment on lines +233 to +235
def _one_body_terms(p, q, coefficient):
r"""Terms of a^\dagger_p a_q + h.c. under Jordan-Wigner, as a dict."""
terms = {}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update _one_body_terms to accept an optional terms dictionary to allow in-place accumulation of terms.

Suggested change
def _one_body_terms(p, q, coefficient):
r"""Terms of a^\dagger_p a_q + h.c. under Jordan-Wigner, as a dict."""
terms = {}
def _one_body_terms(p, q, coefficient, terms=None):
r"""Terms of a^\dagger_p a_q + h.c. under Jordan-Wigner, as a dict."""
if terms is None:
terms = {}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Keeping the dict-return form here, same reason as the _merge_terms thread: the local drop has to happen against zero, not against the running total.

Comment on lines +270 to +272
def _two_body_terms(p, q, r, s, coefficient):
r"""Terms of a^\dagger_p a^\dagger_q a_r a_s + h.c. under JW, as a dict."""
terms = {}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update _two_body_terms to accept an optional terms dictionary to allow in-place accumulation of terms.

Suggested change
def _two_body_terms(p, q, r, s, coefficient):
r"""Terms of a^\dagger_p a^\dagger_q a_r a_s + h.c. under JW, as a dict."""
terms = {}
def _two_body_terms(p, q, r, s, coefficient, terms=None):
r"""Terms of a^\dagger_p a^\dagger_q a_r a_s + h.c. under JW, as a dict."""
if terms is None:
terms = {}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here, keeping the dict return so that first drop stays local.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant