Skip to content

Python API¤

The qontrol Python API consists of the function optimize which performs optimal control, and various models that can be optimized including SESolveModel and MESolveModel. There are additionally various utility functions and classes that help to define an optimization routine (cost functions, file input output, options, etc.)

Optimization¤

optimize(parameters: ArrayLike | dict, costs: Cost, model: Model, *, optimizer: GradientTransformation = optax.adam(0.0001, b1=0.99, b2=0.99), plotter: Plotter | None = None, method: Method = Tsit5(), gradient: Gradient | None = None, dq_options: dq.Options = dq.Options(), opt_options: dict | None = None, filepath: str | None = None) -> Array | dict ¤

Perform gradient descent to optimize Hamiltonian parameters.

This function takes as input parameters which parametrize a model when called performs time-dynamics simulations using Dynamiqs. How to update parameters is encoded in the list of cost functions costs that contains e.g. infidelity contributions, pulse amplitude penalties, etc.

Parameters:

  • parameters (ArrayLike | dict) –

    parameters to optimize over that are used to define the Hamiltonian and control times.

  • costs (Cost) –

    List of cost functions used to perform the optimization.

  • model (Model) –

    Model that is called at each iteration step.

  • optimizer (GradientTransformation, default: adam(0.0001, b1=0.99, b2=0.99) ) –

    optax optimizer to use for gradient descent. Defaults to the Adam optimizer.

  • plotter (Plotter | None, default: None ) –

    Plotter for monitoring the optimization.

  • method (Method, default: Tsit5() ) –

    Method passed to Dynamiqs.

  • gradient (Gradient | None, default: None ) –

    Gradient passed to Dynamiqs.

  • filepath (str | None, default: None ) –

    Filepath of where to save optimization results.

  • dq_options (Options, default: Options() ) –

    Options for the Dynamiqs integrator.

  • opt_options (dict | None, default: None ) –

    Options for grape optimization.

    Detailed opt_options API
    • verbose (bool, default: True): If True, the optimizer will print out the infidelity at each epoch step to track the progress of the optimization.
    • epochs (int, default: 2000): Number of optimization epochs.
    • batch_initial_parameters (bool, default: False): Whether to batch over initial parameters. If True, then the first dimension of parameters defines the number of simulations to batch over (note: not supported for parameters that are dictionaries or lists of dictionaries). If False, then parameters is assumed to not be batched.
    • plot (bool, default: True): Whether to plot the results during the optimization (for the epochs where results are plotted, necessarily suffer a time penalty).
    • plot_period (int, default: 30): If plot is True, plot every plot_period.
    • save_period (int, default: 30): If a filepath is provided, save every save_period.
    • xtol (float, default: 1e-8): Terminate the optimization if the parameters are not being updated.
    • ftol (float, default: 1e-8): Terminate the optimization if the cost function is not changing above this level.
    • gtol (float, default: 1e-8): Terminate the optimization if the norm of the gradient falls below this level.

Returns:

  • Array | dict –

    Optimized parameters from the final timestep.

Cost functions¤

incoherent_infidelity(target_states: list[QArrayLike], cost_multiplier: float = 1.0, target_cost: float = 0.005) -> IncoherentInfidelity ¤

Instantiate the cost function for calculating infidelity incoherently.

This fidelity is defined as $$ F_{\rm incoherent} = \sum_{k}|\langle\psi_{t}^{k}|\psi_{i}^{k}(T)\rangle|^2, $$ where the states at the end of the pulse are \(|\psi_{i}^{k}(T)\rangle\) and the target states are \(|\psi_{t}^{k}\rangle\).

Parameters:

  • target_states (list[QArrayLike]) –

    Shape of QArrays = (s, n, 1) or (s, n, n). Target states for the initial states passed to optimize. If performing master-equation optimization, the target states should be passed as a list of density matrices.

  • cost_multiplier (float, default: 1.0 ) –

    Weight for this cost function relative to other cost functions.

  • target_cost (float, default: 0.005 ) –

    Target value for this cost function. If options.all_costs is True, the optimization terminates early if all cost functions fall below their target values. If options.all_costs is False, the optimization terminates if only one cost function falls below its target value.

Returns:

  • IncoherentInfidelity –

    Callable object that returns the incoherent infidelity and whether the infidelity is below the target value.

coherent_infidelity(target_states: list[QArrayLike], cost_multiplier: float = 1.0, target_cost: float = 0.005) -> CoherentInfidelity ¤

Instantiate the cost function for calculating infidelity coherently.

This fidelity is defined as $$ F_{\rm coherent} = |\sum_{k}\langle\psi_{t}^{k}|\psi_{i}^{k}(T)\rangle|^2, $$ where the states at the end of the pulse are \(|\psi_{i}^{k}(T)\rangle\) and the target states are \(|\psi_{t}^{k}\rangle\).

Parameters:

  • target_states (list[QArrayLike]) –

    Shape = (s, n, 1) or (s, n, n). Target states for the initial states passed to optimize. If performing master-equation optimization, the target states should be passed as a list of density matrices.

  • cost_multiplier (float, default: 1.0 ) –

    Weight for this cost function relative to other cost functions.

  • target_cost (float, default: 0.005 ) –

    Target value for this cost function. If options.all_costs is True, the optimization terminates early if all cost functions fall below their target values. If options.all_costs is False, the optimization terminates if only one cost function falls below its target value.

Returns:

  • CoherentInfidelity –

    Callable object that returns the coherent infidelity and whether the infidelity is below the target value.

propagator_infidelity(target_unitary: QArrayLike, cost_multiplier: float = 1.0, target_cost: float = 0.005) -> PropagatorInfidelity ¤

Instantiate the cost function for calculating infidelity of a propagator.

This fidelity is defined as $$ F_{\rm propagator} = \Tr(U_{t}^{\dagger} U/d)^2, $$ where the propagator at the end of the pulse is \(U\), the dimension of the system is \(d\) and the target unitary \(U_{t}\).

Parameters:

  • target_unitary (QArrayLike) –

    Shape = (n,n). Target unitary for the initial states passed to optimize.

  • cost_multiplier (float, default: 1.0 ) –

    Weight for this cost function relative to other cost functions.

  • target_cost (float, default: 0.005 ) –

    Target value for this cost function. If options.all_costs is True, the optimization terminates early if all cost functions fall below their target values. If options.all_costs is False, the optimization terminates if only one cost function falls below its target value.

Returns:

  • PropagatorInfidelity –

    Callable object that returns the propagator infidelity and whether the infidelity is below the target value.

forbidden_states(forbidden_states_list: list[QArrayLike], cost_multiplier: float = 1.0, target_cost: float = 0.0) -> ForbiddenStates ¤

Instantiate the cost function for penalizing forbidden-state occupation.

This cost function is defined as $$ C=\sum_{k}\sum_{f}\int_{0}^{T}dt|\langle\psi_{f}^{k}|\psi_{i}^{k}(t)\rangle|^2, $$ where \(|\psi_{i}^{k}(t)\rangle\) is the \(k^{\rm th}\) initial state propagated to time \(t\) and |\psi_{f}^{k}\rangle is a forbidden state for the \(k^{\rm th}\) initial state.

Parameters:

  • forbidden_states_list (list[QArrayLike]) –

    Shape of QArrays = (n, 1) or (n, n)_. A list of forbidden states for each initial state (indexed by s). The inner lists need not all be of the same shape, for instance if some initial states have more forbidden states than others. The array is eventually reshaped to have shape (s, 1, f, n, 1) or (s, 1, f, n, n) (for sesolve or mesolve, respectively) where s is the number of initial states, f is the length of the longest forbidden-state list (with zero-padding for meaningless entries) and 1 is an added batch dimension for eventually batching over tsave.

  • cost_multiplier (float, default: 1.0 ) –

    Weight for this cost function relative to other cost functions.

  • target_cost (float, default: 0.0 ) –

    Target value for this cost function. If options.all_costs is True, the optimization terminates early if all cost functions fall below their target values. If options.all_costs is False, the optimization terminates if only one cost function falls below its target value.

Returns:

  • ForbiddenStates –

    Callable object that returns the forbidden-state cost and whether the cost is below the target value.

control_area(threshold: float = 0.0, cost_multiplier: float = 1.0, target_cost: float = 0.0) -> ControlCostArea ¤

Control area cost function.

Penalize the area under the pulse curves according to $$ C = \sum_{j}\int_{0}^{T}\Omega_{j}(t)dt, $$ where the \(\Omega_{j}\) are the individual controls and \(T\) is the pulse time.

Parameters:

  • threshold (float, default: 0.0 ) –

    Threshold to use for penalizing areas above this value in absolute magnitude.

  • cost_multiplier (float, default: 1.0 ) –

    Weight for this cost function relative to other cost functions.

  • target_cost (float, default: 0.0 ) –

    Target value for this cost function. If options.all_costs is True, the optimization terminates early if all cost functions fall below their target values. If options.all_costs is False, the optimization terminates if only one cost function falls below its target value.

Returns:

  • ControlCostArea –

    Callable object that returns the control-area cost and whether the cost is below the target value.

control_norm(threshold: float, cost_multiplier: float = 1.0, target_cost: float = 0.0) -> ControlCostNorm ¤

Control norm cost function.

Penalize the norm of the controls above some threshold according to $$ C = \sum_{j}\int_{0}^{T}ReLU(|\Omega_{j}(t)|-\Omega_{max})dt, $$ where the \(\Omega_{j}\) are the individual controls, \(T\) is the pulse time and \(\Omega_{max}\) is the threshold.

Parameters:

  • threshold (float) –

    Threshold to use for penalizing amplitudes above this value in absolute magnitude.

  • cost_multiplier (float, default: 1.0 ) –

    Weight for this cost function relative to other cost functions.

  • target_cost (float, default: 0.0 ) –

    Target value for this cost function. If options.all_costs is True, the optimization terminates early if all cost functions fall below their target values. If options.all_costs is False, the optimization terminates if only one cost function falls below its target value.

Returns:

  • ControlCostNorm –

    Callable object that returns the control-norm cost and whether the cost is below the target value.

custom_control_cost(cost_fun: callable, cost_multiplier: float = 1.0, target_cost: float = 0.0) -> CustomControlCost ¤

Cost function based on an arbitrary transformation of the controls.

Penalize the controls according to an arbitrary function F $$ C = \sum_{j}\int_{0}^{T}F(\Omega_{j}(t))dt, $$

Parameters:

  • cost_fun (callable) –

    Cost function which must have signature (control_amp: Array) -> Array.

  • cost_multiplier (float, default: 1.0 ) –

    Weight for this cost function relative to other cost functions.

  • target_cost (float, default: 0.0 ) –

    Target value for this cost function. If options.all_costs is True, the optimization terminates early if all cost functions fall below their target values. If options.all_costs is False, the optimization terminates if only one cost function falls below its target value.

Returns:

  • CustomControlCost –

    Callable object that returns the cost for the custom function and whether the cost is below the target value.

Example: Penalize negative drive amplitudes 
def penalize_negative(control_amp: jax.Array) -> jax.Array:
    return jax.nn.relu(-control_amp)


negative_amp_cost = ql.custom_control_cost(penalize_negative)

custom_cost(cost_fun: callable, cost_multiplier: float = 1.0, target_cost: float = 0.0) -> CustomCost ¤

A custom cost function.

In many (most!) cases, the user may want to add a cost function to their optimization that is not included in the hardcoded set of available cost functions.

Parameters:

  • cost_fun (callable) –

    Cost function which must have signature (result: dq.SolveResult, H: dq.TimeQArray, parameters: dict | Array) -> Array.

  • cost_multiplier (float, default: 1.0 ) –

    Weight for this cost function relative to other cost functions.

  • target_cost (float, default: 0.0 ) –

    Target value for this cost function. If options.all_costs is True, the optimization terminates early if all cost functions fall below their target values. If options.all_costs is False, the optimization terminates if only one cost function falls below its target value.

Returns:

  • CustomCost –

    Callable object that returns the cost for the custom function.

Example: Penalize an expectation value

Let's imagine we want to penalize the value of some expectation value at the final time in tsave.

def penalize_expect(
    result: SolveResult, H: TimeQArray, parameters: dict | Array
) -> Array:
    # 0 is the index of the operator, -1 is the time index
    return jnp.sum(jnp.abs(result.expects[0, -1]))


expect_cost = ql.custom_cost(penalize_expect)
Then expect_cost can be added to the other utilized cost functions. The only thing happening under the hood is that the penalize_expect function is passed to jax.tree_util.Partial to enable it to be passed through jitted functions without issue.

Models¤

sesolve_model(H_function: callable, psi0: QArrayLike, tsave_or_function: ArrayLike | callable, *, exp_ops: list[ArrayLike] | None = None) -> SESolveModel ¤

Instantiate sesolve model.

Here we instantiate the model that is called at each step of the optimization iteration, returning a tuple of the result of calling sesolve as well as the Hamiltonian evaluated at the parameter values.

Parameters:

  • H_function (callable) –

    function specifying how to update the Hamiltonian

  • psi0 (QArrayLike) –

    Shape = (..., n, 1). Initial states.

  • tsave_or_function (ArrayLike | callable) –

    If ArrayLike, then Shape = (ntsave,). Either an array of times passed to sesolve or a method specifying how to update the times that are passed to sesolve.

  • exp_ops (list[ArrayLike] | None, default: None ) –

    Operators to calculate expectation values of, in case some of the cost functions depend on the value of certain expectation values.

Returns:

  • SESolveModel –

    Model that when called with the parameters we optimize over as argument returns the results of sesolve as well as the updated Hamiltonian.

Piecewise-constant control 

tsave = jnp.linspace(0.0, 11.0, 10)
psi0 = [dq.basis(2, 0)]
H1s = [dq.sigmax(), dq.sigmay()]


def H_pwc(values: Array) -> dq.TimeQArray:
    H = dq.sigmaz()
    for idx, _H1 in enumerate(H1s):
        H += dq.pwc(tsave, values[idx], _H1)
    return H


sesolve_model = ql.sesolve_model(H_pwc, psi0, tsave)
See for example this tutorial.

Spline control

In more complex cases, we can imagine that the optimized parameters are the control points fed into a spline, and additionally the control times themselves are optimized. 

init_drive_params_topt = {
    'dp': -0.001 * jnp.ones((len(H1s), len(tsave))),
    't': tsave[-1],
}


def H_func_topt(t: float, drive_params_dict: dict) -> dq.TimeQArray:
    drive_params = drive_params_dict['dp']
    new_tsave = jnp.linspace(0.0, drive_params_dict['t'], len(tsave))
    drive_spline = _drive_spline(drive_params, envelope, new_tsave)
    drive_amps = drive_spline.evaluate(t)
    drive_Hs = jnp.einsum('d,dij->ij', drive_amps, H1s)
    return H0 + drive_Hs


def update_H_topt(drive_params_dict: dict) -> dq.TimeQArray:
    new_H = jtu.Partial(H_func_topt, drive_params_dict=drive_params_dict)
    return dq.timecallable(new_H)


def update_tsave_topt(drive_params_dict: dict) -> jax.Array:
    return jnp.linspace(0.0, drive_params_dict['t'], len(tsave))


se_t_opt_Kerr_model = ql.sesolve_model(update_H_topt, psi0, update_tsave_topt)

mesolve_model(H_function: callable, jump_ops: list[QArrayLike | TimeQArray], rho0: QArrayLike, tsave_or_function: ArrayLike | callable, *, exp_ops: list[ArrayLike] | None = None) -> MESolveModel ¤

Instantiate mesolve model.

Here we instantiate the model called at each step of the optimization iteration, returning a tuple of the result of calling mesolve as well as the Hamiltonian evaluated at the parameter values.

Parameters:

  • H_function (callable) –

    function specifying how to update the Hamiltonian

  • jump_ops (list[QArrayLike | TimeQArray]) –

    Each of QArray or TimeQArray is has shape = (...Lk, n, n)). List of jump operators.

  • rho0 (QArrayLike) –

    Shape = (..., n, n). Initial density matrices.

  • tsave_or_function (ArrayLike | callable) –

    If ArrayLike, then Shape = (ntsave,). Either an array of times passed to sesolve or a method specifying how to update the times that are passed to mesolve.

  • exp_ops (list[ArrayLike] | None, default: None ) –

    Operators to calculate expectation values of, in case some of the cost functions depend on the value of certain expectation values.

Returns:

  • MESolveModel –

    Model that when called with the parameters we optimize over as argument returns the results of mesolve as well as the updated Hamiltonian.

Instantiating a MESolveModel

Instantiating a MESolveModel is quite similar to instantiating an SESolveModel, with the two differences being that we need to supply jump operators, and the initial and target states should be specified as density matrices. Continuing the last example from sesolve_model 

jump_ops = [0.03 * dq.sigmax()]
me_initial_states = dq.todm(psi0)
me_Kerr_model = ql.mesolve_model(
    update_H_topt, jump_ops, me_initial_states, update_tsave_topt
)
See this tutorial for example

sepropagator_model(H_function: callable, tsave_or_function: ArrayLike | callable) -> SEPropagatorModel ¤

Instantiate sepropagator model.

Here we instantiate the model that is called at each step of the optimization iteration, returning a tuple of the result of calling sepropagator as well as the Hamiltonian evaluated at the parameter values.

Parameters:

  • H_function (callable) –

    function specifying how to update the Hamiltonian

  • tsave_or_function (ArrayLike | callable) –

    If ArrayLike, then Shape = (ntsave,). Either an array of times passed to sepropagator or a method specifying how to update the times that are passed to sepropagator.

Returns:

  • SEPropagatorModel –

    Model that when called with the parameters we optimize over as argument returns the results of sepropagator as well as the updated Hamiltonian.

Instantiating a SEPropagatorModel 
tsave = jnp.linspace(0.0, 11.0, 10)
H1s = [dq.sigmax(), dq.sigmay()]


def H_pwc(values: Array) -> dq.TimeQArray:
    H = dq.sigmaz()
    for idx, _H1 in enumerate(H1s):
        H += dq.pwc(tsave, values[idx], _H1)
    return H


sepropagator_model = ql.sepropagator_model(H_pwc, tsave)

mepropagator_model(H_function: callable, jump_ops: list[QArrayLike | TimeQArray], tsave_or_function: ArrayLike | callable) -> MEPropagatorModel ¤

Instantiate mepropagator model.

Here we instantiate the model that is called at each step of the optimization iteration, returning a tuple of the result of calling mepropagator as well as the Hamiltonian evaluated at the parameter values.

Parameters:

  • H_function (callable) –

    function specifying how to update the Hamiltonian

  • jump_ops (list[QArrayLike | TimeQArray]) –

    Each of QArray or TimeQArray is has shape = (...Lk, n, n)). List of jump operators.

  • tsave_or_function (ArrayLike | callable) –

    If ArrayLike, then Shape = (ntsave,). Either an array of times passed to mepropagator or a method specifying how to update the times that are passed to mepropagator.

Returns:

  • MEPropagatorModel –

    Model that when called with the parameters we optimize over as argument returns the results of mepropagator as well as the updated Hamiltonian.

Instantiating a MEPropagatorModel

Instantiating a MEPropagatorModel is quite similar to instantiating an SEPropagatorModel, with the difference being that we need to supply jump operators. Continuing the last example from sepropagator_model 

jump_ops = [0.03 * dq.sigmax()]
mep_Kerr_model = ql.mepropagator_model(
    update_H_topt, jump_ops, update_tsave_topt
)

Custom Plotting¤

plot_costs(ax: Axes, costs: Cost, cost_values_over_epochs: list | np.ndarray) -> Axes ¤

Plot the evolution of the cost function values.

get_controls(H: TimeQArray, tsave: np.ndarray) -> list[np.ndarray] ¤

Extract the Hamiltonian prefactors at the supplied times.

plot_controls(ax: Axes, _expects: Array | None, model: Model, parameters: Array | dict) -> Axes ¤

Plot the Hamiltonian prefactors, usually corresponding to controls.

plot_fft(ax: Axes, _expects: Array | None, model: Model, parameters: Array | dict) -> Axes ¤

Plot the fft of the Hamiltonian controls.

plot_expects(ax: Axes, expects: Array | None, model: Model, parameters: Array | dict) -> Axes ¤

Plot the expectation values obtained from the time evolution.

custom_plotter(plotting_functions: list[Callable]) -> Plotter ¤

Instantiate a custom Plotter for tracking results during optimization.

This function returns a Plotter that can be passed to optimize to track the progress of an optimization run. Note that the cost function values are always plotted in the first panel and that there is no limit to the number of plots a user can ask for: if more than four, additional plots will appear in a new row of four, and so on.

Parameters:

  • plotting_functions (list[Callable]) –

    list of functions that each return a plot useful for tracking intermediate results, such as the value of the optimized controls, fft of the controls, expectation values, etc. Each function must have signature example_plot_function(ax, expects, model, parameters) where ax is the matplotlib.pyplot.Axes instance where the results are plotted, expects is of type dq.SolveResult.expects (which could be None), model is of type ql.Model and parameters are the parameters being optimized. Of course, some of these arguments may be unused for a particular plot (for instance if we are plotting expectation values, we don't need access to parameters).

Returns:

  • Plotter –

    Plotter whose update_plots method is repeatedly called during an optimization

  • Plotter –

    run.

Example

We plot the controls as well as the expectation values for two different initial states 

import dynamiqs as dq
import jax.numpy as jnp
import numpy as np
import qontrol as ql
from functools import partial

H1s = [dq.sigmax(), dq.sigmay()]
H1_labels = ['X', 'Y']


def plot_states(
    ax: Axes,
    expects: Array | None,
    model: ql.Model,
    parameters: Array | dict,
    which=0,
) -> Axes:
    ax.set_facecolor('none')
    tsave = model.tsave_function(parameters)
    batch_idxs = np.ndindex(*expects.shape[:-3])
    for batch_idx in batch_idxs:
        ax.plot(tsave, np.real(expects[tuple(batch_idx), which, 0]))
    ax.set_xlabel('time [ns]')
    ax.set_ylabel(
        f'population in $|1\\rangle$ for initial state $|{which}\\rangle$'
    )
    return ax


def plot_controls(
    ax: Axes, _expects: Array | None, model: ql.Model, parameters: Array | dict
) -> Axes:
    ax.set_facecolor('none')
    tsave = model.tsave_function(parameters)
    finer_tsave = jnp.linspace(0.0, tsave[-1], 10 * len(tsave))
    for idx, control in enumerate(parameters):
        H_c = dq.pwc(tsave, control, H1s[idx])
        ax.plot(
            finer_tsave,
            np.real(jax.vmap(H_c.prefactor)(finer_tsave)) / 2 / np.pi,
            label=H1_labels[idx],
        )
    ax.legend(loc='lower right', framealpha=0.0)
    ax.set_ylabel('pulse amplitude [GHz]')
    ax.set_xlabel('time [ns]')
    return ax


plotter = ql.custom_plotter(
    [
        plot_controls,
        partial(plot_states, which=0),
        partial(plot_states, which=1),
    ]
)
See for example this tutorial.

File utilities¤