diff --git a/.github/workflows/nightly_ci.yaml b/.github/workflows/nightly_ci.yaml index 6c1fd0d3..64e82186 100644 --- a/.github/workflows/nightly_ci.yaml +++ b/.github/workflows/nightly_ci.yaml @@ -65,6 +65,7 @@ jobs: cd .. python -m pip install git+https://github.com/qutip/qutip-qip python -m pip install --no-deps git+https://github.com/qutip/qutip-jax + python -m pip install --no-deps git+https://github.com/qutip/qutip-qoc git clone -b master https://github.com/qutip/qutip-qtrl.git cd qutip-qtrl diff --git a/.github/workflows/notebook_ci.yaml b/.github/workflows/notebook_ci.yaml index 58013fba..0e53cfb3 100644 --- a/.github/workflows/notebook_ci.yaml +++ b/.github/workflows/notebook_ci.yaml @@ -67,6 +67,7 @@ jobs: cd .. python -m pip install git+https://github.com/qutip/qutip-qip python -m pip install --no-deps git+https://github.com/qutip/qutip-jax + python -m pip install --no-deps git+https://github.com/qutip/qutip-qoc git clone -b master https://github.com/qutip/qutip-qtrl.git cd qutip-qtrl diff --git a/tutorials-v5/miscellaneous/v5_paper-optimal-control.md b/tutorials-v5/miscellaneous/v5_paper-optimal-control.md new file mode 100644 index 00000000..efd501db --- /dev/null +++ b/tutorials-v5/miscellaneous/v5_paper-optimal-control.md @@ -0,0 +1,329 @@ +--- +jupyter: + jupytext: + text_representation: + extension: .md + format_name: markdown + format_version: '1.3' + jupytext_version: 1.13.8 + kernelspec: + display_name: Python 3 (ipykernel) + language: python + name: python3 +--- + +# QuTiPv5 Paper Example: The Quantum Optimal Control Package + +Authors: Maximilian Meyer-Mölleringhof (m.meyermoelleringhof@gmail.com), Boxi Li (etamin1201@gmail.com), Neill Lambert (nwlambert@gmail.com) + +Quantum systems are naturally sensitive to their environment and external perturbations. +This is great, as it allows for very precise measurements. +However, it also makes handling errors and imprecisions a big challenge. +In the case for quantum computing, finding the optimal parameters to achieve a desired operation is this thus an important problem. +Optimization parameters may include amplitude, frequency, duration, bandwidth, etc. and are generally directly dependent on the considered hardware. + +To find these optimal control parameters, several methods have been developed. +Here, we look at three algorithms: *gradient ascent pulse engineering* (GRAPE) [\[3\]](#References), *chopped random basis* (CRAB) [\[4\]](#References) and *gradient optimization af analytic controls* (GOAT) [\[5\]](#References). +Whereas the former two have been part of the `QuTiP-QTRL` package of QuTiPv4, the latter is a new addition in version 5. +Althogether, these algorithms are now included in the new `QuTiP-QOC` package that also adds `QuTiP-JAX` [\[6\]](#References) integration via the JAX optimization technique (JOPT). + +```python +import matplotlib.pyplot as plt +import numpy as np +from jax import jit, numpy +from qutip import (about, gates, liouvillian, qeye, sigmam, sigmax, sigmay, + sigmaz) +from qutip_qoc import Objective, optimize_pulses + +%matplotlib inline +``` + +## Introduction + +In this example we want to implement a Hadamard gate on a single qubit. +In general, a qubit might be subject to decoherence which can be captured using the Lindblad formalism with the jump operator $\sigma_{-}$. + +For simplicity, we consider a control Hamiltonian parametrized by $\sigma_x$, $\sigma_y$ and $\sigma_z$: + +$H_c(t) = c_x(t) \sigma_x + c_y(t) \sigma_y + c_z(t) \sigma_z$ + +with $c_x(t)$, $c_y(t)$ and $c_z(t)$ as independent control parameters. +Additionally, we model a constant drift Hamiltonian + +$H_d = \dfrac{1}{2} (\omega \sigma_z + \delta \sigma_x)$, + +with associated energy splitting $\omega$ and tunneling rate $\delta$. +The amplitude damping rate for the collapse operator $C = \sqrt{\gamma} \sigma_-$ is denoted as $\gamma$. + +```python +# energy splitting, tunneling, amplitude damping +omega = 0.1 # energy splitting +delta = 1.0 # tunneling +gamma = 0.1 # amplitude damping +sx, sy, sz = sigmax(), sigmay(), sigmaz() + +Hc = [sx, sy, sz] # control operator +Hc = [liouvillian(H) for H in Hc] + +Hd = 1 / 2 * (omega * sz + delta * sx) # drift term +Hd = liouvillian(H=Hd, c_ops=[np.sqrt(gamma) * sigmam()]) + +# combined operator list +H = [Hd, Hc[0], Hc[1], Hc[2]] +``` + +```python +# objectives for optimization +initial = qeye(2) +target = gates.hadamard_transform() +fid_err = 0.01 +``` + +```python +# pulse time interval +times = np.linspace(0, np.pi / 2, 100) +``` + +## Implementation + +### GRAPE Algorithm + +The GRAPE algorithm works by minimizing an infidelity loss function that measures how close the final state or unitary tranformation is to the desired target. +Starting from the provided `guess` control pulse, it optimizes evenly spaced piecewise constant pulse amplitudes. +In the end, it strives to achieve the desired target infidelity, sepcified by the `fid_err_targ` keyword. + +```python +res_grape = optimize_pulses( + objectives=Objective(initial, H, target), + control_parameters={ + "ctrl_x": {"guess": np.sin(times), "bounds": [-1, 1]}, + "ctrl_y": {"guess": np.cos(times), "bounds": [-1, 1]}, + "ctrl_z": {"guess": np.tanh(times), "bounds": [-1, 1]}, + }, + tlist=times, + algorithm_kwargs={"alg": "GRAPE", "fid_err_targ": fid_err}, +) +``` + +### CRAB Algorithm + +This algorithm is based on the idea of expanding the control fields in a random basis and optimizing the expansion coefficients $\vec{\alpha}$. +This has the advantage of using analytical control functions $c(\vec{\alpha}, t)$ on a continuous time interval, and is by default a Fourier expansion. +This reduces the search space to the function parameters. +Typically, these parameters can efficiently be calculated through direct search algorithms (like Nelder-Mead). +The basis function is only expanded for some finite number of summands and the initial basis coefficients are usually picked at random. + +```python +n_params = 3 # adjust in steps of 3 +alg_args = {"alg": "CRAB", "fid_err_targ": fid_err, "fix_frequency": False} +``` + +```python +res_crab = optimize_pulses( + objectives=Objective(initial, H, target), + control_parameters={ + "ctrl_x": { + "guess": [1 for _ in range(n_params)], + "bounds": [(-1, 1)] * n_params, + }, + "ctrl_y": { + "guess": [1 for _ in range(n_params)], + "bounds": [(-1, 1)] * n_params, + }, + "ctrl_z": { + "guess": [1 for _ in range(n_params)], + "bounds": [(-1, 1)] * n_params, + }, + }, + tlist=times, + algorithm_kwargs=alg_args, +) +``` + +### GOAT Algorithm + +Similar to CRAB, this method also works with analytical control functions. +By constructing a coupled system of equations of motion, the derivative of the (time ordered) evolution operator with respect to the control parameters can be calculated after numerical forward integration. +In unconstrained settings, GOAT was found to outperform the previous described methods in terms of convergence and fidelity achievement. +The QuTiP implementation allows for arbitrary control functions provided together with their respective derivatives in a common python manner. + +```python +def sin(t, c): + return c[0] * np.sin(c[1] * t) + + +# derivatives +def grad_sin(t, c, idx): + if idx == 0: # w.r.t. c0 + return np.sin(c[1] * t) + if idx == 1: # w.r.t. c1 + return c[0] * np.cos(c[1] * t) * t + if idx == 2: # w.r.t. time + return c[0] * np.cos(c[1] * t) * c[1] +``` + +```python +H = [Hd] + [[hc, sin, {"grad": grad_sin}] for hc in Hc] + +bnds = [(-1, 1), (0, 2 * np.pi)] +ctrl_param = {id: {"guess": [1, 0], "bounds": bnds} for id in ["x", "y", "z"]} +``` + +For even faster convergence QuTiP extends to original algorithm with the option to optimize controls with +respect to the overall time evolution, which can be enabled by specifying the additional time keyword +argument: + +```python +# treats time as optimization variable +ctrl_param["__time__"] = { + "guess": times[len(times) // 2], + "bounds": [times[0], times[-1]], +} +``` + +```python +# run the optimization +res_goat = optimize_pulses( + objectives=Objective(initial, H, target), + control_parameters=ctrl_param, + tlist=times, + algorithm_kwargs={ + "alg": "GOAT", + "fid_err_targ": fid_err, + }, +) +``` + +### JOT Algorithm - JAX integration + +QuTiP's new JAX backend provides automatic differentiation capabilities that can be directly be used with the new control framework. +As with QuTiP’s GOAT implementation, any analytically defined control function can be handed to the algorithm. +However, in this method, JAX automatic differentiation abilities take care of calculating the derivative throughout the whole system evolution. +Therefore we don't have to provide any derivatives manually. +Compared to the previous example, this simply means to swap the control functions with their just-in-time compiled version. + +```python +@jit +def sin_y(t, d, **kwargs): + return d[0] * numpy.sin(d[1] * t) + + +@jit +def sin_z(t, e, **kwargs): + return e[0] * numpy.sin(e[1] * t) + + +@jit +def sin_x(t, c, **kwargs): + return c[0] * numpy.sin(c[1] * t) +``` + +```python +H = [Hd] + [[Hc[0], sin_x], [Hc[1], sin_y], [Hc[2], sin_z]] +``` + +```python +res_jopt = optimize_pulses( + objectives=Objective(initial, H, target), + control_parameters=ctrl_param, + tlist=times, + algorithm_kwargs={ + "alg": "JOPT", + "fid_err_targ": fid_err, + }, +) +``` + +## Comparison of Results + +After running the global and local optimization, one can compare the results obtained by the various +algorithms through a `qoc.Result` object, which provides common optimization metrics along with the `optimized_controls`. + +### Pulse Amplitudes + +```python +fig, ax = plt.subplots(1, 3, figsize=(13, 5)) + +goat_range = times < res_goat.optimized_params[-1] +jopt_range = times < res_jopt.optimized_params[-1] + +for i in range(3): + ax[i].plot(times, res_grape.optimized_controls[i], ":", label="GRAPE") + ax[i].plot(times, res_crab.optimized_controls[i], "-.", label="CRAB") + ax[i].plot( + times[goat_range], + np.array(res_goat.optimized_controls[i])[goat_range], + "-", + label="GOAT", + ) + ax[i].plot( + times[jopt_range], + np.array(res_jopt.optimized_controls[i])[jopt_range], + "--", + label="JOPT", + ) + + ax[i].set_xlabel(r"Time $t$") + +ax[0].legend(loc=0) +ax[0].set_ylabel(r"Pulse amplitude $c_x(t)$", labelpad=-5) +ax[1].set_ylabel(r"Pulse amplitude $c_y(t)$", labelpad=-5) +ax[2].set_ylabel(r"Pulse amplitude $c_z(t)$", labelpad=-5) + +plt.show() +``` + +### Infidelities and Processing Time + +```python +print("GRAPE: ", res_grape.fid_err) +print(res_grape.total_seconds, " seconds") +print() +print("CRAB : ", res_crab.fid_err) +print(res_crab.total_seconds, " seconds") +print() +print("GOAT : ", res_goat.fid_err) +print(res_goat.total_seconds, " seconds") +print() +print("JOPT : ", res_jopt.fid_err) +print(res_jopt.total_seconds, " seconds") +``` + +## References + +[1] [QuTiP 5: The Quantum Toolbox in Python](https://arxiv.org/abs/2412.04705) + +[2] [QuTiP-QOC Repository](https://github.com/qutip/qutip-qoc) + +[3] [Khaneja, et. al, Journal of Magnetic Resonance (2005)](https://www.sciencedirect.com/science/article/pii/S1090780704003696) + +[4] [Caneva, et. al, Phys. Rev. A (2011)](https://link.aps.org/doi/10.1103/PhysRevA.84.022326) + +[5] [Machnes, et. al, Phys. Rev. Lett. (2018)](https://link.aps.org/doi/10.1103/PhysRevLett.120.150401) + +[6] [QuTiP-JAX Repository](https://github.com/qutip/qutip-jax) + + + +## About + +```python +about() +``` + +## Testing + +```python +assert ( + res_grape.fid_err < fid_err +), f"GRAPE did not reach the target infidelity of < {fid_err}." +assert ( + res_crab.fid_err < fid_err +), f"CRAB did not reach the target infidelity of < {fid_err}." +assert ( + res_goat.fid_err < fid_err +), f"GOAT did not reach the target infidelity of < {fid_err}." +assert ( + res_jopt.fid_err < fid_err +), f"JOPT did not reach the target infidelity of < {fid_err}." +```