Dynamiqs: Optimal control

Quantum control of a cavity coupled to a qubit

Hosted by

Alice & Bob

Quantum control of a cavity coupled to a qubit

If this is your first competition on Aqora, we highly recommend you follow the H2 Groundstate Energy Tutorial to get familar with the platform and the CLI.

In this document, we are going to perform GRadient Ascent Pulse Engineering (GRAPE) in order to prepare arbitrary states in a quantum harmonic oscillator. The harmonic oscillator alone is a linear system, so it can only access gaussian states (coherent states). To access non-classical states it must be coupled to a non linear system, such as a qubit.

Downloading the template

To download the template for this use case, you can run the following command in the terminal

aqora template dynamiqs-2025-2

This will download the template into a folder called test.

You can then open the folder in Visual Studio Code by running the following command in the terminal

aqora lab -p dynamiqs-2025-2

This should open the folder in Visual Studio Code. If you receive a prompt, you can click on "Yes, I trust the authors".

Uploading the submission

You can find a template notebook in submission/solution.ipynb. Fill in your solution. You can run the notebook locally to test your solution by running the following in the terminal

aqora test

And when you are ready to submit run

aqora upload

Getting Started

Your goal is to find the right controls to apply to both the qubit and the oscillator in order to prepare these states. We split the time interval in small steps. At each time step you can choose a drive on the oscillator and the qubit (

\varepsilon_{a}

and

\varepsilon_{q}

), on both

x

and

y

quadratures.

The full Hamiltonian of the system reads:

\begin{aligned} H_{\text{total}} &= H_{\text{interaction}} + H_{\text{control}}\\ H_{\text{interaction}} &= \chi a^{\dagger} a \sigma_z \\ H_{\text{control}}&=\varepsilon_{y, a}(t)\left(a+a^{\dagger}\right)+i \varepsilon_{x, a}(t)\left(a-a^{\dagger}\right)+\varepsilon_{x, q}(t) \sigma_x+\varepsilon_{y, q}(t) \sigma_y \end{aligned}

where

\chi

is the coupling between the oscillator and the qubit.

a

is the destruction operator of the harmonic oscillator and

\sigma_{x, y, z}

are the Pauli matrices associated to the qubit. The parameter space you have to work with is the 4 drives for the 100 time steps so 400 parameters in total.

In this tutorial we are going to start from the oscillator in vacuum and the qubit in ground state (

\ket{0}_a \otimes \ket{g}_q

) and prepare the oscillator in a state with 2 photons, with the qubit still in ground state (

\ket{2}_a \otimes \ket{g}_q

Note: running the following code in full takes ~30 minutes on a 4090 GPU. You should use a GPU too in order to finish the computation in reasonable time. Or maybe there are more clever ways to find the sequences?

import dynamiqs as dq
import jax
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import optax
dq.plot.mplstyle(dpi=150)

tsave, target = input

# Declare useful helpers to handle units
MHz = 2 * jnp.pi
kHz = 2 * jnp.pi * 1e-3
us = 1.0
ns = 1.0e-3

# Instanciate our bosonic annihilation operator
Na = 15 # bosonic mode truncation
a = dq.destroy(Na) & dq.eye(2)

# Qubit Pauli gate operators
X = dq.eye(Na) & dq.sigmax() # X gate
Y = dq.eye(Na) & dq.sigmay() # Y gate
Z = dq.eye(Na) & dq.sigmaz() # Z gate

Definition of the model

# coupling between the memory and the qubit
chi_aq = 2.0 * MHz 

tsave = np.linspace(0, 1 * us, 100)
target = dq.fock(Na, 2) & dq.basis(2, 0)

def simulate(drives, progress_meter=True):
    # Redefine Hamiltonian with the new g2 value
    H = dq.pwc(tsave, drives["memory"]["x"], a + a.dag()) # memory drive along x
    H += dq.pwc(tsave, drives["memory"]["y"], 1j * (a - a.dag())) # memory drive along y
    H += dq.pwc(tsave, drives["qubit"]["x"], X) # qubit drive, pauli X
    H += dq.pwc(tsave, drives["qubit"]["y"], Y) # qubit drive, pauli Y
    H += dq.constant(chi_aq * a.dag() @ a @ Z) # memory-qubit coupling

    # initial state: memory in vacuum and qubit in |0>
    psi0 = dq.coherent(Na, 0) & dq.basis(2, 0)

    options = dq.Options(
        progress_meter=
        dq.progress_meter.TqdmProgressMeter() if progress_meter 
        else dq.progress_meter.NoProgressMeter(),
    )
    
    # Perform the master equation simulation
    return dq.sesolve(H, psi0, tsave, options=options)

def cost_function(drive, do_sum=True):
    final_states = simulate(drive, progress_meter=False).states[:, -1]
    loss = (1 - abs(target.dag() @ final_states) ** 2)
    if do_sum:
        loss = loss.sum()
    return loss

Running the optimization

We are going to use the batching functionnality of dynamiqs. The batching dimension is given by BATCH_SIZE. Here it's set at 200, meaning that we are going to optimize on 200 candidate set of control parameters at the time

BATCH_SIZE = 200
NUM_TRAIN_STEPS = 500
MAX_AMP = 2.0 # maximum drive strength

key = jax.random.PRNGKey(42) # random number generation seed
keys = jax.random.split(key, 4)

draw = lambda key: dq.random.real(key, (BATCH_SIZE, len(tsave)-1), max=0.1)
memory_drive = dict(x=draw(keys[0]), y=draw(keys[1]))
qubit_drive  = dict(x=draw(keys[2]), y=draw(keys[3]))

params = dict(memory=memory_drive, qubit=qubit_drive)
optimizer = optax.chain(
    optax.adam(learning_rate=1e-2),
    optax.clip(1.0)
)
opt_state = optimizer.init(params)

@jax.jit
def step(params, opt_state):
    cost = cost_function(params, do_sum=False)[..., 0, 0] # log all costs
    cost_grad = jax.grad(cost_function)(params) # compute costs gradient
    updates, opt_state = optimizer.update(cost_grad, opt_state, params)
    params = optax.apply_updates(params, updates)

    return cost, params, opt_state

params_history, costs_history = [], []
for i in tqdm(range(NUM_TRAIN_STEPS), disable=False):
    cost, params, opt_state = step(params, opt_state)

    params_history.append(params)
    costs_history.append(cost)

costs_history = np.array(costs_history)

plt.plot(costs_history.min(axis=1), label="min")
plt.plot(costs_history.mean(axis=1), label="mean")
plt.plot(costs_history.max(axis=1), label="max")
plt.legend()

# this is the index of the pulse sequence that minimizes the cost
idx_best_sequence = np.argmin(costs_history[:, -1])
res = simulate(params_history[-1]).states[idx_best_sequence]
dq.plot.wigner_mosaic(res.ptrace(0), n=15, nrows=3)

fig, axs = plt.subplots(2, 1)

memory_pulse = params_history[-1]["memory"]["x"][idx_best_sequence] + 1j * params_history[-1]["memory"]["y"][idx_best_sequence]
qubit_pulse = params_history[-1]["qubit"]["x"][idx_best_sequence] + 1j * params_history[-1]["qubit"]["y"][idx_best_sequence]

dq.plot.pwc_pulse(tsave, memory_pulse, ax=axs[0])
dq.plot.pwc_pulse(tsave, qubit_pulse, ax=axs[1])

On the top graph we have the control applied on the memory, and on the qubit on the bottom. Each color represents a quadrature.

Submitting the solution

We return our parameters for submission by writing

output = dict(
  memory=dict(
    x=params_history[-1]["memory"]["x"][idx_best_sequence],
    y=params_history[-1]["memory"]["y"][idx_best_sequence],
  ),
  qubit=dict(
    x=params_history[-1]["qubit"]["x"][idx_best_sequence],
    y=params_history[-1]["qubit"]["y"][idx_best_sequence],
  ),
)

And then by running the command

$ aqora upload