Find optimal pulses with automated optimization

In this tutorial you will optimize a control pulse which implements a single-qubit gate on a simulated realistic experiment. You will set up a cross-entropy closed-loop optimizer which will keep suggesting new pulses to try in the experiment, in order to find one which minimizes the gate infidelity.

The closed-loop optimization toolkit in Boulder Opal executes a closed optimization loop where the optimizer communicates with the experimental apparatus without your direct involvement. In this kind of setting, your experimental apparatus produces an initial set of results, which it sends to the optimizer. Using this information, the optimizer produces a set of improved test points (which in this case are controls) that it recommends back to the experimental apparatus. The results corresponding to these test points are sent back to the optimizer, and the cycle repeats itself until any of the results has a sufficiently low cost function value, or until it meets any other ending condition that you imposed. This setup is illustrated in the figure below.

The Boulder Opal Toolkits are currently in beta phase of development. Breaking changes may be introduced.

Obtaining an optimal pulse to implement a single qubit gate

Your task is to find an optimal pulse which implements an X gate on a qubit by using Boulder Opal's closed-loop optimizer. You would usually do this for an experiment, but in this case you will interact with a simulated qubit implemented by the run_experiments function below.

The system simulated in run_experiments is a qubit with leakage to a higher level, thus it's a three level system. The qubit starts in state $|0\rangle$ and is evolved by the controls. The function returns measurements of what level the qubit is in, either 0, 1 or 2. Our goal of implementing an X gate corresponds to evolving the qubit from state $|0\rangle$ to state $|1\rangle$.

Do not worry about the exact details of this function, just think of it as a function that interacts with the experiment for which you want to find the optimal gate. It takes in an array of controls, a duration, and a shot_count, and returns an array of measurements on the qubit associated with each control and shot.

def run_experiments(controls, duration, shot_count):
    """
    Simulates a single qubit experiment using Boulder Opal
    with the given piecewise-constant controls.

    Parameters
    ----------
    controls : np.ndarray
        The controls to simulate.
        A 2D NumPy array of shape (control_count, segment_count)
        with the per-segment values of each control.
    duration : float
        The duration (in nanoseconds) of the controls.
    shot_count : int
        The number of shots for which to generate measurements.

    Returns
    -------
    np.ndarray
        The qubit measurement results (either 0, 1, or 2) associated to each control
        and shot, as an array of shape (len(controls), shot_count).
    """

    # Create Boulder Opal graph.
    graph = qctrl.create_graph()

    # Define simulation parameters and operators.
    filter_cutoff_frequency = 2 * np.pi * 0.3  # GHz
    segment_count = 128
    max_drive_amplitude = 2 * np.pi * 0.1  # GHz
    delta = -0.33 * 2 * np.pi  # GHz
    big_delta = 0.01 * 2 * np.pi  # GHz
    number_operator = graph.number_operator(3)
    drive_operator = 0.5 * graph.annihilation_operator(3)
    confusion_matrix = np.array(
        [[0.99, 0.01, 0.01], [0.01, 0.98, 0.01], [0.0, 0.01, 0.98]]
    )

    # Retrieve control information.
    control_values = controls * max_drive_amplitude

    # Define initial state.
    initial_state = graph.fock_state(3, 0)

    # Construct constant Hamiltonian terms.
    frequency_term = big_delta * number_operator
    anharmonicity_term = (
        delta * (number_operator @ number_operator - number_operator) / 2
    )

    # Construct filtered drive.
    filtered_drive = graph.convolve_pwc(
        pwc=graph.pwc_signal(duration=duration, values=control_values),
        kernel=graph.sinc_convolution_kernel(filter_cutoff_frequency),
    )
    drive = graph.discretize_stf(filtered_drive, duration, segment_count)
    drive_term = graph.hermitian_part(drive * drive_operator)

    # Build Hamiltonian and calculate unitary evolution operators.
    hamiltonian = drive_term + frequency_term + anharmonicity_term
    unitary = graph.time_evolution_operators_pwc(
        hamiltonian=hamiltonian, sample_times=np.array([duration])
    )
    unitary = unitary[:, -1]

    # Evolve initial state and calculate final (normalized) populations.
    final_state = unitary @ initial_state[:, None]
    populations = graph.abs(final_state) ** 2

    # Apply the confusion matrix to the populations.
    populations = confusion_matrix @ populations
    populations.name = "populations"

    # Execute graph and retrieve the populations.
    result = qctrl.functions.calculate_graph(
        graph=graph, output_node_names=["populations"]
    )
    populations = result.output["populations"]["value"].squeeze(axis=2)

    # Simulate measurements for each control.
    measurements_list = []
    for p in populations:
        # Sample and store measurements.
        measurements = np.random.choice(3, size=shot_count, p=p / np.sum(p))
        measurements_list.append(measurements)

    return np.array(measurements_list)

1. Import libraries and start a Boulder Opal session

Before doing any calculation with Boulder Opal, you always need to import the necessary libraries and start a session.

In this case, import the numpy, matplotlib.pyplot, qctrlvisualizer, and qctrl packages. To learn more about installing Boulder Opal and starting a session, see the Get started guide.

As each iteration in the closed-loop optimizer is a separate job, to obtain a neater output from the optimizer, we mute messsages and progress bars by passing verbosity="QUIET" to the Qctrl object constructor.

# Import packages.
import numpy as np
import matplotlib.pyplot as plt
import qctrlvisualizer
from qctrl import Qctrl

# Apply Q-CTRL style to plots created in pyplot.
plt.style.use(qctrlvisualizer.get_qctrl_style())

# Start a Boulder Opal session.
qctrl = Qctrl(verbosity="QUIET")  # Mute status messages.

2. Configure closed-loop optimization

Define the optimizer

First, you will set up a cross-entropy optimizer to carry out your optimization. You can read our Choosing a control-design (optimization) strategy in Boulder Opal topic to learn more about the optimization routines available in Boulder Opal. Each one of them has their own strengths, and will require a slightly different setup in this step.

Create a qctrl.closed_loop.CrossEntropy object with the elite_fraction of best pulses you want to keep at each iteration, which should be a value smaller than 1. You can also pass it an integer seed for the random number generator that the optimizer will use internally if you want the optimizer to generate deterministic results.

# Define the state object for the closed-loop optimizer.
optimizer = qctrl.closed_loop.CrossEntropy(elite_fraction=0.25, seed=1)

Create a function to calculate the cost from the controls

The optimizer will aim to find the minimum of a cost function that takes in a set of parameters (controls) and returns their associated costs. In this case, the cost function must be suitable for the task of performing an X gate on the qubit. As this is equivalent to maximizing the probability of finding the qubit in state $|1\rangle$, define the cost for each control as one minus the probability of finding the qubit in state $|1\rangle$.

This can be calculated by calling the experiment function (run_experiments) which returns an array with the qubit measurements associated to each control. For each control, the measurements are a list of 0, 1, or 2, corresponding to the measured qubit state after the control is applied.

# Calculate cost from experiment results.
def cost_function(controls, duration, shot_count):
    """
    Accepts an array of controls and returns their associated costs.
    """
    measurements = run_experiments(controls, duration, shot_count)

    costs = []
    for shots in measurements:
        shots_in_one = np.count_nonzero(shots == 1)
        fidelity = shots_in_one / len(shots)
        costs.append(1 - fidelity)

    return costs

Create an initial set of results to seed the optimizer

Next, you need to create a set of initial controls. If you had some pulses that you want to refine you could use those, but since you don't, you can generate a random set of pulses.

Start by defining the parameters for the controls that you want to obtain, the number of piecewise-constant segments and the duration. Next we define the the experimental parameters, the number of projective measurements (shots) you want to take after each control and the number of controls (test points) you want to test at each optimization iteration. For this last one, make sure that elite_fraction × test_point_count is at least two so the cross-entropy optimizer can advance.

With these parameters, create an array of shape (test_point_count, segment_count) random values for the controls to start the optimization with.

# Number of segments in each control.
segment_count = 10

# Duration of each control.
duration = 100  # ns

# Number of controls to retrieve from the optimizer each run.
test_point_count = 32

# Number of projective measurements to take after the control is applied.
shot_count = 1000

# Define parameters as a set of controls with piecewise-constant segments.
rng = np.random.default_rng(seed=1)
initial_controls = rng.uniform(0.0, 1.0, size=(test_point_count, segment_count))

Define the bounds

Finally we define the per-parameter bounds on the test points, which limit the values the controls generated by the optimizer can take. As there's no reason to treat any of the pulse segments differently, we use uniform bounds, with a lower bound of 0 and an upper one of 1.

bounds = np.repeat([[0.0, 1.0]], segment_count, axis=0)

3. Execute optimization loop

You now have all of the elements to create a closed optimization loop to obtain the best controls.

To run the closed-loop optimization, use the function qctrl.closed_loop.optimize. This starts an iterative cycle, where on each iteration the optimizer suggests new test points predicted to have a low cost. This cycle keeps going until the best cost is lower than the target cost, target_cost.

It is a good idea to set a maximum number of iterations in case the problem is hard for the optimizer to converge.

The function will return a dictionary with the results of the optimization, namely, the best parameters best_parameters, their associated cost best_cost, and the history of best cost values best_cost_history.

max_iteration_count = 20
target_cost = 0.02

# Run the optimization loop until the cost (infidelity) is sufficiently small.
results = qctrl.closed_loop.optimize(
    cost_function=lambda controls: cost_function(
        controls, duration=duration, shot_count=shot_count
    ),
    initial_test_parameters=initial_controls,
    optimizer=optimizer,
    bounds=bounds,
    target_cost=target_cost,
    max_iteration_count=max_iteration_count,
)

Running closed loop optimization
----------------------------------------
  Optimizer            : cross entropy
  Number of test points: 32
  Number of parameters : 10
----------------------------------------

Calling cost function…
  Initial best cost: 0.167

Running optimizer…
Calling cost function…
  Best cost after 1 iterations: 0.030

Running optimizer…
Calling cost function…
  Best cost after 2 iterations: 0.030

Running optimizer…
Calling cost function…
  Best cost after 3 iterations: 0.027

Running optimizer…
Calling cost function…
  Best cost after 4 iterations: 0.026

Running optimizer…
Calling cost function…
  Best cost after 5 iterations: 0.019

Target cost reached. Stopping the optimization.

You can see that as the iterations progress, the best cost found by the optimizer gets smaller until it reaches the target cost and the optimization stops.

4. Retrieve the best results obtained by the optimizer

You can now print the best cost and controls the optimizer has achieved and visualize them with the plot_controls function from the Q-CTRL Visualizer.

# Print final best cost.
print(f"Best cost reached: {results['best_cost']:.3f}")

# Print and plot controls that correspond to the best cost.
print(f"Best control values: {np.round(results['best_parameters'], 3)}")

qctrlvisualizer.plot_controls(
    {
        r"$\Omega(t)$": qctrl.utils.pwc_arrays_to_pairs(
            duration, results["best_parameters"]
        )
    },
    two_pi_factor=False,
    unit_symbol=" a.u.",
)
plt.suptitle("Best control")

plt.show()

Best cost reached: 0.019
Best control values: [0.21  0.702 0.506 0.426 0.59  0.687 0.283 0.021 0.718 0.435]

Notice that your result for the best control will look different to the one plotted here due to a different random initialization of the optimizer.

5. Visualize the convergence of the optimizer

You can also plot the cost history during the optimization, and see its decrease until it crosses the target threshold. You can easily do this with the plot_cost_histories function from the Q-CTRL Visualizer.

qctrlvisualizer.plot_cost_histories(results["best_cost_history"])

6. Compare initial and converged parameters

Finally we can compare the initial and optimized distributions of measurements. For this plot, you will contrast the optimized controls with the best initial controls (the controls which have the highest probability of finding the qubit in $|1\rangle$).

As expected the converged controls perform substantially better than the random initial controls.

# Obtain a set of initial experimental results.
measurements = run_experiments(
    controls=initial_controls, duration=duration, shot_count=shot_count
)

# Find the best initial (random) control.
best_initial_control = np.argmax(np.count_nonzero(measurements == 1, axis=1))
initial_best_counts = np.unique(
    measurements[best_initial_control], return_counts=True, axis=0
)[1]
initial_best_probability = initial_best_counts / shot_count
print(f"Best initial probabilities: {initial_best_probability}")


# Obtain a set of converged experimental results.
measurements = run_experiments(
    controls=np.array([results["best_parameters"]]),
    duration=duration,
    shot_count=shot_count,
)
optimized_counts = np.unique(measurements, return_counts=True)[1]
optimized_probability = optimized_counts / shot_count
print(f"Optimized probabilities: {optimized_probability}")

# Plot distribution of measurements for the initial & converged sets of pulses.
fig, ax = plt.subplots()
ax.bar(np.arange(3) - 0.1, optimized_probability, width=0.2, label="Optimized controls")
ax.bar(
    np.arange(3) + 0.1,
    initial_best_probability,
    width=0.2,
    label="Best initial controls",
)
ax.legend()
ax.set_ylabel("Probability")
ax.set_xlabel("Measurement results")
ax.set_xticks(np.arange(3))
ax.set_xticklabels([rf"$|{k}\rangle$" for k in np.arange(3)])

plt.show()

Best initial probabilities: [0.149 0.841 0.01 ]
Optimized probabilities: [0.008 0.982 0.01 ]

This concludes the tutorial. Congratulations on running a closed-loop optimization!

You can now try to change the optimization loop you are running here by changing the optimization parameters or using a completely different optimizer. You can also play around with the function simulating the experiment and, for instance, make it harder or easier for the optimizer to find solutions by changing the confusion matrix or the anharmonicity in the system.

Our user guides offer more examples and problems to which you can apply this automated optimization tool. In particular, you might be interested in reading about obtaining optimal controls with other closed-loop optimization routines, calibrating control hardware, or using the open-source M-LOOP package to integrate the experiment to Boulder Opal closed-loop optimizer.

If you want to learn more about Boulder Opal and its capabilities, visit our topics page.

This notebook was run using the following package versions. It should also be compatible with newer versions of the Q-CTRL Python package.