autodiff.AutodiffPropagation
autodiff.AutodiffPropagation(
setup,
*,
step='ab',
strategy='s3',
checkpoints=None,
rtol=None,
atol=None,
max_steps=100000,
n_steps=None,
solver=None,
)A differentiable version of a :class:~clerq.PropagationSolver problem.
.. versionadded:: 0.14.0
Build it with :meth:clerq.PropagationSolver.to_autodiff. It reproduces PropagationSolver.solve as a pure function of a parameter pytree, so the propagated field and any loss of it can be differentiated with respect to the parameters, with reverse- or forward-mode autodiff, and composed with jax.jit, jax.vmap and optimisers.
Parameters are a dict; every key is optional and defaults to the value in the problem (see :meth:default_params):
pulse_args list, one dict per field, of the arguments of its analytic input pulse (ampl, centre, fwhm/width, n_pi, on, off, order, tau). The edge time at which a pulse is discontinuous (the on/off of square, the on of exp_decay, the off of exp_rise) can be set but not differentiated (the field jumps there, so its derivative is a delta function that autodiff would return as 0); value_and_grad raises if asked, and holds them fixed by default. To differentiate the edge times use a pulse that is continuous there (ramp_on, ramp_off, ramp_onoff, ramp_offon or sin2). Calling jax.grad yourself through :meth:propagate with those traced gives exactly that silent 0, so keep them out of the differentiated parameters. rabi_freq array (n_fields,), cyclic rabi_freq multiplying each pulse. interaction_strengths array (n_fields,), as the interaction_strengths argument. detunings array (n_fields,), cyclic fixed detuning of each field. decay_scale scalar multiplying every decay rate. decay_rates array (n_decays,), the cyclic rate of each entry of the system’s decays (spontaneous decay, ground-state decoherence, or pure dephasing through a [a, a] channel). The effective rate of a channel is decay_scale * decay_rates[i]. Differentiable at any value, including zero, so a rate can be fitted or optimised from nothing. Rates must be non-negative: a concrete negative value raises, but an optimiser can still step below zero (a negative rate is gain, and the field runs away), so optimise an unconstrained variable and map it, for example with :func:clerq.autodiff.controls.bounded. Empty for a system without decays. b_field scalar, the DC B-field in natural units (mu_B B / h in the frequency unit of the energies), so the Zeeman shift of the levels can be differentiated: H_0 = H_base + 2 pi B Z. Requires the system to have a b_field block giving g_F and m_F (it is then in the default parameters; without one it raises ValueError). omega0 array (n_fields, n_t), the input field samples on tlist as a cyclic Rabi frequency, the same unit as rabi_freq. If given, the first slice is driven by their cubic spline instead of the analytic pulse (and pulse_args / rabi_freq are unused).
All parameters are cyclic (like the problem definition), and so is the propagated field. It is returned in the layout and units of PropagationSolver.rabi_freqs_zt: shape (n_fields, z_steps + 1, n_t), complex, cyclic Rabi frequency, the same unit as rabi_freq and omega0.
Attributes
| Name | Type | Description |
|---|---|---|
| GRAD_CACHE_SIZE | how many compiled value_and_grad programs (one per loss function and mode) are kept; the least recently used is dropped. |
Note
Counter-propagating fields are not supported. The result matches PropagationSolver for every z_steps_inner; versions of clerq before #328 stored Adams-Bashforth slices one inner step early when z_steps_inner > 1 and will not match there.
Precision
The backend computes in the precision JAX is configured for: float32 by default, float64 with jax.config.update("jax_enable_x64", True). Float64 is the reference. Float32 agrees with it to 1e-6 to 2e-5 of the peak on the fields and 6e-7 to 5e-5 on gradients on the benchmark problems, with the default tolerances (rtol=1e-6, atol=1e-9, against 1e-8 and 1e-10 in float64). Use float64 for finite-difference checks, for tolerances tighter than about 1e-6, and when the loss is very close to zero, where float32 gradients are noisier.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| setup | PropagationSetup | the problem, from setup_from_solver. |
required |
| step | str | "ab" (Adams-Bashforth, default) or "euler". |
'ab' |
| strategy | str | "s3" (default) recomputes a slice on the backward pass; "s2" stores every slice; "s4" additionally limits the solver’s inner checkpoints to checkpoints. |
's3' |
| checkpoints | int | None | inner checkpoint count, required (a positive integer) for strategy="s4" and not allowed with the other strategies. |
None |
| rtol | float | None | relative tolerance of the adaptive step-size control; default 1e-8 in float64 and 1e-6 in float32. | None |
| atol | float | None | absolute tolerance of the adaptive step-size control; default 1e-10 in float64 and 1e-9 in float32. | None |
| max_steps | int | the most steps one slice solve may take (the solve errors if it needs more); reverse mode sizes its default checkpoint count from it, so a smaller value lowers memory use for large problems. | 100000 |
| n_steps | int | None | fixed-step solver with this many steps per slice (useful for finite-difference reference gradients); must not be commensurate with the time grid. | None |
| solver | Any | a diffrax solver, default Tsit5. |
None |
Methods
| Name | Description |
|---|---|
| clear_cache | Drop the compiled programs, releasing the loss functions (and the arrays |
| default_params | The parameters of the problem this object was built from. |
| from_solver | Build from a :class:~clerq.PropagationSolver. |
| minimize | Minimise loss(Omega_zt) over params with an optax optimiser. |
| propagate | Propagate the fields; a pure, jit-compiled function of params. |
| value_and_grad | Value and gradient of loss(Omega_zt) with respect to params. |
clear_cache
autodiff.AutodiffPropagation.clear_cache()Drop the compiled programs, releasing the loss functions (and the arrays they capture) that value_and_grad keeps alive.
default_params
autodiff.AutodiffPropagation.default_params()The parameters of the problem this object was built from.
Every value is a float, whatever literal the problem definition used. b_field is included only if the system has a b_field block.
from_solver
autodiff.AutodiffPropagation.from_solver(solver, **options)Build from a :class:~clerq.PropagationSolver.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| solver | Any | the problem to differentiate. | required |
| **options | Any | the constructor options of this class (step, strategy, checkpoints, rtol, atol, max_steps, n_steps, solver). |
{} |
minimize
autodiff.AutodiffPropagation.minimize(
loss,
params=None,
*,
fixed=None,
optimizer=None,
learning_rate=None,
steps=100,
gtol=None,
)Minimise loss(Omega_zt) over params with an optax optimiser.
Each step is one :meth:value_and_grad call and one optimiser update, so the loss is compiled once, on the first step (plus one final evaluation). Needs optax, which the jax extra installs.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| loss | Callable | scalar-valued function of the propagated field, as for :meth:value_and_grad; reuse one function object. |
required |
| params | dict | None | starting point and the parameters to optimise (any subset of the keys). Default: every differentiable parameter of the problem. | None |
| fixed | dict | None | parameters held constant, as for :meth:value_and_grad. |
None |
| optimizer | Any | an optax.GradientTransformation (for example optax.adam(1e-2) or optax.lbfgs()). Pass the same object to reuse its compiled update between calls. Default optax.adam(learning_rate). Parameters differ in scale (a pulse area is of order 1, an interaction strength of order 1e3), so the right learning_rate depends on what is being optimised; rescale the parameters inside loss if needed. |
None |
| learning_rate | float | None | step size of the default Adam optimiser (0.05 if not given); not allowed together with optimizer. |
None |
| steps | int | the most updates to take. | 100 |
| gtol | float | None | stop when the gradient’s global norm is below this (default: always take steps updates). |
None |
Returns
| Name | Type | Description |
|---|---|---|
| OptimizeResult | The An :class:~clerq.autodiff.OptimizeResult, at the last point reached. That is the |
|
| OptimizeResult | best one for a descending optimiser but need not be for one that | |
| OptimizeResult | overshoots; losses shows the path. |
Raises
| Name | Type | Description |
|---|---|---|
| ImportError | if optax is not installed. |
|
| RuntimeError | if JAX is not in 64-bit mode. | |
| ValueError | as for :meth:value_and_grad, or if steps is not a positive integer, or both optimizer and learning_rate are given. |
propagate
autodiff.AutodiffPropagation.propagate(params=None)Propagate the fields; a pure, jit-compiled function of params.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| params | dict | None | parameter pytree (see the class docstring); defaults to the problem’s. | None |
Returns
| Name | Type | Description |
|---|---|---|
| Any | Omega(z, t), shape (n_fields, z_steps + 1, n_t), complex. |
Raises
| Name | Type | Description |
|---|---|---|
| RuntimeError | if JAX is not in 64-bit mode. | |
| ValueError | for unknown keys or wrong shapes. |
value_and_grad
autodiff.AutodiffPropagation.value_and_grad(
loss,
params=None,
*,
fixed=None,
mode='reverse',
)Value and gradient of loss(Omega_zt) with respect to params.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| loss | Callable | scalar-valued function of the propagated field Omega_zt, shape (n_fields, z_steps + 1, n_t), complex. Reuse the same function object between calls to reuse the compiled program: it is compiled once per function object, and only the last GRAD_CACHE_SIZE (8) are kept, so a new lambda on every call recompiles every time. |
required |
| params | dict | None | the parameters to differentiate (any subset of the keys; the gradient has the same structure). Default: every differentiable parameter of the problem, i.e. default_params() without the edge times (on, off) of square pulses, which stay fixed because the field’s derivative with respect to a discontinuity’s location is not defined. |
None |
| fixed | dict | None | parameters held constant at the given values, with the same structure as params (for example the on/off of a square pulse, or a different target interaction strength). They are not differentiated. |
None |
| mode | str | "reverse" (cost ~5x one propagation regardless of the number of parameters) or "forward" (cost grows with the number of parameters; useful as a cross-check or for very few parameters). |
'reverse' |
Returns
| Name | Type | Description |
|---|---|---|
| tuple[Any, dict] | (value, grads) with grads of the same structure as params. |
Raises
| Name | Type | Description |
|---|---|---|
| RuntimeError | if JAX is not in 64-bit mode. | |
| ValueError | for an unknown mode, unknown keys, wrong shapes, or a request for the gradient of a square pulse’s edge times. |