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 optimal control using the qruise-toolset.
The tutorial consists of the following:
- Defining the Hamiltonian
- Defining the time-dependent Hamiltonian with external drive
- Solving the Schrödinger equation
- Plotting the qubit dynamics
- Performing optimal control
- Optimal control with gate fidelity
1. Defining the Hamiltonian¶
The system we are interested in here 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$.
As we only have one spin in this system, the only term we have in our stationary (time-independent) Hamiltonian, $H_0$ (also known as the drift term), corresponds to the Zeeman effect:
$$ H_0 =\frac{\gamma}{2} \vec{B} \cdot \vec{\sigma} , $$where $\vec{B}$ is the external magnetic field vector, and $\vec{\sigma} = (\sigma_x,\sigma_y,\sigma_z)$ is the Pauli matrix vector. 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.
Note: This Hamiltonian does not yet include an external drive.
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 spins 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 qubit frequency, $\omega_q$, can then be extracted as $\omega_q=\gamma B_z$. This is also know as the Larmor frequency, which is the frequency at which the spin precesses around the magnetic field.
Okay, now we understand our stationary Hamiltonian, we can start coding it. First, we need to define some initial qubit parameters: for an electron, $\gamma = 2.8 \, \text{Hz} \, \text{T}^{-1}$, and we'll set $B_z = 0.1 \, \text{T}$.
Tip: Make sure that your environment is correctly set to handle float64 precision by setting JAX_ENABLE_X64=True or add the following codeblock to your scripts' preamble:
import jax
jax.config.update("jax_enable_x64", True)
# Spin qubit constants
gamma = -2.8024951386169e10 # gyromagnetic ratio of an electron (Hz/T)
Bz = 0.1 # magnetic field strength (T)
freq = gamma * Bz # Larmor or qubit frequency (Hz)
Then we can define the stationary part of the Hamiltonian.
from qruise.toolset import Hamiltonian
from qutip import sigmaz
# define stationary Hamiltonian
H = Hamiltonian(freq * sigmaz() / 2)
2. Defining the time-dependent Hamiltonian with external drive¶
To construct our time-dependent Hamiltonian, $H$, we simply take our stationary Hamiltonian and add a term describing the effect of the external drive, $c(t)$. The time-dependent Hamiltonian is then given by
$$ H = \frac{\gamma}{2} B_z \sigma_z + c(t)\sigma_x . $$For this system, we’ll use a Gaussian pulse modulated by a cosine:
$$ c(t; \{a,\sigma,\mu\}) = \frac{a}{\sigma\sqrt{2\pi}} \mathrm{exp}\left[-\frac{(t - \mu)^2}{2\sigma^2}\right] \cos(\omega_{d}t), $$which is characterised by four parameters:
(i) $a$, a scalar that adjusts the amplitude of the pulse
(ii) $\sigma$, the pulse variance
(iii) $\mu$, the time at which the pulse reaches its maximum amplitude
(iv) $\omega_{d}$, the local oscillator frequency, which is often set to $\omega_q=\gamma B_z$.
We can use this equation to define our drive pulse.
import jax.numpy as jnp
def drive(t, params):
"""Envelope: Gaussian pulse modulated by a cosine"""
amp = params["a"] # Gaussian pulse amplitude
sigma = params["sigma"] # Std dev of Gaussian pulse (controls pulse width)
lo_freq = params["omega"] # Carrier (local oscillator) frequency
factor = amp / jnp.sqrt(2 * jnp.pi) / sigma
gaussian = factor * jnp.exp(-((t - tfinal / 2) ** 2) / (2 * sigma**2))
# Note: mu is set to tfinal/2
return gaussian * jnp.cos(lo_freq * t)
Now we need to define the time and pulse parameters for the simulation.
# time parameters for simulation
t0 = 0.0 # initial time (s)
tfinal = 20e-9 # final time (s)
grid_size = 1000 # simulation grid size (number of time points)
ts = jnp.linspace(t0, tfinal, grid_size)
# drive pulse parameters
params = {"a": 2.0, "sigma": 0.2 * tfinal, "omega": freq}
It’s a good idea to check your drive behaves as desired before proceeding, so let’s quickly plot it.
from qruise.toolset import PlotUtil
canvas = PlotUtil(x_axis_label="t [ns]", y_axis_label="Amplitude", notebook=True)
canvas.plot(ts / 1e-9, drive(ts, params), labels=["Pulse"])
canvas.show_canvas()
Great, our drive looks as expected!
Now, we simply add the drive term to our Hamiltonian.
from qutip import sigmax
# Add the drive term to the Hamiltonian
H.add_term(sigmax(), drive) # H(t) = H0 + sigma_x * drive
3. Solving the Schrödinger equation¶
Now we can solve the Schrödinger equation using a piecewise constant solver.
We’ll start by specifying our initial qubit state, which we choose as the ground state.
# define initial qubit state
y0 = jnp.array([1.0, 0.0])
We then define our Problem, which takes the Hamiltonian, the initial qubit state, and the time range as inputs. We then use Qruise’s PWCSolver to solve the Problem and calculate the wavefunction of the system at each timestamp (n=grid_size).
from qruise.toolset import Problem, PWCSolver
# define Problem
prob = Problem(H, y0, params, (0.0, tfinal))
# uses PWCSolver to solve Schrödinger equation
solver = PWCSolver(n=grid_size, store=True)
solver.set_system(prob.sepwc())
_, res = solver.evolve(*prob.problem())
Note: Setting store=True in the PWCSolver ensures that the results are stored at each timestamp during the simulation.
4. Plotting the qubit dynamics¶
Finally, we can view the results of our simulation by calculating and plotting the expectation value of the transmon state populations against time.
from qruise.toolset import get_population
# calculates the expectation value of each qubit state
pop = get_population(res)
# plots populations against time
canvas.init_canvas(x_axis_label="t [ns]", y_axis_label="Population")
canvas.plot(ts / 1e-9, pop, labels=["Ground state", "1st excited state"])
canvas.show_canvas()
Congratulations! You’ve just simulated your first spin qubit. You’ll notice, however, that the population exchange between the ground and first excited states is not complete. We can improve this by performing quantum optimal control.
5. Performing optimal control¶
As we saw in the previous plot, the drive pulse we defined caused a certain degree of population exchange between the ground and excited states. However, if we want to perform an X gate on our qubit, we need to optimise the drive pulse parameters such that, at the end of the simulation, the population of the ground and first excited states is fully inverted.
To perform the optimisation, we first need to define a loss function, $\mathcal{L}$, to quantify how far we are from our desired target. We can use the infidelity, which is given by $1-\mathcal{F}$, with $\mathcal{F}$ being the fidelity. The loss function is then given by
$$ \mathcal{L}=1-\mathcal{F}=1− |\langle \psi(t=t_\text{final})|\psi_t \rangle|^2, $$where $\mathcal{F}$ is the gate fidelity, $|\psi(t=t_\text{final})\rangle$ is the wavefunction of the system at the end time of the pulse, and $|\psi_t\rangle$ is the target (desired) wavefunction. The objective is to optimise the parameters of the drive pulse by minimising $\mathcal{L}$.
We can define the loss function as follows:
def loss(x, y):
"""
Returns the infidelity (1 - |<x|y>|^2)
of two wavefunctions x and y.
"""
o = jnp.matmul(
x.conj().T, y
) # Calculates the inner product (overlap) of the two wavefunctions
return jnp.real(1.0 - o.conj() * o) # Returns the real part of the infidelity
We are now ready to define a quantum optimal control problem (QOCProblem). Similar to the definition of a Problem for simulating the dynamics of a quantum system, QOCProblem further defines the optimisation protocol. As inputs, it takes
- the Hamiltonian,
H - the initial qubit state,
y0 - the pulse parameters,
params - the time interval of interest,
t0totfinal - the desired qubit state,
yt, i.e., 1st excited state - the loss function,
loss
Note: We defined the first four inputs earlier in this tutorial.
from qruise.toolset import QOCProblem
yt = jnp.array([0.0, 1.0]) # define desired state
opt_prob = QOCProblem(H, y0, params, (t0, tfinal), yt, loss)
The whole workflow now reduces to solving the optimal control problem using our initial guess for the pulse 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 gradient values, the parameter values are updated iteratively until they converge and minimise the loss function.
The Optimiser class is provided for this specific task. As a user, you need to specify both the solver used to simulate the dynamics and the optimisation algorithm itself. In this case, we again use a piecewise constant solver (PWCSolver), together with the Broyden–Fletcher–Goldfarb–Shanno (BFGS) method from the optimistix library.
import optimistix as optx
from qruise.toolset import OptimistixMinimiser, Optimiser
minimiser = OptimistixMinimiser(optx.BFGS, tol=1e-3)
opt = Optimiser(minimiser, PWCSolver(n=grid_size, eq=prob.sepwc()))
# Set the loss function for the optimiser to minimise (infidelity)
opt.set_optimisation(opt_prob.loss)
# Run optimisation on problem
opt_params, opt_summary = opt.optimise(*opt_prob.problem())
Note: Since qruise-toolset v2 is fully written in JAX, we use other libraries within the same ecosystem, such as diffrax for solving differential equations and optimistix for minimisation problems.
We can now compare the initial values of the parameters with those the optimisation yielded.
from pprint import pprint
from qruise.toolset import PlotUtil
canvas = PlotUtil(x_axis_label="t [ns]", y_axis_label="Amplitude", notebook=True)
# Print initial (unoptimised) and optimised parameters
pprint(f"Unoptimised {params}")
pprint(f"Optimised {opt_params}")
# Plot initial (unoptimised) and optimised parameters
canvas.plot(ts / 1e-9, drive(ts, params), labels=["Unoptimised"])
canvas.plot(ts / 1e-9, drive(ts, opt_params), labels=["Optimised"])
canvas.show_canvas()
"Unoptimised {'a': 2.0, 'sigma': 4e-09, 'omega': -2802495138.6169}"
("Optimised {'a': Array(3.14144187, dtype=float64), 'omega': "
"Array(-2.80249514e+09, dtype=float64), 'sigma': Array(1.72359068e-09, "
'dtype=float64)}')
We can see that the optimised pulse has increased amplitude and a wider envelope (greater duration) compared to the unoptimised pulse.
Let’s now simulate and plot the state populations to see if our optimisation worked.
# Check the simulation with optimised parameters
_, opt_res = solver.evolve(*opt_prob.problem(params=opt_params)[:-1])
# Get the population of the ground and excited states before optimisation
pop = get_population(res)
# Get the population of the ground and excited states after optimisation
opt_pop = get_population(opt_res)
# Plot unoptimised and optimised populations on same graph
canvas.init_canvas(x_axis_label="t [ns]", y_axis_label="Population")
canvas.plot(ts / 1e-9, pop, labels=["Unopt. ground state", "Unopt. 1st excited state"])
canvas.plot(ts / 1e-9, opt_pop, labels=["Opt. ground state", "Opt. 1st excited state"])
canvas.show_canvas()
We can see here that the population exchange between the ground and first excited states is almost perfect. The final populations are around 0.03% and 99.7%, respectively, compared to 30% and 70% before optimisation. Our optimisation was successful!
We can quantify this success by printing the fidelity before and after optimisation. (Remember: we defined our loss function as the infidelity = 1 - fidelity).
fidelity = [1 - loss(res[-1, :], yt), 1 - loss(opt_res[-1, :], yt)]
pprint(f"Unoptimised: {fidelity[0]*100:.1f} %")
pprint(f"Optimised: {fidelity[1]*100:.1f} %")
'Unoptimised: 69.6 %' 'Optimised: 99.7 %'
6. Optimal control with unitary (gate) fidelity¶
So far, we have optimised the drive pulse using state fidelity: the overlap between the final wavefunction and a target state vector. This approach works well when we care about preparing a specific state from a known initial state.
However, in quantum computing we often want to implement a specific quantum gate — a unitary transformation that should work correctly regardless of the input state. For this, we use gate fidelity, which measures how close the achieved propagator (unitary) $U(t_\text{final})$ is to the target unitary $U_t$.
Computing the propagator¶
The propagator $U(t)$ satisfies the same Schrödinger equation as the state, but with the identity matrix $\mathbf{I}$ as the initial condition:
$$ i\hbar \dot{U}(t) = H(t)\, U(t), \quad U(0) = \mathbf{I}. $$By passing U0 = jnp.eye(2) as the initial condition to the solver, we obtain the full propagator at each time step.
Gate fidelity¶
The gate fidelity between the achieved propagator $U$ and the target $U_t$ in a $d$-dimensional Hilbert space is:
$$ \mathcal{F}_\text{gate} = \frac{|\mathrm{Tr}(U_t^\dagger U)|^2}{d^2}. $$When the system Hilbert space is larger than the target gate subspace (e.g., a leakage subspace exists), we use project_op from qruise.toolset.utils to project $U$ onto the subspace spanned by the columns of $U_t$ before computing the fidelity.
Lab-frame target¶
Our system Hamiltonian includes a drift term $H_0 = \frac{\omega_q}{2}\sigma_z$, which continuously rotates the qubit in the lab frame. This means the lab-frame propagator always accumulates a z-rotation on top of the gate. To correctly target an X gate, the target unitary must account for this accumulated rotation:
$$ U_t^\text{lab} = e^{-i \frac{\omega_q}{2} \sigma_z\, t_\text{final}} \cdot X. $$Following the same approach as in Section 5, we now define a loss function based on gate infidelity. Notice the close analogy to the state infidelity loss we defined earlier — the key difference is that we now compare full propagators rather than state vectors, and normalise by the Hilbert space dimension $d$:
$$ \mathcal{L}_\text{gate} = 1 - \frac{|\mathrm{Tr}(U_t^\dagger U)|^2}{d^2}. $$Note: When the system Hilbert space is larger than the target gate subspace (for example, when leakage levels are present), project_op from qruise.toolset.utils can be used to project $U$ onto the subspace spanned by the columns of $U_t$ before computing the fidelity.
We can define the gate fidelity loss function as follows:
from jax import Array
def loss_gate_fidelity(U: Array, Ut: Array) -> Array:
"""Returns the gate infidelity (1 - |Tr(Ut† U)|² / d²)."""
d = U.shape[-1]
overlap = jnp.trace(jnp.matmul(Ut.conj().T, U))
return jnp.real(1.0 - overlap.conj() * overlap / d**2)
Before defining the optimal control problem, we need to specify two key inputs: the initial condition for the propagator and the target unitary.
As described above, the propagator starts as the identity matrix, so we set U0 = jnp.eye(2). The target is an X gate (a spin flip), but because our Hamiltonian includes the drift term $H_0 = \frac{\omega_q}{2}\sigma_z$, the qubit continuously rotates in the lab frame throughout the pulse. We must therefore account for this accumulated z-rotation in our target:
# propagator initial condition: identity matrix
U0 = jnp.eye(2, dtype=jnp.complex128)
# X gate
Xgate = jnp.array([[0.0, 1.0], [1.0, 0.0]], dtype=jnp.complex128)
# z-rotation accumulated over the pulse duration
U_z = jnp.diag(
jnp.array(
[jnp.exp(-1j * freq / 2 * tfinal), jnp.exp(1j * freq / 2 * tfinal)],
dtype=jnp.complex128,
)
)
# lab-frame target: Ut = exp(-i * freq/2 * sigma_z * tfinal) @ X
Ut = jnp.matmul(U_z, Xgate)
We are now ready to define the gate QOCProblem. The setup mirrors Section 5 exactly, with two changes: we pass U0 as the initial condition (so the solver evolves the propagator rather than a state vector) and Ut as the target. Rather than starting from scratch, we use the state-optimised parameters opt_params as an initial guess — this warm start significantly speeds up convergence, since the pulse is already close to performing an X rotation.
For this optimisation we use ScipyMinimiser with the L-BFGS-B algorithm, which is well suited to smooth, gradient-based problems and offers fine-grained convergence control.
from qruise.toolset import ScipyMinimiser
# gate QOC problem, warm-started from state-optimised params
gate_opt_prob = QOCProblem(H, U0, opt_params, (t0, tfinal), Ut, loss_gate_fidelity)
# L-BFGS-B minimiser
gate_minimiser = ScipyMinimiser("L-BFGS-B", maxiter=2000, tol=1e-12)
gate_opt = Optimiser(gate_minimiser, PWCSolver(n=grid_size, eq=gate_opt_prob.sepwc()))
gate_opt.set_optimisation(gate_opt_prob.loss)
# run optimisation
gate_opt_params, gate_opt_summary = gate_opt.optimise(*gate_opt_prob.problem())
Let's compare the state-optimised and gate-optimised pulse shapes.
canvas = PlotUtil(x_axis_label="t [ns]", y_axis_label="Amplitude", notebook=True)
pprint(f"Initial (state-optimised) {opt_params}")
pprint(f"Gate-optimised {gate_opt_params}")
canvas.plot(ts / 1e-9, drive(ts, opt_params), labels=["Initial (state-optimised)"])
canvas.plot(ts / 1e-9, drive(ts, gate_opt_params), labels=["Gate-optimised"])
canvas.show_canvas()
("Initial (state-optimised) {'a': Array(3.14144187, dtype=float64), 'omega': "
"Array(-2.80249514e+09, dtype=float64), 'sigma': Array(1.72359068e-09, "
'dtype=float64)}')
("Gate-optimised {'a': Array(3.14144187, dtype=float64), 'omega': "
"Array(-2.80249514e+09, dtype=float64), 'sigma': Array(3.61435632e-09, "
'dtype=float64)}')
We can also visualise how the gate-optimised pulse affects the state populations when starting from the ground state.
# state-vector solver
state_solver = PWCSolver(n=grid_size, store=True)
state_solver.set_system(opt_prob.sepwc())
# simulate with state-optimised and gate-optimised params
_, state_opt_state_res = state_solver.evolve(
y0.astype(jnp.complex128), opt_params, (t0, tfinal)
)
_, gate_opt_state_res = state_solver.evolve(
y0.astype(jnp.complex128), gate_opt_params, (t0, tfinal)
)
# get populations
init_pop = get_population(state_opt_state_res)
gate_opt_pop = get_population(gate_opt_state_res)
canvas.init_canvas(x_axis_label="t [ns]", y_axis_label="Population")
canvas.plot(
ts / 1e-9,
init_pop,
labels=["State-opt. ground state", "State-opt. 1st excited state"],
)
canvas.plot(
ts / 1e-9,
gate_opt_pop,
labels=["Gate-opt. ground state", "Gate-opt. 1st excited state"],
)
canvas.show_canvas()
Finally, let's quantify the improvement by computing the gate fidelity before and after optimisation.
# unitary solver (separate from state-vector solver to avoid type inconsistency)
unitary_solver = PWCSolver(n=grid_size, store=False)
unitary_solver.set_system(gate_opt_prob.sepwc())
# compute propagator with original and gate-optimised parameters
_, U_unopt = unitary_solver.evolve(U0, params, (t0, tfinal))
_, U_gate_opt = unitary_solver.evolve(
*gate_opt_prob.problem(params=gate_opt_params)[:-1]
)
gate_fidelity_before = 1 - loss_gate_fidelity(U_unopt, Ut)
gate_fidelity_after = 1 - loss_gate_fidelity(U_gate_opt, Ut)
pprint(
f"Gate fidelity (unoptimised, original params): {gate_fidelity_before*100:.1f} %"
)
pprint(f"Gate fidelity (gate-optimised): {gate_fidelity_after*100:.1f} %")
'Gate fidelity (unoptimised, original params): 69.5 %' 'Gate fidelity (gate-optimised): 99.8 %'
We can clearly see that our optimisation has once again significantly enhanced the fidelity!
Congratulations, you’ve just performed optimal control on a spin qubit, using both state and gate fidelity! You can now use these methods to start optimising your qubit performance.