Gate Optimization with Autodiff

Simphony is a Python package designed for simulating the spin dynamics of point defects, in particular the nitrogen-vacancy (NV) center, which is surrounded by nuclear spins and serves as a central-spin quantum register.

This tutorial demonstrates how to leverage Simphony’s automatic differentiation (autodiff) capabilities to optimize a quantum gate at the pulse level.

We will:

• Define a physical model based on NV-center spins

• Set up an objective function using gate infidelity

• Use autodiff to compute gradients

• Apply gradient-based optimization to improve gate fidelity

Import the packages

First, enable automatic differentiation in Simphony, set the computation to run on the GPU — highly recommended for autodiff — and import the required packages:

import simphony
simphony.Config.set_platform('gpu')
simphony.Config.set_autodiff_mode(True)
simphony.Config.set_matplotlib_format('retina')

import numpy as np
import jax.numpy as jnp

from qiskit import QuantumCircuit

from matplotlib import pyplot as plt

Define the model

We use a predefined NV-center model from Simphony, which includes the electron spin and the nitrogen nuclear spin, but no carbon nuclear spins:

model = simphony.default_nv_model(nitrogen_isotope = 14,
                                  static_field_strength = 0.005)
model.plot_levels()
model.spin_names

['e', 'N']

Infidelity as the Objective Function

To perform autodiff, we first define an objective function as a Python function. Infidelity quantifies how far the implemented gate deviates from the target ideal gate, making it a natural choice for the objective function.

We start by implementing an \(RX(\pi/2)\) gate on the electron spin using a resonant pulse with a rectangular envelope. Since the electron spin splitting depends on the state of the nitrogen nuclear spin, we set the driving frequency to the average electron spin splitting across nuclear spin states. The pulse duration will be treated as the optimization variable, while the amplitude is set to achieve a \(\pi/2\) rotation — corresponding to a quarter of a Rabi oscillation:

frequency = model.splitting_qubit(spin_name='e')
qc = QuantumCircuit(2)
qc.rx(np.pi/2,0)

def infidelity(duration, model):

    period_time = 4 * duration # 4 -> pi/2 gate
    amplitude = model.rabi_cycle_amplitude_qubit(driving_field_name = 'MW_x',
                                                 spin_name = 'e', 
                                                 period_time = period_time)
    
    model.remove_all_pulses()
    model.driving_field('MW_x').add_rectangle_pulse(frequency = frequency,
                                                    amplitude = amplitude,
                                                    duration = duration,
                                                    phase = 0)

    result = model.simulate_time_evolution(n_eval=2)
    result.ideal = qc
    return 1 - result.average_gate_fidelity()

We can compute the gate fidelity for a given pulse duration:

duration = 0.03
infidelity(duration, model)

Array(0.01549852, dtype=float64)

Alternatively, we can scan the duration parameter over a specified range:

durations = np.linspace(0.003, 0.03, 251)
infidelities = np.array([infidelity(duration, model) for duration in durations])

We can visualize the infidelity as a function of the pulse duration:

plt.plot(durations, infidelities, '-')
plt.xlabel('duration [$\mu$s]')
plt.ylabel('infidelity')
plt.ylim(0)
plt.show()

At short pulse durations, the infidelity is high due to leakage out of the computational subspace. At longer durations—corresponding to weaker amplitude—frequency-averaged driving cannot achieve resonance regardless of the nuclear spin state, leading to a gradual increase in infidelity. Bloch–Siegert oscillations appear due to the presence of the counter-rotating term, which is retained since the rotating wave approximation (RWA) is not applied.

To observe the leakage explicitly, we can calculate it as follows:

def leakage(duration, model):

    period_time = 4 * duration
    amplitude = model.rabi_cycle_amplitude_qubit(driving_field_name = 'MW_x',
                                                 spin_name = 'e', 
                                                 period_time = period_time)
    
    model.remove_all_pulses()
    model.driving_field('MW_x').add_rectangle_pulse(frequency = frequency,
                                                    amplitude = amplitude,
                                                    duration = duration,
                                                    phase = 0)

    result = model.simulate_time_evolution(n_eval=2)
    return result.leakage()

leakages = np.array([leakage(duration, model) for duration in durations])

plt.plot(durations, leakages, '-')
plt.xlabel('duration [$\mu$s]')
plt.ylabel('leakage')
plt.ylim(0)
plt.show()

The optimal pulse duration is:

durations[np.argmin(infidelities)]

np.float64(0.00732)

Autodiff with JAX and Simple Gradient Descent

To compute the gradient, we use JAX’s value_and_grad function:

from jax import value_and_grad

value_and_grad_infidelity = value_and_grad(infidelity)

The value_and_grad_infidelity function works similarly to the infidelity function, but returns a tuple:

• The first element is the infidelity value.

• The second element is the gradient with respect to the first argument of the infidelity function, returned as a JAX Array:

duration = 0.03
value_and_grad_infidelity(duration, model)

(Array(0.01549852, dtype=float64),
 Array(1.304514, dtype=float64, weak_type=True))

We illustrate gate optimization starting from an initial pulse duration using gradient descent. The procedure involves iteratively computing the gradient and updating the duration according to:

\[ \text{duration}_{i+1} = \text{duration}_{i} - \epsilon \cdot \text{grad}_i, \]

where \(\epsilon\) is the learning rate:

duration = 0.03
epsilon = 1e-3
n_step = 25

infidelities_optimized = []
durations_optimized = []

for _ in range(n_step):
    
    value, grad = value_and_grad_infidelity(duration, model)
    infidelities_optimized.append(value)
    durations_optimized.append(duration)

    duration -= epsilon*grad

We can visualize the optimization process:

fig, axs = plt.subplots(1, 2, figsize=(12,4))

axs[0].plot(infidelities_optimized, 'o-', c='#ff7f0e')
axs[0].set_xlabel('steps')
axs[0].set_ylabel('infidelity')
axs[0].set_ylim(0)

axs[1].plot(durations, infidelities, '-')
axs[1].plot(durations_optimized, infidelities_optimized, '.-')
axs[1].set_xlabel('duration [$\mu$s]')
axs[1].set_ylabel('infidelity')
axs[1].set_ylim(0)

plt.show()

In the left plot, the infidelity is shown as a function of the number of optimization steps. In the right plot, we show the trajectory of the optimization on the landscape.

Note: For oscillatory landscapes, the optimization problem becomes non-trivial and may require a more sophisticated optimization method.

Optimize a Pulse Sequence

As a second example, we consider a more complex scenario: optimizing a pulse sequence with a piecewise-constant envelope. The optimization variables are real numbers, where:

• even-indexed parameters correspond to the real parts of the envelope segments

• odd-indexed parameters correspond to the imaginary parts

def infidelity_2(parameters, dt, model):

    samples = parameters[::2] + 1j*parameters[1::2]
    
    model.remove_all_pulses()
    model.driving_field('MW_x').add_discrete_pulse(frequency = frequency,
                                                   samples = samples,
                                                   dt = dt)

    result = model.simulate_time_evolution(n_eval=2)
    result.ideal = qc
    return 1 - result.average_gate_fidelity()

value_and_grad_infidelity_2 = value_and_grad(infidelity_2)

Gradient descent is initialized from a randomly generated pulse sequence. The sequence consists of 8 constant-envelope pulse segments with a total duration of 0.01 μs. The optimization runs for 100 iterations.

Note: Executing this code may take several minutes, depending on system performance:

np.random.seed(42)

N = 8
duration = 0.01
dt = duration / N
parameters = 0.1 * (np.random.random(2 * N) - 0.5)

epsilon = 2e-5
n_step = 100

infidelities_optimized = []

for i in range(n_step):
    
    value, grads = value_and_grad_infidelity_2(parameters, dt, model)
    infidelities_optimized.append(value)

    parameters -= epsilon*grads
    print(i, end=' ')

0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 

Infidelity is plotted as a function of the number of optimization steps. Starting from an initial gate with infidelity near unity, the optimization gradually converges toward zero infidelity using a highly conservative learning rate.

While a higher learning rate could accelerate convergence, it often leads to suboptimal gate performance. In such cases, more advanced optimization methods are needed to achieve better gate fidelities and faster convergence.

plt.plot(infidelities_optimized, 'o-')
plt.ylabel('infidelity')
plt.xlabel('steps')
plt.yscale('log')
plt.show()

The resulting optimized pulse sequence is shown below:

model.driving_field('MW_x').plot_pulses(function='complex_envelope')

Note: This optimization process is conceptually similar to the Gradient Ascent Pulse Engineering (GRAPE) method, but does not rely on the rotating wave approximation (RWA).