How to automate complex calibration tasks with Boulder Opal

pip install qctrl-experiment-scheduler

import numpy as np
import matplotlib.pyplot as plt
import qctrlvisualizer as qv
import boulderopal as bo
import qctrlexperimentscheduler as es
from qctrlexperimentscheduler.predefined import (
    SimulationAmplitudeCalibrationBackend,
    create_amplitude_calibration_graph,
)

plt.style.use(qv.get_qctrl_style())

# Mute messages from Boulder Opal calls.
bo.cloud.set_verbosity("QUIET")

# Define a backend, that is, a way for the graph to communicate with
# the hardware that performs the calibrations. In this case, instead of
# hardware we define a simulation with high noise.
backend = SimulationAmplitudeCalibrationBackend(
    maximum_rabi_rate=5e7, pulse_duration=320e-9, time_step=0.25e-9, dephasing_noise=4e5
)

# Create an instance of the predefined calibration graph.
calibration_graph = create_amplitude_calibration_graph(backend=backend)

# Visualize the graph.
calibration_graph.visualize()

# Ask the scheduler to execute the node on which all the other nodes
# depend. As the other calibrations haven't been performed yet, this will
# trigger their execution as well.
final_node = calibration_graph.get_node("fine_pi_pulse_amplitude_calibration")

calibration_graph.execute(final_node)

# Print final value of the amplitude variables.
for variable in ["pi_pulse_amplitude", "pi_2_pulse_amplitude"]:
    print(f"{variable:20s}: {calibration_graph.get_variable(variable).get()}")

original_pi_2_amplitude = calibration_graph.get_variable("pi_2_pulse_amplitude").get()

# Check the results by plotting the inifidelity after
# several repetitions of the experiment.

repetitions = np.array([4, 16, 24, 36, 48])
pi_pulse_data = []
pi_2_pulse_data = []

for repetition in repetitions:
    pi_2_pulse_data.append(
        backend.pi_2_pulse_experiment(
            amplitude=original_pi_2_amplitude,
            drag=np.array([0]),
            repetition_count=repetition,
            shot_count=100000,
        )
    )
    pi_pulse_data.append(
        backend.pi_pulse_experiment(
            amplitude=calibration_graph.get_variable("pi_pulse_amplitude").get(),
            drag=np.array([0]),
            repetition_count=repetition,
            shot_count=100000,
        )
    )

plt.xlabel("Gate repetitions")
plt.ylabel("Infidelity")

plt.plot(repetitions, pi_2_pulse_data, label=r"$\pi$/2 pulse")
plt.plot(repetitions, pi_pulse_data, label=r"$\pi$ pulse")

plt.legend()
plt.show()

pi_2_pulse_amplitude = calibration_graph.get_variable("pi_2_pulse_amplitude")
drag_bound = 320e-9
pi_2_pulse_drag = calibration_graph.create_variable(
    name="pi_2_pulse_drag",
    initial_value=np.array([0]),
    lower_bound=-drag_bound,
    upper_bound=drag_bound,
)

# Define custom calibration function.
def sample_closed_loop_calibration(_calibration_graph):
    sample_point_count = 51

    # Gets the current value of the pi/2 pulse amplitude, use it to
    # generate a series of test points.
    initial_parameters = np.array(
        [
            np.clip(np.linspace(0.85, 1.15, sample_point_count), 0, 1)
            * pi_2_pulse_amplitude.get(),
            np.linspace(-drag_bound, drag_bound, sample_point_count),
        ]
    ).T

    # Define a simulated experiment applying the pi/2 pulse
    # different numbers of repetitions.
    def _experiment(input_array):
        amplitudes = input_array[:, 0]
        drag_parameters = input_array[:, 1]
        experiment_data = {
            count: backend.pi_2_pulse_experiment(
                amplitude=amplitudes,
                drag=drag_parameters,
                shot_count=10000,
                repetition_count=count,
            )
            for count in [5, 6, 7, 8]
        }
        return np.squeeze(
            2 * np.abs(0.5 - experiment_data[5])
            + np.abs(1 - experiment_data[6])
            + 2 * np.abs(0.5 - experiment_data[7])
            + np.abs(experiment_data[8])
        )

    # Run a simple closed-loop optimization of the pi/2 pulse.
    amplitude_bounds = np.clip(
        np.array([0.85, 1.15]) * pi_2_pulse_amplitude.get(), 0, 1
    )

    bounds = np.array([amplitude_bounds, [-drag_bound, drag_bound]])

    result = bo.closed_loop.optimize(
        cost_function=_experiment,
        initial_parameters=initial_parameters,
        optimizer=bo.closed_loop.GaussianProcess(
            bounds=bo.closed_loop.Bounds(bounds), seed=7
        ),
        max_iteration_count=10,
    )

    # Gather the best parameter of the calibration.
    best_parameters = result["parameters"]
    pi_2_pulse_amplitude.set(np.array([best_parameters[0]]))
    pi_2_pulse_drag.set(np.array([best_parameters[1]]))
    return es.CalibrationStatus.PASS


# Add a new node to the end of the previously-defined graph.
new_node = calibration_graph.create_node(
    name="closed_loop_calibration",
    dependencies=[final_node],
    calibration_function=sample_closed_loop_calibration,
)

# Visualize the updated graph.
calibration_graph.visualize()

# Run graph this time with the new node.
# As the previous calibrations have already been performed,
# the scheduler only runs the new node.
calibration_graph.execute(new_node)

# Print updated values of the pulse variables.
for variable in ["pi_pulse_amplitude", "pi_2_pulse_amplitude", "pi_2_pulse_drag"]:
    print(f"{variable:20s}: {calibration_graph.get_variable(variable).get()}")

# Check the results by plotting the inifidelity after
# several repetitions of the experiment.

repetitions = np.array([4, 16, 24, 36, 48])

original_pi_2_pulse_data = []
new_pi_2_pulse_data = []

for repetition in repetitions:
    original_pi_2_pulse_data.append(
        backend.pi_2_pulse_experiment(
            amplitude=original_pi_2_amplitude,
            drag=np.array([0]),
            repetition_count=repetition,
            shot_count=100000,
        )
    )
    new_pi_2_pulse_data.append(
        backend.pi_2_pulse_experiment(
            amplitude=calibration_graph.get_variable("pi_2_pulse_amplitude").get(),
            drag=calibration_graph.get_variable("pi_2_pulse_drag").get(),
            repetition_count=repetition,
            shot_count=100000,
        )
    )

plt.xlabel("Gate repetitions")
plt.ylabel("Infidelity")

plt.plot(repetitions, original_pi_2_pulse_data, label=r"Original $\pi$/2 pulse")
plt.plot(repetitions, new_pi_2_pulse_data, label=r"New $\pi$/2 pulse")

plt.legend()
plt.show()