Simulation of a single NV centre qubit¶

In this tutorial, we’ll show you, step-by-step, how to simulate the dynamics of a single spin in an NV centre and perform optimal control using the qruise-toolset.

The tutorial consists of the following:

1. Defining the Hamiltonian (theory)
2. Defining the drive
3. Simulating NV spin without noise
4. Simulating NV spin with noise
5. Performing optimal control

1. Defining the Hamiltonian (theory)¶

The system we are interested in here is an NV centre. Its Hamiltonian, $H$, is given by:

$$ H = \underbrace{D S_z^2}_{\text{zero-field splitting}} + \underbrace{\gamma_e B_0 S_z}_{\text{Zeeman interaction}} + \underbrace{\gamma_e B_1(t) \left[\cos(\omega t + \phi) S_x + \sin(\omega t + \phi) S_y\right]}_{\text{driving field interaction}}, \tag{1} $$

where $S_{x,y,z}$ are the $S=1$ spin operators.

To better understand our Hamiltonian, let’s break it down into the different terms.

• Zero-field splitting: represents the energy difference between the $m_s=0$ and $m_s = \pm 1$ spin states in the absence of an external magnetic field, with $D$ being the zero-field splitting constant.
• Zeeman interaction: describes the interaction of the electron spin with a static magnetic field, $B_0$. $\gamma_e$ is the gyromagnetic ratio of the electron, which quantifies its sensitivity to the field.
• Driving field interaction: describes the effect on the electron of the applied magnetic field, which is used to control the spin state of the NV centre. This field oscillates in time $t$, with $B_1(t)$, $\omega$, and $\phi$ being the amplitude, frequency, and relative phase of the magnetic field, respectively.

We can approximate our system to a two-level system by choosing our drive frequency to be resonant with the $|0\rangle \rightarrow |-1\rangle$ transition:

$$ \omega = D -\gamma_e B_0 . \tag{2} $$

If the condition $D+\gamma_e B_0 \gg \gamma_e B_1$ is met, the $m_s=0 \rightarrow 1$ transition frequency will be much greater than that of the $m_s=0 \rightarrow -1$ transition, approximating a two-level system, with $m_s=0$ denoting the ground state $|0\rangle$ and $m_s=-1$ denoting the $|1\rangle$ state.

Substituting the Pauli matrices, $\sigma_{x,y,z}$, in for the spin operators, we then get

$$ H = D \frac{\sigma_z}{2} - \gamma_e B_0 \frac{\sigma_z}{2} + \gamma_e B_1(t) \left[\cos(\omega t + \phi) \frac{\sigma_x}{2} + \sin(\omega t + \phi) \frac{\sigma_y}{2}\right]. \tag{3} $$

We can greatly simplify our Hamiltonian by moving into the rotating frame. Doing this, we absorb the drive into the complex Rabi rate $\Omega(t) = \gamma_e B_1(t)e^{i\phi(t)}$, giving

$$ H= \Omega(t) \frac{\sigma_x}{2}. \tag{4} $$

Taking into account the effect of noise, our final Hamiltonian is given by

$$ H = \delta \frac{\sigma_z}{2} + (1+ \beta)\Omega(t)\frac{\sigma_x}{2}, \tag{5} $$

where $\delta$ accounts for inhomogeneities in the static magnetic field ($B_0$) and hyperfine interactions between the NV centre and nearby nuclear spins, and $\beta$ represents the variation in Rabi rate due to inhomogeneities in the applied microwave field ($B_1$).

Now that we have a simplified version of our Hamiltonian that takes into account the relevant sources of noise, we can start coding it. Let's start by defining our external drive.

---

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 script's preamble:

In [1]:

import jax

jax.config.update("jax_enable_x64", True)

import jax jax.config.update("jax_enable_x64", True)

---

2. Defining the 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 it reaches its maximum amplitude

We can use this equation to define our drive pulse, which takes the time and pulse parameters as inputs.

In [2]:

import jax.numpy as jnp


def drive(t, params):
    """Gaussian"""

    amp = params["a"]  # Gaussian pulse amplitude
    sigma = params["sigma"]  # Std dev of Gaussian pulse (controls pulse width)

    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

import jax.numpy as jnp def drive(t, params): """Gaussian""" amp = params["a"] # Gaussian pulse amplitude sigma = params["sigma"] # Std dev of Gaussian pulse (controls pulse width) 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.

In [3]:

# 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,
}

# 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, }

It’s a good idea to check your drive behaves as desired before proceeding, so let’s quickly plot it.

In [4]:

from qruise.toolset.plots import PlotUtil

canvas = PlotUtil(x_axis_label="t [s]", y_axis_label="Amplitude", notebook=True)
canvas.plot(ts, drive(ts, params), labels=["Pulse"])
canvas.show_canvas()

from qruise.toolset.plots import PlotUtil canvas = PlotUtil(x_axis_label="t [s]", y_axis_label="Amplitude", notebook=True) canvas.plot(ts, drive(ts, params), labels=["Pulse"]) canvas.show_canvas()

Great, our drive looks as expected! Now we can start simulating our NV centre spin, firstly without including the effects of noise.

3. Simulating NV spin without noise¶

We'll start by defining our Hamiltonian. Without noise, $\delta=0$ and $\beta=0$, so we have only the drive term in our Hamiltonian, i.e. our Hamiltonian is given by Equation 4. We'll use Qruise's Hamiltonian function, with the first input, which represents the stationary term, set to None, and the second input being our drive term.

In [5]:

from qruise.toolset import Hamiltonian
from qutip import sigmax

H = Hamiltonian(None, [(sigmax() / 2, drive)])

from qruise.toolset import Hamiltonian from qutip import sigmax H = Hamiltonian(None, [(sigmax() / 2, drive)])

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, here the ground state)
• the pulse parameters (params)
• the time interval of the simulation (t0 to tfinal).

In [6]:

from qruise.toolset import Problem

# define initial qubit state
y0 = jnp.array([1.0, 0.0], dtype=jnp.complex128)

# define Problem
prob = Problem(H, y0, params, (t0, tfinal))

from qruise.toolset import Problem # define initial qubit state y0 = jnp.array([1.0, 0.0], dtype=jnp.complex128) # 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.

In [7]:

from qruise.toolset import PWCSolver

# uses PWCSolver to solve Schrödinger equation
solver = PWCSolver(dt=dt, store=True)
solver.set_system(prob.sepwc())
res = solver.evolve(*prob.problem())

from qruise.toolset import PWCSolver # 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.

---

Finally, we can view the results of our simulation by calculating and plotting the expectation value of the spin state populations against time.

In [8]:

from qruise.toolset.utils 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 [s]", y_axis_label="Population")
canvas.plot(ts, pop, labels=["Ground state", "1st excited state"])
canvas.show_canvas()

from qruise.toolset.utils 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 [s]", y_axis_label="Population") canvas.plot(ts, pop, labels=["Ground state", "1st excited state"]) canvas.show_canvas()

Congratulations! You’ve just simulated your first NV centre. You’ll notice, however, that the population exchange between the ground and first excited states is incomplete. We can improve this by performing quantum optimal control, which we'll do in Section 5 of this tutorial, but first we'll simulate our system again, this time including noise.

4. Simulating NV spin with noise¶

For this simulation, we need to start by defining our noise parameters, $\delta$ and $\beta$ (see Equation 5).

In [9]:

# define noise parameters
delta = 1e5  # accounts for inhomogeneities in B0 and hyperfine interactions (Hz)
beta = 0.05  # variation in Omega due to inhomogeneities in B1

# define noise parameters delta = 1e5 # accounts for inhomogeneities in B0 and hyperfine interactions (Hz) beta = 0.05 # variation in Omega due to inhomogeneities in B1

Then we can code our Hamiltonian, firstly by defining our stationary Hamiltonian, then by adding the drive term.

In [10]:

from qutip import sigmaz

# define stationary Hamiltonian
H = Hamiltonian(delta * sigmaz() / 2)

# add time-dependent term
H.add_term((1 + beta) * sigmax() / 2, drive)

from qutip import sigmaz # define stationary Hamiltonian H = Hamiltonian(delta * sigmaz() / 2) # add time-dependent term H.add_term((1 + beta) * sigmax() / 2, drive)

To solve the Schrödinger equation for this new Hamiltonian, we simply need to redefine our Problem using the new Hamiltonian, H, and solve again using the PWCSolver.

In [11]:

# define Problem with new Hamiltonian
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())

# define Problem with new Hamiltonian 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())

We can view the results of our new simulation by again calculating and plotting the expectation value of the transmon state populations against time.

In [12]:

# calculates the expectation value of each qubit state
pop = get_population(res)

# plots populations against time
canvas.init_canvas(x_axis_label="t [s]", y_axis_label="Population")
canvas.plot(ts, pop, labels=["Ground state", "1st excited state"])
canvas.show_canvas()

# calculates the expectation value of each qubit state pop = get_population(res) # plots populations against time canvas.init_canvas(x_axis_label="t [s]", y_axis_label="Population") canvas.plot(ts, pop, labels=["Ground state", "1st excited state"]) canvas.show_canvas()

Excellent! Now that we know how to simulate a spin in an NV centre both without and with noise, we can try performing optimal control to see if we can get a better fidelity.

5. Optimal control¶

As we saw in our population plots, 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. We'll perform the optimisation on the noisy system, but the process is identical for the system without noise.

We start by defining 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⟩||, \tag{7} $$

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:

In [13]:

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

def loss(x, y): """ Returns the infidelity (1 - ||^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, t0 to tfinal
• 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.

---

In [14]:

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)

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 our ScipyMinimiser library for optimisation.

In [15]:

from qruise.toolset import Optimiser, ScipyMinimiser

# define parameter bounds
params_bounds = {"sigma": (1e-12, tfinal), "a": (0.01, 20.0)}

# find optimal parameters
minimiser = ScipyMinimiser("L-BFGS-B", tol=1e-4, bounds=params_bounds)
opt = Optimiser(minimiser, PWCSolver(dt=dt, eq=prob.sepwc()))
opt.set_optimisation(opt_prob.loss)
result, minimiser_summary = opt.optimise(*opt_prob.problem())

from qruise.toolset import Optimiser, ScipyMinimiser # define parameter bounds params_bounds = {"sigma": (1e-12, tfinal), "a": (0.01, 20.0)} # find optimal parameters minimiser = ScipyMinimiser("L-BFGS-B", tol=1e-4, bounds=params_bounds) opt = Optimiser(minimiser, PWCSolver(dt=dt, eq=prob.sepwc())) opt.set_optimisation(opt_prob.loss) result, minimiser_summary = opt.optimise(*opt_prob.problem())

We can now compare the initial values of the parameters with those the optimisation yielded.

In [16]:

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 {result}")

# Plot initial (unoptimised) and optimised parameters
canvas.plot(ts, drive(ts, params), labels=["Unoptimised"])
canvas.plot(ts, drive(ts, result), labels=["Optimised"])
canvas.show_canvas()

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 {result}") # Plot initial (unoptimised) and optimised parameters canvas.plot(ts, drive(ts, params), labels=["Unoptimised"]) canvas.plot(ts, drive(ts, result), labels=["Optimised"]) canvas.show_canvas()

"Unoptimised {'a': 2.0, 'sigma': 4e-09}"
("Optimised {'a': np.float64(2.9669300803146355), 'sigma': "
 'np.float64(6.620637250427527e-10)}')

We can see that the optimised pulse has increased amplitude and a slightly larger variance compared to the unoptimised pulse.

Let’s now simulate and plot the state populations to see if our optimisation worked.

In [17]:

# 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 [s]", y_axis_label="Population")
canvas.plot(ts, pop, labels=["Unopt. ground state", "Unopt. 1st excited state"])
canvas.plot(ts, opt_pop, labels=["Opt. ground state", "Opt. 1st excited state"])
canvas.show_canvas()

# 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 [s]", y_axis_label="Population") canvas.plot(ts, pop, labels=["Unopt. ground state", "Unopt. 1st excited state"]) canvas.plot(ts, 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 almost 0% and 100%, respectively, compared to 26% and 74% 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).

In [18]:

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} %")

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: 74.10818 %'
'Optimised: 99.98269 %'

We can clearly see that our optimisation has significantly enhanced the fidelity!

Congratulations, you’ve just performed optimal control on a spin in an NV centre! You can now use these methods to start optimising your qubit performance.