Diffrax solver options¶
JAXQuantum uses Diffrax for Schrödinger and Lindblad equations. SolverOptions accepts native Diffrax objects, while sesolve_result and mesolve_result expose the complete diffrax.Solution. Legacy string-based options remain available with a FutureWarning during migration.
In [1]:
Copied!
import diffrax
import jax
import jax.numpy as jnp
import jaxquantum as jqt
psi0 = jqt.basis(2, 0)
H = 0.5 * jqt.sigmax()
times = jnp.linspace(0.0, 2.0, 41)
import diffrax
import jax
import jax.numpy as jnp
import jaxquantum as jqt
psi0 = jqt.basis(2, 0)
H = 0.5 * jqt.sigmax()
times = jnp.linspace(0.0, 2.0, 41)
/opt/hostedtoolcache/Python/3.11.16/x64/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
Defaults and native Diffrax configuration¶
The defaults are Tsit5, an adaptive PID controller with rtol=1e-7 and atol=1e-9, the first tlist interval as dt0, and at most 100,000 steps. Pass native objects to change any policy.
In [2]:
Copied!
options = jqt.SolverOptions(
solver=diffrax.Dopri5(),
stepsize_controller=diffrax.PIDController(rtol=1e-8, atol=1e-10),
dt0=None, # ask Diffrax to select the initial step
)
states = jqt.sesolve(H, psi0, times, solver_options=options)
states.data.shape
options = jqt.SolverOptions(
solver=diffrax.Dopri5(),
stepsize_controller=diffrax.PIDController(rtol=1e-8, atol=1e-10),
dt0=None, # ask Diffrax to select the initial step
)
states = jqt.sesolve(H, psi0, times, solver_options=options)
states.data.shape
0% |░░░░░░░░░░| [00:00<?, ?%/s]
0% |░░░░░░░░░░| [00:00<00:00, 426.16%/s]
100% |██████████| [00:00<00:00, 42577.44%/s]
Out[2]:
(41, 2)
Full solutions and custom saving¶
Use a result function for solver statistics, events, continuation, dense interpolation, or a SaveAt function that does not return quantum states.
In [3]:
Copied!
final_result = jqt.sesolve_result(
H,
psi0,
times,
solver_options=jqt.SolverOptions(saveat=diffrax.SaveAt(t1=True)),
)
final_result.stats, final_result.ys.shape
final_result = jqt.sesolve_result(
H,
psi0,
times,
solver_options=jqt.SolverOptions(saveat=diffrax.SaveAt(t1=True)),
)
final_result.stats, final_result.ys.shape
0% |░░░░░░░░░░| [00:00<?, ?%/s]
2% |░░░░░░░░░░| [00:00<00:00, 4354.55%/s]
100% |██████████| [00:00<00:00, 63996.09%/s]
Out[3]:
({'max_steps': 100000,
'num_accepted_steps': Array(8, dtype=int64, weak_type=True),
'num_rejected_steps': Array(0, dtype=int64, weak_type=True),
'num_steps': Array(8, dtype=int64, weak_type=True)},
(1, 2))
In [4]:
Copied!
population_result = jqt.sesolve_result(
H,
psi0,
times,
solver_options=jqt.SolverOptions(
saveat=diffrax.SaveAt(
ts=times, fn=lambda t, y, args: jnp.abs(y[1]) ** 2
)
),
)
population_result.ys.shape
population_result = jqt.sesolve_result(
H,
psi0,
times,
solver_options=jqt.SolverOptions(
saveat=diffrax.SaveAt(
ts=times, fn=lambda t, y, args: jnp.abs(y[1]) ** 2
)
),
)
population_result.ys.shape
0% |░░░░░░░░░░| [00:00<?, ?%/s]
2% |░░░░░░░░░░| [00:00<00:00, 3986.98%/s]
100% |██████████| [00:00<00:00, 63879.13%/s]
Out[4]:
(41,)
Events and progress meters¶
event accepts a native diffrax.Event. Existing calls retain JAXQuantum's progress bar; use None to disable it, "default" for the JAXQuantum bar, or pass any diffrax.AbstractProgressMeter instance.
In [5]:
Copied!
event_result = jqt.sesolve_result(
H,
psi0,
times,
solver_options=jqt.SolverOptions(
event=diffrax.Event(lambda t, y, args, **kwargs: t >= 0.75),
saveat=diffrax.SaveAt(t1=True),
),
)
assert event_result.ts[-1] < times[-1]
jaxquantum_bar = jqt.SolverOptions(progress_meter="default")
text_bar = jqt.SolverOptions(progress_meter=diffrax.TextProgressMeter())
event_result = jqt.sesolve_result(
H,
psi0,
times,
solver_options=jqt.SolverOptions(
event=diffrax.Event(lambda t, y, args, **kwargs: t >= 0.75),
saveat=diffrax.SaveAt(t1=True),
),
)
assert event_result.ts[-1] < times[-1]
jaxquantum_bar = jqt.SolverOptions(progress_meter="default")
text_bar = jqt.SolverOptions(progress_meter=diffrax.TextProgressMeter())
0% |░░░░░░░░░░| [00:00<?, ?%/s]
2% |░░░░░░░░░░| [00:00<00:00, 4781.47%/s]
100% |██████████| [00:00<00:00, 68422.58%/s]
Differentiation¶
jax.grad requests a derivative. Leaving adjoint=None lets Diffrax choose its recommended reverse-mode method. Specify ForwardMode only for jax.jvp or jax.jacfwd.
In [6]:
Copied!
def excited_population(scale, adjoint=None):
result = jqt.sesolve_result(
scale * jqt.sigmax(),
psi0,
times,
solver_options=jqt.SolverOptions(
adjoint=adjoint, saveat=diffrax.SaveAt(t1=True)
),
)
return jnp.abs(result.ys[-1, 1]) ** 2
reverse_derivative = jax.grad(excited_population)(0.5)
forward_derivative = jax.jacfwd(
lambda scale: excited_population(scale, diffrax.ForwardMode())
)(0.5)
reverse_derivative, forward_derivative
def excited_population(scale, adjoint=None):
result = jqt.sesolve_result(
scale * jqt.sigmax(),
psi0,
times,
solver_options=jqt.SolverOptions(
adjoint=adjoint, saveat=diffrax.SaveAt(t1=True)
),
)
return jnp.abs(result.ys[-1, 1]) ** 2
reverse_derivative = jax.grad(excited_population)(0.5)
forward_derivative = jax.jacfwd(
lambda scale: excited_population(scale, diffrax.ForwardMode())
)(0.5)
reverse_derivative, forward_derivative
0% |░░░░░░░░░░| [00:00<?, ?%/s]
2% |░░░░░░░░░░| [00:00<00:00, 3953.91%/s]
100% |██████████| [00:00<00:00, 49362.17%/s]
0% |░░░░░░░░░░| [00:00<?, ?%/s]
2% |░░░░░░░░░░| [00:00<00:00, 4843.31%/s]
100% |██████████| [00:00<00:00, 68579.20%/s]
Out[6]:
(Array(1.81859483, dtype=float64, weak_type=True), Array(1.81859483, dtype=float64))
Continuing a solve¶
Ask SaveAt for the internal states, then pass them into the next solve. This avoids reinitializing the solver and controller.
In [7]:
Copied!
continuation_saveat = diffrax.SaveAt(
t1=True, solver_state=True, controller_state=True, made_jump=True
)
first = jqt.sesolve_result(
H,
psi0,
jnp.array([0.0, 1.0]),
solver_options=jqt.SolverOptions(saveat=continuation_saveat),
)
continued_state = jqt.Qarray.create(first.ys[-1], dims=psi0.dims)
second = jqt.sesolve_result(
H,
continued_state,
jnp.array([1.0, 2.0]),
solver_options=jqt.SolverOptions(
saveat=diffrax.SaveAt(t1=True),
solver_state=first.solver_state,
controller_state=first.controller_state,
made_jump=first.made_jump,
),
)
assert jnp.allclose(second.ys[-1], final_result.ys[-1], atol=1e-7)
continuation_saveat = diffrax.SaveAt(
t1=True, solver_state=True, controller_state=True, made_jump=True
)
first = jqt.sesolve_result(
H,
psi0,
jnp.array([0.0, 1.0]),
solver_options=jqt.SolverOptions(saveat=continuation_saveat),
)
continued_state = jqt.Qarray.create(first.ys[-1], dims=psi0.dims)
second = jqt.sesolve_result(
H,
continued_state,
jnp.array([1.0, 2.0]),
solver_options=jqt.SolverOptions(
saveat=diffrax.SaveAt(t1=True),
solver_state=first.solver_state,
controller_state=first.controller_state,
made_jump=first.made_jump,
),
)
assert jnp.allclose(second.ys[-1], final_result.ys[-1], atol=1e-7)
0% |░░░░░░░░░░| [00:00<?, ?%/s]
0% |░░░░░░░░░░| [00:00<?, ?%/s]
100% |██████████| [00:00<00:00, 67606.45%/s]
0% |░░░░░░░░░░| [00:00<?, ?%/s]
0% |░░░░░░░░░░| [00:00<?, ?%/s]
100% |██████████| [00:00<00:00, 67923.95%/s]
Open-system evolution¶
mesolve_result has the same interface and returns Diffrax diagnostics alongside the density matrices.
In [8]:
Copied!
collapse_ops = jqt.Qarray.from_list([0.2 * jqt.sigmam()])
master_result = jqt.mesolve_result(
0.5 * jqt.sigmaz(),
jqt.basis(2, 1).to_dm(),
times,
c_ops=collapse_ops,
)
master_result.stats, jnp.trace(master_result.ys[-1])
collapse_ops = jqt.Qarray.from_list([0.2 * jqt.sigmam()])
master_result = jqt.mesolve_result(
0.5 * jqt.sigmaz(),
jqt.basis(2, 1).to_dm(),
times,
c_ops=collapse_ops,
)
master_result.stats, jnp.trace(master_result.ys[-1])
0% |░░░░░░░░░░| [00:00<?, ?%/s]
2% |░░░░░░░░░░| [00:00<00:00, 4182.59%/s]
100% |██████████| [00:00<00:00, 61212.84%/s]
Out[8]:
({'max_steps': 100000,
'num_accepted_steps': Array(3, dtype=int64, weak_type=True),
'num_rejected_steps': Array(0, dtype=int64, weak_type=True),
'num_steps': Array(3, dtype=int64, weak_type=True)},
Array(1.+0.j, dtype=complex128))