1c Quantum control.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantum control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# file: seeq/control.py\n",
    "import numpy as np\n",
    "from seeq.evolution import evolve\n",
    "from numbers import Number\n",
    "import scipy.optimize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem statement"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to use the parameterized pulse approach to quantum control. Our model assumes a Hamiltonian of the form\n",
    "$$H = H_0 + g(t;x) H_1$$\n",
    "where $g(t)$ is a time-dependent control that will be eventually expanded in a basis of functions\n",
    "$$g(t;x) = \\sum_n x_n f_n(t),$$\n",
    "or which may have some other more complicated dependency, such as\n",
    "$$g(t;x) = x_0 \\cos(x_1 t + x_2).$$\n",
    "We will denote this fact by writing $H(t;x).$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our goal is to control $N$ states, that must suffer a unitary operation $U_g$ after a time $T.$ If the unitary evolution with $H(t;x)$ is given by $U(t),$ we will have\n",
    "$$U(T) |\\psi_n\\rangle \\simeq U |\\psi_n\\rangle,\\;n=1,\\ldots,N.$$\n",
    "The unitary operator satisfies the Schrödinger equation\n",
    "$$i \\frac{d}{dt} U(t;x) = H(t;x) U(t;x),$$\n",
    "with $U(0) = 1.$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We measure the quality of our control by studying the \"fidelity\"\n",
    "$$F = \\frac{1}{N}\\mathrm{Re}\\left(\\sum_n \\langle\\psi_n | U_g^\\dagger U(T)|\\psi_n\\rangle\\right).$$\n",
    "When the protocol is exact, $F\\simeq 1.$ When the unitary is not exactly the one we want, or even when the phases are different, this quantity decreases towards 0.\n",
    "\n",
    "We can therefore reformulate our problem as computing the minimum of the cost function\n",
    "$$x_{opt} = \\mathrm{argmin}_{x} [-F(x)].$$\n",
    "subject to any additional constraints that we may have on $x.$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## State gradient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The optimization of our cost function is simpler if know the cost function changes when we perturb our control. In other words, given an infinitesimal change $x\\to x+\\epsilon,$ we want to know whether $F(x+\\epsilon) \\simeq F(x) + \\epsilon \\nabla F(x)$ increases or decreases, so as to choose a new control that is more optimal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As explained in the appendices of [O. Romero-Isart and J. J. García-Ripoll\n",
    "Phys. Rev. A 76, 052304 (2007)](https://journals.aps.org/pra/abstract/10.1103/PhysRevA.76.052304), we can estimate $\\nabla F(x)$ by analyzing how the unitary operator $U(t;x)$ changes under small perturbations in the control. To first order in perturbation theory, this unitary operator will satisfy\n",
    "$$i \\frac{d}{dt}U(t;x+\\epsilon) = \\left[H(t) + \\epsilon \\frac{\\partial H}{\\partial x}\\right] U(t;x+\\epsilon).$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We introduce an interaction picture to compute the deviation $W$ induced by the perturbation\n",
    "$$U(t;x+\\epsilon) = U(t;x) W(t;x).$$\n",
    "The new operator satisfies\n",
    "$$i\\frac{d}{dt}W(t;x) = \\epsilon U(t;x)^\\dagger \\frac{\\partial H}{\\partial x}(t;x) U(t;x) W(t;x).$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Keeping terms to $\\mathcal{O}(\\epsilon^2),$\n",
    "$$W(t;x) = 1 - i \\epsilon \\int_{t_0}^t \\mathrm{d}\\tau\\, U(\\tau;x)^\\dagger \\frac{\\partial H}{\\partial x}(\\tau;x) U(\\tau;x) + \\mathcal{O}(\\epsilon^2).$$\n",
    "Or in the original operator, \n",
    "$$U(t;x+\\epsilon) = U(t;x) - i \\epsilon U(t;x) \\int_{t_0}^t \\mathrm{d}\\tau\\, U(\\tau;x)^\\dagger \\frac{\\partial H}{\\partial x}(\\tau;x) U(\\tau;x) + \\mathcal{O}(\\epsilon^2).$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This implies a formula for the derivative of the unitary operator\n",
    "$$\\frac{\\partial}{\\partial x} U(t;x) = - i U(t;x) \\int_{t_0}^t \\mathrm{d}\\tau\\, U(\\tau;x)^\\dagger \\frac{\\partial H}{\\partial x}(\\tau;x) U(\\tau;x) + \\mathcal{O}(\\epsilon^2).$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And from here it follows a formula for the gradient of the cost function\n",
    "$$\\frac{\\partial}{\\partial x}F = \\frac{1}{N}\\mathrm{Re}\\sum_n (-i)\\int_0^T \\left\\langle U_g^\\dagger U(T;x) U(\\tau;x)^\\dagger \\frac{\\partial H}{\\partial x}(\\tau;x) U(\\tau;x) \\right\\rangle_{\\psi_n}\\mathrm{d}\\tau.$$\n",
    "Note that all formulas extend trivially for any dimensionality of the variable $x.$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fidelity gradient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These integral formulations may be reformulated in a friendlier way, that amounts to solving a couple of differential equations:\n",
    "\n",
    "1. Compute the backwards evolved states $|\\xi_n(0)\\rangle = U(T;x)^\\dagger U_g |\\psi_n\\rangle.$\n",
    "\n",
    "2. Solve the system of equations\n",
    "\\begin{eqnarray}\n",
    "i\\frac{d}{dt}|\\psi_n(t;x)\\rangle &=& H(t;x)|\\psi_n(t;x)\\rangle,\\\\\n",
    "i\\frac{d}{dt}|\\xi_n(t;x)\\rangle &=& H(t;x)|\\xi_n(t;x)\\rangle,\\\\\n",
    "\\frac{d}{dt}f_{n,i} &=& \\mathrm{Im}\\langle \\xi_n(t;x)|\\frac{\\partial H}{\\partial x_i}(t;x)|\\psi_n(t;x)\\rangle\n",
    "\\end{eqnarray}\n",
    "\n",
    "3. Estimate the gradient $$\\frac{\\partial F}{\\partial x_i} = \\frac{1}{N}\\mathrm{Im}\\sum_{n,i} f_{n,i}(T).$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other measures: overlap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of computing the fidelity, we may want to transfer a set of states to a different subset, without regard for the phases that are achieved. In this case we can use a different measure\n",
    "$$O = \\frac{1}{N}\\sum_n |\\langle \\psi_n |U_g^\\dagger U(T) |\\psi_n\\rangle|^2 = \\sum_n F_n F_n^*.$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The quantities $F_n$ are such that\n",
    "$$F = \\frac{1}{N}\\mathrm{Re}\\sum_n F_n.$$\n",
    "The gradient of this function is slightly different. We just need to realize that\n",
    "$$\\frac{\\partial F_n}{\\partial x_i} =(-i) \\sum_i f_{n,i}(T).$$\n",
    "Hence,\n",
    "$$\\frac{\\partial O}{\\partial x_i} = \\frac{1}{N}\\sum_{n,i} [ F_n^* (-i)f_{n,i}(T) + \\mathrm{c.c}]\n",
    "=\\frac{2}N\\mathrm{Im}\\sum_{n,i} F_n^* f_{n,i}(T).$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code below implements two versions of a parameterized quantum control. In the simplest incantation, it requires a function `H(t,x,ψ)` that returns the product $H(t;x)\\psi.$ It also requires the target gate `Ug`, the states that are to be controlled as the columns of `ψ0` and a vector of times `T` used to solve the evolution. In this simple form, the algorithm does not compute $\\nabla F,$ and uses Scipy's algoritms for [global optimization](https://docs.scipy.org/doc/scipy/reference/optimize.html) to find the optimal control."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A more refined version of the code is activated when you pass the optional argument `dH(t,x,ψ)`, which is a function that returns a list of matrices, `[dH1, dH2,...]` with `dHn` being the result of computing $\\frac{\\partial H}{\\partial x_n}(t;x) \\psi.$ When `dH` is passed, the function uses the algorithm above to estimate $\\nabla F$ and use this to speed up Scipy's optimization algorithms."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function `parametric_control()` relies on the routine `evolve()` from [SeeQ](Time%20evolution.ipynb). It defaults to the Chebyshev algorithm, but you may choose other algorithms that will be faster for smaller problems, such as `'eig'` or `'expm'`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As mentioned above, our code uses Scipy for computing the optimal control with or without gradients. This function has different solvers or optimizers, such as `'BFGS'`, `'COBYLA'`, etc. You can change the default optimizer by supplying the argument `optimizer` to `parametric_control()`. If the optimizer you select requires additional options, they can be passed as keyword arguments at the end of your call to `parametric_control()`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the accuracy of the time evolution is critical for the control problem to succeed. The number of time steps affects both the estimates of $\\psi(t_n)$ and $\\xi(t_n)$ as well as the estimate of $\\nabla F,$ which is computed using a Simpson integration formula. To select the number of time steps you may either supply a vector of real times `T=[t0,t1,...,tN]`, or make `T` a real number and pass also the number of `steps` as a keyword argument.\n",
    "\n",
    "Hint: a crude way to check convergence is to double the number of steps and inspect whether the results changed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.optimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# file: seeq/control.py\n",
    "\n",
    "def parametric_control(x0, H, ψ0, T, Ug=None, ξT=None, dH=None, check_gradient=False,\n",
    "                       steps=100, tol=1e-10, method='chebyshev', target='fidelity',\n",
    "                       debug=False, optimizer='BFGS', dispiter=0, **kwdargs):\n",
    "    \"\"\"Solve the quantum control problem for a Hamiltonian H acting on\n",
    "    a basis of states ψ.\n",
    "    \n",
    "    Arguments:\n",
    "    ----------\n",
    "    H      - Callable object H(t,x,ψ) that applies H(t,x) on ψ\n",
    "    ψ0     - N x d object with N wavefunctions\n",
    "    Ug, ξT - Desired quantum operation or desired quantum states\n",
    "    T      - Either a time or a vector of times\n",
    "    steps  - # time steps in the algorithm (if T is a number)\n",
    "    tol    - Optimization tolerance\n",
    "    target - 'fidelity' or 'overlap', depending on the cost function\n",
    "    method - Solution method for the time evolution\n",
    "    debug  - Return cost functions\n",
    "    optimizer - Method for scipy.optimize.minimize\n",
    "    \n",
    "    Output:\n",
    "    -------\n",
    "    x      - Optimal control\n",
    "    F      - Fidelity\"\"\"\n",
    "\n",
    "    if isinstance(T, Number):\n",
    "        times = np.linspace(0, T, steps)\n",
    "    else:\n",
    "        times = np.array(T)\n",
    "        steps = len(T)\n",
    "    if ψ0.ndim == 1:\n",
    "        ψ0 = ψ0.reshape(ψ0.size, 1)\n",
    "    N = ψ0.shape[1]\n",
    "    if ξT is None:\n",
    "        if Ug is None:\n",
    "            raise Exception('Must provide either Ug or ξT')\n",
    "        ξT = Ug @ ψ0\n",
    "    \n",
    "    def compute_ψT(x):\n",
    "        for _, ψT in evolve(ψ0, lambda t,ψ: H(t,x,ψ), times, method=method):\n",
    "            pass\n",
    "        return renormalize(ψT)\n",
    "    \n",
    "    def renormalize(ψ):\n",
    "        # Remove errors in norm due to numerical dissipation\n",
    "        return ψ / np.sqrt(np.sum(ψ.conj() * ψ, 0).real)\n",
    "\n",
    "    def compute_Fn_dFndx(x):\n",
    "        Hx = lambda t, ψ: H(t, x, ψ)\n",
    "        ξ = evolve(ξT, Hx, times[-1::-1], method=method)\n",
    "        ψ = evolve(ψ0, Hx, times, method=method)\n",
    "        f = np.zeros((len(times), len(x), N), dtype=np.complex128)\n",
    "        for (i, ((t, ψt), (t, ξt))) in enumerate(zip(ψ, reversed(list(ξ)))):\n",
    "            ξt = renormalize(ξt)\n",
    "            ψt = renormalize(ψt)\n",
    "            f[i,:,:] = np.sum(ξt.conj() * dH(t, x, ψt), 1)\n",
    "        Fn = np.sum(ξT.conj() * ψt, 0)\n",
    "        dFndx = -1j * scipy.integrate.simps(f, x=times, axis=0)\n",
    "        return Fn, dFndx\n",
    "    \n",
    "    global iterations\n",
    "    iterations = (0,0,0)\n",
    "\n",
    "    def maybe_display(F,x,gradF=None):\n",
    "        global iterations\n",
    "        if dispiter and iterations[0] % dispiter == 0:\n",
    "            dF = F - iterations[1] if iterations[0] else 0\n",
    "            dx = np.linalg.norm(x - iterations[2]) if iterations[0] else 0\n",
    "            gradF = np.linalg(gradF) if gradF is not None else 0\n",
    "            print(f'evaluation # {iterations[0]:6} - error={F:+10.4e} - variation={dF:+10.4e} - |dx|={dx:10.4e} - |grad|={gradF:10.4e}')\n",
    "        iterations = (iterations[0]+1, F, x)\n",
    "        return F\n",
    "\n",
    "    def fidelity(x):\n",
    "        F = np.vdot(ξT, compute_ψT(x)).real/N\n",
    "        return maybe_display(1.0-F,x)\n",
    "\n",
    "    def fidelity_and_gradient(x):\n",
    "        Fn, dFndx = compute_Fn_dFndx(x)\n",
    "        F = np.sum(Fn).real/N\n",
    "        dFdx = np.sum(dFndx, axis=-1).real/N\n",
    "        return maybe_display(1.0-F,x), -dFdx\n",
    "    \n",
    "    def overlap(x):\n",
    "        Fn = np.sum(ξT.conj() * compute_ψT(x), 0)\n",
    "        return maybe_display(1.0-np.vdot(Fn, Fn).real/N,x)\n",
    "\n",
    "    def overlap_and_gradient(x, verbose=False):\n",
    "        Fn, dFndx = compute_Fn_dFndx(x)\n",
    "        O = np.vdot(Fn, Fn).real/N\n",
    "        dOdx = 2*np.sum(Fn.conj() * dFndx, axis=-1).real/N\n",
    "        return maybe_display(1.0-O,x), -dOdx\n",
    "    \n",
    "    if target == 'fidelity':\n",
    "        fn = fidelity\n",
    "        dfn = fidelity_and_gradient\n",
    "    else:\n",
    "        fn = overlap\n",
    "        dfn = overlap_and_gradient\n",
    "    if dH is None:\n",
    "        r = scipy.optimize.minimize(fn, x0, tol=tol, method=optimizer, **kwdargs)\n",
    "    else:\n",
    "        if check_gradient:\n",
    "            _, grad = dfn(x0)\n",
    "            grad2 = scipy.optimize.approx_fprime(x0, fn, 1e-6)\n",
    "            err = np.max(np.abs(grad2 - grad))\n",
    "            print(f'max gradient error:   {err}')\n",
    "            print(f'Finite diff gradient: {grad2}')\n",
    "            print(f'Our estimate:         {grad}')\n",
    "        r = scipy.optimize.minimize(dfn, x0, tol=tol, jac=True, method=optimizer, **kwdargs)\n",
    "    if debug:\n",
    "        return r, fn\n",
    "    else:\n",
    "        return r"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now provide some examples of application of this algorithm with and without gradients."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### a) Qubit flip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a simple control. We have a Hamiltonian\n",
    "$$H = \\Omega \\sigma_y$$\n",
    "and we want to implement a spin flip\n",
    "$$U = \\exp(-i (\\pi/2) \\sigma_y) = -i \\sigma_y = \\left(\\begin{matrix} 0 & -1 \\\\ 1 & 0 \\end{matrix}\\right).$$\n",
    "The solution to achieve this in time $T$ is to use $\\Omega = \\pi / 2T.$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def test_qubit_flip():\n",
    "    # Operators that we use in the Hamiltonian\n",
    "    σz = np.diag([1., -1.])\n",
    "    σy = np.array([[0., -1.j],[1.j, 0.]])\n",
    "    # Desired gate\n",
    "    Ug = -1j * σy\n",
    "    # Desired time to execute the gate\n",
    "    T = 1.0\n",
    "    # Some initial guess of the control\n",
    "    x0 = [1.0]\n",
    "    # We want to control all states\n",
    "    ψ0 = np.eye(2)\n",
    "    # A function that implements the control\n",
    "    H = lambda t, x, ψ: x * (σy @ ψ)\n",
    "    # Execute\n",
    "    r = parametric_control([1.0], H, ψ0, T=1.0, Ug=Ug, method='eig')\n",
    "    # Plot the ouptut\n",
    "    sz = np.array([[t, obs[0,0], obs[0,1]]\n",
    "                   for t, obs in evolve(ψ0, r.x * σy, np.linspace(0, T, 31),\n",
    "                                        observables=[σz], method='eig')])\n",
    "    fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(12,4))\n",
    "    ax1.plot(sz[:,0], sz[:,1], '-', label='$\\\\uparrow$')\n",
    "    ax1.plot(sz[:,0], sz[:,2], '--', label='$\\\\downarrow$')\n",
    "    ax2.plot(sz[:,0], sz[:,0]*0+r.x, '--', label='control')\n",
    "    ax1.legend()\n",
    "    ax1.set_xlabel('$t$')\n",
    "    ax1.set_ylabel('$\\\\langle\\\\sigma_z\\\\rangle$')\n",
    "    ax2.set_xlabel('$t$')\n",
    "    ax2.set_ylabel('$x(t)$')\n",
    "\n",
    "test_qubit_flip()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To illustrate the role of derivatives, we introduce an extra degree of difficulty, by incorporating a time-dependent control with a constant amplitude `x`, which we optimize\n",
    "$$H(t,x) = x \\cos(t) \\sigma_y.$$\n",
    "In this example, we use the estimate of the derivative of $H$ to speed up the optimization\n",
    "$$\\frac{\\partial H}{\\partial x} = \\cos(t)\\sigma_y.$$\n",
    "This derivative is created in the function `dH(t,x,ψ)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def test_qubit_flip():\n",
    "    # Operators that we use in the Hamiltonian\n",
    "    π = np.pi\n",
    "    σz = np.diag([1., -1.])\n",
    "    σy = np.array([[0., -1.j],[1.j, 0.]])\n",
    "    # Desired gate\n",
    "    Ug = -1j * σy\n",
    "    # Desired time to execute the gate\n",
    "    T = 1.0\n",
    "    # Some initial guess of the control\n",
    "    x0 = [1.0]\n",
    "    # We want to control all states\n",
    "    ψ0 = np.eye(2)\n",
    "    # A function that implements the control\n",
    "    H = lambda t, x, ψ: x[0] * np.cos(t) * (σy @ ψ)\n",
    "    dH = lambda t, x, ψ: np.array([np.cos(t) * (σy @ ψ)])\n",
    "    # Execute\n",
    "    r = parametric_control([1.0], H, ψ0, T=1.0, Ug=Ug, dH=dH, target='overlap', check_gradient=True, method='eig')\n",
    "    xopt = r.x\n",
    "    # Plot the output\n",
    "    sz = np.array([[t, obs[0,0], obs[0,1]]\n",
    "                   for t, obs in evolve(ψ0, lambda t, ψ: H(t, xopt, ψ), np.linspace(0, T, 31),\n",
    "                                        observables=[σz], method='eig')])\n",
    "    fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(12,4))\n",
    "    t = sz[:,0]\n",
    "    ax1.plot(t, sz[:,1], '-', label='$\\\\uparrow$')\n",
    "    ax1.plot(t, sz[:,2], '--', label='$\\\\downarrow$')\n",
    "    ax2.plot(t, xopt*np.cos(t), '--', label='control')\n",
    "    ax1.legend()\n",
    "    ax1.set_xlabel('$t$')\n",
    "    ax1.set_ylabel('$\\\\langle\\\\sigma_z\\\\rangle$')\n",
    "    ax2.set_xlabel('$t$')\n",
    "    ax2.set_ylabel('$x(t)$')\n",
    "\n",
    "test_qubit_flip()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## b) Qubit drive"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to excite a qubit that has a fixed gap $\\Delta=1$ using a model of the form\n",
    "$$H(t) = \\frac{\\Delta}{2}\\sigma^z + \\epsilon \\cos(x t) \\sigma^y.$$\n",
    "If we apply an interaction picture,\n",
    "$$\\psi(t) = e^{-i\\Delta \\sigma^z t /2} \\xi(t) = U_0(t)\\xi(t),$$\n",
    "the new state evolves according to the Schrödinger equation\n",
    "$$i \\partial_t \\xi = U_0(t)^\\dagger H(t) U_0(t) \\xi(t).$$\n",
    "In other words\n",
    "$$i \\partial_t \\xi = \\frac{\\epsilon}{2} (e^{i x t} + e^{-i x t})(-i\\sigma^+ e^{i\\Delta t} + i\\sigma^- e^{-i\\Delta t}).$$\n",
    "If we apply a rotating wave approximation whereby $x\\simeq \\Delta$\n",
    "$$i \\partial_t \\xi \\simeq \\frac{\\epsilon}{2} (-i\\sigma^+  + i\\sigma^-) = \\frac{\\epsilon}{2} \\sigma^y,$$\n",
    "which achieves a total rotation in a time $T=\\pi/\\epsilon.$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "\n",
    "def test_qubit_drive():\n",
    "    # Operators that we use in the Hamiltonian\n",
    "    π = np.pi\n",
    "    σz = np.diag([1., -1.])\n",
    "    σx = np.array([[0, 1.0],[1.0, 0.]])\n",
    "    σy = np.array([[0., -1.j],[1.j, 0.]])\n",
    "    # Desired time to execute the gate\n",
    "    ϵ = 0.15\n",
    "    T = π/ϵ\n",
    "    steps = 300\n",
    "    times = np.linspace(0, T, steps)\n",
    "    # Desired gate\n",
    "    Ug = -1j * σy @ scipy.linalg.expm(-1j * T * σz)\n",
    "    # Some initial guess of the control\n",
    "    x0 = [1.13]\n",
    "    # We want to control all states\n",
    "    ψ0 = np.eye(2)\n",
    "    # A function that implements the control\n",
    "    H = lambda t, x, ψ: 0.5*(σz @ ψ) + ϵ * math.cos(x[0] * t) * (σy @ ψ)\n",
    "    dH = lambda t, x, ψ: np.array([-ϵ * t * math.sin(x[0] * t) * (σy @ ψ)])\n",
    "    # Execute\n",
    "    r, cost = parametric_control(x0, H, ψ0, T=times, Ug=Ug, dH=dH, debug=True, method='eig', check_gradient=True)\n",
    "    xopt = r.x\n",
    "    # Plot the output\n",
    "    sz = np.array([[obs[0,0], obs[0,1]]\n",
    "                   for _, obs in evolve(ψ0, lambda t, ψ: H(t, xopt, ψ), times, observables=[σz], method='eig')])\n",
    "    fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(12,4))\n",
    "    ax1.plot(times, sz[:,0], '-', label='$\\\\uparrow$')\n",
    "    ax1.plot(times, sz[:,1], '--', label='$\\\\downarrow$')\n",
    "    ax1.legend()\n",
    "    ax1.set_xlabel('$t$')\n",
    "    ax1.set_ylabel('$\\\\langle\\\\sigma_z\\\\rangle$')\n",
    "    xrange = np.linspace(0, 2, 41)\n",
    "    costs = np.array([cost([x]) for x in xrange])\n",
    "    ax2.plot(xrange, costs)\n",
    "    ax2.plot([xopt, xopt], [np.min(costs), np.max(costs)], 'r-.')\n",
    "    ax2.plot(xrange, r.fun + 0*xrange, 'r-.')\n",
    "    ax2.set_xlabel('$x/\\\\pi$')\n",
    "    ax2.set_ylabel('cost')\n",
    "\n",
    "test_qubit_drive()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following code is only required for the automated testing of the library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# file: seeq/test/test_parametric_control.py\n",
    "\n",
    "from seeq.control import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# file: seeq/test/test_parametric_control.py\n",
    "\n",
    "import unittest\n",
    "\n",
    "class TestQControl(unittest.TestCase):\n",
    "    π = np.pi\n",
    "    σz = np.array([[1., 0.],[0., -1.]])\n",
    "    σx = np.array([[0., 1.],[1., 0.]])\n",
    "    σy = np.array([[0., -1.j],[1.j, 0.]])\n",
    "    ψ0 = np.eye(2)\n",
    "\n",
    "    def test_nothing(self):\n",
    "        \"\"\"For a qubit to remain the same, we do nothing.\"\"\"\n",
    "        Ug = np.eye(2)\n",
    "        H = lambda t, x, ψ: x * (self.σx @ ψ)\n",
    "        r = parametric_control([1.0], H, self.ψ0, T=1.0, Ug=Ug, tol=1e-8, method='expm')\n",
    "        self.assertEqual(len(r.x), 1)\n",
    "        self.assertAlmostEqual(r.x[0], 0.0, delta=1e-7)\n",
    "\n",
    "    def test_nothing2(self):\n",
    "        \"\"\"For a qubit to remain the same, we cancel the frequency.\"\"\"\n",
    "        Ug = np.eye(2)\n",
    "        H = lambda t, x, ψ: x[0] * (self.σx @ ψ) + (1.0 - x[1]) * (self.σz @ ψ)\n",
    "        r = parametric_control([1.0, 0.1], H, self.ψ0, T=1.0, Ug=Ug, tol=1e-8, method='expm')\n",
    "        self.assertEqual(len(r.x), 2)\n",
    "        self.assertAlmostEqual(r.x[0], 0.0, delta=1e-7)\n",
    "        self.assertAlmostEqual(r.x[1], 1.0, delta=1e-7)\n",
    "\n",
    "    def test_qubit_flip(self):\n",
    "        \"\"\"Construct a π/2 pulse.\"\"\"\n",
    "        Ug = -1j*self.σy\n",
    "        H = lambda t, x, ψ: (x * self.σy) @ ψ\n",
    "        r = parametric_control([1.0], H, self.ψ0, T=1.0, Ug=Ug, tol=1e-9, method='expm')\n",
    "        self.assertEqual(len(r.x), 1)\n",
    "        self.assertAlmostEqual(r.x[0], self.π/2., delta=1e-7)\n",
    "\n",
    "    def test_nothing_derivative(self):\n",
    "        \"\"\"For a qubit to remain the same, we do nothing (with gradients).\"\"\"\n",
    "        Ug = np.eye(2)\n",
    "        H = lambda t, x, ψ: x * (self.σx @ ψ)\n",
    "        dH = lambda t, x, ψ: [self.σx @ ψ]\n",
    "        r = parametric_control([1.0], H, self.ψ0, T=1.0, Ug=Ug, dH=dH, tol=1e-8, method='expm')\n",
    "        self.assertEqual(len(r.x), 1)\n",
    "        self.assertAlmostEqual(r.x[0], 0.0, delta=1e-7)\n",
    "\n",
    "    def test_qubit_flip_derivative(self):\n",
    "        \"\"\"Construct a π/2 pulse (with gradients).\"\"\"\n",
    "        Ug = -1j*self.σy\n",
    "        H = lambda t, x, ψ: (x * self.σy) @ ψ\n",
    "        dH = lambda t, x, ψ: [self.σy @ ψ]\n",
    "        r = parametric_control([1.0], H, self.ψ0, T=1.0, Ug=Ug, dH=dH, tol=1e-9, method='expm')\n",
    "        self.assertEqual(len(r.x), 1)\n",
    "        self.assertAlmostEqual(r.x[0], self.π/2., delta=1e-7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "suite1 = unittest.TestLoader().loadTestsFromNames(['__main__.TestQControl'])\n",
    "unittest.TextTestRunner(verbosity=2).run(suite1);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}