{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculating filter functions of optimized gates\n",
    "In this example we want to optimize a Quantum Fourier Transform gate using QuTiP and compute the filter functions for the optimized gate for different initial pulse amplitudes.\n",
    "\n",
    "For more information on the QuTiP optimization, see the tutorials [here] and, for QFT in particular, [this] notebook.\n",
    "\n",
    "[here]: https://qutip.org/tutorials.html\n",
    "[this]: https://nbviewer.jupyter.org/github/qutip/qutip-notebooks/blob/master/examples/control-pulseoptim-QFT.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import qutip as qt\n",
    "from qutip.control import grape, pulseoptim\n",
    "from qutip.qip import algorithms\n",
    "\n",
    "import filter_functions as ff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use widget for interactive mode\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use GRAPE to optimize the pulse for an IQ toy model with $\\sigma_{x,i}$ and $\\sigma_{y,i}$ as well as an exchange coupling $\\sigma_z\\otimes\\sigma_z$ as controls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# No constant Hamiltonian\n",
    "H_d = qt.tensor(qt.sigmaz(), qt.sigmaz())*0.0\n",
    "\n",
    "H_c = [\n",
    "    qt.tensor(qt.sigmax(), qt.qeye(2)),\n",
    "    qt.tensor(qt.sigmay(), qt.qeye(2)),\n",
    "    qt.tensor(qt.qeye(2), qt.sigmax()),\n",
    "    qt.tensor(qt.qeye(2), qt.sigmay()),\n",
    "    qt.tensor(qt.sigmaz(), qt.sigmaz())\n",
    "]\n",
    "n_ctrls = len(H_c)\n",
    "# start point for the gate evolution\n",
    "U_0 = qt.identity(4)\n",
    "# Target for the gate evolution - Quantum Fourier Transform gate\n",
    "U_tar = algorithms.qft(2)\n",
    "\n",
    "evo_time = 1\n",
    "n_ts = 50\n",
    "\n",
    "# Some optimization parameters\n",
    "fid_err_targ = 1e-10\n",
    "max_iter = 10**5\n",
    "max_wall_time = 10**3\n",
    "#min_grad = 1e-20\n",
    "\n",
    "# Initial pulse amplitudes, 'RND'|'LIN'|'ZERO'|'SINE'|'SQUARE'|'TRIANGLE'|'SAW\n",
    "p_types = ['LIN', 'SINE', 'RND']\n",
    "results = {\n",
    "    p_type: pulseoptim.optimize_pulse_unitary(\n",
    "        H_d, H_c, U_0, U_tar, n_ts, evo_time, fid_err_targ=fid_err_targ,\n",
    "        max_iter=max_iter, max_wall_time=max_wall_time, \n",
    "        out_file_ext=None, alg='GRAPE', init_pulse_type=p_type, gen_stats=True,\n",
    "        phase_option='PSU',  # ignore global phase\n",
    "    ) for p_type in p_types\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now easily set up ``PulseSequence`` instances from the final amplitudes and compute the filter functions of each pulse:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from filter_functions import plotting\n",
    "\n",
    "identifiers = ['XI', 'YI', 'IX', 'IY', 'ZZ']\n",
    "pulses = {\n",
    "    p_type: ff.PulseSequence(\n",
    "        list(zip(H_c, result.final_amps.T, identifiers)),\n",
    "        list(zip(H_c, np.ones((n_ctrls, n_ts)), identifiers)),\n",
    "        [evo_time/n_ts]*n_ts\n",
    "    ) for p_type, result in results.items()\n",
    "}\n",
    "\n",
    "fig, ax = plt.subplots(2, 3, figsize=(12, 6))\n",
    "for i, (p_type, pulse) in enumerate(pulses.items()):\n",
    "    # Plot the pulse train\n",
    "    *_, = plotting.plot_pulse_train(pulse, fig=fig, axes=ax[0, i])\n",
    "    ax[0, i].set_title(p_type)\n",
    "    # Plot the filter functions\n",
    "    omega = ff.util.get_sample_frequencies(pulse, spacing='log')\n",
    "    *_, = plotting.plot_filter_function(pulse, omega, fig=fig, axes=ax[1, i])\n",
    "    \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evidently, the pulses with linear and random initial amplitudes filter out DC noise better than those with sinusoidal initial amplitudes, in particular for $\\sigma_y$ noise on the first qubit. There, the filter function has DCG-like character to an extent."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (Spyder)",
   "language": "python3",
   "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.12.3"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}