Skip to content

Support time dependent collapse operators in mesolve - #43

Open
Noro-Official wants to merge 11 commits into
EQuS:mainfrom
Noro-Official:feature/time-dependent-collapse
Open

Noro-Official wants to merge 11 commits into
EQuS:mainfrom
Noro-Official:feature/time-dependent-collapse

Conversation

@Noro-Official

Copy link
Copy Markdown

JAXQuantum already allows the Hamiltonian in mesolve to 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 mesolve and mesolve_result to accept a callable c_ops(t) that returns a Qarray stack of collapse operators.

The implementation follows the existing time-dependent Hamiltonian interface. When c_ops is 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 mesolve API? And is there a cleaner way to handle the initial operator evaluation without complicating the static path?

@Noro-Official
Noro-Official marked this pull request as ready for review November 9, 2026 06:32
@Noro-Official
Noro-Official marked this pull request as draft November 9, 2026 06:34
@Noro-Official
Noro-Official marked this pull request as ready for review November 10, 2026 06:07

This branch has not been deployed

No deployments
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