{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Barren Plateau Analysis\n",
    "This notebook focuses on the analysis of the Barren Plateau problem. \n",
    "The code is heavily based on [Tensorflow's Quantum tutorial](https://www.tensorflow.org/quantum/tutorials/barren_plateaus)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# add the parent directory to the path so that we can import the module\n",
    "import sys\n",
    "import os\n",
    "\n",
    "sys.path.append(os.path.dirname(os.path.dirname(os.getcwd())))\n",
    "import qandle\n",
    "\n",
    "qandle.__reimport()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install missing libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "skip-execution"
    ]
   },
   "outputs": [],
   "source": [
    "!pip install pandas seaborn matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for simulation\n",
    "import qandle \n",
    "import torch\n",
    "import qw_map\n",
    "\n",
    "# for plotting\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# helper libraries\n",
    "import random\n",
    "import dataclasses\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the problem\n",
    "We use a fixed seed for reproducibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEPTH = 50\n",
    "NUM_REPEATS = 10\n",
    "MAX_QUBITS = 5\n",
    "\n",
    "qandle.config.DRAW_SHIFT_LEFT = True\n",
    "qandle.config.DRAW_SHOW_VALUES = True\n",
    "\n",
    "random.seed(42)\n",
    "torch.random.manual_seed(42)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate random circuits\n",
    "Our circuits will first have a fixed layer, rotating each qubit by $R_y\\left(\\frac{\\pi}{4}\\right)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pi8 = torch.tensor(torch.pi / 8)\n",
    "# the matrix representation of the RY(pi/4) gate\n",
    "ry_pi_4 = torch.tensor(\n",
    "    [[torch.cos(pi8), -torch.sin(pi8)], [torch.sin(pi8), torch.cos(pi8)]]\n",
    ")\n",
    "\n",
    "\n",
    "def generate_random_circuit(num_qubits, remapping=None):\n",
    "    layers = []\n",
    "    # initialize the circuit with RY(pi/4) gates on all qubits\n",
    "    for i in range(num_qubits):\n",
    "        layers.append(\n",
    "            qandle.CustomGate(\n",
    "                i, matrix=ry_pi_4, num_qubits=num_qubits, self_description=\"RY(pi/4)\"\n",
    "            )\n",
    "        )\n",
    "    for _ in range(DEPTH):\n",
    "        for qubit in range(num_qubits):\n",
    "            # randomly choose a gate to apply\n",
    "            gate = random.choice([qandle.RX, qandle.RY, qandle.RZ])\n",
    "            layers.append(\n",
    "                gate(\n",
    "                    qubit=qubit,\n",
    "                    remapping=remapping,\n",
    "                )\n",
    "            )\n",
    "        # add CNOT gates, with even qubits as controls and odd qubits as targets\n",
    "        for control in range(0, num_qubits, 2):\n",
    "            target = control + 1\n",
    "            if target < num_qubits:\n",
    "                layers.append(qandle.CNOT(control=control, target=target))\n",
    "        # add CNOT gates, with odd qubits as controls and even qubits as targets\n",
    "        for control in range(1, num_qubits, 2):\n",
    "            target = control + 1\n",
    "            if target < num_qubits:\n",
    "                layers.append(qandle.CNOT(control=control, target=target))\n",
    "\n",
    "    # automatically split the circuit into smaller circuits if it is too large\n",
    "    return qandle.Circuit(layers, split_max_qubits=6)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test a circuit\n",
    "We generate a random circuit with 5 qubits and print it. We also execute it by calling `example_circuit()` and print the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_circuit = generate_random_circuit(5)\n",
    "print(example_circuit.draw())\n",
    "\n",
    "print(example_circuit())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Measure the gradients\n",
    "We create `NUM_REPEATS` random circuits, and measure their gradients after a forward and backward pass. The mean and standard deviation of the gradients are saved in a pandas DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclasses.dataclass\n",
    "class GradData:\n",
    "    \"\"\"Helper class for neatly storing gradient data before converting to a DataFrame.\"\"\"\n",
    "\n",
    "    num_qubits: int\n",
    "    remapping: str\n",
    "    grad_std: float\n",
    "    grad_mean: float\n",
    "\n",
    "\n",
    "circ_grads = []\n",
    "combinations = [[r, q] for r in [None, qw_map.tanh] for q in range(2, MAX_QUBITS + 1)]\n",
    "\n",
    "for remapping, num_qubits in combinations:\n",
    "    remap_str = remapping.__name__ if remapping is not None else \"None\"\n",
    "    for _ in range(NUM_REPEATS):\n",
    "        c = generate_random_circuit(num_qubits, remapping=remapping)\n",
    "        # forward pass and backward pass using a dummy loss\n",
    "        c().sum().abs().backward()\n",
    "        c_grads = [p.grad.item() for p in c.parameters() if p.grad is not None]\n",
    "        c_grads = torch.tensor(c_grads)\n",
    "        # standard deviation of the gradients\n",
    "        grad_std = c_grads.std().item()\n",
    "        # we use the absolute value of the mean gradient\n",
    "        grad_mean = c_grads.abs().mean().item()\n",
    "        circ_grads.append(GradData(num_qubits, remap_str, grad_std, grad_mean))\n",
    "\n",
    "df = pd.DataFrame([dataclasses.asdict(gd) for gd in circ_grads])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the results\n",
    "Remapping increases both mean absolute gradient and standard deviation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "pltargs = dict(data=df, x=\"num_qubits\", hue=\"remapping\")\n",
    "sns.lineplot(**pltargs, ax=ax, y=\"grad_std\", linestyle=\"-\")\n",
    "sns.lineplot(**pltargs, ax=ax.twinx(), y=\"grad_mean\", linestyle=\":\", legend=False)\n",
    "# set major x ticks to integer values\n",
    "ax.xaxis.set_major_locator(plt.MaxNLocator(integer=True))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qc",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}