Simulation & optimal control of a single spin qubit¶
In this tutorial, we'll show you, step-by-step, how to simulate the dynamics of a single spin qubit and perform quantum optimal control using qruise-ml.
The tutorial consists of the following:
1. Defining the Hamiltonian¶
1.1. Physical model¶
The system we are interested in is a single electron in a single quantum dot. With spin qubits, spin-up corresponds to the ground state, i.e. $|0\rangle = |\uparrow \rangle$, and spin-down corresponds to the excited state, i.e., $|1 \rangle = |\downarrow \rangle$.
Our stationary (time-independent) Hamiltonian, $H_0$ (also known as the drift term), corresponds to the Zeeman interaction between the electron spin and an external static magnetic field:
$$ H_0 =\frac{\gamma}{2} \vec{B} \cdot \vec{\sigma} . $$Here, $\vec{B}$ is the external magnetic field vector, and $\vec{\sigma} = (\sigma_x,\sigma_y,\sigma_z)$ is the vector of Pauli matrices. The gyromagnetic ratio, $\gamma$, quantifies the response of an electron to a magnetic field and can be expressed by $\gamma = \frac{-g\mu_B}{\hbar}$, where $g$ is the $g$-factor ($\sim 2$ for a free electron), $\mu_B$ is the Bohr magneton quantifying the magnetic moment of an electron, and $\hbar$ is the reduced Planck's constant.
Let's assume we only want to apply our external magnetic field in the $z$-direction, i.e., $B_x=B_y=0$. Our Hamiltonian then simplifies to
$$ H_0 = \frac{\gamma}{2} B_z \sigma_z. $$Substituting in the eigenvalues of $\sigma_z$ (±1), we see that the two spin states have an energy (Zeeman) splitting, $\Delta E$, given by
$$ \Delta E = \frac{\gamma}{2} B_z (1) - \frac{\gamma}{2} B_z (-1) = \gamma B_z . $$The corresponding qubit transition frequency is
$$ \omega_q=\gamma B_z $$which is also known as the Larmor frequency. This is the frequency at which the spin precesses about the magnetic field. In the simulations that follow, we use $\gamma = 2.8 \times 10^{10}~\mathrm{Hz~T^{-1}}$ and an external magnetic field of $B_z = 0.1~\mathrm{T}$.
To manipulate the qubit, we apply a time-dependent drive $c(t)$ that couples to the spin through the $\sigma_x$ operator. The full Hamiltonian is therefore
$$ H(t) = \frac{\gamma}{2} B_z \sigma_z + c(t)\sigma_x. $$The first term determines the natural evolution of the qubit, while the second term drives transitions between the $|0\rangle$ and $|1\rangle$ states.
1.2. Time-dependent drive function¶
For the drive, we’ll use a Gaussian pulse modulated by a cosine:
$$ c(t; \{a,\sigma,\mu,\omega_d\}) = \frac{a}{\sigma\sqrt{2\pi}} \mathrm{exp}\left[-\frac{(t - \mu)^2}{2\sigma^2}\right] \cos(\omega_{d}t). $$This is characterised by four parameters:
(i) $a$, a scalar that adjusts the amplitude of the pulse
(ii) $\sigma$, which controls the pulse width
(iii) $\mu$, the time at which the pulse reaches its maximum amplitude
(iv) $\omega_{d}$, the local oscillator frequency, which is often chosen to match $\omega_q$.
In QruiseML, drive functions are defined as annotated Python functions. This means that each input parameter, as well as the function output, must have a type annotation. These annotations must use the types defined in the types module in QruiseML. Here, we'll set $\mu=\frac{t_1}{2}$ and use type Float64 for all the inputs and the output.
import numpy as np
from qruise.toolset.types import Float64
def drive(t: Float64, amp: Float64, sigma: Float64, omega_d: Float64) -> Float64:
"""Envelope: Gaussian pulse modulated by a cosine"""
factor = amp / np.sqrt(2 * np.pi) / sigma
gaussian = factor * np.exp(-((t - t1 / 2) ** 2) / (2 * sigma**2))
# Note: mu is set to t1/2
return gaussian * np.cos(omega_d * t)
Now we need to define the time, qubit, and initial drive parameters for the simulation.
# time parameters
t0 = 0.0 # initial time (s)
t1 = 20e-9 # final time (s)
N = 1000 # number of time points
ts = np.linspace(t0, t1, N) # time array
# spin qubit parameters
gamma = -2.8024951386169e10 # gyromagnetic ratio of an electron (Hz/T)
Bz = 0.1 # magnetic field strength (T)
omega_q = gamma * Bz # Larmor (qubit) frequency (Hz)
# drive parameters
amp = 2.0 # amplitude
sigma = 0.2 * t1 # width
omega_d = omega_q # frequency (Hz)
It’s a good idea to check your drive behaves as desired before proceeding, so let’s quickly plot it.
import xarray as xr
import hvplot.xarray
xr.DataArray(
drive(ts, amp, sigma, omega_d),
coords={"time": ("time", ts, {"units": "s"})},
dims=["time"],
name="Amplitude (arb. u.)",
).hvplot()
As expected, we have a Gaussian pulse modulated by a cosine — great!
1.3. Drive parameters and channels¶
Next, we'll organise the drive parameters into a drive channel. A drive channel groups together the parameters associated with a particular drive, providing a convenient way to keep related parameters organised. When working with multiple drives, this makes it easy to identify which parameters belong to which drive.
We use Parameters to construct a parameter space from the channel parameters. Each parameter is specified as a tuple of the form (value, lower_bound, upper_bound). If you're only performing a simulation, the bounds can be omitted. However, they are required for optimal control, so we'll include them here.
from qruise.toolset import Parameters, Drive
# associate drive parameters with channel and construct parameter space
ps = Parameters(
{
"d1": {
# "parameter name": (value, lower_bound, upper_bound)
"amp": (amp, 1.0, 5.0),
"sigma": (sigma, 1e-9, 8e-9),
"omega_d": (omega_d, omega_q - 5e8, omega_q + 5e8),
}
}
)
Note: When working with multiple drives, drive channels allow different drives to reuse parameter names — such as amp, sigma, and omega_d — without conflicts. In this example, we only have a single drive, but it's good practice to create a drive channel from the outset so it's straightforward to add additional drives later.
Finally, we use the Drive class to associate the drive function with the drive channel, which itself is associated to the relevant parameters:
drv1 = Drive("d1", drive)
1.4. Constructing the Hamiltonian¶
The Hamiltonian is constructed in a similar way to QuTiP. We instantiate the Hamiltonian class by passing the stationary part, H0, together with a list of (drive, operator) pairs defining the time-dependent terms:
import qutip as qt
from qruise.toolset import Hamiltonian
# Pauli operators from QuTiP
sx = qt.sigmax()
sz = qt.sigmaz()
# define Hamiltonian
H = Hamiltonian(omega_q * sz / 2, [(drv1, sx)])
2. Defining the quantum simulation problem¶
Now that we've defined our Hamiltonian, we can start setting up a quantum simulation problem. In QruiseML, we use the Problem class to group together the equations and parameters that govern the system. To instantiate the Problem, we need:
- the equation that governs the dynamics (in this example, we'll use the
"Schroedinger"equation) - the Hamiltonian (
H) - the initial qubit state (
u0, here the ground state $|0\rangle$) - the time interval of the simulation (
t0totfinal) - the pulse parameters (
ps)
from qruise.toolset import Problem
u0 = qt.basis(2, 0) # initial vector state, |0>
prob = Problem("Schroedinger", H, u0, (t0, t1), ps)
3. Starting a session and solving the problem¶
The Problem object defines the quantum system we want to simulate,
but it does not run the simulation itself. To execute the time evolution, we need to create a Session, then use qsprob_init to initialise the Problem within the session. The session prepares the problem for numerical simulation and manages the solver execution.
from qruise.toolset.session import Session
sess = Session()
sess.qsprob_init(prob)
[juliapkg] Found dependencies: /home/ci/.venv/lib/python3.11/site-packages/juliacall/juliapkg.json [juliapkg] Found dependencies: /home/ci/.venv/lib/python3.11/site-packages/juliapkg/juliapkg.json [juliapkg] Locating Julia ~1.9, ^1.10.3 [juliapkg] Using Julia 1.11.6 at /home/ci/.juliaup/bin/julia [juliapkg] Using Julia project at /__w/qruise-ml/qruise-ml/julia/python_env
To compute the system time evolution, we use the Session.evolve method to solve the quantum simulation problem numerically.
In the implementation below, "Tsit5" specifies the ODE algorithm, reltol
and abstol the numerical tolerances, and saveat the times at which the result is stored. Other solver options can also be passed as keyword arguments.
ds = sess.evolve("Tsit5", reltol=1e-3, abstol=1e-4, saveat=ts)
ds
ODE [ ] 0% ODE [########################################] 100%
<xarray.Dataset> Size: 40kB
Dimensions: (time: 1000, row: 2)
Coordinates:
* time (time) float64 8kB 0.0 2.002e-11 4.004e-11 ... 1.998e-08 2e-08
Dimensions without coordinates: row
Data variables:
qstates (time, row) complex128 32kB (1+0j) ... (-0.20567110345887604+0.8...This returns an xarray.Dataset object containing coordinates corresponding to the time axis and quantum-state labels (here, $|0\rangle$ and $|1\rangle$), as well as the variable qstates, which stores either the state vector or density matrix of the system.
4. Plotting the qubit dynamics¶
Let's take a look at the effect of the drive on the qubit populations. Since xarray.Dataset supports additional derived variables, we can compute the state populations and store them alongside the original simulation results. We can also rename the row dimension to state and assign the basis-state labels $|0\rangle$ and $|1\rangle$ for clearer indexing and visualisation.
# calculate population of each state
ds["population"] = np.abs(ds["qstates"]) ** 2
ds = ds.rename({"row": "state"}) # rename the dimension "row" to "states"
ds = ds.assign_coords(state=[r"|0>", r"|1>"]) # assign labels to the dimensions
xarray provides different backends to conveniently plot the result. For example, we can use hvplot to plot the population of each state:
# plot populations
ds.hvplot.line(
x="time",
y="population",
by="state",
xlabel="time (s)",
ylabel="population (arb. u.)",
)
Congratulations! You’ve just simulated your first spin qubit.
5. Performing quantum optimal control¶
As we can see in the population plot above, the drive pulse does not fully invert the qubit population. If we want to implement an $X$ gate, we need to adjust the pulse parameters so that the qubit is driven from $|0\rangle$ to $|1\rangle$ by the end of the simulation. To achieve this, we can perform quantum optimal control.
For this, we use QOCProblem, which is almost identical to Problem, but with two additional inputs:
- the desired (target) qubit state (
ut, here the excited state $|1\rangle$), which must be compatible with the type and shape ofu0 - a loss function, $\mathcal{L}$, that we want to minimise. Here, we'll use
"svo", which uses the state vector overlap to compute the infidelity:
where $u_f(t=t_1)$ is the simulated qubit state at the final time $t_1$, and $u_t$ is the desired target state. The objective is to optimise the parameters of the drive pulse by minimising $\mathcal{L}$, bringing the final qubit state as close as possible to the target state.
from qruise.toolset import QOCProblem
# define target state
ut = qt.basis(2, 1)
# define QOCProblem
qocprob = QOCProblem(
"Schroedinger", H, u0, (t0, t1), ps, ut, "svo"
) # the Quantum Optimal Control problem
The whole workflow now reduces to solving the optimal control problem using our initial guess for the drive parameters. At each step, the system is simulated with the current parameters to obtain the final wavefunction and compute the infidelity with respect to the target state. Based on the resulting gradients, the parameters are updated iteratively until they converge and minimise the loss function.
Similar to Problem, QOCProblem defines the optimal control problem we want to solve, but it does not perform the optimisation itself. We first create a Session, then use qocprob_init to initialise the QOCProblem within the session, before calling Session.optimise to perform the optimisation. Here, we'll use the L-BFGS (limited-memory Broyden–Fletcher–Goldfarb–Shanno) algorithm.
# initialise session with QOCProblem
sess.qocprob_init(qocprob)
# perform optimisation using L-BFGS algorithm
optparams = sess.optimise("LBFGS")
loss: 0.30420622234238104 loss: 0.6352895949387161 loss: 0.005802778138724651 loss: 0.02613618729904643 loss: 0.0016626924216938122 loss: 0.00011674909770675868 loss: 2.0748237039569517e-7 loss: 2.0670808498213233e-7
We can see the loss function is progressively reduced as the optimisation converges. Let's now compare the initial values of the parameters with those the optimisation yielded.
# create parameter space for optimised parameters
optps = Parameters(optparams)
# print initial and optimised parameters
print("Initial parameters:", *ps.ps)
print("Optimised parameters:", *optps.ps)
# create xarray Dataset to store initial and optimised drive parameters
pulse_shapes = xr.Dataset(
{
"Unoptimised": ("time", drive(ts, *ps.ps)),
"Optimised": ("time", drive(ts, *optps.ps)),
},
coords={"time": ts / 1e-9},
)
# plot initial and optimised drives
pulse_shapes.hvplot.line(
x="time",
xlabel="Time (ns)",
ylabel="Amplitude (arb. u.)",
)
Initial parameters: 2.0 4e-09 -2802495138.6169 Optimised parameters: 3.1620023558947286 3.59989480878582e-09 -2809284981.8920965
We can see that the optimisation primarily increases the drive amplitude, while making minor adjustments to the width and frequency.
To determine whether these changes improve the population inversion, let's simulate the dynamics using the optimised parameters. We then create a new simulation problem, initialise it with Session.qsprob_init, and solve it as before.
# create a new quantum simulation problem with the optimised parameters
optprob = prob.remake(ps=optps)
# instantiate the new problem with the optimised parameters
sess.qsprob_init(optprob)
# evolve the system
optds = sess.evolve("Tsit5", saveat=ts)
optds
ODE [ ] 0% ODE [########################################] 100%
<xarray.Dataset> Size: 40kB
Dimensions: (time: 1000, row: 2)
Coordinates:
* time (time) float64 8kB 0.0 2.002e-11 4.004e-11 ... 1.998e-08 2e-08
Dimensions without coordinates: row
Data variables:
qstates (time, row) complex128 32kB (1+0j) ... (-0.18048959203517007+0.9...To compare the dynamics, we calculate the state populations after optimisation and plot them together with the unoptimised populations from earlier.
optds["population"] = np.abs(optds["qstates"]) ** 2
optds = optds.rename({"row": "state"})
optds = optds.assign_coords(state=[r"|0>", r"|1>"])
(
ds["population"].isel(state=0).hvplot.line(label="Unoptimised |0⟩")
* ds["population"].isel(state=1).hvplot.line(label="Unoptimised |1⟩")
* optds["population"].isel(state=0).hvplot.line(label="Optimised |0⟩")
* optds["population"].isel(state=1).hvplot.line(label="Optimised |1⟩")
).opts(
xlabel="Time (s)",
ylabel="Population",
title="",
)
We can see here that the population exchange between the ground and first excited states is almost complete. Our optimisation was successful!
We can quantify this success by printing the fidelity before and after optimisation:
def fidelity(uf, ut):
"""State fidelity between two pure states."""
uf = np.asarray(uf).squeeze()
ut = ut.full().squeeze()
return np.abs(np.vdot(ut, uf)) ** 2
fid_unopt = fidelity(ds["qstates"].isel(time=-1), ut)
fid_opt = fidelity(optds["qstates"].isel(time=-1), ut)
print(f"Unoptimised: {100*fid_unopt:.6f} %")
print(f"Optimised: {100*fid_opt:.6f} %")
Unoptimised: 69.775021 % Optimised: 99.999979 %
We can clearly see that our optimisation has significantly enhanced the fidelity!
Congratulations, you’ve just performed optimal control on a spin qubit! You can now use these methods to start optimising your qubit performance.