Dynamiqs: Experimental data fitting background cover

Dynamiqs: Experimental data fitting

Fitting a cat qubit experiment with real-world data

Alice & Bob

Hosted by

Alice & Bob

Overview
Docs
Code
Rules
Discussion
Leaderboard

Fitting a cat qubit experiment with real-world data

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.
dynamiqs enables fast simulations of quantum systems with automatic gradient computation, making it highly suitable for data fitting tasks. This tutorial will guide you through the process of simulating and fitting real-world experimental data, step by step.

Downloading the template

To download the template for this use case, you can run the following command in the terminal
aqora template dynamiqs-2025-1
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-1
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

Context of the experiment

In the experiment of Réglade, Bocquet et al., "Quantum control of a cat-qubit with bit-flip times exceeding ten seconds" (2023), arxiv:2307.06617, the authors aim to perform quantum control of a cat qubit while preserving its exponential bias in bit-flip errors. One challenging aspect of the experiment is to calibrate the two-photon exchange rate g2g_2 between the cat qubit oscillator (the memory) and the ancillary reservoir mode (the buffer). The object of this tutorial is to perform this calibration with dynamiqs.
The experiment can be modelled by the following master equation:
dρ^dt=−i[H^,ρ^]+κ∗a(1+n∗th,a)D[a^]ρ+κ∗an∗th,aD[a^†]ρ+κ∗b(1+n∗th,b)D[b^]ρ+κ∗bn∗th,bD[b^†]ρ, \begin{aligned} \frac{d\hat\rho}{dt} = -i[\hat H, \hat \rho] &+ \kappa*a (1 + n*\mathrm{th, a})\mathcal{D}[\hat a] \rho + \kappa*a n*\mathrm{th, a} \mathcal{D}[\hat a^\dagger] \rho \\ &+ \kappa*b (1 + n*\mathrm{th, b})\mathcal{D}[\hat b] \rho + \kappa*b n*\mathrm{th, b} \mathcal{D}[\hat b^\dagger] \rho, \end{aligned}
with
H^=g2a^2b^†+g2∗a^†2b^ \hat H = g_2 \hat a^2 \hat b^\dagger + g_2^* \hat a^{\dagger 2} \hat b
where a^\hat a and b^\hat b denote the annihilation operators of the memory and the buffer, respectively. Both the memory and buffer dissipation rates and their average thermal populations are assumed to be known from previous experimental calibrations.
At the beginning of the experiment the memory is populated with a coherent state and the buffer is in thermal vacuum. We then turn on the two-photon exchange mechanism, such that pairs of photons leave the memory and are converted into single buffer photons that are quickly dissipated into the environment. The number of photons in the memory, ⟨a†a⟩\langle a^\dagger a \rangle is measured at different times and saved into the data variable. We then infer g2g_2 from this data.
Pulse sequence
In this tutorial, we will load the previously described experimental data (or generate synthetic data), encode the two-photon exchange model in dynamiqs, simulate the system evolution and measure the number of photons at different times. We will then use gradient descent to fit the g2,nth,a,nth,b,κa,κbg_2, n_{th,a}, n_{th,b}, \kappa_a, \kappa_b parameters from the data.
We start by importing all necessary libraries, and defining the physical parameters of the system.
import dynamiqs as dq
import jax
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
dq.plot.mplstyle(dpi=150)

# 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 operators
Na, Nb = 15, 7
a, b = dq.destroy(Na, Nb)

Loading the data

# Load the data
def load(file):
    x = np.loadtxt(".aqora/data/data/%s" % file)
    x = jnp.array(x)
    return x

if isinstance(input, tuple):
    data, tsave = input
else:
    data = load("short.npy")
    tsave = load("time_short.npy")
alphas = jnp.sqrt(data[0])

# Plot experimental data
plt.figure()
plt.plot(tsave / us, data, "+")
plt.grid()
plt.xlabel("Time [us]")
plt.ylabel("Photon number")
output

Definition of the model

# expectation operator: number of photons
exp_ops = [dq.dag(a) @ a]

def simulate(g2, kappa_a, kappa_b, nth_a, nth_b, progress_meter=True):
    # Redefine Hamiltonian with the new g2 value
    H = g2 * a @ a @ dq.dag(b)
    H += dq.dag(H)

    jump_ops = [
        jnp.sqrt(kappa_a * (1 + nth_a)) * a,
        jnp.sqrt(kappa_a * nth_a) * dq.dag(a),
        jnp.sqrt(kappa_b * (1 + nth_b)) * b,
        jnp.sqrt(kappa_b * nth_b) * dq.dag(b)
    ]

    # intialize buffer state with some thermal population
    psi0_b = dq.unit(dq.fock(Nb, 0) + nth_b ** 0.5 * dq.fock(Nb, 1)) 
    psi0 = dq.stack([dq.tensor(dq.coherent(Na, alpha), psi0_b) for alpha in alphas])

    options = dq.Options(progress_meter=
        dq.progress_meter.TqdmProgressMeter() if progress_meter 
        else dq.progress_meter.NoProgressMeter()
    )
    
    # Perform the master equation simulation
    return dq.mesolve(
        H,
        jump_ops,
        psi0,
        tsave,
        exp_ops=exp_ops,
        options=options
    )
We choose an initial guess that is reasonable with respect to the data we have. We converge to a steady state in roughly 1 μs1\:\mu\mathrm{s} so g2g_2 should be of the order of 1 MHz1\:\mathrm{MHz}. Let us begin from this estimated value. Other parameters guesses come from previous physical measurements.
# Initial guess
g2 = jnp.array(1.0 * MHz)
nth_a, nth_b = 0.1, 0.011
kappa_a = 5 * kHz
kappa_b = 10 * MHz

# Run the simulation with the initial guess
result = simulate(alphas, tsave, g2, kappa_a, kappa_b, nth_a, nth_b)

# Plot the initial guess simulation and the data
plt.figure()
for i in range(len(alphas)):
    plt.plot(tsave / us, data[:, i].real, f"+C{i}", alpha=0.6)
    plt.plot(tsave / us, result.expects[i, 0].real, f"C{i}")
plt.grid()
plt.xlabel("Time [us]")
plt.ylabel("Photon number")
output

Running the optimization

Next, we run the optimization, with a straightforward cost function (distance between the simulated and experimental data), and with a plain gradient descent. There are many smarter ways to perform gradient descent, for instance the ones provided in Optax. The optimization loop is defined as follow
@jax.jit
def cost_fn(x, data):
    """Simulate the system with the given the parameters and return the cost."""
    g2, kappa_a, kappa_b, nth_a, nth_b = x
    res = simulate(*x, progress_meter=False)
    return jnp.sum((data - res.expects[:, 0].real.T) ** 2)

physical_parameters =  [jnp.array([g2, kappa_a, kappa_b, nth_a, nth_b])]
costs = [cost_fn(physical_parameters[-1], data, tsave)] 
lr = 1e-3 # learning rate
for _ in tqdm(range(40)):
    # Compute the cost and its gradient
    parameter = physical_parameters[-1]
    cost, cost_grad = jax.value_and_grad(cost_fn)(parameter, data, tsave)

    # Clip the gradient to avoid numerical instability
    cost_grad = jnp.clip(cost_grad, -100, 100)

    # Perform the gradient descent step
    parameter = parameter - lr * cost_grad
    # all parameters in the problem are bounded by some lower values
    parameter = np.maximum(
        parameter,
        jnp.array([0.1 * MHz, 0.1 * kHz, 0.1 * MHz, 0.01, 1e-5])
    )

    # Save the cost and the parameter values
    costs.append(cost)
    physical_parameters.append(parameter)

# Save the final cost and convert to jnp arrays
costs, physical_parameters = jnp.array(costs), jnp.array(physical_parameters)

Plotting the results

Once the optimization is finished, we can plot the loss to ensure we converged correctly.
fig, axs = plt.subplots(2, 3, tight_layout=True)

ax = axs[0, 0]
ax.plot(costs)
ax.set_xlabel("Epoch")
ax.set_ylabel("Cost")

labels = ["g2", "kappa_a", "kappa_b", "nth_a", "nth_b"]
ref_unit = [MHz, kHz, MHz, 1, 1]
for i, label in enumerate(labels):
    ax = axs.ravel()[i+1]
    ax.plot(physical_parameters[:, i] / ref_unit[i])
    ax.set_xlabel("Epoch")
    ax.set_ylabel(label)

plt.tight_layout()
output
Finally we plot the fit and compare with the data
g2, kappa_a, kappa_b, nth_a, nth_b  = physical_parameters[-1]
# Print the final g2 value
print(f"g2 = {g2 / MHz:.3f} MHz")
print(f"kappa_a = {kappa_a / kHz:.3f} kHz")
print(f"kappa_b = {kappa_b / MHz:.3f} MHz")
print(f"nth_a = {nth_a:.3f}")
print(f"nth_b = {nth_b:.3f}")

# Run a final simulation with the fitted g2 value
result = simulate(alphas, tsave, g2, kappa_a, kappa_b, nth_a, nth_b)

# Plot the fit result
plt.figure()
for i in range(len(alphas)):
    plt.plot(tsave / us, data[:, i].real, f"+C{i}", alpha=0.6)
    plt.plot(tsave / us, result.expects[i, 0].real, f"C{i}")
plt.grid()
plt.xlabel("Time [us]")
plt.ylabel("Photon number")
output

Submitting the solution

We return our parameters for submission by writing
output = g2, kappa_a, kappa_b, nth_a, nth_b
And then by running the command
$ aqora upload