Simulation & optimal control of a single-qubit gate in an NV quantum processor¶
In this tutorial, we’ll show you how to simulate and optimise a single-qubit gate in an NV centre quantum processor using the qruise-toolset
. The system consists of an NV centre coupled to two nuclear spins: ${}^{13}\text{C}$ and ${}^{15}\text{N}$.
The tutorial consists of the following:
- Defining the Hamiltonian (theory)
- Defining the stationary Hamiltonian
- Defining the external drive
- Solving the Schrödinger equation
- Performing quantum optimal control
1. Defining the Hamiltonian (theory)¶
The general Hamiltonian for an NV centre coupled to $i$ nuclear spins (where $i$ is the index of each spin) is:
$$ H = \underbrace{D \left(S_z^2 - \tfrac{3}{2}\right)\vphantom{\sum}}_{\text{zero-field splitting}} + \underbrace{\gamma_e B_0 S_z \vphantom{\sum}}_{\text{NV Zeeman interaction}} + \underbrace{\gamma_e B_1^e(t) S_x \vphantom{\sum}}_{\text{NV drive interaction}} + \sum\limits_{i} \Big[\underbrace{\gamma_i B_0 J_{iz} \vphantom{\sum}}_{\text{nuclear Zeeman interaction}} + \underbrace{\gamma_i B_{1}^n(t) J_{ix} \vphantom{\sum}}_{\text{nuclear drive interaction}} + \underbrace{ S \cdot A_i \cdot J_i \vphantom{\sum}}_{\text{hyperfine interaction}} + \underbrace{Q_i \left(J_{iz}^2 - \tfrac{3}{2}\right) \vphantom{\sum}}_{\text{quadrupole interaction}} \Big]. \tag{1} $$
Here, $S_{x,y,z}$ are the spin-$1$ operators for the NV electronic spin, and $J_{i,x,y,z}$ are the spin-$\frac{1}{2}$ operators for the $i^\text{th}$ nuclear spin.
If you're not familiar with this Hamiltonian, it might look overwhelming, so let’s break it down into the different terms to better understand what's going on.
Zero-field splitting (NV): represents the energy difference between the $m_s=0$ and $m_s = \pm 1$ spin states of the NV centre in the absence of an external magnetic field, with $D$ being the zero-field splitting constant.
Zeeman interaction (NV and nuclear): describes the interaction of a spin with a static magnetic field, $B_0$. $\gamma_e$ and $\gamma_i$ are the gyromagnetic ratios of the NV electron spin and the $i^\text{th}$ nuclear spin, respectively, a parameter that quantifies their sensitivity to the field.
Drive interaction (NV and nuclear): describes the effect of time-dependent magnetic fields used to control the spin states. $B_1^e(t)$ is the amplitude of the oscillating microwave drive field applied to the NV electronic spin and $B_1^n(t)$ is the amplitude of the oscillating radio-frequency drive field applied to the nuclear spins.
Hyperfine interaction (NV-nuclear): represents the coupling between the NV spin and the nearby nuclear spins, where $A_i$ is the hyperfine coupling tensor describing the interaction between the NV spin and the $i^\text{th}$ nuclear spin.
Nuclear quadrupole interaction (nuclear): describes the interaction between nuclear spins with $J > \frac{1}{2}$ (e.g., ${}^{14}\mathrm{N}$) and the local electric field gradient, where $Q_i$ is the quadrupole tensor of the $i^\text{th}$ nucleus.
Now let's consider the quantum system we're working with: one NV centre spin coupled to two nuclear spins, a ${}^{13}\text{C}$ and a ${}^{15}\text{N}$. We can simplify our Hamiltonian for a single qubit gate with the following simplifications & approximations:
- Single-qubit gates and nuclear spin control occur in the $|m_s=-1\rangle$ subspace (computational subspace) as shown in the figure below. The zero-field splitting and the NV Zeeman interaction terms therefore can be ignored as they don't affect the system's dynamics.
[Taken from Chen et. al.]
$B_1^e(t) = 0$: Since the NV spin is not actively driven, the applied microwave field and therefore the NV drive interaction term are equal to zero.
$J=\frac{1}{2}$: ${}^{13}\text{C}$ and ${}^{15}\text{N}$ both have spin equal to $\frac{1}{2}$, so the quadrupole interaction term is equal to zero.
Secular approximation: In the presence of a strong static magnetic field, $B_0$, along the $z$-axis, the $S_x$, $S_y$, $J_{ix}$, and $J_{iy}$ spin components can be ignored under secular approximation, leaving only the $A_{i,zz} S_z J_{iz}$ term in the Hamiltonian.
We'll limit ourselves to nuclei that are nearly aligned with the NV axis. Due to this, diagonalisation of the nuclear spin Hamiltonian combines the hyperfine and nuclear Zeeman terms into an effective nuclear transition frequency: $$ \omega_i = \sqrt{(A_{i,zz} + \gamma_i B_0)^2 + A_{i,xz}^2 + A_{i,yz}^2}. \tag{2} $$
This leaves us with a significantly simplified Hamiltonian for our single-qubit gate:
$$ H_{\text{single}} = \sum_i \left(-\omega_i J_{iz} - \gamma_i B_{1}^n(t) J_{ix} \right). \tag{3} $$
We can simplify further by moving into the rotating frame, where the drive is absorbed into the complex Rabi rate $\Omega_i(t) = \gamma_i B_1^n(t)e^{i\phi}$, and applying the rotating wave approximation. In this notebook, we'll implement an X gate, which sets the phase $\phi=0$. Our Hamiltonian becomes:
$$ H = \sum_{i} \left(-\Delta \omega_i \frac{\sigma_{iz}}{2} - \Omega_i(t) \frac{\sigma_{ix}}{2} \right). \tag{4} $$
Here, $\Delta \omega_i = \omega_i - \omega_d$ is the offset of the drive frequency, $\omega_d$, from the transition frequency of the $i^{\text{th}}$ nuclear spin, and $\sigma_{iz}$ and $\sigma_{ix}$ are the Pauli operators of the $i^{\text{th}}$ nuclear spin.
Now that we have a sufficiently simplified Hamiltonian, we can start coding our NV system.
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)
2. Defining the stationary Hamiltonian¶
First, we’ll set up some general parameters for the simulation of our qubit. These include the strength of the static magnetic field ($B_0$), and the nuclear hyperfine coupling constants ($A_i$) and gyromagnetic ratios ($\gamma_i$).
# static magnetic field in z-direction (Tesla)
B0 = 0.62
# hyperfine coupling constants for 13C and 15N (Hz)
A_zz_C13 = 2.281e6
A_xz_C13 = 0.24e6
A_yz_C13 = 0.24e6
A_zz_N15 = 3.03e6
A_xz_N15 = 0.0
A_yz_N15 = 0.0
# gyromagnetic ratios for 13C and 15N (Hz/T)
gamma_C13 = 10.7084e6
gamma_N15 = -4.316e6
gamma_ratio = gamma_N15 / gamma_C13
We can then use these parameters to define the effective nuclear transition frequency, $\omega_i$, which is given by Equation 2 above.
import jax.numpy as jnp
# effective nuclear transition frequencies (Hz)
omega_C13 = jnp.sqrt((A_zz_C13 + gamma_C13 * B0) ** 2 + A_xz_C13**2 + A_yz_C13**2)
omega_N15 = jnp.sqrt((A_zz_N15 + gamma_N15 * B0) ** 2 + A_xz_N15**2 + A_yz_N15**2)
Before we can define our drive or our Hamiltonian, we need to define our operators. For this, we import the Pauli $\sigma$-operators and the identity operator from the QuTip library. In a multi-spin system, it's essential to ensure each operator acts only on the intended spin while leaving the other spins unaffected. To do this, we use tensor products, where the relevant $\sigma$-operator acts on the target spin and the identity acts on the other two spins.
from qutip import tensor, sigmax, sigmaz, qeye
# define identity operator using Hilbert space dimension
hilbert_dim = 2
Id = qeye(hilbert_dim)
# construct Pauli operators for NV spin using tensor products
sigmax_NV = tensor(sigmax(), Id, Id)
sigmaz_NV = tensor(sigmaz(), Id, Id)
# construct Pauli operators for 13C spin using tensor products
sigmax_C13 = tensor(Id, sigmax(), Id)
sigmaz_C13 = tensor(Id, sigmaz(), Id)
# construct Pauli operators for 15N spin using tensor products
sigmax_N15 = tensor(Id, Id, sigmax())
sigmaz_N15 = tensor(Id, Id, sigmaz())
Now we can define the stationary part of our Hamiltonian. In this case, we'll choose ${}^{13}\text{C}$ as the target spin, so we'll set the drive frequency equal to the transition frequency of ${}^{13} \text{C}$, i.e., $\Delta \omega_{13\text{C}} = 0$ and $\Delta \omega_{15\text{N}} = \omega_{15\text{N}} - \omega_{13\text{C}}$.
from qruise.toolset import Hamiltonian
# define offset from drive frequency (Hz)
delta_omega_C13 = (
omega_C13 - omega_C13
) # Set to zero because 13C is target qubit for gate operation
delta_omega_N15 = omega_N15 - omega_C13 # Approximately -8.27 MHz
# define stationary Hamiltonian
H = Hamiltonian(-(delta_omega_C13 * sigmaz_C13 / 2 + delta_omega_N15 * sigmaz_N15 / 2))
3. Defining the external drive¶
For our external drive, we'll use a Gaussian pulse given by:
$$ \Omega(t; \{a,\sigma,\mu\}) = \frac{a}{\sigma\sqrt{2\pi}} \mathrm{exp}\left[-\frac{(t - \mu)^2}{2\sigma^2}\right], \tag{6} $$
which is characterised by three parameters:
$a$, a scalar that adjusts the amplitude of the pulse
$\sigma$, the pulse variance
$\mu$, the time at which the pulse reaches its maximum amplitude
We can use this equation to define a Gaussian drive pulse on the two nuclei, which takes time and pulse parameters as inputs.
def drive_nuclei(t, params):
"""Gaussian drive on nuclei"""
amp = params["amp"]
sigma = params["sigma"]
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
Now we need to define the time and pulse parameters for the simulation.
# time parameters for simulation (s)
t0 = 0.0 # initial time
tfinal = 1e-6 # final time
grid_size = 1000 # simulation grid size (number of time points)
ts = jnp.linspace(t0, tfinal, grid_size)
# drive pulse parameters
v2hz = 1 # Volts-to-Hertz conversion factor
params = {"amp": 1.0, "sigma": 0.2 * tfinal} # 13C drive parameters
It’s a good idea to check your drive behaves as desired before proceeding, so let’s quickly plot it.
from qruise.toolset.plots import PlotUtil
canvas = PlotUtil(x_axis_label="t [a.u.]", y_axis_label="Amplitude [Hz]", notebook=True)
canvas.plot(ts, drive_nuclei(ts, params), labels=["Nuclei drive"])
canvas.show_canvas()
Great! As expected, we see a Gaussian pulse shape for the drive. Now we can add the time-dependent (drive) terms to the Hamiltonian.
# add time-dependent terms with drive to the Hamiltonian
H.add_term(v2hz * (sigmax_C13 + gamma_ratio * sigmax_N15), drive_nuclei)
4. Solving the Schrödinger equation¶
To solve the Schrödinger equation, we need to specify the equations and parameters that govern the system. In the qruise-toolset
, we combine these using a Problem
. To instantiate the Problem
, we need:
- the Hamiltonian (
H
) - the initial qubit state (
y0
) - the pulse parameters (
params
) - the time interval of the simulation (
t0
totfinal
).
Before we can define our initial qubit state, we need to define define and label our basis states. As we have three two-level systems, we have $2^3=8$ basis states: $|000\rangle, |001\rangle, |010\rangle, |011\rangle, |100\rangle, |101\rangle, |110\rangle, |111\rangle$. We'll set our initial qubit state as the ground state: $|000\rangle$.
from qruise.toolset.utils import computational_basis
from qruise.toolset import Problem
# define basis states and label
num_subsystems = 3 # number of subsystems (spins)
basis, labels = computational_basis([hilbert_dim] * num_subsystems)
# define initial qubit state (ground state)
y0 = basis[labels.index((0, 0, 0))]
# define Problem
prob = Problem(H, y0, params, (t0, tfinal))
The user can then define the type of solver they want to use and the equation they want to solve (for example, the Schrödinger equation, the master equation, or the Lindblad equation).
In this case, we’ll choose to use the piece-wise constant solver (PWCSolver
) to solve the Schrödinger equation (sepwc
). We can then calculate the wavefunction of the system at each timestamp (n=grid_size
).
from qruise.toolset import PWCSolver
# 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.
Finally, we can view the results of our simulation by calculating and plotting the expectation value of the spin state populations against time.
from qruise.toolset import get_population
from qruise.toolset.plots import qruise_colour_palette
# calculate the expectation value of each qubit state
pop = get_population(res)
# plot populations against time
canvas.init_canvas(
x_axis_label="t [s]",
y_axis_label="Population",
colour_palette=qruise_colour_palette,
)
canvas.plot(ts, pop, labels=[str(label) for label in labels])
canvas.show_canvas()
Nice! You've just simulated a single-qubit gate in an NV quantum processor. We observe a decrease in the population of the $|000\rangle$ basis state and an increase in the population of the $|010\rangle$ basis state, i.e., the excited state of the ${}^{13}\text{C}$ nuclear spin. However, we also see some undesired population in other states.
For an ideal X gate, we would expect a full population inversion between the $|000\rangle$ and $|010\rangle$ states with the population in all other states remaining at zero. We can perform quantum optimal control to improve the gate fidelity.
5. Performing quantum optimal control¶
To perform optimal control, we need a metric that quantifies how far we are from the target state. For this, we define a loss function, $\mathcal{L}$; for example, the infidelity:
$$ \mathcal{L}=1-\mathcal{F}=1− |\langle \psi(t=t_\text{final})|\psi_t \rangle|^2, \tag{7} $$
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}$.
# define loss function (infidelity)
def loss(x, y):
"""
Returns the infidelity
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 (
t0
totfinal
) - the desired qubit state (
yt
, i.e., $|010\rangle$) - the loss function (
loss
)
Note: We defined the first four inputs earlier in the tutorial.
from qruise.toolset import QOCProblem
# define target state |010>
yt = basis[labels.index((0, 1, 0))]
# define optimal control problem
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 L-BFGS-B algorithm provided by the ScipyMinimiser
class from qruise-toolset
. L-BFGS-B is a variant of the Broyden–Fletcher–Goldfarb–Shanno (BFGS) optimisation method that supports parameter bounds, allowing us to constrain parameters within physically meaningful ranges during the optimisation.
from qruise.toolset import Optimiser, ScipyMinimiser
# define parameter bounds
params_bounds = {
"sigma": (1e-12, 0.4 * tfinal),
"amp": (0.01, 20.0),
}
# find optimal parameters
minimiser = ScipyMinimiser("L-BFGS-B", tol=1e-4, bounds=params_bounds)
opt = Optimiser(minimiser, PWCSolver(n=grid_size, eq=prob.sepwc()))
opt.set_optimisation(opt_prob.loss)
opt_params, minimiser_summary = opt.optimise(*opt_prob.problem())
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 [s]", 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, drive_nuclei(ts, params), labels=["Unoptimised"])
canvas.plot(ts, drive_nuclei(ts, opt_params), labels=["Optimised"])
canvas.show_canvas()
"Unoptimised {'amp': 1.0, 'sigma': 2e-07}" ("Optimised {'amp': Array(1.72517793, dtype=float64), 'sigma': " 'Array(2.94733233e-07, dtype=float64)}')
We can see that the optimised pulse has increased amplitude and variance compared to the unoptimised pulse.
Let’s now simulate and plot the state populations to see if our optimisation worked.
from qruise.toolset import get_population
# solve with optimised parameters
_, opt_res = solver.evolve(*prob.problem(params=opt_params))
# 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 [s]", y_axis_label="Population")
canvas.plot(ts, pop[:, 0], labels=["Unopt. (0, 0, 0)"])
canvas.plot(ts, pop[:, 2], labels=["Unopt. (0, 1, 0)"])
canvas.plot(ts, opt_pop[:, 0], labels=["Opt. (0, 0, 0)"])
canvas.plot(ts, opt_pop[:, 2], labels=["Opt. (0, 1, 0)"])
canvas.show_canvas()
Note: For simplicity, we have only plotted the $|000\rangle$ and $|010\rangle$ basis states, but it's a good idea to check all the states to make sure nothing unexpected occurs with the optimised drive parameters.
We can see here that the population exchange between the $|000\rangle$ and $|010\rangle$ states is almost perfect. The final populations are almost 0% and 100%, respectively, compared to 30% and 69% 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:.5f} %")
pprint(f"Optimised: {fidelity[1]*100:.5f} %")
'Unoptimised: 69.12456 %' 'Optimised: 99.99954 %'
We can see that our optimisation has significantly enhanced the fidelity!
Congratulations, you’ve just performed optimal control on single-qubit gate in an NV quantum processor! You can now use these methods to start optimising your qubit performance.