Autodiff recipes

Author

Thomas Ogden

PropagationSolver.to_autodiff() returns a differentiable copy of a propagation problem: the propagated field becomes a pure function of a dictionary of parameters, so any loss of the field can be differentiated exactly with JAX. This page is a set of short recipes for the parts of the API you will reach for most. It needs the jax extra:

pip install 'clerq[jax]'

Setup: 64-bit precision

JAX computes in 32-bit by default, which is not enough for the master equation. clerq never changes JAX’s global configuration for you; it raises a RuntimeError if 64-bit mode is off. Enable it before creating any JAX arrays:

import jax
jax.config.update("jax_enable_x64", True)

or set the environment variable JAX_ENABLE_X64=1.

import json

import jax
import jax.numpy as jnp
import numpy as np

jax.config.update("jax_enable_x64", True)

from clerq import PropagationSolver

We’ll use a small ladder problem with two fields: a weak Gaussian probe on \(|0\rangle \leftrightarrow |1\rangle\) and a square coupling pulse on \(|1\rangle \leftrightarrow |2\rangle\).

problem = {
    "system": {
        "num_states": 3,
        "energies": [0.0, 0.0, 0.0],
        "decays": [{"channels": [[0, 1]], "rate": 1.0}],
        "fields": [
            {
                "coupled_levels": [[0, 1]],
                "rabi_freq": 0.5,
                "rabi_freq_t_func": "gaussian",
                "rabi_freq_t_args": {"ampl": 1.0, "centre": 0.0, "fwhm": 0.5},
            },
            {
                "coupled_levels": [[1, 2]],
                "rabi_freq": 1.0,
                "rabi_freq_t_func": "square",
                "rabi_freq_t_args": {"ampl": 1.0, "on": -1.0, "off": 1.0},
            },
        ],
        "initial_state": [1.0, 0.0, 0.0],
    },
    "t_min": -2.0,
    "t_max": 2.0,
    "t_steps": 80,
    "z_min": 0.0,
    "z_max": 1.0,
    "z_steps": 20,
    "interaction_strengths": [10.0, 10.0],
}

solver = PropagationSolver.from_json_str(json.dumps(problem))
ad = solver.to_autodiff()

ad is an AutodiffPropagation. Options such as step, strategy, rtol and max_steps are passed to to_autodiff() and are covered below.

The parameter dictionary

Parameters are a dictionary; every key is optional and defaults to the value in the problem. default_params() shows them:

ad.default_params()

{'pulse_args': [{'ampl': 1.0, 'centre': 0.0, 'fwhm': 0.5},
  {'ampl': 1.0, 'on': -1.0, 'off': 1.0}],
 'rabi_freq': array([0.5, 1. ]),
 'interaction_strengths': array([10., 10.]),
 'detunings': array([0., 0.]),
 'decay_scale': 1.0,
 'decay_rates': array([1.])}

Key	Shape	Meaning
`pulse_args`	list, one dict per field	arguments of each field’s analytic pulse (`ampl`, `centre`, `fwhm`/`width`, `n_pi`, `on`, `off`)
`rabi_freq`	`(n_fields,)`	the `rabi_freq` multiplying each pulse
`interaction_strengths`	`(n_fields,)`	the coupling strength of each field to the medium
`detunings`	`(n_fields,)`	the fixed detuning of each field
`decay_scale`	scalar	multiplies every decay rate
`decay_rates`	`(n_decays,)`	the `rate` of each entry of the system’s `decays` (non-negative); the effective rate is `decay_scale * decay_rates[i]`
`omega0`	`(n_fields, n_t)`	sampled input fields on `tlist`, used instead of the analytic pulses

Propagating

propagate(params) returns the field \(\Omega(z, t)\) with shape (n_fields, z_steps + 1, n_t). Pass nothing for the problem’s own parameters, or any subset of keys to override them; everything else keeps its default.

omega = ad.propagate()
omega_strong = ad.propagate({"rabi_freq": jnp.array([0.5, 2.0])})
print(omega.shape, omega.dtype)

(2, 21, 81) complex128

Pulse arguments are merged per field, so you only name the ones you change. A field’s alternative arguments replace one another: giving n_pi drops the default ampl.

wider = ad.propagate({"pulse_args": [{"fwhm": 1.0}, {}]})

Gradients with respect to any subset

value_and_grad(loss, params) differentiates loss(Omega_zt) with respect to the parameters you pass, and returns gradients with the same structure. Here, the gradient of the output probe intensity with respect to the coupling rabi_freq and the probe interaction_strengths only:

def loss(omega):
    """Negative probe energy at the end of the medium."""
    return -jnp.sum(jnp.abs(omega[0, -1]) ** 2)


params = {
    "rabi_freq": jnp.array([0.5, 1.0]),
    "interaction_strengths": jnp.array([10.0, 10.0]),
}
value, grads = ad.value_and_grad(loss, params)
print(value)
grads

-0.0024447873079130617

{'interaction_strengths': Array([4.52445709e-04, 4.38689474e-06], dtype=float64),
 'rabi_freq': Array([-0.00932889, -0.00800311], dtype=float64)}

With no params, every differentiable parameter of the problem is used.

Holding values constant with `fixed`

fixed= takes parameters with the same structure but does not differentiate them. Use it to evaluate at a different operating point, or to hold a parameter the optimiser should not touch:

value, grads = ad.value_and_grad(
    loss,
    {"rabi_freq": jnp.array([0.5, 1.0])},
    fixed={"detunings": jnp.array([0.0, 0.3]), "decay_scale": 2.0},
)
print(value)
grads

-0.0010771536965100022

{'rabi_freq': Array([-0.00415884, -0.00493268], dtype=float64)}

Units

All parameters are cyclic, exactly as in the problem definition (rabi_freq in \(2\pi\,\mathrm{MHz}\) means you pass the cyclic number, and detunings and omega0 follow suit). The returned field is in the layout and units of PropagationSolver.Omegas_zt: angular frequency. Divide by \(2\pi\) to get the cyclic Rabi frequency:

solver.solve()
print(np.abs(omega - solver.Omegas_zt).max() / np.abs(solver.Omegas_zt).max())  # agrees with solve() to solver tolerance
rabi_cyclic = omega / (2 * np.pi)

0.8408450498108503

/var/folders/41/hkxt57911pzcfzrv2p66w1z80000gn/T/ipykernel_39902/795264175.py:1: UserWarning: The z-step may be too coarse: estimated accumulated relative error is 12.2% (tolerance 10.0%). Try z_steps_inner=4 (currently 2) or z_steps_inner='auto'.
  solver.solve()
/var/folders/41/hkxt57911pzcfzrv2p66w1z80000gn/T/ipykernel_39902/795264175.py:1: UserWarning: The time grid may be too coarse: the fields or coherences are interpolated from every other time sample with a relative error of 48.3% (tolerance 10.0%). Try t_steps=387 (currently 80), and repeat if the warning remains; suggest_t_steps() gives an estimate from the problem's frequency scales.
  solver.solve()
/var/folders/41/hkxt57911pzcfzrv2p66w1z80000gn/T/ipykernel_39902/795264175.py:2: DeprecationWarning: PropagationSolver.Omegas_zt is deprecated and will be removed in a future release; use rabi_freqs_zt (cyclic) instead.
  print(np.abs(omega - solver.Omegas_zt).max() / np.abs(solver.Omegas_zt).max())  # agrees with solve() to solver tolerance

Reverse and forward mode

The default mode="reverse" costs roughly five propagations however many parameters you differentiate. mode="forward" costs more as parameters are added, so it is mainly useful as a cross-check or when there are very few parameters.

v_rev, g_rev = ad.value_and_grad(loss, params, mode="reverse")
v_fwd, g_fwd = ad.value_and_grad(loss, params, mode="forward")
print(np.allclose(g_rev["rabi_freq"], g_fwd["rabi_freq"]))

True

Memory strategies

Reverse mode has to remember something about the forward pass. The strategy option, given to to_autodiff(), trades memory for recomputation:

`strategy`	Stores	Cost
`"s2"`	every \(z\)-slice’s residuals	least recomputation, most memory
`"s3"` (default)	slice boundaries; recomputes inside a slice on the backward pass	a moderate recomputation overhead
`"s4"`	as `s3`, and limits the ODE solver’s inner checkpoints to `checkpoints`	least memory; needs `checkpoints=N`

max_steps is the most steps one slice solve may take: the solve raises if it needs more. In reverse mode the default checkpoint count is sized from it, so a smaller max_steps lowers memory use on large problems.

ad_s2 = solver.to_autodiff(strategy="s2")
ad_s4 = solver.to_autodiff(strategy="s4", checkpoints=16, max_steps=2000)

g_s2 = ad_s2.value_and_grad(loss, params)[1]
g_s4 = ad_s4.value_and_grad(loss, params)[1]
print(np.allclose(g_s2["rabi_freq"], g_s4["rabi_freq"]))

True

All strategies give the same gradients; only the memory and run time differ. checkpoints is an error with any strategy other than "s4".

`jax.jit` and reusing the loss function

propagate and value_and_grad are already jit-compiled, and compose with jax.jit, jax.vmap and optimisers such as Optax. Two things to know:

value_and_grad compiles once per loss function object (and mode). Define the loss once and reuse it. Only the last 8 are kept (an LRU, AutodiffPropagation.GRAD_CACHE_SIZE), so a new lambda on every call recompiles every time.
Captured arrays stay alive as long as the compiled function does. Call ad.clear_cache() to drop the compiled programs and release them.

# A loss with a captured target is compiled once, then reused for every step.
target = jnp.abs(omega[0, -1]) ** 2


def make_loss(target):
    def mismatch(omega):
        return jnp.sum((jnp.abs(omega[0, -1]) ** 2 - target) ** 2)

    return mismatch


mismatch = make_loss(target)  # build once...
for rabi in (0.8, 0.9, 1.0):  # ...and reuse
    p = {"rabi_freq": jnp.array([0.5, rabi])}
    v, _ = ad.value_and_grad(mismatch, p)
    print(rabi, float(v))

ad.clear_cache()

0.8 3.089429785615958e-07
0.9 9.875228538778031e-08
1.0 7.11648226219785e-22

A gradient step is then a short loop around value_and_grad.

Limits

Co-propagating fields only. to_autodiff() raises NotImplementedError for a problem with counter-propagating fields.
Discontinuous edge times are held fixed. The field jumps at the on and off of square, the on of exp_decay and the off of exp_rise, so its derivative with respect to them is a delta function that autodiff would silently return as 0. value_and_grad refuses to differentiate them: pass them in fixed=, or use a continuous pulse. The ramp_on, ramp_off, ramp_onoff, ramp_offon and sin2 pulses are continuous, so their on and off can be differentiated.
Supported pulses are those ported to JAX; others raise NotImplementedError from to_autodiff().
Float64 is required (see above).

For example, asking for a square pulse’s edge gives a clear error rather than a silent zero:

try:
    ad.value_and_grad(loss, {"pulse_args": [{}, {"on": -1.0}]})
except ValueError as e:
    print(e)

pulse_args[1] contains ['on'], which place the edges of a 'square' pulse. The field is discontinuous there, so its derivative with respect to them is a delta function that autodiff cannot represent (it would return 0). Hold them fixed by passing them in fixed=..., or use a ramp pulse ('ramp_on', 'ramp_off', 'ramp_onoff' or 'ramp_offon') or 'sin2', pulses with smooth edges whose on/off times are differentiable.