grape.py`¶

Because the simulator is written in JAX, the gradient of \(J_T\) with respect to every pulse sample comes for free: jax.grad differentiates straight through the eigh-based propagators and the closed-form Weyl coordinates. grape.py uses that for three things:

Use	Call
Plain GRAPE from random guesses; the QSL scan (paper Fig. 3)	`optimise(random_guesses(...), sector_hamiltonian(model), u_max, cfg)`, `qsl_scan`
Robust (ensemble) GRAPE over detuned Hamiltonians	`optimise(u0, ensemble_hamiltonian(model, omegas), u_max, cfg)`
Refinement of an RL pulse	`optimise(rl_pulse[None], ham, u_max, cfg)`

The optimiser¶

def optimise(u0, ham, u_max, cfg: GrapeConfig):
    """Optimise a batch of initial pulses u0 (n, n_samples) within |u| <= u_max.

    ``ham`` is one SectorHamiltonian, or an ensemble with a leading axis (see
    ``ensemble_hamiltonian``), in which case the loss is log10 of the mean J_T over
    the ensemble (robust GRAPE). Returns the best pulses, their (mean) J_T, C, U
    and the per-iteration history of the (mean) J_T.
    """
    ensemble = ham.drift1.ndim == 3
    theta0 = jnp.arctanh(jnp.clip(u0 / u_max, -0.99, 0.99))     # u = u_max tanh(theta) keeps |u| <= u_max
    tx = optax.adam(optax.cosine_decay_schedule(cfg.lr, cfg.n_iter, cfg.lr_final_frac))

    def cost(u):
        if not ensemble:
            return final_cost(u, ham, cfg)
        JT, (C, U) = jax.vmap(lambda h: final_cost(u, h, cfg))(ham)
        return JT.mean(), (C.mean(), U.mean())

    def loss(theta):
        JT, aux = cost(u_max * jnp.tanh(theta))
        return jnp.log10(jnp.maximum(JT, 1e-12)), (JT, aux)

    grad_fn = jax.value_and_grad(loss, has_aux=True)

    def run_one(theta):
        def step(carry, _):
            theta, opt_state, best = carry
            (l, (JT, _)), g = grad_fn(theta)
            g = jnp.nan_to_num(g)          # sqrt(abcd) in the Weyl eigenvalues has an infinite slope at 0
            updates, opt_state = tx.update(g, opt_state)
            best_theta, best_JT = best
            better = JT < best_JT
            best = (jnp.where(better, theta, best_theta), jnp.where(better, JT, best_JT))
            return (optax.apply_updates(theta, updates), opt_state, best), JT

        init = (theta, tx.init(theta), (theta, jnp.inf))
        (_, _, (best_theta, _)), hist = jax.lax.scan(step, init, None, cfg.n_iter)
        u = u_max * jnp.tanh(best_theta)
        JT, (C, U) = cost(u)
        return u, JT, C, U, hist

    return jax.jit(jax.vmap(run_one))(theta0)

How it works:

Bounded by construction. The pulse is \(u = u_{\max}\tanh\theta\) and the optimiser moves \(\theta\), so every iterate respects \(|u| \le u_{\max}\) with no clipping or penalty.
Loss. \(\log_{10} J_T\) rather than \(J_T\): it keeps gradients useful when \(J_T\) spans 1e-1 to 1e-6.
Ensemble. If ham has a leading axis, the cost is the mean \(J_T\) over it (computed with vmap). Minimising \(\log_{10}\) of the mean penalises the worst members most, which is what makes the pulse robust.
Best iterate kept. Adam with a cosine-decayed step is not monotone, so the loop tracks the best \(\theta\) seen.
Restarts in parallel. vmap(run_one) optimises a batch of initial pulses at once.

Robust GRAPE in three lines¶

from rlquantopt.jx import grape, physics
model = physics.ModelParams()
omegas = grape.detuning_grid(model, half_width_mhz=10.0, n_per_axis=5)   # 25 detuned systems
hams = grape.ensemble_hamiltonian(model, omegas)                          # stacked sector Hamiltonians
u, JT_mean, C, U, history = grape.optimise(u0, hams, u_max=20.0, cfg=grape.GrapeConfig())

Refining an RL pulse¶

import pandas as pd, jax.numpy as jnp
rl = jnp.asarray(pd.read_csv("rlquantopt/paper_plots/final_pulses/RL_pulse.csv")["amplist"].to_numpy()[1:346])
cfg = grape.GrapeConfig(n_iter=2000, lr=0.01)       # small steps: the start is already good
u, JT, *_ = grape.optimise(rl[None], physics.sector_hamiltonian(model), 20.0, cfg)

On the paper's RL pulse (J_T = 9.5e-5 at 17.25 ns), 200 iterations take 2 s on the CPU and reach J_T = 7.3e-6. Results of the full comparison are in RL vs GRAPE vs robust GRAPE.

Cost¶

One GRAPE iteration is one forward and one backward pass through the pulse per ensemble member and restart. A 345-sample pulse costs about 0.1 ms per member per iteration on the CPU: 2000 iterations of a 25-member ensemble with 4 restarts take ~25 minutes.

Limits¶

GRAPE sees the whole Hamiltonian and its gradients; RL sees only rewards. Comparisons of robustness are fair only at the same gate time, amplitude bound and model, which is how the experiment is set up. Comparisons of sample cost between the two are not like for like.
Gradients near \(abcd = 0\) are replaced with 0 (see Gate metrics); starting from \(u = 0\) therefore does not move. Random guesses always include a component at the qubit-qubit detuning (0.86 GHz) to start inside the useful region.

GRAPE: rlquantopt/jx/grape.py¶