{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Canonical Quantum Phase estimation\n",
    "\n",
    "**Download this notebook - {nb-download}`canonical-qpe.ipynb`**\n",
    "\n",
    "\n",
    "\n",
    "Based on this [stack exchange post](https://quantumcomputing.stackexchange.com/questions/32594/how-would-you-draw-the-phase-estimation-circuit-for-the-eigenvalues-of-u-mat/32598#32598). See also the old [pytket example notebook on QPE](https://docs.quantinuum.com/tket/user-guide/examples/algorithms_and_protocols/phase_estimation.html).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from guppylang import guppy\n",
    "from guppylang.std.angles import pi\n",
    "from guppylang.std.quantum import qubit, h, crz, x, measure_array, discard_array\n",
    "from guppylang.std.builtins import result, array, mem_swap\n",
    "\n",
    "from pytket.circuit import DiagonalBox, QControlBox\n",
    "from pytket.passes import AutoRebase\n",
    "from pytket import OpType\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Phase estimation is an important quantum algorithm for estimating the eigenvalues of a unitary operator $U$ to some precision. Quantum phase estimation appears as an important subroutine in Shor's algorithm and various fault tolerant approaches to quantum chemistry. In this notebook we will consider the \"canonical\" QPE variant which is implemented by a pure unitary circuit.\n",
    "\n",
    "\n",
    "\n",
    "If $U$ is a unitary matrix, its eigenvalues must lie on the unit circle.\n",
    "\n",
    "\\begin{equation*}\n",
    "U |\\psi \\rangle = e^{2 \\pi i \\theta}|\\psi\\rangle\\,, \\quad \\theta \\in [0, 1) \n",
    "\\end{equation*}\n",
    "\n",
    "Here $|\\psi\\rangle$ is an eigenstate of $U$.\n",
    "\n",
    "We estimate the eigenvalues by approximating $\\theta$ in the equation above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will consider a very simplified version of phase estimation wherein $U$ is a diagonal matrix. This means the true eigenvalues can be read off the diagonal. This will allow us to clearly see that our implementation is correct."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "U =\n",
    "\\begin{pmatrix}\n",
    "1 & 0 & 0 & 0 \\\\\n",
    "0 & 1 & 0 & 0 \\\\\n",
    "0 & 0 & e^{ i \\frac{\\pi}{4}} & 0 \\\\\n",
    "0 & 0 & 0 & e^{ i \\frac{\\pi}{8}}\n",
    "\\end{pmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trivial State Preparation\n",
    "\n",
    "Clearly the matrix $U$ has the eigenvalue $e^{i \\frac{\\pi}{8}}$ corresponding to the eigenstate $|11\\rangle = (0, 0, 0, 1)^T$.\n",
    "\n",
    "We can prepare this trivial eigenstate with two Pauli $X$ gates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def prepare_trivial_eigenstate() -> array[qubit, 2]:\n",
    "    q0, q1 = qubit(), qubit()\n",
    "    x(q0)\n",
    "    x(q1)\n",
    "    return array(q0, q1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Controlled-$U$ Operations\n",
    "\n",
    "Next we need to create a subroutine applying a Controlled-$U$ operation. This will be repeatedly applied and will kick back phase factors of $e^{i \\frac{\\pi}{8}}$ onto the ancilla qubits. We can generate this using a pytket [DiagonalBox](https://docs.quantinuum.com/tket/api-docs/circuit.html#pytket.circuit.DiagonalBox) where we apply an additional control. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_box = DiagonalBox(np.array([1, 1, np.exp(1j * np.pi / 4), np.exp(1j * np.pi / 8)]))\n",
    "controlled_u_op = QControlBox(d_box, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Circuits created using pytket can be loaded into Guppy as functions using the `load_pytket` function (or alternatively by specifying a function stub with the same signature as the circuit and annotating it with `@pytket(circ)`). However, the circuit being loaded can only contain gates that exist in the Guppy `quantum` library, so we need to rebase the generated circuit to remove a `TK1` gate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "rebase = AutoRebase({OpType.CX, OpType.Rz, OpType.H, OpType.CCX})\n",
    "circ = controlled_u_op.get_circuit()\n",
    "rebase.apply(circ)\n",
    "\n",
    "controlled_u = guppy.load_pytket(\"controlled_u_circuit\", circ, use_arrays=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inverse Quantum Fourier Transform\n",
    "\n",
    "The final subroutine we need is the inverse quantum fourier transform (IQFT). This has the effect of inducing destructive interference at the end of our circuit. This means that we are more likely to measure a single basis state (or a small set of basis states).\n",
    "\n",
    "We can define a generalised Guppy function over $n$ qubits. We do this by defining a natural number variable $n$. Our Guppy program for the IQFT is then polymorphic over $n$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = guppy.nat_var(\"n\")\n",
    "\n",
    "@guppy\n",
    "def inverse_qft(qs: array[qubit, n]) -> None:\n",
    "    # Reverse qubit order with swaps\n",
    "    for k in range(n // 2):\n",
    "        mem_swap(qs[k], qs[n - k - 1])\n",
    "\n",
    "    for i in range(n):\n",
    "        h(qs[n - i - 1])\n",
    "        for j in range(n - i - 1):\n",
    "            crz(qs[n - i - 1], qs[n - i - j - 2], -pi / 2 ** (j + 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we can have a IQFT subroutine for any number of qubits that we like by adjusting the size of the input array."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The QPE Program\n",
    "\n",
    "Now that we have defined subroutines for the state preparation, controlled unitaries and IQFT steps, we can combine these into a single function to perform quantum phase estimation.\n",
    "\n",
    "First we define an array of measurement qubits of size $m$. The more measurement qubits we have the more precise our estimate of the phase $\\theta$ will be."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can define a function to implement phase estimation using our `inverse_qft` and `controlled_u` functions. Here we will have $n$ measurement qubits. We fix the size of the initial state to have only two qubits. \n",
    "\n",
    "The QPE construction can  be generalised in a similar manner to the inverse QFT function. A larger value of $n$ will mean that we can estimate the eigenphase of $U$ to greater precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def phase_estimation(measured: array[qubit, n], state: array[qubit, 2]) -> None:\n",
    "    for i in range(n):\n",
    "        h(measured[i])\n",
    "\n",
    "    # Add 2^n - 1 controlled unitaries sequentially\n",
    "    for n_index in range(n):\n",
    "        control_index: int = n - n_index - 1\n",
    "        for _ in range(2**n_index):\n",
    "            controlled_u(measured[control_index], state[0], state[1])\n",
    "\n",
    "    inverse_qft(measured)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Execution on Selene\n",
    "\n",
    "Let's execute this QPE program on the Selene emulator for 500 shots."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define a `main` function which includes our six qubit phase estimation subroutine and measurements. This can then be compiled for execution on the Selene simulator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def main() -> None:\n",
    "    state = prepare_trivial_eigenstate()\n",
    "    measured = array(qubit() for _ in range(4))\n",
    "    phase_estimation(measured, state)\n",
    "\n",
    "    # state qubits are not measured so have to be explicitly discarded\n",
    "    discard_array(state)\n",
    "\n",
    "    # Create a result from the measured array\n",
    "    result(\"c\", measure_array(measured))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_shots = 500\n",
    "sim_result = main.emulator(n_qubits=6).with_seed(5).with_shots(n_shots).run()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have executed our phase estimation instance on the Selene emulator we can analyse our results.\n",
    "\n",
    "Let's look at our measurement outcomes. In this highly idealised QPE instance we expect all of our measurement outcomes to be $|0001\\rangle$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counter({'0001': 500})\n"
     ]
    }
   ],
   "source": [
    "result_counter = sim_result.register_counts()[\"c\"]\n",
    "assert result_counter[\"0001\"] == n_shots\n",
    "print(result_counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that all of our measurements yield the basis state $|0001\\rangle$ which encodes the integer $j=1$ in four bits.\n",
    "\n",
    "$$\n",
    "\\theta = \\frac{j}{2^m}\n",
    "$$\n",
    "\n",
    "Here $n$ is the number of evaluation qubits (4 in our case). The value of $j$ is given by the decimal value of the most frequent measurement outcome.\n",
    "\n",
    "$$\n",
    "\\theta = \\frac{1}{2^4} = \\frac{1}{16}\n",
    "$$"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}