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
- Optimal control
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 total pulse duration
(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)
dt = 1e-11 # time step (s)
ts = jnp.arange(t0, tfinal, dt)
# 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], dtype=jnp.complex128)
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.
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(dt=dt, 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. 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−⟨\psi(t=T)|\psi_t⟩||, $$
where $|\psi(t=t_{final})\rangle$ is the wavefunction of the system at the end of the pulse duration, 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], dtype=jnp.complex128) # define desired state
opt_prob = QOCProblem(H, y0, params, (t0, tfinal), yt, loss)
The whole workflow now reduces to simulating the dynamics using the initial guess values for the parameters we defined in the last tutorial. At the end of the simulation, we get the new wavefunction and calculate the infidelity with respect to the target wavefunction. We then calculate the gradients of the loss function for the initial value of the parameters. Based on the gradient values, we update the parameter values iteratively until they converge and minimise the loss function.
The Optimiser class is provided for this specific task. As a user, you choose the solver you want to use to simulate the dynamics and which algorithm to use to carry out the optimisation. Here, we use a piecewise constant solver (PWCOptimiser) and the Broyden–Fletcher–Goldfarb–Shannon (BFGS) method from the optimistix library for optimisation.
import optimistix as optx
from qruise.toolset import OptimistixMinimiser, Optimiser
minimiser = OptimistixMinimiser(optx.BFGS, rtol=1e-2, atol=1e-3)
opt = Optimiser(minimiser, PWCSolver(dt=dt, eq=prob.sepwc()))
# Set the loss function for the optimiser to minimise (infidelity)
opt.set_optimisation(opt_prob.loss)
# Run optimisation on problem
result, opt_summary = opt.optimise(*opt_prob.problem())
# Check the simulation with optimised parameters
res_opt = solver.evolve(y0, result, (t0, tfinal))
Note: Since qruise-simple 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 {result}")
# Plot initial (unoptimised) and optimised parameters
canvas.plot(ts / 1e-9, drive(ts, params), labels=["Unoptimised"])
canvas.plot(ts / 1e-9, drive(ts, result), labels=["Optimised"])
canvas.show_canvas()
"Unoptimised {'a': 2.0, 'sigma': 4e-09, 'omega': -2802495138.6169}"
("Optimised {'a': Array(3.1414308, dtype=float64), 'omega': "
"Array(-2.80249514e+09, dtype=float64), 'sigma': Array(1.73108914e-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.
# simulate with new set of optimised parameters
from qruise.toolset import get_population
opt_res = solver.evolve(y0, result, (t0, tfinal))
# 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(res_opt[-1, :], yt)]
pprint(f"Unoptimised: {fidelity[0]*100:.1f} %")
pprint(f"Optimised: {fidelity[1]*100:.1f} %")
'Unoptimised: 69.6 %' 'Optimised: 99.7 %'
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.