How to optimize controls robust to strong noise sources

Design controls that are robust against strong time-dependent noise sources with stochastic optimization

In this notebook we demonstrate how to apply stochastic optimization to design robust pulses using our highly-flexible optimization engine.

The stochastic optimizer is well suited to treat strong and non-static noises, allowing us to go beyond the weak noise regime required for using filter functions in robust optimization, as illustrated in this user guide. For stronger noise fields, typically neither the first-order Magnus approximation of the toggling frame Hamiltonian nor the second-order approximation of the time evolution operator hold.

Summary Workflow

1. Define the system, variables to be optimized and cost function in the computational graph

The flexible Q-CTRL optimization engine expresses all optimization problems as data flow graphs, which describe how optimization variables (variables that can be tuned by the optimizer) are transformed into the cost function (the objective that the optimizer attempts to minimize).

You can define physical constraints, decide on which variables to optimize and provide a value range, include noise terms, and specify the cost function, which can be based on the target operator and the infidelity.

2. Execute graph-based optimization

With the graph object created, an optimization can be run by providing the name of the graph node representing the cost function to be minimized, the names of the desired output nodes and the graph itself to the calculate_stochastic_optimization function.

Optionally, you can also set the number of iterations to perform, after which the results with the lowest cost are returned (defaults to 1000), a target cost which early-stops the optimization when/if reached (else the optimizer runs until all iterations are completed), and an optimization algorithm (with Adam as the default).

The function returns the results of the optimization, which include the lowest cost achieved across all iterations, as well as the specified nodes evaluated at the optimized variables.

You can then use the qctrlvisualizer package to plot the optimized pulses, which are available in the result object.

Worked example: Robust control of a qubit under strong noise

In this example we show how to find robust controls with the stochastic optimizer. The single-qubit system is represented by the following Hamiltonian:

\begin{align*} H(t) &= (1 + \beta(t))(\gamma_x(t) \sigma_x + \gamma_y(t) \sigma_y) \end{align*}

where $\gamma_x(t)$ and $\gamma_y(t)$ are the real time pulses and $\beta(t)$ is a stochastic noise process.

To perform a stochastic optimization, we assume that the stochastic process can be decomposed into 10 Fourier components

\begin{align*} \beta(t) &= \sum_{i=1}^{10} a_i \cos(\omega_i t) + b_i \sin(\omega_i t) \end{align*}

where $a_i$ and $b_i$ are amplitude samples of a normal distribution and $\omega_i$ are frequency samples of a uniform distribution. At each iteration step a new set of samples is taken.

This example uses nodes random_uniform and random_normal to generate the stochastic process. These nodes provide random values sampled from uniform and normal probability distributions, respectively. Other ways of generating a stochastic process include using the random_choices node to shuffle your own set of numerical values (which you can sample from any distribution you want) or random_colored_noise_stf_signal to create a stochastic signal from a power spectral density.

The cost is calculated by averaging the infidelities corresponding to a batch of samples of the stochastic process. Averaging over a batch can lead to more reliable cost function estimates and gradients, which in turn can result in better overall optimization performance.

import matplotlib.pyplot as plt
import numpy as np
from qctrlvisualizer import get_qctrl_style, plot_controls

plt.style.use(get_qctrl_style())

from qctrl import Qctrl

# Starting a session with the API
qctrl = Qctrl()
# Define standard matrices
sigma_x = np.array([[0, 1], [1, 0]])
sigma_y = np.array([[0, -1j], [1j, 0]])
sigma_z = np.array([[1, 0], [0, -1]])

# Define physical constraints
gamma_min = -np.pi
gamma_max = np.pi

segment_count = 20
duration = 2
sample_times = (0.5 + np.arange(segment_count)) * duration / segment_count

batch_size = 200

# Create graph object
with qctrl.create_graph() as graph:

    var_x = qctrl.operations.optimization_variable(
        count=segment_count,
        lower_bound=gamma_min,
        upper_bound=gamma_max,
    )
    var_y = qctrl.operations.optimization_variable(
        count=segment_count,
        lower_bound=gamma_min,
        upper_bound=gamma_max,
    )

    gamma_x = qctrl.operations.pwc_signal(var_x, duration=duration, name="gamma_x")
    gamma_y = qctrl.operations.pwc_signal(var_y, duration=duration, name="gamma_y")

    noise_x_term = []
    noise_y_term = []

    # a cos(wt) + b sin(wt)
    for _ in range(10):
        a = qctrl.operations.random_normal(
            (batch_size, 1), mean=0.0, standard_deviation=0.05
        )
        b = qctrl.operations.random_normal(
            (batch_size, 1), mean=0.0, standard_deviation=0.05
        )
        omega = qctrl.operations.random_uniform(
            (batch_size, 1), lower_bound=np.pi, upper_bound=2 * np.pi
        )
        sampled_fourier = a * qctrl.operations.cos(
            omega * sample_times[None]
        ) + b * qctrl.operations.sin(omega * sample_times[None])

        noise_x_term.append(
            qctrl.operations.pwc_signal(
                var_x[None] * sampled_fourier, duration=duration
            )
            * sigma_x
        )
        noise_y_term.append(
            qctrl.operations.pwc_signal(
                var_y[None] * sampled_fourier, duration=duration
            )
            * sigma_y
        )

    # Define Hamiltonian
    total_noise_terms = qctrl.operations.pwc_sum(
        noise_x_term
    ) + qctrl.operations.pwc_sum(noise_y_term)
    hamiltonian = gamma_x * sigma_x + gamma_y * sigma_y + total_noise_terms

    # Create target
    target_operator = qctrl.operations.target(operator=sigma_x)

    # Create infidelity
    infidelities = qctrl.operations.infidelity_pwc(
        hamiltonian,
        target_operator,
        name="infidelities",
    )
    cost = qctrl.operations.sum(infidelities / batch_size, name="cost")
# Run the optimization

optimization_result = qctrl.functions.calculate_stochastic_optimization(
    graph=graph,
    cost_node_name="cost",
    output_node_names=["gamma_x", "gamma_y", "infidelities"],
    iteration_count=10000,
    target_cost=1e-6,
)

print(f"\nOptimized cost:\t {optimization_result.best_cost:.3e}\n")
Your task calculate_stochastic_optimization has started.
Your task calculate_stochastic_optimization has completed.

Optimized cost:	 2.470e-05

# Plot histogram of infidelities evaluated at the optimized variables

print(
    f'Batch mean: {optimization_result.best_output["infidelities"]["value"].mean():.2e}, \
standard deviation: {optimization_result.best_output["infidelities"]["value"].std():.2e}'
)

plt.title("Batch of Optimized Infidelities")
plt.xlabel("Infidelities")
plt.ylabel("Count")
plt.hist(
    optimization_result.best_output["infidelities"]["value"], bins="auto", ec="black"
)
plt.xlim([0.0, 1e-4])
plt.show()
Batch mean: 2.69e-05, standard deviation: 3.94e-05
# Plot the optimized controls

plot_controls(
    plt.figure(),
    controls={
        "$gamma_x$": optimization_result.best_output["gamma_x"],
        "$gamma_y$": optimization_result.best_output["gamma_y"],
    },
)
plt.show()