improved Lanczos iteration and introduced numba usage leading to significant speedup#310
Merged
Conversation
Contributor
📝 WalkthroughSummary by CodeRabbit
WalkthroughReplaces fixed Lanczos API with an adaptive inlined Krylov expm, adds Numba-accelerated Lanczos and TDVP kernels, updates tests for the new paths, lowers a dense-threshold, adjusts CI/macOS support, adds numba deps, and minor test/fixture tuning across the suite. Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant matrix_expm as expm_krylov
participant LanczosNumba as lanczos_numba
participant LinearAlgebra as Diagonalize/EXPM
Caller->>matrix_expm: call expm_krylov(op, vec, dt, max_lanczos_iterations, tol)
matrix_expm->>LanczosNumba: (optional) use orthogonalize_step/normalize_and_store in loop
alt numba available & size >= NUMBA_THRESHOLD
LanczosNumba-->>matrix_expm: updated alpha, beta, v, norm
else fallback pure-Python
matrix_expm-->>matrix_expm: perform Python Lanczos steps (in-place)
end
matrix_expm->>LinearAlgebra: form tridiagonal, diagonalize/check residual
LinearAlgebra-->>matrix_expm: eigenpairs / convergence flag
alt converged or max reached
matrix_expm->>LinearAlgebra: compute Krylov projection via _compute_krylov_result
LinearAlgebra-->>Caller: projected vector (expm action)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/core/methods/test_matrix_exponential.py (1)
8-29:⚠️ Potential issue | 🟡 MinorUpdate the module docstring to reflect removal of
lanczos_iterationtests.✏️ Suggested fix
-This module provides unit tests for the internal functions `lanczos_iteration` and `expm_krylov`, -which are utilized in YAQS for efficient computation of matrix exponentials. +This module provides unit tests for the internal function `expm_krylov`, +which is utilized in YAQS for efficient computation of matrix exponentials. @@ -The tests verify that: -- The Lanczos iteration correctly generates orthonormal bases and respects expected shapes. -- Early termination of the Lanczos iteration occurs appropriately when convergence conditions are met. -- Krylov subspace approximations to matrix exponentials match exact computations when the subspace dimension - equals the full space, and remain within acceptable error bounds for smaller subspace dimensions. +The tests verify that: +- Krylov subspace approximations to matrix exponentials match exact computations when the subspace dimension + equals the full space, and remain within acceptable error bounds for smaller subspace dimensions.
🤖 Fix all issues with AI agents
In `@CHANGELOG.md`:
- Line 14: Update the CHANGELOG.md entry by removing the backticks around the PR
reference so the PR link resolves correctly; change the line "- improved Lanczos
iteration and introduced numba usage leading to significant speedup ([`#310`])
([**@aaronleesander**])" to use an unescaped PR reference like "([`#310`])" (and
keep the author token unchanged) so the entry matches the project changelog
template and the PR number resolves.
In `@pyproject.toml`:
- Line 52: The pyproject.toml dependency "numba>=0.59.0" is incompatible with
Python 3.13/3.14; update the dependency declaration to either bump the minimum
to a release that provides wheels for newer Pythons (e.g., change
"numba>=0.59.0" to "numba>=0.63.0") or make it conditional with an environment
marker (e.g., keep "numba>=0.59.0; python_version<\"3.13\"" and treat 3.13+ as
optional/conditional), ensuring the change is applied to the dependency entry
named "numba" in pyproject.toml.
In `@src/mqt/yaqs/core/methods/matrix_exponential.py`:
- Around line 40-64: The docstring header "Adaptive Step Size Control" is
misleading because the function matrix_exponential adapts the number of Lanczos
iterations, not any time step; update the header to something like "Adaptive
Iteration Control" (or "Adaptive Iteration Count") and adjust the one-sentence
description below it to state that the algorithm adapts the iteration count
until the estimated error is below tol or max_lanczos_iterations is reached so
the docstring for matrix_exponential accurately reflects its behavior.
- Around line 74-96: The Lanczos basis array v is always allocated C-contiguous
but the numba kernel expects column-contiguous (Fortran) for best performance;
change the allocation to pick order = "F" when use_numba is True (keep "C"
otherwise) and then create v = np.zeros((vec.size, m_max), dtype=np.complex128,
order=order) so that the lanczos_numba path gets an F-ordered basis; ensure
subsequent uses (v0 = vec / vec_norm; v[:, 0] = v0) remain valid with either
memory order.
In `@src/mqt/yaqs/core/methods/tdvp_numba.py`:
- Around line 69-86: The code contains unnecessary noqa comments suppressing
PLR1702 on the outer loops (the lines with "for o in range(o_dim): # noqa:
PLR1702" and the later "for o in range(o_dim): # noqa: PLR1702"); remove these
unused "# noqa: PLR1702" directives so ruff no longer sees suppressed
rules—leave the loops and their prange annotations (e.g., "for a in
prange(a_in)" and "for aa in prange(a_out)") intact and only delete the trailing
noqa comments.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Aaron Sander <61705296+aaronleesander@users.noreply.github.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Aaron Sander <61705296+aaronleesander@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@src/mqt/yaqs/core/methods/tdvp_numba.py`:
- Around line 145-156: The p and w loops should be fused and parallelized
instead of using prange on the inner u loop: replace the nested for p in
range(p_dim)/for w in range(w_dim) with a single prange over fused_row in
prange(p_dim * w_dim) (or similar) and compute p = fused_row // w_dim and w =
fused_row % w_dim to derive row_idx; keep the inner u/v/a loops sequential and
remove prange from u. Update references to row_idx, left_env, right_env, out,
and any numba typing/annotations to reflect the new loop shape so the outer
iteration is parallelized for better thread utilization.
Signed-off-by: Aaron Sander <61705296+aaronleesander@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In @.pre-commit-config.yaml:
- Around line 60-64: The pre-commit hook with id "uv-lock" and entry "uv lock"
is failing because pre-commit appends matched filenames by default; update the
hook configuration to include pass_filenames: false so pre-commit will not pass
file paths to the "uv lock" command (which expects no filename arguments) and
will run the lock operation for the whole repo.
- Around line 65-69: Update the pre-commit hook with id "ty-check" (name "ty
check", entry "uv run ty check") to use the forward-compatible runner by
replacing language: system with language: unsupported; keep require_serial: true
and other fields unchanged so the hook behavior remains identical while avoiding
future deprecation.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/digital/test_digital_tjm.py (1)
454-520:⚠️ Potential issue | 🟡 MinorUpdate long‑range test docstring and consider preserving a full‑coverage variant.
The docstring still states 10 layers (0..9) and tolerance 0.1, but the test now uses 2 layers and tol 0.15. Also, reducing layers/trajectories can weaken regression sensitivity; consider a
pytest.mark.slow(or parametrized fast/slow) variant to keep the full‑coverage version in CI/nightly.✏️ Proposed docstring fix
- - 4 qubits, periodic Ising single timestep dt=0.1 composed into 10 layers (with SAMPLE_OBSERVABLES barriers) + - 4 qubits, periodic Ising single timestep dt=0.1 composed into 2 layers (with SAMPLE_OBSERVABLES barriers) @@ - - Compare per-layer Z expectations (qubits 0..3, layers 0..9) against hardcoded Qiskit density-matrix - - Tolerance 0.1 + - Compare per-layer Z expectations (qubits 0..3, layers 0..1) against hardcoded Qiskit density-matrix + - Tolerance 0.15
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@tests/core/methods/test_matrix_exponential.py`:
- Line 168: The statement np.ones(size, dtype=complex) in
tests/core/methods/test_matrix_exponential.py is unused and should be removed;
locate the stray np.ones(size, dtype=complex) line (referencing the size
variable in the test) and delete it so the test no longer contains an unused
array allocation.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Aaron Sander <61705296+aaronleesander@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@tests/core/methods/test_matrix_exponential.py`:
- Around line 166-192: Remove the duplicate assignment to size (the repeated
"size = 10") so size is only defined once before constructing lanczos_mat, and
delete the unused "# noqa: SLF001" comment on the call to
matrix_exponential._compute_krylov_result; ensure the test still calls
matrix_exponential._compute_krylov_result(alpha, beta, lanczos_mat, nrm, dt) and
keeps the mock_eigh assertions unchanged.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Aaron Sander <61705296+aaronleesander@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@tests/core/methods/test_matrix_exponential.py`:
- Around line 41-49: Rename the local variable lanczos_iterations to
max_lanczos_iterations to match the new expm_krylov parameter name; update the
assignment (previously "lanczos_iterations = 2") and the call site where
expm_krylov is invoked (expm_krylov(...,
max_lanczos_iterations=lanczos_iterations)) so it passes
max_lanczos_iterations=max_lanczos_iterations, ensuring the variable name
matches the parameter used by matrix_free_operator, v, and dt in
test_matrix_exponential.py.
- Around line 100-129: Update the test_expm_krylov_numba_execution to explicitly
assert the Numba path runs by spying on a Numba-specific helper: use
pytest.importorskip("mqt.yaqs.core.methods.matrix_exponential.lanczos_numba")
(or skip if unavailable), then patch the helper function used by expm_krylov
(e.g., lanczos_numba.orthogonalize_step or the actual helper symbol imported in
matrix_exponential) with unittest.mock.patch as a Mock and assert it was called
when calling expm_krylov with NUMBA_THRESHOLD patched to 50; keep the existing
numerical assertion but add the mock call assert to ensure the Numba branch
executed.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/core/methods/test_matrix_exponential.py (1)
10-11:⚠️ Potential issue | 🟡 MinorOutdated docstring reference to removed function.
The docstring mentions
lanczos_iterationas a tested function, but according to the PR changes, this function has been removed from the public API and replaced by the adaptive inline logic inexpm_krylov. Update the docstring to reflect the current API.📝 Suggested fix
-This module provides unit tests for the internal functions `lanczos_iteration` and `expm_krylov`, +This module provides unit tests for the `expm_krylov` function and its internal helpers, which are utilized in YAQS for efficient computation of matrix exponentials.
🤖 Fix all issues with AI agents
In `@tests/core/methods/test_matrix_exponential.py`:
- Around line 103-106: Remove the duplicated comment and the unused noqa
directive in the test import block: delete the repeated "Verify Numba module is
available" line and remove the "# noqa: PLC0415" trailing the "from
mqt.yaqs.core.methods.lanczos_numba import orthogonalize_step" import so the
block only keeps a single descriptive comment, the
pytest.importorskip("mqt.yaqs.core.methods.lanczos_numba") call, and the clean
import of orthogonalize_step.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This pull request makes several changes to how the Lanczos iteration works, particularly using an adaptive stopping, as well as introduces numba for JIT computations at critical points.
This was benchmarked across several instances and leads to speedups between 30% and 30x for different applications.
Due to numba's requirements, testing of x86 macOS was dropped in this PR as well.
Checklist: