Simulation & optimal control of a single superconducting qubit¶
In this tutorial, we’ll show you, step-by-step, how to simulate the dynamics of a single superconducting qubit and perform optimal control using the qruise-toolset.
The tutorial consists of the following:
- Defining the stationary Hamiltonian
- Defining the external drive
- Solving the Schrödinger equation
- Plotting the qubit dynamics
- Optimal control
1. Defining the stationary Hamiltonian¶
The system we are interested in here is a single three-level transmon. The transmon circuit (shown in the top panel of the figure below) consists of a Josephson junction (JJ), with inductance $L_J$ and capacitance $C_J$, connected in series to a capacitor of capacitance $C_s$. The transmon energy levels as a function of the phase difference, $\phi$, across the JJ are plotted in the bottom panel.
The Hamiltonian, $H$, of our system is given by
$$ H=\omega_q a^\dagger a + \frac{\alpha}{2} a^\dagger a^\dagger a a , $$
where $\omega_q=\omega_{01}$ is the qubit resonance frequency, $\alpha = \omega_q - \omega_{12}$ is the anharmonicity, where $\omega_{12}$ is the qubit resonance frequency between the 1st and 2nd excited states. $a^\dagger$ and $a$ are the creation and annihilation operators, respectively.
First, we’ll set up some general parameters for the simulation of our three-level transmon. We’ll use $\omega_q$ = 5 GHz and $\alpha$ = -210 MHz, which lie in the standard range of such a system.
Tip: Make sure that your environment is correctly set to handle float64 precision by setting JAX_ENABLE_X64=True or add
import jax
jax.config.update("jax_enable_x64", True)
to your scripts' preamble.
# Transmon constants
freq = 5e9 # qubit frequency, omega_q (Hz)
anhar = -210e6 # anharmonicity, alpha (Hz)
qubit_lvls = 3 # number of qubit energy levels
We also need to define $a$ and $a^{\dagger}$ by using the destroy and create functions, which can be imported from the qutip library.
from qutip import destroy, create
# define operators
a = destroy(qubit_lvls)
a_dag = create(qubit_lvls)
We then have all we need to define the stationary (time-independent) part of the Hamiltonian.
from qruise.toolset import Hamiltonian
# define stationary Hamiltonian
H = Hamiltonian(freq * a_dag @ a + anhar / 2 * a_dag @ a_dag @ a @ a)
2. Defining the external drive¶
To add the time-dependent part of the Hamiltonian, we need to define our drive function, $c(t)$. For this system, 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), $$
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 it reaches its maximum amplitude
(iv) $\omega_{d}$, the local oscillator frequency, which is often slightly detuned from $\omega_q$.
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 parameters for the simulation and a few pulse parameters.
# 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) # defines time axis values
# drive pulse parameters
sideband = 50e6 # frequency shift between qubit and drive pulse
v2hz = 1e9 # pulse amplitude conversion factor (Volts to Hz)
params = {
"a": 3e-9,
"sigma": 0.5 * tfinal,
"omega": freq + sideband,
}
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 function and an additional interaction term to the Hamiltonian to account for the effects of our drive pulse. The interaction term, v2hz * (a_dag + a), represents the interaction between the qubit and the external driving pulse.
# Add the drive term to the Hamiltonian
H.add_term(v2hz * (a_dag + a), drive) # H(t) = (a_dag + a) * drive(t, params)
3. 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, here the ground state) - the pulse parameters (
params) - the time interval of the simulation (
t0totfinal).
from qruise.toolset import Problem
# define initial qubit state (ground state)
y0 = jnp.array([1.0, 0.0, 0.0], dtype=jnp.complex128)
# H, params, t0, tfinal already defined
# 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.
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.
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", "2nd excited state"]
)
canvas.show_canvas()
Congratulations! You’ve just simulated your first three-level transmon. 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¶
We often want the drive function to implement a specific quantum gate, e.g. an X gate. For this to happen, we need to first determine which pulse parameters result in that specific gate operation.
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, with no population in the second excited state.
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, 0.0], dtype=jnp.complex128) # define desired state
# y0 = initial state was defined above
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
from qruise.toolset import Optimiser
minimiser = OptimistixMinimiser(optx.BFGS, rtol=1e-2, atol=1e-3)
opt = Optimiser(minimiser, PWCSolver(dt=dt, eq=prob.sepwc()))
# set the loss function and the time interval
opt.set_optimisation(opt_prob.loss)
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
canvas = PlotUtil(x_axis_label="t [ns]", y_axis_label="Amplitude", notebook=True)
pprint(f"Unoptimised {params}")
pprint(f"Optimised {result}")
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': 3e-09, 'sigma': 1e-08, 'omega': 5050000000.0}"
("Optimised {'a': Array(1.85645497e-08, dtype=float64), 'omega': "
"Array(5.05e+09, dtype=float64), 'sigma': Array(4.34472589e-08, "
'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",
"Unopt. 2nd excited state",
],
)
canvas.plot(
ts / 1e-9,
opt_pop,
labels=["Opt. ground state", "Opt. 1st excited state", "Unopt. 2nd excited state"],
)
canvas.show_canvas()
We can see here that the population exchange between the ground and first excited states is much more complete with the optimised parameters. The final populations are around 4% and 88%, respectively, compared to 33% and 62% before optimisation. Our optimisation was successful!
Note: You’ll notice that we still have a non-zero population in the second excited state. This is due to unintended excitation by our drive pulse to higher states. In order to eliminate this unwanted excitation, we need to instead use a DRAG pulse.
We can quantify our 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: 62.4 %' 'Optimised: 90.0 %'
We can clearly see that our optimisation enhanced the fidelity!
Congratulations, you’ve just performed optimal control on a three-level transmon! You can now use these methods to start optimising your qubit performance.