{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The SWAP test and loading `pytket` circuits\n",
    "\n",
    "**Download this notebook - {nb-download}`swap_test.ipynb`**\n",
    "\n",
    "In this example we will use the SWAP test to demonstrate the construction of a simple Guppy program and show how pytket circuits with multiple registers can be loaded as Guppy functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from guppylang import guppy\n",
    "from guppylang.std.quantum import qubit, h, cx, toffoli, measure, discard_array\n",
    "from guppylang.emulator import EmulatorResult\n",
    "from guppylang.std.builtins import result, array\n",
    "\n",
    "\n",
    "from pytket import Circuit\n",
    "from pytket.circuit import StatePreparationBox\n",
    "from pytket.passes import DecomposeBoxes\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Background"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The swap test is a simple quantum circuit that allows us to approximate the squared overlap between two quantum states. \n",
    "\n",
    "\n",
    "![](./img/swap_test_circuit.png){w=240px align=center}\n",
    "\n",
    "\n",
    "We apply controlled swap operation to swap the qubits in the $|\\psi\\rangle$ and $|\\phi\\rangle$ registers. This swaps the two target qubits if the control qubit is in the $|1\\rangle$ state and acts as the identity otherwise. \n",
    "\n",
    "\n",
    "Notice how only the first qubit is measured, giving $|0\\rangle$ or $|1\\rangle$. The value of our squared inner product is related to the probability of measuring $|0\\rangle$ and $|1\\rangle$ as follows.\n",
    "\n",
    "$$\n",
    "P_0 = \\frac{1}{2} + \\frac{1}{2}|\\langle\\psi|\\phi\\rangle|^2\\,, \\quad \n",
    "P_1 = \\frac{1}{2} - \\frac{1}{2}|\\langle\\psi|\\phi\\rangle|^2\n",
    "$$\n",
    "\n",
    "To approximate $|\\langle\\psi|\\phi\\rangle|^2$ within an error $\\epsilon$ we will need to run the program for $\\mathcal{O}(\\frac{1}{\\epsilon^2})$ shots. For more background on the swap test, see this [Wikipedia article](https://en.wikipedia.org/wiki/Swap_test).\n",
    "\n",
    "\n",
    "In this example we will consider the following three qubit states\n",
    "\n",
    "$$\n",
    "|W\\rangle = \\frac{1}{\\sqrt{3}} \\big(|001\\rangle + |010\\rangle + |100\\rangle \\big)\n",
    "$$\n",
    "\n",
    "$$\n",
    "|S\\rangle = \\frac{1}{\\sqrt{7}}\\big(|001\\rangle+|010\\rangle+|011\\rangle+|100\\rangle+|101\\rangle+|110\\rangle+|111\\rangle \\big)\\,.\n",
    "$$\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Controlled SWAP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Firstly lets define a Guppy function to implement our CSWAP gate in terms of Toffoli and CX gates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def cswap(control: qubit, q1: qubit, q2: qubit) -> None:\n",
    "    cx(q1, q2)\n",
    "    toffoli(control, q2, q1)\n",
    "    cx(q1, q2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we are doing the swap test for two states with three qubits each, we will need to use three CSWAP gates to swap all of the qubits in both state preparation registers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def cswap_layer(ancilla: qubit, arr0: array[qubit, 3], arr1: array[qubit, 3]) -> None:\n",
    "    for i in range(3):\n",
    "        cswap(ancilla, arr0[i], arr1[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## State Preparation using pytket\n",
    "\n",
    "Now for the state preparation phase. Here we will leverage the [StatePreparationBox](https://docs.quantinuum.com/tket/api-docs/circuit.html#pytket.circuit.StatePreparationBox) construct from pytket.\n",
    "\n",
    "We will create a pytket `Circuit` with three registers. One qubit is used for the ancilla, and three qubits each for the two state preparation registers.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[StatePreparationBox s[0], s[1], s[2]; StatePreparationBox w[0], w[1], w[2]; ]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Define two StatePreparationBox(es) by passing the amplitudes for |S> and |W> as numpy arrays\n",
    "w_state = 1 / np.sqrt(3) * np.array([0, 1, 1, 0, 1, 0, 0, 0])\n",
    "s_state = 1 / np.sqrt(7) * np.array([0] + [1] * 7)\n",
    "\n",
    "w_state_box = StatePreparationBox(w_state)\n",
    "s_state_box = StatePreparationBox(s_state)\n",
    "\n",
    "\n",
    "# Build a pytket Circuit with 7 qubits\n",
    "pytket_circ = Circuit()\n",
    "ancilla = pytket_circ.add_q_register(\"a\", 1)\n",
    "w_qubits = pytket_circ.add_q_register(\"w\", 3)\n",
    "s_qubits = pytket_circ.add_q_register(\"s\", 3)\n",
    "\n",
    "\n",
    "# Append the state preparation subroutines to the empty circuit\n",
    "pytket_circ.add_gate(w_state_box, list(w_qubits))\n",
    "pytket_circ.add_gate(s_state_box, list(s_qubits))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before loading the state preparation circuit into Guppy, we can use the [DecomposeBoxes](https://docs.quantinuum.com/tket/api-docs/passes.html#pytket.passes.DecomposeBoxes) pass to decompose the pytket circuit into CX and Ry gates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DecomposeBoxes().apply(pytket_circ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now load in the pytket circuit and create a corresponding Guppy function which takes the ancilla, $|S\\rangle$ and $|W\\rangle$ registers as inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "pytket_state_prep = guppy.load_pytket(\"pytket_state_prep\", pytket_circ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, the separate quantum registers in the circuit are treated as distinct Guppy arrays. The loaded circuit can now be invoked as a `pytket_state_prep` Guppy function which takes three arrays (corresponding to `\"a\"` `\"s\"` and `\"w\"`) as input. Note that as with pytket the qubit ordering here is lexicographic. So the arrays are arranged in alphabetical order with `\"a\"` first followed by `\"s\"` and finally `\"w\"`.\n",
    "\n",
    "The `load_pytket` function has a `use_arrays` flag which is set to `True` by default. Setting this argument to `False` means that Guppy will not create arrays for the distinct registers but instead will treat each named qubit as a separate function argument to `pytket_state_prep`. In this case we would have seven distinct qubit arguments to `pytket_state_prep` instead of three arrays."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Execution on Selene"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can compose the different parts of our program together into a `main` function. Here we only need to measure the first qubit so we discard the two state preparation registers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "@guppy\n",
    "def main() -> None:\n",
    "    w_qubits = array(qubit() for _ in range(3))\n",
    "    s_qubits = array(qubit() for _ in range(3))\n",
    "    ancilla_reg = array(qubit())\n",
    "    # The pytket function only acts on arrays\n",
    "    pytket_state_prep(ancilla_reg, s_qubits, w_qubits)\n",
    "\n",
    "    (ancilla,) = ancilla_reg\n",
    "\n",
    "    h(ancilla)\n",
    "    cswap_layer(ancilla, w_qubits, s_qubits)\n",
    "    h(ancilla)\n",
    "\n",
    "    result(\"c\", measure(ancilla))\n",
    "\n",
    "    # We are only interested in measuring the first qubit\n",
    "    # Discard all the of |W> and |S> qubits to avoid linearity violation.\n",
    "    discard_array(w_qubits)\n",
    "    discard_array(s_qubits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now compile and execute our program on Selene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_shots = main.emulator(n_qubits=7).with_seed(91919).with_shots(2000).run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's look at the frequencies of the $|0\\rangle$ and $|1\\rangle$ measurement outcomes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counter({'0': 1414, '1': 586})\n"
     ]
    }
   ],
   "source": [
    "print(my_shots.register_counts()[\"c\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, lets get an estimate for the value of $|\\langle W| S \\rangle|^2$ by rearranging our expression for $P_0$\n",
    "\n",
    "$$\n",
    "|\\langle W | S\\rangle|^2 = 2P_0 -1\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_squared_overlap(shots: EmulatorResult) -> float:\n",
    "    counter = shots.register_counts()[\"c\"]\n",
    "    p0 = counter[\"0\"] / (counter[\"0\"] + counter[\"1\"])\n",
    "    return 2 * p0 - 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.4139999999999999"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_squared_overlap(my_shots)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets compare with the classical calculation. We can just compute the inner product between the $|S\\rangle$ and $|W\\rangle$ statevectors with numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.42857142857142866\n"
     ]
    }
   ],
   "source": [
    "print(abs(np.vdot(w_state, s_state)) ** 2)"
   ]
  }
 ],
 "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
}