Optimization for a Kerr oscillator¤
In this example we will try to implement DRAG-like pulses on a Kerr oscillator, penalizing population going to leakage states. We will additionally use a slightly more complicated control scheme than in the qubit example. There we performed traditional grape, with piecewise-constant controls. Here we fit a spline to control knot points: it is these knot points that are optimized.
This technique has the benefit of ensuring smooth controls if the control points are spread far enough apart in time, which is not guaranteed in the case of piece-wise constant controls. It is important to emphasize here that the control points are in no way related to the points chosen by the numerical integrator for solving the differential equation! This is a significant difference from typical QOC/GRAPE implementations that perform numerical integration of the Schrödinger equation by calculating step propagators with sufficiently small time steps.
This example is available as a Jupyter notebook here.
import diffrax as dx
import dynamiqs as dq
import jax.numpy as jnp
import optax
from jax import Array
import qontrol as ql
time = 20.0
control_dt = 2.0
ramp_nts = 4
ntimes = int(time // control_dt) + 1
tsave = jnp.linspace(0, time, ntimes)
optimizer = optax.adam(learning_rate=0.001, b1=0.999, b2=0.999)
opt_options = {'verbose': False, 'epochs': 4000, 'plot': True, 'plot_period': 40}
dq_options = dq.Options(save_states=True, progress_meter=None)
We define our Kerr oscillator to have a 100 MHz Kerr. We optimize the drive envelopes on the I and Q quadratures.
dim = 5
Kerr = -2.0 * jnp.pi * 0.100
a = dq.destroy(dim)
H0 = -0.5 * Kerr * dq.dag(a) @ dq.dag(a) @ a @ a
H1s = [a + dq.dag(a), 1j * (a - dq.dag(a))]
H1_labels = ['I', 'Q']
Here we attempt to perform an X gate, and penalize population leaking to the higher-lying states
initial_states = [dq.basis(dim, 0), dq.basis(dim, 1)]
final_states = [dq.basis(dim, 1), dq.basis(dim, 0)]
_forbidden_states = [dq.basis(dim, idx) for idx in range(2, dim)]
forbidden_states_list = len(initial_states) * [_forbidden_states]
We initialize the guess for the controls and moreover define an envelope of cosine ramps with a flat top to ensure that the control turns on slowly.
init_drive_params = {'dp': -0.001 * jnp.ones((len(H1s), ntimes))}
cos_ramp = (1 - jnp.cos(jnp.linspace(0.0, jnp.pi, ramp_nts))) / 2
envelope = jnp.concatenate(
(cos_ramp, jnp.ones(ntimes - 2 * ramp_nts), jnp.flip(cos_ramp))
)
def _drive_spline(
drive_params: Array, envelope: Array, ts: Array
) -> dx.CubicInterpolation:
drive_w_envelope = jnp.einsum('t,...t->t...', envelope, drive_params)
drive_coeffs = dx.backward_hermite_coefficients(ts, drive_w_envelope)
return dx.CubicInterpolation(ts, drive_coeffs)
def H_func(drive_params_dict: dict) -> dq.TimeArray:
drive_params = drive_params_dict['dp']
H = H0
for H1, drive_param in zip(H1s, drive_params):
drive_spline = _drive_spline(drive_param, envelope, tsave)
H += dq.modulated(drive_spline.evaluate, H1)
return H
exp_ops = [dq.basis(dim, idx) @ dq.dag(dq.basis(dim, idx)) for idx in range(dim)]
se_Kerr_model = ql.sesolve_model(H_func, initial_states, tsave, exp_ops=exp_ops)
costs = ql.coherent_infidelity(target_states=final_states, target_cost=0.005)
costs += 40.0 * ql.forbidden_states(
forbidden_states_list=forbidden_states_list, target_cost=0.4
)
opt_params = ql.optimize(
init_drive_params,
costs,
model=se_Kerr_model,
optimizer=optimizer,
opt_options=opt_options,
dq_options=dq_options,
)
Try playing with the cost_multiplier of the forbidden states cost function to see how it modifies the pulse! Specifically, try setting it to 0 so that the optimizer is happy to populate the 2nd excited state without penalty.
Master-equation optimization¤
We can also instead perform the optimization directly on the master equation. Continuing the above example, we need now to initialize an mesolve_model, which calls dq.mesolve, and to correctly specify the initial and target states as density matrices.
jump_ops = [0.01 * a]
me_initial_states = dq.todm(initial_states)
me_final_states = dq.todm(final_states)
me_forbidden_states_list = len(initial_states) * [dq.todm(_forbidden_states)]
me_costs = ql.coherent_infidelity(target_states=me_final_states, target_cost=0.005)
me_costs += 40.0 * ql.forbidden_states(
forbidden_states_list=me_forbidden_states_list, target_cost=0.4
)
me_Kerr_model = ql.mesolve_model(
H_func, jump_ops, me_initial_states, tsave, exp_ops=exp_ops
)
opt_params_me = ql.optimize(
init_drive_params,
me_costs,
me_Kerr_model,
optimizer=optimizer,
opt_options=opt_options,
dq_options=dq_options,
)
