Designing a pulse with automatic differentiation

Author

Thomas Ogden

Requires clerq 0.14.0 or later (and the jax extra: pip install 'clerq[jax]').

PropagationSolver.to_autodiff() turns a propagation problem into a function that JAX can differentiate. Any scalar quantity you can compute from the propagated field then has an exact gradient with respect to the pulse, the medium and the detunings, for about the cost of five to six propagations however many parameters there are. That makes gradient-based design practical: instead of scanning a parameter and looking at the result, you follow the gradient.

This notebook designs the input pulse of a resonant two-level medium. The goal is to get as much of the pulse energy as possible through the medium. The self-induced transparency area theorem says that a pulse of area \(2\pi\) passes without loss, and the optimiser should end up near that on its own.

This needs the jax extra: pip install 'clerq[jax]'. JAX must run in 64-bit mode, which clerq.autodiff checks and refuses to change for you.

import jax
import jax.numpy as jnp
import numpy as np
import plotly.graph_objects as go

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

from clerq import PropagationSolver

The problem

A sech pulse enters a resonant two-level medium. The pulse is described, as in any clerq problem, by its shape (sech) and its parameters. Here the pulse area is given by n_pi, in units of \(\pi\), which is the number we will optimise. The grid is smaller than in the SIT notebook, because the optimiser propagates the pulse many times.

t_width = 0.001

def make_solver(n_pi):
    return PropagationSolver.from_json_str(f'''
    {{
      "system": {{
        "fields": [{{"coupled_levels": [[0, 1]], "rabi_freq_t_func": "sech",
          "rabi_freq_t_args": {{"n_pi": {n_pi}, "centre": 0.0, "width": {t_width}}}}}],
        "num_states": 2,
        "decays": [{{"channels": [[0, 1]], "rate": 1.0}}]
      }},
      "t_min": {-10 * t_width}, "t_max": {40 * t_width}, "t_steps": 200,
      "z_min": 0.0, "z_max": 1.0, "z_steps": 30,
      "interaction_strengths": [{0.6 / t_width}]
    }}
    ''')

solver = make_solver(1.6)
solver.solve(recalc=True);

Making it differentiable

to_autodiff() returns an object built from the same problem definition. Its propagate() method reproduces solve() as a pure function of a dictionary of parameters, in the same layout and units as rabi_freqs_zt (a cyclic Rabi frequency): field, then \(z\), then time.

ad = solver.to_autodiff()
omegas = ad.propagate()  # default parameters: the ones in the problem
print(omegas.shape, solver.rabi_freqs_zt.shape)
rel = np.abs(np.asarray(omegas) - solver.rabi_freqs_zt).max() / np.abs(solver.rabi_freqs_zt).max()
print(f"largest difference from solve(): {rel:.0e} of the peak field")
(1, 31, 201) (1, 31, 201)
largest difference from solve(): 2e-05 of the peak field

The two agree to a few parts in \(10^{5}\) of the peak field. The remaining difference is on the solve() side: the autodiff result does not change when its own tolerances are tightened further.

The parameters you can differentiate with respect to are the arguments of the input pulse (here n_pi, width, centre), rabi_freq, interaction_strengths, detunings, a decay_scale that multiplies every decay rate, and omega0, the sampled input field. Any subset works, and the gradient has the same structure as the parameters you pass.

A loss function

The loss is any scalar function of the field Omega[field, z, t]. We want the fraction of the pulse energy that gets through to the end of the medium, so we minimise minus the energy at the last slice divided by the energy at the first.

ts = jnp.asarray(solver.tlist)

def transmission(omegas):
    energy = jnp.trapezoid(jnp.abs(omegas[0]) ** 2, ts, axis=-1)  # per z slice
    return energy[-1] / energy[0]

def loss(omegas):
    return -transmission(omegas)

params = {"pulse_args": [{"n_pi": 1.6}]}
value, grads = ad.value_and_grad(loss, params)
print(f"transmission at n_pi = 1.6: {-float(value):.3f}")
print(f"d(loss)/d(n_pi) = {float(grads['pulse_args'][0]['n_pi']):+.3f}")
transmission at n_pi = 1.6: 0.865
d(loss)/d(n_pi) = -0.531

Checking the gradient

Check a gradient against finite differences before trusting an optimisation to it. Finite differences through an adaptive solver are noisy, because the step sizes change with the parameter, so the check uses the n_steps option, which takes a fixed number of steps per \(z\) slice. Choose a count that does not divide the output time grid evenly: clerq raises if it does, because of a known diffrax problem with step boundaries that coincide with the output times.

ad_fixed = solver.to_autodiff(n_steps=2003)

def fixed_loss(n_pi):
    return float(loss(ad_fixed.propagate({"pulse_args": [{"n_pi": n_pi}]})))

h = 1e-4
fd = (fixed_loss(1.6 + h) - fixed_loss(1.6 - h)) / (2 * h)
_, g_fixed = ad_fixed.value_and_grad(loss, params)
print(f"finite difference: {fd:+.6f}")
print(f"autodiff:          {float(g_fixed['pulse_args'][0]['n_pi']):+.6f}")
finite difference: -0.531123
autodiff:          -0.531123

What the loss looks like

With one parameter we can afford to see the whole landscape. Weak pulses are absorbed, so transmission is poor below about \(1.4\pi\). It then rises sharply, and peaks on a broad plateau, above 95% from about \(1.9\pi\) to \(2.65\pi\), before slowly falling again.

n_pi_grid = np.linspace(0.4, 3.2, 57)
T_grid = np.array([
    float(transmission(ad.propagate({"pulse_args": [{"n_pi": float(n)}]})))
    for n in n_pi_grid
])

fig = go.Figure(go.Scatter(x=n_pi_grid, y=T_grid, mode="lines", name="transmission"))
fig.update_layout(
    xaxis_title="input pulse area (units of π)",
    yaxis_title="transmitted energy fraction",
    title="Transmission through the medium",
)
fig.show(renderer="notebook_connected")

Following the gradient

ad.minimize(loss, params) runs an optax optimiser, Adam by default, on the loss. Each step calls value_and_grad once. The first call compiles the function for these array shapes and later calls reuse it, so loss is defined once above and reused. Pass your own optimiser with optimizer=, for example optax.adam(0.05) or optax.lbfgs(); optax is installed with the jax extra.

result = ad.minimize(loss, {"pulse_args": [{"n_pi": 1.6}]}, steps=40, learning_rate=0.05)
best = float(result.params["pulse_args"][0]["n_pi"])
T_start, T_best = -result.losses[0], -result.value
print(f"start: n_pi = 1.60, transmission {T_start:.3f}")
print(f"end:   n_pi = {best:.2f}, transmission {T_best:.3f}")
start: n_pi = 1.60, transmission 0.865
end:   n_pi = 2.33, transmission 0.968

The optimiser climbs the rise and settles on the plateau, a little above \(2\pi\). The area theorem’s \(2\pi\) is the result for an idealised, lossless medium, so the best area for this finite medium with decay need not be exactly \(2\pi\). result.losses holds the loss at the start and after every step, if you want to plot the convergence.

fig = go.Figure()
fig.add_trace(go.Scatter(x=n_pi_grid, y=T_grid, mode="lines", name="transmission"))
fig.add_trace(go.Scatter(x=[1.6, best], y=[T_start, T_best], mode="markers+text",
                         text=["start", "optimised"], textposition="bottom right",
                         name="optimiser", marker=dict(size=9)))
fig.update_layout(
    xaxis_title="input pulse area (units of π)",
    yaxis_title="transmitted energy fraction",
    title="The optimiser on the transmission curve",
)
fig.show(renderer="notebook_connected")

Checking the result with the ordinary solver

The optimised pulse is just a clerq problem with a different n_pi, so build it and solve it with PropagationSolver itself.

check = make_solver(best)
check.solve(recalc=True)
energy = np.trapezoid(np.abs(check.rabi_freqs_zt[0]) ** 2, check.tlist, axis=-1)
print(f"PropagationSolver transmission at n_pi = {best:.3f}: {energy[-1] / energy[0]:.4f}")
print(f"autodiff transmission:                      {T_best:.4f}")
PropagationSolver transmission at n_pi = 2.329: 0.9685
autodiff transmission:                      0.9685

What else you can differentiate

The same pattern works for the other parameters:

  • interaction_strengths, detunings and decay_scale (the medium and the detuning of each field);
  • rabi_freq and every pulse argument, for several co-propagating fields at once;
  • omega0, the input field sampled on tlist, which lets the optimiser shape the pulse freely. It is interpolated with the same cubic spline that clerq uses, so the result still matches PropagationSolver.

value_and_grad(..., mode="forward") uses forward-mode autodiff, useful as a cross-check or when there are very few parameters. propagate() is a pure function, so it composes with jax.jit, jax.vmap and optimisers such as optax.

Limits

  • Fields must be co-propagating; counter-propagating problems raise NotImplementedError.
  • The on and off edge times of a square pulse can be set but not differentiated, because the field is discontinuous there. value_and_grad holds them fixed and raises if you ask for their gradient; pass them through fixed=... to change them. To differentiate edge times, use a ramp pulse (ramp_onoff and friends), whose edges are smooth.
  • JAX must be in 64-bit mode. Gradients are validated on the CPU; the GPU path is untested.