{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Grover and quantum search\n",
        "\n",
        "The purpose of this notebook is to briefly demonstrate the programming and execution of a simple Quantum search algorithm in the QLM.\n",
        "\n",
        "## Grover's algorithm\n",
        "\n",
        "Grover's algorithm rely on a two main ingredients:\n",
        "- a diffusion operator $\\mathcal{D} = 2 |s\\rangle\\langle s| - I$, where $|s\\rangle = \\frac{1}{\\sqrt{2^n}}\\sum |i\\rangle$\n",
        "- and an oracle \"marking\" some basis states by flipping their phases: $O_f = \\sum_i (-1)^{f(i)}|i\\rangle\\langle i|$ for some boolean predicate $f$\n",
        "\n",
        "\n",
        "### Diffusion operator\n",
        "\n",
        "Lets first start by programming a piece of pyAQASM code that will generate a diffusion operator.\n",
        "$\\mathcal{D}$ can be obtained by:\n",
        "- performing a basis change from the computation basis to the diagonal basis (i.e a cascade of H gates)\n",
        "- performing a collectiong of bit flips in order to send the $|0^n\\rangle$ state onto $|1^n\\rangle$\n",
        "- performing a \"global\" controlled-Z gate in order to flip the phase of $|0^n\\rangle\n",
        "- undoing the bit flips\n",
        "- undoing the basis change\n",
        "\n",
        "\n",
        "This can be easily expressed in pyAQASM as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from qat.lang.AQASM import QRoutine, Z, H, X, CNOT\n",
        "from qat.lang.AQASM.misc import build_gate\n",
        "\n",
        "@build_gate(\"D\", [int], arity=lambda n: n)\n",
        "def diffusion(nbqbits):\n",
        "    rout = QRoutine()\n",
        "    wires = rout.new_wires(nbqbits)\n",
        "    with rout.compute():\n",
        "        for wire in wires:\n",
        "            H(wire)\n",
        "            X(wire)\n",
        "    Z.ctrl(nbqbits - 1)(wires)\n",
        "    rout.uncompute()\n",
        "    return rout\n",
        "\n",
        "example = diffusion(4)\n",
        "example.display()\n",
        "example.display(depth=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now we have a function that produces diffusion operators of the correct size. Moreover, this function is wrapped in a \"box\" called \"D\" for diffusion (this is mostly for visualization purpose).\n",
        "\n",
        "### Oracle\n",
        "\n",
        "In order to produce a self-contained example, we will assume that we want to run a Grover instance for a particular predicate characterizing clean colorations of a graph $G$.\n",
        "\n",
        "Given some graph $G(V, E)$, a clean coloration is a function $c: V \\rightarrow \\mathcal{C}$ for some finite set $\\mathcal{C}$ such that for all edge $\\{i, j\\}\\in E$, we have $c(i) \\neq c(j)$.\n",
        "\n",
        "#### Representing colorations:\n",
        "First, we need to represent graph coloration using only qubits. To do so, we will consider that $\\mathcal{C}= [0, ...,k -1]$ and use $log(k)$ qubits to represent each color.\n",
        "To keep things simple, we will assume that the number of colors is a power of two: $k = 2^m$. That way, we will use $m$ qubits to store the color of each vertex.\n",
        "\n",
        "This brings up the number of qubits to $n.m$ where $n = |V|$.\n",
        "\n",
        "#### Checking colorations\n",
        "\n",
        "In order to check that a coloration is clean, we need to check that for each edge $\\{i, j\\}$ we have $c(i) \\neq c(j)$.\n",
        "This can be easily checked by:\n",
        "- xoring the two colors $c(i)$ and $c(j)$\n",
        "- checking that the result is not $|0\\rangle$\n",
        "- computing the logical AND of all these conditions\n",
        "\n",
        "Lets first write a small routine that will check a single edge:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "@build_gate(\"CHECK\", [int], arity=lambda m:2 * m + 1)\n",
        "def check_edge(m):\n",
        "    rout = QRoutine()\n",
        "    color1 = rout.new_wires(m)\n",
        "    color2 = rout.new_wires(m)\n",
        "    output = rout.new_wires(1)\n",
        "    with rout.compute():\n",
        "        for wire1, wire2 in zip(color1, color2):\n",
        "            CNOT(wire1, wire2)\n",
        "        for wire in color2:\n",
        "            X(wire)\n",
        "    X.ctrl(m)(color2, output)\n",
        "    rout.uncompute()\n",
        "    X(output)\n",
        "    return rout\n",
        "\n",
        "example = check_edge(3)\n",
        "example.display()\n",
        "example.display(depth=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now that we have a routine checking color for a single edge, we can easily write a routine check the coloration of a full graph.\n",
        "\n",
        "We need to compute a large AND of $|E|$ clauses. This can be done using $|E|$ ancillae and a single, large Toffoli gate with $|E|$ controls.\n",
        "This works, but is quite expensive. In order to slightly optimize the number of qubits, we will use the following trick:\n",
        "- initialize some quantum register to $L = |0\\rangle$ of size $log |E|$\n",
        "- for each edge, if the corresponding clause is verified, increment $L$ by one\n",
        "- finally check that $L = |E|$ and flip the phase of the classical state. If $L \\neq |E|$ , some clause was violated and the coloration is not clean!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import networkx as nx\n",
        "# we need an adder:\n",
        "from qat.lang.AQASM.qftarith import add\n",
        "\n",
        "@build_gate(\"CHECK_GRAPH\", [nx.Graph, int], arity=lambda g, m: g.number_of_nodes() * m)\n",
        "def check_graph(graph, m):\n",
        "    rout = QRoutine()\n",
        "    # Colors array\n",
        "    colors = [rout.new_wires(m) for _ in graph.nodes()]\n",
        "    # Our counter L\n",
        "    size_l = graph.number_of_edges().bit_length()\n",
        "    L = rout.new_wires(size_l)\n",
        "    # a work qubit to store the result of a check_edge\n",
        "    tmp = rout.new_wires(1)\n",
        "    # some routines (check_edge and an adder)\n",
        "    check_routine = check_edge(m)\n",
        "    adder = add(size_l, 1)\n",
        "    with rout.compute():\n",
        "        for a, b in graph.edges():\n",
        "            # checking the edge\n",
        "            with rout.compute():\n",
        "                check_routine(colors[a], colors[b], tmp)\n",
        "            # adding the result in our counter\n",
        "            adder(L, tmp)\n",
        "            # uncomputing 'tmp'\n",
        "            rout.uncompute()\n",
        "    # checking if l = |E|\n",
        "    E = graph.number_of_edges()\n",
        "    with rout.compute():\n",
        "        for i in range(size_l):\n",
        "            if ((E >> i) & 1) == 0:\n",
        "                X(L[i])\n",
        "    Z.ctrl(size_l - 1)(L)\n",
        "    # uncomputing the X's\n",
        "    rout.uncompute()\n",
        "    # uncomputing L\n",
        "    rout.uncompute()\n",
        "    # tmp and L are work qubits and should be flagged to that they can be re-used\n",
        "    rout.set_ancillae(L, tmp)\n",
        "    return rout\n",
        "    \n",
        "graph = nx.generators.path_graph(3)\n",
        "example = check_graph(graph, 2)\n",
        "\n",
        "example.display()\n",
        "example.display(depth=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Ok, we're almost there! We now have our oracle that inverts the phase of a classical state if and only if it describes a clean coloring.\n",
        "We can now write our full Grover algorithm:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from qat.lang.AQASM import Program\n",
        "from qat.lang.AQASM.qint import QInt\n",
        "from qat.qpus import get_default_qpu\n",
        "\n",
        "def graph_coloration_grover(graph, m, nbsteps):\n",
        "    prog = Program()\n",
        "    # Notice that each color is stored in a different register with type \"int\"\n",
        "    # This is so it is easy to display final measurement results in the next cells\n",
        "    data = [prog.qalloc(m, QInt) for _ in range(graph.number_of_nodes())]\n",
        "    # Initializing a superposition of all possible coloring\n",
        "    for color in data:\n",
        "        for qbit in color:\n",
        "            H(qbit)\n",
        "    # Our diffusion routine and our oracle\n",
        "    O = check_graph(graph, m)\n",
        "    D = diffusion(graph.number_of_nodes() * m)\n",
        "    # The main loop\n",
        "    for _ in range(nbsteps):\n",
        "        O(data)\n",
        "        D(data)\n",
        "    # We also return the list of registers carying the colors to simplify data retrieval\n",
        "    return prog.to_circ(), data\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Lets run a Grover algorithm on our path graph with 4 possible colors per nodes (completely overkill) and for 2 iterations.\n",
        "\n",
        "We will iterate over the possible measurement outcomes and compute the success probability of the algorithm:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "color, reg = graph_coloration_grover(graph, 2, 2)\n",
        "print(\"Using\", color.nbqbits, \"qubits\")\n",
        "qpu = get_default_qpu()\n",
        "# Here we specify that we want to measure the color registers\n",
        "# the results will be formatted nicely thanks to this (using the .value field below)\n",
        "result = qpu.submit(color.to_job(qubits=reg))\n",
        "\n",
        "success_proba = 0\n",
        "for sample in result:\n",
        "    c1, c2, c3 = sample.state.value\n",
        "    if c1 != c2 != c3:\n",
        "        success_proba += sample.probability\n",
        "print(\"Overall success probability:\", success_proba)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "It seems to work quite alright. Lets try with a larger graph. \n",
        "We will:\n",
        "- generate a random graph\n",
        "- run a grover to look for clean 4-colorations for several number of iterations\n",
        "- print the success probability for each number of iterations\n",
        "\n",
        "The main reason why we need this for loop is because we can't predict the proportion of correct colorations among the full set of colorations.\n",
        "\n",
        "In practice, the correct way of tackling this issue would be to use a slightly more advanced quantum search algorithm by Brassard et al."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "graph = nx.generators.erdos_renyi_graph(5, 0.5)\n",
        "\n",
        "color, reg = graph_coloration_grover(graph, 2, 0)\n",
        "print(\"Using\", color.nbqbits, \"qubits\")\n",
        "qpu = get_default_qpu()\n",
        "\n",
        "for nbsteps in range(2, 8):\n",
        "    result = qpu.submit(color.to_job(qubits=reg))\n",
        "    color, reg = graph_coloration_grover(graph, 2, nbsteps)\n",
        "    success_proba = 0\n",
        "    for sample in result:\n",
        "        colors = sample.state.value\n",
        "        if all(colors[a] != colors[b] for a, b in graph.edges()):\n",
        "            success_proba += sample.probability\n",
        "    print(nbsteps, success_proba)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Going further: QSearch\n",
        "\n",
        "As mentionned before, the main issue of our current algorithm is that we can't predict the optimal number of iterations required to find a clean coloration with high probability.\n",
        "\n",
        "In 2000, Brassard et al proposed a solution for that particular problem (which is quite recurrent when dealing with quantum search).\n",
        "All the details can be found [here](https://arxiv.org/abs/quant-ph/0005055).\n",
        "\n",
        "The algorithm is a classical loop consisting of many calls to a Grover search with number of steps growing exponentially.\n",
        "\n",
        "For simplicity, lets first define a wrapping function that will:\n",
        "- build a Grover search for a given number of steps\n",
        "- sample a single bitstring out of the final quantum state and return the corresponding coloration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def single_search(graph, m, nbsteps):\n",
        "    grover, registers = graph_coloration_grover(graph, m, nbsteps)\n",
        "    job = grover.to_job(qubits=registers, nbshots=1)\n",
        "    result = qpu.submit(job)\n",
        "    return result[0].state.value    "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can now declare a `q_search` function implementing the classical loop as in the paper. For clarity we will try to use the same notations as the one in the reference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def q_search(graph, m):\n",
        "    l = 1\n",
        "    c = 1.5\n",
        "    M = int(np.ceil(c ** l))\n",
        "    while True:\n",
        "        j = np.random.randint(1, M)\n",
        "        print(\"Trying with\", j, \"steps\")\n",
        "        coloration = single_search(graph, m, j)\n",
        "        if all(coloration[a] != coloration[b] for a, b in graph.edges()):\n",
        "            return coloration\n",
        "        l += 1\n",
        "        M = int(np.ceil(c ** l))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "coloration = q_search(graph, 2)\n",
        "print(\"Found a valid coloration:\", coloration)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Of course all this is just an brief introduction to quantum search algorithms!\n",
        "There are plenty of alternatives and variants that would deserve hundreds of notebooks.\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.9.16"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}