GitHub

Differentiable Quantum Optimal Control

Gradient-based quantum optimal control means shaping a control pulse, evolving the state, measuring how far the final state is from a target, and descending the gradient of that error with respect to the pulse parameters. Two packages divide the work:

  • SecondQuantizedAlgebra is the symbolic layer. You write the Hamiltonian and operators symbolically, with the control entering through a time_parameter value function (p, t) -> value of the solver parameter vector p and time t. to_numeric compiles that symbolic expression into the numeric operator (a QobjEvo) the solver consumes. Because p is threaded into its coefficients, that operator is differentiable with respect to p. No manual operator assembly is needed.
  • QuantumToolbox is the numeric engine. It solves the dynamics and, through the SciMLSensitivity adjoint machinery with a reverse-mode backend such as Mooncake or Enzyme, differentiates the solve with respect to p. This is specific to the QuantumToolbox backend; the QuantumOpticsBase backend has no solver parameter vector to differentiate against.

We solve the same task twice: first as a closed system with sesolve, then, after adding photon loss, as an open system with mesolve. The task is single-photon state preparation in a Kerr resonator. The Kerr nonlinearity shifts the $1 \to 2$ transition out of resonance (photon blockade), so a resonant drive can pump $0 \to 1$ but not beyond. We optimize the drive pulse to reach the Fock state $|1\rangle$ from vacuum.

Setup

The control Hamiltonian in the frame rotating at the drive frequency is

\[H(t) = \Delta\, a^\dagger a + \frac{K}{2}\, a^\dagger a^\dagger a a + \varepsilon(t)\,(a^\dagger + a),\]

with detuning $\Delta$, Kerr strength $K$, and a real control $\varepsilon(t)$.

using SecondQuantizedAlgebra
using QuantumToolbox
using Symbolics

using SciMLSensitivity
import Mooncake                       # loads the Mooncake reverse-mode backend
using DifferentiationInterface: AutoMooncake
using Optimization, OptimizationOptimJL

h = FockSpace(:c)
@qnumbers a::Destroy(h)
@variables Δ K ε κ

\[ \begin{equation} \left[ \begin{array}{c} \Delta \\ K \\ \varepsilon \\ \kappa \\ \end{array} \right] \end{equation} \]

Physical constants and control discretization. The pulse is a Fourier sine series, which is smooth and vanishes at both ends of the interval, and its amplitudes are the parameters we optimize.

N = 10                       # Fock cutoff (Hilbert dimension N + 1)
Δ0, K0 = 0.0, -2.0           # on resonance for 0 -> 1; Kerr blockade detunes 1 -> 2
κ0 = 0.05                    # photon-loss rate (used only for the open-system part)
T = 6.0                      # control duration
tlist = collect(range(0, T, 61))
M = 6                        # number of control parameters

pulse(p, t) = sum(p[k] * sinpi(k * t / T) for k in eachindex(p))

The differentiable control Hamiltonian

The static part is fixed with parameter (substituted eagerly). The drive is registered through time_parameter as the two-argument (p, t) function, whose arity marks the coefficient as solver-parameter aware and makes the returned QobjEvo differentiable with respect to p. The Hamiltonian is built once and reused by both solvers; each solver call supplies the current p.

Hc = to_numeric(
    Δ * a' * a + (K / 2) * a' * a' * a * a + ε * (a' + a), h, N;
    backend = QuantumToolboxBackend(),
    parameter = Dict(Δ => Δ0, K => K0),
    time_parameter = Dict(ε => (p, t) -> pulse(p, t)),
)

ψ0 = fock(N + 1, 0)
proj1 = fock(N + 1, 1) * fock(N + 1, 1)'
sensealg = BacksolveAdjoint(autojacvec = SciMLSensitivity.MooncakeVJP())

Optimization.jl is the idiomatic SciML interface: wrap the cost in an OptimizationFunction carrying the AD backend, build an OptimizationProblem, and solve. For a smooth, deterministic, low-dimensional control problem with exact gradients, the quasi-Newton LBFGS converges in far fewer solver evaluations than a first-order method such as Adam. This small helper optimizes a cost from a flat initial pulse and records the loss trace.

function optimize_pulse(cost)
    losses = Float64[]
    optf = OptimizationFunction((u, _p) -> cost(u), AutoMooncake())
    prob = OptimizationProblem(optf, fill(0.3, M))
    res = solve(prob, OptimizationOptimJL.LBFGS(); maxiters = 100, callback = (s, l) -> (push!(losses, l); false))
    return res.u, losses
end

Closed system (sesolve)

With no dissipation the state stays pure. The cost is the infidelity $1 - |\langle 1 | \psi(T) \rangle|^2$, one minus the population in $|1\rangle$ at the final time. sesolve receives params and the adjoint chosen for correct gradients here, BacksolveAdjoint with the Mooncake vector-Jacobian product; Enzyme works too, while Zygote returns silently wrong gradients for time evolution.

infidelity_closed(params) =
    1 - real(expect(proj1, sesolve(Hc, ψ0, tlist; params, sensealg, progress_bar = Val(false)).states[end]))

p_closed, losses_closed = optimize_pulse(infidelity_closed)
1 - infidelity_closed(p_closed)   # near unity: closed dynamics conserve probability
0.9999361731679142

The optimized pulse drives the mean photon number from 0 to 1, and the final population lands almost entirely in $|1\rangle$; the Kerr blockade suppresses $|2\rangle$.

using CairoMakie

n_op = create(N + 1) * destroy(N + 1)
sol_closed = sesolve(Hc, ψ0, tlist; params = p_closed, e_ops = [n_op], progress_bar = Val(false))
nt_closed = real.(sol_closed.expect[1, :])
εt_closed = pulse.(Ref(p_closed), tlist)
finalpops_closed = [real(expect(fock(N + 1, n) * fock(N + 1, n)', sol_closed.states[end])) for n in 0:4]

fig_closed = Figure(size = (1050, 320))

ax1 = Axis(fig_closed[1, 1]; xlabel = "iteration", ylabel = "infidelity 1 − F", yscale = log10, title = "LBFGS convergence")
lines!(ax1, 1:length(losses_closed), losses_closed; color = :crimson)

ax2 = Axis(fig_closed[1, 2]; xlabel = "t", ylabel = "⟨n⟩ / ε(t)", title = "optimized pulse & dynamics")
lines!(ax2, tlist, nt_closed; label = "⟨n⟩(t)", color = :navy)
lines!(ax2, tlist, εt_closed; label = "ε(t)", color = :darkorange, linestyle = :dash)
axislegend(ax2; position = :lt)

ax3 = Axis(fig_closed[1, 3]; xlabel = "Fock state n", ylabel = "population", title = "final state", xticks = 0:4)
barplot!(ax3, 0:4, finalpops_closed; color = :seagreen)

fig_closed
Example block output

Open system (mesolve)

Now add photon loss through a single collapse operator $c = \sqrt{\kappa}\, a$, giving the Lindblad master equation $\dot\rho = -i[H(t), \rho] + \kappa\,\mathcal{D}[a]\rho$. The collapse operator is assembled the same way as the Hamiltonian, from sqrt(κ) * a.

c_op = to_numeric(sqrt(κ) * a, h, N; backend = QuantumToolboxBackend(), parameter = Dict(κ => κ0))

The loss rate is a control too: writing time_parameter = Dict(κ => (p, t) -> ...) would make $\kappa$ a differentiable coefficient as well, so a tunable dissipation could be optimized alongside the drive.

The cost is now $1 - \langle 1 | \rho(T) | 1 \rangle$ on the final density matrix. mesolve builds its Liouvillian from the time-dependent operators as a lazy superoperator sum of a single concrete type, so the adjoint solve sees no dynamic dispatch and the same sensealg applies.

infidelity_open(params) =
    1 - real(expect(proj1, mesolve(Hc, ψ0, tlist, [c_op]; params, sensealg, progress_bar = Val(false)).states[end]))

p_open, losses_open = optimize_pulse(infidelity_open)
1 - infidelity_open(p_open)   # below unity: decoherence-limited
0.9065668158944354

Open-system results

Photon loss caps the achievable fidelity below the closed-system value: the single excitation decays throughout the pulse, so the final $|1\rangle$ population sits below one and the state is left mixed.

sol_open = mesolve(Hc, ψ0, tlist, [c_op]; params = p_open, e_ops = [n_op], progress_bar = Val(false))
nt_open = real.(sol_open.expect[1, :])
εt_open = pulse.(Ref(p_open), tlist)
finalpops_open = [real(expect(fock(N + 1, n) * fock(N + 1, n)', sol_open.states[end])) for n in 0:4]

fig_open = Figure(size = (1050, 320))

bx1 = Axis(fig_open[1, 1]; xlabel = "iteration", ylabel = "infidelity 1 − F", yscale = log10, title = "LBFGS convergence")
lines!(bx1, 1:length(losses_open), losses_open; color = :crimson)

bx2 = Axis(fig_open[1, 2]; xlabel = "t", ylabel = "⟨n⟩ / ε(t)", title = "optimized pulse & dynamics")
lines!(bx2, tlist, nt_open; label = "⟨n⟩(t)", color = :navy)
lines!(bx2, tlist, εt_open; label = "ε(t)", color = :darkorange, linestyle = :dash)
axislegend(bx2; position = :lt)

bx3 = Axis(fig_open[1, 3]; xlabel = "Fock state n", ylabel = "population", title = "final state", xticks = 0:4)
barplot!(bx3, 0:4, finalpops_open; color = :seagreen)

fig_open
Example block output

This page was generated using Literate.jl.