Repository navigation
Support time dependent collapse operators in mesolve - #43
Open
Noro-Official wants to merge 11 commits into
Open
Noro-Official wants to merge 11 commits into
Noro-Official wants to merge 11 commits into
Conversation
Noro-Official
marked this pull request as ready for review
November 9, 2026 06:32
Noro-Official
marked this pull request as draft
November 9, 2026 06:34
Noro-Official
marked this pull request as ready for review
November 10, 2026 06:07
This branch has not been deployed
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
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.
JAXQuantum already allows the Hamiltonian in
mesolveto vary with time. Collapse operators, however, have to remain fixed. This is limiting when modeling transient relaxation or dephasing, where the dissipative rates can change substantially during a single evolution.This PR extends
mesolveandmesolve_resultto accept a callablec_ops(t)that returns aQarraystack of collapse operators.The implementation follows the existing time-dependent Hamiltonian interface. When
c_opsis callable, the solver evaluates the operators and their adjoints inside the Diffrax ODE function. Static operators continue through the original path, with their adjoints computed only once. I wanted to keep that distinction rather than introduce additional work into calls that do not need time dependence.The important constraint is that the operator structure must remain fixed. A callable can change the coefficients, but not the number of operators, their dimensions or their batch shapes. The implementation checks the initial return type and Hilbert-space dimensions. For sparse BCOO operators, the stack must also be constructed outside the traced callback; otherwise, reconstructing its sparsity pattern can trigger a JAX concretization error.
I added ten regression tests. Several compare the numerical evolution directly against analytical solutions for time-dependent amplitude damping and dephasing, including a case where both the Hamiltonian and collapse operators vary. The remaining tests cover static/callable equivalence, JIT compilation,
vmap, automatic differentiation, batching, sparse operators and invalid inputs.The full repository checks pass locally on macOS with Python 3.11.18 and JAX 0.10.2: 510 standard tests, 26 sharding tests, Ruff checks and the documentation build. I also added a runnable example to the Diffrax tutorial.
I was concerned that evaluating collapse operators within the ODE function might introduce unnecessary overhead, so I benchmarked the static path against an equivalent constant callable. The warmed-runtime ratio changed from 1.046 in a 20-iteration run to 0.942 over 100 iterations. These measurements do not establish a meaningful speedup or regression, but they also give no clear reason to restructure the solver for performance.
Two points remain open for review. Does a fixed-shape callable fit the intended
mesolveAPI? And is there a cleaner way to handle the initial operator evaluation without complicating the static path?