Simulating a Mølmer-Sørensen gate¶
This tutorial shows you step by step how to simulate a Mølmer-Sørensen (MS) gate in a trapped-ion system using the qruise-toolset. The MS gate creates maximally entangled states by coupling two qubits through a shared motional mode. We'll first simulate an ideal gate using a rectangular pulse, and then include a finite rise time to model a more realistic, imperfect gate.
The tutorial consists of the following:
- Defining the Hamiltonian (theory)
- Defining the external drive
- Defining the Hamiltonian
- Solving the Schrödinger equation
- Adding Gaussian filter (imperfect gate)
1. Defining the Hamiltonian (theory)¶
The system we're interested in consists of two trapped ion qubits coupled through a shared motional mode. Its Hamiltonian, $\hat{H}$, in the interaction frame, within the Lamb-Dicke limit, is given by
$$ \hat{H}=-2 \hbar \eta \Omega \hat{J}_y\left(\hat{a} e^{i \delta t} + \hat{a}^\dagger e^{-i \delta t}\right) . $$
Here, $\eta$ is the Lamb-Dicke parameter, $\Omega$ is the Rabi frequency associated with the laser-ion interaction, $\delta$ is the detuning from the sideband transition frequency, and $\hat{a}$ and $\hat{a}^{\dagger}$ are the annihilation and creation operators of the shared motional mode, respectively. $\hat{J}_y$ is the global spin-$y$ operator in the two-qubit subspace and is given by
$$ \hat{J}_y=\frac{1}{2}(\hat{I} \otimes \hat{\sigma}_y+\hat{\sigma}_y \otimes \hat{I}) ,$$
where $\hat{I}$ is the identity operator and $\hat{\sigma}_y$ is the Pauli $y$-matrix acting on the qubit. You can read more about where this Hamiltonian comes from here.
Note: Within the Lamb-Dicke limit, the spatial extent of the ion's motion is much smaller than the wavelength of the driving laser field and $\eta \ll 1$.
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 external drive¶
For our external drive, we'll use a piecewise constant (PWC) pulse. This type of pulse is a good choice for this exercise as it allows us to independently control the amplitude at each time step, rather than being constrained by a predefined function. To start with, we'll use a rectangular pulse with a finite number of time steps, $N$, of duration $dt = \frac{t_{\text{final}}}{N}$, where $t_\text{final}$ is the total duration of our pulse.
Let's begin by defining the relevant system parameters.
import jax.numpy as jnp
# Hamiltonian parameters
Omega = 2 * jnp.pi * 8e5 # Rabi frequency (Hz)
eta = 0.1 # Lamb-Dicke parameter (Note: eta<<1 as required for the Lamb-Dicke limit)
delta = 4 * eta * Omega # detuning from sideband transition frequency (Hz)
# pulse parameters
amp = 1.0 # pulse amplitude, envelope for the drive
t0 = 0.0 # start time of simulation (s)
tfinal = jnp.pi / (2 * eta * Omega) # pulse length (s)
N = 1000 # number of steps
ts = jnp.linspace(0.0, tfinal, N) # time values
dt = float(ts[1] - ts[0]) # time step, dt (s)
# define dictionary of pulse parameters
parameters_rect = {
"amp": amp, # pulse amplitude
}
We now need to define a rectangular shape for the pulse envelope.
# define rectangular pulse
def rect_pulse(t, params):
"""Rectangular pulse. Returns amplitude at every time step."""
amp = params["amp"]
return jnp.ones_like(t) * amp
We'll use Qruise's PWCPulse class to create the drive pulse. For this, we need to define an amplitude array for the envelope using rect_pulse(). Then we can define a parameter dictionary for the drive pulse, consisting of the start time, the time step, and the envelope, and use it to create our pulse.
import qruise.toolset as qr
# define pulse amplitude array using rect_pulse
amplitude_array = jnp.asarray(rect_pulse(ts, parameters_rect))
# define parameters for piecewise constant pulse
pwc_params = {
"t0": (t0, False), # start time (s)
"dt": (dt, False), # time step (s)
"env": (amplitude_array, True), # pulse envelope (rectangle)
}
# create piecewise constant pulse using pwc_params
rabi_drive_pwc = qr.PWCPulse("pwc", pwc_params)
We can merge the parameter dictionaries to make them easier to work with. In this way, instead of calling parameters_rect for the envelope and rabi_drive_pwc.params for the PWC pulse, we can simply call params and the correct parameters are extracted automatically.
# merge parameter dictionaries
params = parameters_rect | rabi_drive_pwc.params
Let's quickly plot the pulse to see if it looks how we'd expect.
canvas = qr.PlotUtil(
x_axis_label="t [µs]", y_axis_label="Amplitude [a.u.]", notebook=True
)
canvas.plot(ts / 1e-6, rabi_drive_pwc(ts, params), labels=["Drive (PWC) pulse"])
canvas.show_canvas()
Great, we have a pulse envelope with amplitude equal to $1.0$ starting at $t=0.0$ ns and ending at $t=t_\text{final} = 2.5~\mu \text{s}$.
3. Defining the Hamiltonian¶
We can now start coding our Hamiltonian. First we'll define the time-dependent envelopes that appear in it using our PWC pulse. Each envelope sets the time dependence for one of the terms in the Hamiltonian, and includes the prefactor $-\eta \Omega$ and the oscillating term $e^{\pm i \delta t}$.
# define envelopes
def env1(t, params):
return rabi_drive_pwc(t, params) * (-eta * Omega * jnp.exp(1j * delta * t))
def env2(t, params):
return rabi_drive_pwc(t, params) * (-eta * Omega * jnp.exp(-1j * delta * t))
Let's quickly plot these to see how they look.
canvas = qr.PlotUtil(
x_axis_label="t [µs]", y_axis_label="Amplitude [a.u.]", notebook=True
)
canvas.plot(ts / 1e-6, jnp.imag(env1(ts, params)), labels=["Envelope 1"])
canvas.plot(ts / 1e-6, jnp.imag(env2(ts, params)), labels=["Envelope 2"])
canvas.show_canvas()
Great! We can now define the operators needed to construct the Hamiltonian. To do this, we'll first specify the Hilbert space dimensions of the system. For the qubits, this is 2 since they're two-level systems. The motional mode, being a harmonic oscillator, has an infinite-dimensional space, but for computational purposes we'll truncate it to 8 Fock states.
import qutip as qt
# Hilbert dimensions
q_dim = 2 # 2-level system
mot_dim = 8 # truncate to 8 Fock states
# qubit operators
sigma_y = qt.sigmay() # sigma_y
Id = qt.qeye(q_dim) # identity
Jy = qt.tensor(Id, sigma_y) + qt.tensor(sigma_y, Id) # J_y, global spin-y operator
# motional mode operators
a = qt.destroy(mot_dim) # a
a_dag = qt.create(mot_dim) # a-dagger
# spin-motion interaction operators
Ja = qt.tensor(Jy, a)
Ja_dag = qt.tensor(Jy, a_dag)
Now we can combine our drive envelopes and operators to create our time-dependent Hamiltonian.
# create time-dependent Hamiltonian
H = qr.Hamiltonian(None, [(Ja, env1), (Ja_dag, env2)])
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 (
t0totfinal).
Before we can define our initial qubit state, we need to define and label our basis states. Since we have 2 two-level systems (qubits) and a motional mode truncated to 8 levels, the total system has $ 2^2 \times 8 = 32 $ basis states:
- Qubits: $ |00\rangle, |01\rangle, |10\rangle, |11\rangle $
- Motional mode: $ |0\rangle, |1\rangle, |2\rangle, |3\rangle, |4\rangle, |5\rangle, |6\rangle, |7\rangle $
Thus, the full basis states of the system are:
$$ |000\rangle, |001\rangle, |002\rangle, |003\rangle, |004\rangle, |005\rangle, |006\rangle, |007\rangle, \\ |010\rangle, |011\rangle, |012\rangle, |013\rangle, |014\rangle, |015\rangle, |016\rangle, |017\rangle, \\ |100\rangle, |101\rangle, |102\rangle, |103\rangle, |104\rangle, |105\rangle, |106\rangle, |107\rangle, \\ |110\rangle, |111\rangle, |112\rangle, |113\rangle, |114\rangle, |115\rangle, |116\rangle, |117\rangle . $$
We'll start with both qubits and the motional mode in the ground states, so our initial state, y0, is $ |000\rangle $. We can then define our Problem.
# define basis states and label
basis, labels = qr.utils.computational_basis(
[q_dim, q_dim, mot_dim]
) # Q1, Q2, motional mode
# define states
state_000 = basis[labels.index((0, 0, 0))] # state |000>
state_110 = basis[labels.index((1, 1, 0))] # state |110>
state_010 = basis[labels.index((0, 1, 0))] # state |010>
state_100 = basis[labels.index((1, 0, 0))] # state |100>
# define initial qubit and motional mode state (ground state)
y0 = state_000 # initial state
# define Problem
problem = qr.Problem(H, y0, params, (t0, tfinal))
In the qruise-toolset, you can define the type of solver you want to use and the equation you want to be solved — for example, the Schrödinger equation, the master equation, or the Lindblad equation.
In this case, we’ll choose to use the piecewise constant solver (PWCSolver) to solve the Schrödinger equation (sepwc). We can then calculate the wavefunction of the system at each timestamp.
# use PWCSolver to solve Schrödinger equation
solver = qr.PWCSolver(n=N, store=True)
solver.set_system(problem.sepwc())
time, res = solver.evolve(*problem.problem())
Note: Setting store=True in the PWCSolver ensures that the results are stored at each timestamp during the simulation.
We can then calculate the expectation value of the basis state populations at each timestamp.
from qruise.toolset.utils import get_population
# calculate expectation value of each state
pop = get_population(res)
To keep the plot clear and readable, we’ll exclude basis states whose population remains zero for the entire gate duration.
import numpy as np
# Find states that are zero across all time points
zero_mask = np.all(pop == 0, axis=0)
# keep only non-zero states
nonzero_mask = ~zero_mask
filtered_pop = pop[:, nonzero_mask]
filtered_labels = [str(label) for i, label in enumerate(labels) if not zero_mask[i]]
# plot non-zero states
canvas.init_canvas(x_axis_label="t [µs]", y_axis_label="Population")
canvas.plot(ts / 1e-6, filtered_pop, labels=filtered_labels)
canvas.canvas.legend.spacing = -4
canvas.show_canvas()
As expected, the gate transforms the initial state into an equal superposition of $|000\rangle$ and $|110\rangle$, each with 50% population, while all other states remain unpopulated.
To check our gate is perfect, we can calculate the state fidelity, $\mathcal{F}$, which is given by
$$ \mathcal{F}=|\langle \psi(t=t_\text{final})|\psi_t \rangle|^2, $$
where $|\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.
We just need to define the desired final state, and then use it in the fidelity() function. To do so, we will consider all initial states and their respective final state and average the result, based on the MS gate truth table:
$$ \ket{gg} \rightarrow (\ket{gg} + i \ket{ee})/ \sqrt{2} \\ \ket{eg} \rightarrow (\ket{eg} - i \ket{ge})/ \sqrt{2} \\ \ket{ge} \rightarrow (\ket{ge} - i \ket{eg})/ \sqrt{2} \\ \ket{ee} \rightarrow (\ket{ee} + i \ket{gg})/ \sqrt{2} \\ $$
from pprint import pprint
initial_states = [state_000, state_100, state_010, state_110] # initial states
# final final states
yf_000 = (state_000 + 1j * state_110) / jnp.sqrt(2)
yf_100 = (state_100 - 1j * state_010) / jnp.sqrt(2)
yf_010 = (state_010 - 1j * state_100) / jnp.sqrt(2)
yf_110 = (state_110 + 1j * state_000) / jnp.sqrt(2)
final_states = [yf_000, yf_100, yf_010, yf_110] # final states
def get_fidelity_over_sates(H):
"""Calculates the average fidelity of the final states over the initial states."""
states_fidelity = [] # list to store fidelity results
# define fidelity function
def fidelity(x, y):
"""Returns the fidelity (|<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(o.conj() * o) # Returns the real part of the fidelity
for i, state in enumerate(initial_states):
# define the problem for the different initial states and solve them
problem = qr.Problem(H, state, params, (t0, tfinal))
solver = qr.PWCSolver(n=N, store=True)
solver.set_system(problem.sepwc())
time, res = solver.evolve(*problem.problem())
state_fidelity = fidelity(res[-1, :], final_states[i])
states_fidelity.append(state_fidelity) # append fidelity result
return sum(states_fidelity) / len(
states_fidelity
) # return average fidelity of all states
states_fidelity_average = get_fidelity_over_sates(H)
pprint(f"Ideal gate: {states_fidelity_average*100:.3f} %")
'Ideal gate: 100.000 %'
We can simulate four consecutive identical pulses to illustrate the full coherent evolution of the system, showing how the population dynamics evolve through entanglement and eventually return to the initial state.
We can do this simply by updating the pulse parameters and the Problem and again solving it using the PWCSolver.
from copy import deepcopy
sequence_time = 4 * tfinal # 4 pulses
# define time and amplitude arrays (4N-3 time steps ensures dt remains the same)
N_seq = 4 * N - 3 # number of time steps for full sequence
sequence_ts = jnp.linspace(0.0, sequence_time, N_seq) # time array
amplitude_array_seq = jnp.asarray(
rect_pulse(sequence_ts, parameters_rect)
) # amplitude array
# update pulse parameters for 4-gate sequence (use deepcopy to avoid modifying original params)
params_full_seq = deepcopy(params)
params_full_seq["pwc/env"] = amplitude_array_seq
# define and solve Problem for 4-gate sequence
problem_seq = qr.Problem(H, y0, params_full_seq, (0.0, sequence_time))
solver_seq = qr.PWCSolver(n=N_seq, store=True)
solver_seq.set_system(problem_seq.sepwc())
time_seq, res_seq = solver_seq.evolve(*problem_seq.problem())
# calculate expectation value of each state for 4-gate sequence
pop_seq = get_population(res_seq)
We can then plot the state populations to see how they look for the four-pulse cycle.
# keep only non-zero states
filtered_pop_seq = pop_seq[:, nonzero_mask]
# plot populations for MS gate (4 identical PWC pulses)
canvas.init_canvas(x_axis_label="t [µs]", y_axis_label="Population")
canvas.plot(sequence_ts / 1e-6, filtered_pop_seq, labels=filtered_labels)
canvas.canvas.legend.spacing = -4
canvas.show_canvas()
You can see the populations evolve coherently, with the system returning to the $|000\rangle$ state after four pulses, confirming the cyclic behaviour of the gate.
5. Adding Gaussian filter (imperfect gate)¶
In reality, inherent limitations of the control electronics prevent us from generating a perfect rectangular pulse. To account for this, we can apply a Gaussian filter to simulate a finite rise time.
Here, we'll use a $4^{\text{th}}$ order filter, with a 150 $n \text{s}$ rise time. We can then compute the filter coefficients based on these parameters.
Note: You can read more about simulating a Gaussian rise time in our rise time tutorial.
from qruise.toolset import TransferFunc
from qruise.toolset.utils import compute_rise_time_gaussian_coeffs
rise_time = 150e-9 # rise time (s)
# compute Gaussian filter coefficients to simulate finite rise time
abs_coeffs = compute_rise_time_gaussian_coeffs(rise_time, n_order=4)
b = jnp.array([1.0])
a = jnp.array(abs_coeffs)
# define scalar version of PWC pulse for time differentiation
scalar_pwc = lambda t, p: rabi_drive_pwc(jnp.array([t]), p)[0]
dsource = jax.grad(scalar_pwc, argnums=0)
# define transfer function that adds Gaussian rise time filter to original PWC pulse
pwc_with_rise = TransferFunc(
"tf_pwc",
{
"b": (b, False),
"a": (a, False),
"t0": (0.0, False),
"t1": (tfinal, False),
},
rabi_drive_pwc,
dsource=dsource,
)
As we did earlier, we can merge the parameter dictionaries to make them easier to work with. To do this, we just need to add pwc_with_rise.params to params.
We can then evaluate the PWC pulse with the Gaussian rise time.
# merge new parameters with existing ones
params = params | pwc_with_rise.params
# evaluate PWC pulse with rise time
pwc_with_rise_vec = jax.vmap(pwc_with_rise, (0, None))
pwc_with_rise_result = pwc_with_rise_vec(ts, params)
Let's quickly plot the pulses without and with rise time side-by-side to see the difference.
canvas = qr.PlotUtil(
x_axis_label="t [µs]",
y_axis_label="Amplitude [a.u.]",
notebook=True,
x_range=[0, tfinal / 1e-6],
)
canvas.plot(ts / 1e-6, rabi_drive_pwc(ts, params), labels=["Drive pulse"])
canvas.plot(
ts / 1e-6,
pwc_with_rise_result,
labels=["Drive pulse with rise time"],
line_dash="dashed",
)
canvas.show_canvas()
Great, we see that without the rise time, the pulse envelope is flat with amplitude $1.0$, starting at $t=0.0$ ns and ending at $t = t_\text{final}$. Once the rise time is included, its effect becomes clearly visible.
We can then re-define our Hamiltonian and Problem using the drive pulse with rise time, and simulate the resulting state population dynamics as before.
# define envelopes
def env1_rise(t, params):
return pwc_with_rise(t, params) * (-eta * Omega * jnp.exp(1j * delta * t))
def env2_rise(t, params):
return pwc_with_rise(t, params) * (-eta * Omega * jnp.exp(-1j * delta * t))
# create time-dependent Hamiltonian
H = qr.Hamiltonian(None, [(Ja, env1_rise), (Ja_dag, env2_rise)])
# define Problem
problem = qr.Problem(H, y0, params, (0.0, tfinal))
# use PWCSolver to solve Schrödinger equation
solver = qr.PWCSolver(n=N, store=True)
solver.set_system(problem.sepwc())
time, res = solver.evolve(*problem.problem())
# calculate expectation value of each state
pop = get_population(res)
We'll again plot only the states with non-zero population for readability, and calculate the state fidelity at the end of the gate duration.
# Find states that are zero across all time points
zero_mask = np.all(pop == 0, axis=0)
# keep only non-zero states
nonzero_mask = ~zero_mask
filtered_pop = pop[:, nonzero_mask]
filtered_labels = [str(label) for i, label in enumerate(labels) if not zero_mask[i]]
# plot non-zero states
canvas.init_canvas(x_axis_label="t [µs]", y_axis_label="Population")
canvas.plot(ts / 1e-6, filtered_pop, labels=filtered_labels)
canvas.canvas.legend.spacing = -4
canvas.show_canvas()
# calculate state fidelity
states_fidelity_average = get_fidelity_over_sates(H)
pprint(f"Non-ideal gate: {states_fidelity_average*100:.3f} %")
'Non-ideal gate: 97.947 %'
The difference in the plot is somewhat subtle, but if you look carefully (try zooming in!), you'll notice that the $|101\rangle$ state doesn't return to zero and the populations of the $|000\rangle$ and $|110\rangle$ states are less than 50%. This is clearer when we see that the fidelity has dropped to just 97.947%.
We can again simulate the dynamics for four MS gates, to see more clearly how the rise time affects the system dynamics. This follows the same process as for the ideal gate.
from copy import deepcopy
# update pulse parameters for 4-gate sequence
params_full_seq = deepcopy(params)
params_full_seq["pwc/env"] = amplitude_array_seq
# define and solve Problem for 4-gate sequence
problem_seq = qr.Problem(H, y0, params_full_seq, (0.0, sequence_time))
solver_seq = qr.PWCSolver(
n=N_seq,
store=True,
)
solver_seq.set_system(problem_seq.sepwc())
time_seq, res_seq = solver_seq.evolve(*problem_seq.problem())
# calculate expectation value of each state for 4-gate sequence
pop_seq = get_population(res_seq)
# keep only non-zero states
filtered_pop_seq = pop_seq[:, nonzero_mask]
# plot populations for MS gate (4 identical PWC pulses)
canvas.init_canvas(x_axis_label="t [µs]", y_axis_label="Population")
canvas.plot(sequence_ts / 1e-6, filtered_pop_seq, labels=filtered_labels)
canvas.canvas.legend.spacing = -4
canvas.show_canvas()
We can see that the population of the $|000\rangle$ state does not return to 1, with residual population in the $|101\rangle$ state, demonstrating how the rise time negatively impacts the gate operation.
Great, you just simulated a Mølmer-Sørensen gate and explored how rise time affects the state fidelity.