lecture-3.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "58a834a2",
   "metadata": {},
   "source": [
    "# Lecture 3: Programming and Using a Matching Decoder\n",
    "### James R. Wootton, IBM Quantum\n",
    "\n",
    "## Making an error graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df52e05e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qiskit_qec.utils import Decodoku\n",
    "from qiskit_qec.decoders import DecodingGraph\n",
    "from retworkx.visualization import mpl_draw"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92667289",
   "metadata": {},
   "source": [
    "In this lecture we are going to look at how to write a decoder. We'll start with one of the great workhorses of decoding algorithms: minimum weight perfect matching. This has quite broad applicability, but is especially well-suited to things like repetition codes, surface codes and the Decodoku puzzles we've seen so far.\n",
    "\n",
    "So let's make something to decode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1569d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = 9\n",
    "p = 0.02\n",
    "\n",
    "game = Decodoku(d=d, p=p)\n",
    "\n",
    "game.draw_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96884d3c",
   "metadata": {},
   "source": [
    "Our decoding is, of course, based around the decoding graph. So let's rename it something short and snappy: `dg`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47dadae9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dg = game.decoding_graph.graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c80b1f8",
   "metadata": {},
   "source": [
    "Most of the decoding graph is actually blank space. The only parts that really interest us are the highlighted nodes. So let's create a new graph that contains only them. We'll call this the 'error graph'.\n",
    "\n",
    "First, we'll just make a blank `DecodingGraph` object. We'll also give this the short and snappy name of `eg`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43eed19e",
   "metadata": {},
   "outputs": [],
   "source": [
    "error_graph = DecodingGraph(None)\n",
    "eg = error_graph.graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfecf24b",
   "metadata": {},
   "source": [
    "Next we go through all the nodes of `dg` and add only the highlighted ones to `eg`.\n",
    "\n",
    "We'll also note down their positions. We could use these to help us calculate weights for edges, but here we'll actually just use them for plotting the error graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5992aaa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "error_pos = []\n",
    "for n, node in enumerate(dg.nodes()):\n",
    "    if node['highlighted']:\n",
    "        eg.add_node(node)\n",
    "        error_pos.append(game.graph_pos[n])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91e5a2e5",
   "metadata": {},
   "source": [
    "Let's take a look at our error graph so far."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bbb729a",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl_draw(eg, node_color='red', pos=error_pos, with_labels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2115b9d",
   "metadata": {},
   "source": [
    "We have points but no edges. So what edges shall we add? And how to we account for the different distances?\n",
    "\n",
    "For the latter point, we need to know what those distances are. We could determine these from the positions of the nodes, but this is not always so easy when we have more complex decoding graphs. Instead we can define the distance between two nodes as the number of edges requires to connect them. This makes sense from the perspective of decoding, since it corresponds to the number of errors required to create that pair of highlighted nodes.\n",
    "\n",
    "Packages made to do graph theory usually have a way of figuring out these distances, and `retworkx` is no different. But we are going to do it manually using the coordinates. This is done in the following function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd81fe91",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_manhattan_distance(graph, n0, n1):\n",
    "    '''Takes two nodes of a decoding graph and returns the Manhattan distance'''\n",
    "\n",
    "    node0 = graph[n0]\n",
    "    node1 = graph[n1]\n",
    "    \n",
    "    # x distance\n",
    "    dx = abs(node0['x']-node1['x'])\n",
    "    \n",
    "    # y distance (accounting for boundary)\n",
    "    if node0['is_boundary'] or node1['is_boundary']:\n",
    "        dy = 0\n",
    "    else:\n",
    "        dy = abs(node0['y']-node1['y'])\n",
    "    \n",
    "    return dx + dy\n",
    "\n",
    "get_distance = get_manhattan_distance"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fdcc8fd",
   "metadata": {},
   "source": [
    "We don't actually need a fully connected graph. So to experiment with how to keep it simple, we can try putting an upper limit on the distance for which we assign an edge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caa9b519",
   "metadata": {},
   "outputs": [],
   "source": [
    "dist_max = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0a43e37",
   "metadata": {},
   "source": [
    "Now let's put in those edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68c064af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# these lines just repeat what we already did above, to save some effort if we want to rerun the process\n",
    "error_graph = DecodingGraph(None)\n",
    "eg = error_graph.graph\n",
    "error_pos = []\n",
    "for n, node in enumerate(dg.nodes()):\n",
    "    if node['highlighted']:\n",
    "        eg.add_node(node)  \n",
    "        error_pos.append(game.graph_pos[n])\n",
    "\n",
    "dg_nodes = dg.nodes()\n",
    "# go through all pairs of nodes in eg\n",
    "for n0, node0 in enumerate(eg.nodes()):\n",
    "    for n1, node1 in enumerate(eg.nodes()):\n",
    "        if n0<n1:\n",
    "        \n",
    "            # determine which nodes these are in dg\n",
    "            dg_n0 = dg_nodes.index(node0)\n",
    "            dg_n1 = dg_nodes.index(node1)\n",
    "            # get the distance\n",
    "            dist = get_distance(dg, dg_n0, dg_n1)\n",
    "            # and add the edge\n",
    "            if dist<=dist_max:\n",
    "                error_graph.graph.add_edge(n0, n1, {'distance':dist})\n",
    "        \n",
    "mpl_draw(eg, node_color='red', pos=error_pos)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60ef7913",
   "metadata": {},
   "source": [
    "Now we've made a new graph out of just the highlighted nodes. Before we move on, let's note down how many nodes this has."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "557d143a",
   "metadata": {},
   "outputs": [],
   "source": [
    "bulk_node_num = len(eg.nodes())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b997fa9",
   "metadata": {},
   "source": [
    "## Decoding with matching\n",
    "\n",
    "### Why matching?\n",
    "\n",
    "Let's think a bit more about what we know of decoding.\n",
    "\n",
    "Errors create sets of highlighted nodes. Collective properties of these sets of errors also affect certain logical variables that we are interested in. The job of decoding is to figure out how these variables have been affected. Since this cannot be done with certainty, ideally we want to determine *the most likely way* they were affected.\n",
    "\n",
    "There are often many possible sets of errors that could have caused any given set of higlighted nodes. We can categorize these according to the affect they have on the logical variables. All the sets of errors that cause the same logical effect are said to belong to the same equivalence class. Our job is therefore not to determine the most likely specific set of errors to have occurred, but to determine the most likely equivalence class.\n",
    "\n",
    "This can be done, but it also can be difficult. Sometimes we need to make compromises. Rather than doing the best possible decoding, we'll do an approximation. And the most obvious approximation is to assume that the most likely equivalence class is the one with the most likely error.\n",
    "\n",
    "Finding the most likely error is also a pretty complex optimization problem. Many graph-theoretic optimization problems are inefficient to solve. But there are some that do admit an efficient solution: minimum and maximum weight perfect matching.\n",
    "\n",
    "Put simply, a **matching** is a pairing of the nodes of a graph. The nodes that are paired must be connected by an edge, and no node can be part of more than one pair.\n",
    "\n",
    "A **perfect matching** is when the nodes become fully paired up, with every node part of a pair.\n",
    "\n",
    "a **weighted matching** is where we assign a number as a weight to each edge. The total weight of the matching is then the sum of the weights for all the edges that correspond to the pairs.\n",
    "\n",
    "The minimum and maximum weight matchings are then the matchings with the minimum and maximum total weight, given a particular graph. An algorithm to compute one can also compute the other, because we can simply multiply the weights by $-1$ to turn a minimum weight problem to a maximum weight problem.\n",
    "\n",
    "This is the precisely the problem we need to solve in order to figure out the most likely set of errors in certain decoding problems. Specifically, those for which errors create pairs of highlighted edges.\n",
    "\n",
    "Here we will consider the simple case that each error type occurs with equal probability. The most likely set of errors is therefore the one with the least errors. The weight we use is then simply the distance.\n",
    "\n",
    "# Adding to the error graph\n",
    "\n",
    "For codes like the Toric code, or repetition codes implemented on a loop, the periodic boundary conditions mean that *all* errors affect pairs of nodes. This means there are no weird boundaries to deal with. But in the case we are looking at now, a there is indeed a weird boundary. This means that isolated highlighted nodes can indeed exist, for which there is obviously no pairing possible. So how are we going to deal with this? The answer is to add more nodes.\n",
    "\n",
    "For each node we add one additional 'virtual' twin. We associate this with the nearest boundary.\n",
    "\n",
    "To begin, it'll be useful to know the index of the boundary nodes in `dg`. We'll use `dg_nr` for the right boundary and `dg_nl` for the left."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c0d7366",
   "metadata": {},
   "outputs": [],
   "source": [
    "for n, node in enumerate(dg.nodes()):\n",
    "    if node['x']==-1:\n",
    "        dg_nl = n\n",
    "    if node['x']==d-1:\n",
    "        dg_nr = n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "823ffd10",
   "metadata": {},
   "source": [
    "Now we will add in a virtual node for each bulk node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "010bb59b",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_color = ['red']*len(eg.nodes())\n",
    "virtuals = []\n",
    "\n",
    "for n, node in enumerate(eg.nodes()):\n",
    "    \n",
    "    # copy each bulk node and flag up the copy as being on the boundary\n",
    "    vnode = node.copy()\n",
    "    vnode['is_boundary'] = True\n",
    "    \n",
    "    if node['x']<(d-1)/2:\n",
    "        # if the node is on the left, also note this\n",
    "        vnode['boundary'] = 'l'\n",
    "        \n",
    "        # determine the distance to the left boundary\n",
    "        dg_n = dg_nodes.index(node)\n",
    "        dist = get_distance(dg,dg_n, dg_nl)\n",
    "\n",
    "        # and we place these on the left of any plot\n",
    "        error_pos.append((node['x']-(d+1)/2,-node['y']-0.5))\n",
    "        node_color.append('green')\n",
    "    else:\n",
    "        # similarly for the right\n",
    "        vnode['boundary'] = 'r'\n",
    "        dg_n = dg_nodes.index(node)\n",
    "        dist = get_distance(dg,dg_n, dg_nr)\n",
    "        error_pos.append((node['x']+(d+1)/2,-node['y']-0.5))\n",
    "        node_color.append('blue')\n",
    "\n",
    "    # add it to the graph\n",
    "    nv = eg.add_node(vnode)\n",
    "    \n",
    "    # along with an edge, whose distance is of the straight line to the boundary\n",
    "    eg.add_edge(n, nv, {'distance':d})\n",
    "    \n",
    "    # it will be useful to keep track of which virtual nodes are where\n",
    "    virtuals.append(nv)\n",
    "        \n",
    "mpl_draw(eg, node_color=node_color, pos=error_pos)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4734553",
   "metadata": {},
   "source": [
    "Now if a highlighted node is associated with the boundary, the matching can pair it with its virtual boundary counterpart.\n",
    "\n",
    "But what about when the virtual nodes don't need to be used in this way? What will they pair with? To answer this, we connect them up with each other, using edges with weight zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "108b584d",
   "metadata": {},
   "outputs": [],
   "source": [
    "for nv0 in virtuals:\n",
    "    for nv1 in virtuals:\n",
    "        if nv0<nv1:\n",
    "            eg.add_edge(nv0, nv1, {'distance':0})\n",
    "\n",
    "mpl_draw(eg, node_color=node_color, pos=error_pos, with_labels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b09db50",
   "metadata": {},
   "source": [
    "Now finally we can find our favourite implementation of the algorithm that find minimum (or maximum) weight matchings and use it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7da5906c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from retworkx import max_weight_matching\n",
    "\n",
    "def weight_fn(edge):\n",
    "    return -int(edge[\"distance\"])\n",
    "\n",
    "matching = max_weight_matching(eg, max_cardinality=True, weight_fn=weight_fn)\n",
    "\n",
    "print(matching)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87a622c3",
   "metadata": {},
   "source": [
    "Here is the chosen matching visualized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfba9939",
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_color = []\n",
    "for edge in eg.edge_list():\n",
    "    \n",
    "    if edge in matching or edge[::-1] in matching:\n",
    "        edge_color.append('red')\n",
    "    else:\n",
    "        edge_color.append((0,0,0,0.1))\n",
    "\n",
    "mpl_draw(eg, node_color=node_color, pos=error_pos, edge_color=edge_color, with_labels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d0eb05a",
   "metadata": {},
   "source": [
    "The decoders we write for Decodoku ask us to separate the nodes into clusters, which we think are connected by the same subset of errors. For the matching decoding, these are the pairs that correspond to the minimum weight matching.\n",
    "\n",
    "We also need to find the boundary correction. This is the parity of the number of times that each boundary (i.e. the virtuals of each boundary) are connected to. This is done by the following (where we use `bc` rather than `boundary_corrections`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad97409d",
   "metadata": {},
   "outputs": [],
   "source": [
    "bc = [0,0]\n",
    "eg_nodes = eg.nodes()\n",
    "for pair in matching:\n",
    "    if pair[0]<bulk_node_num or pair[1]<bulk_node_num:\n",
    "        for ne in pair:\n",
    "            node = eg_nodes[ne]\n",
    "            if node not in dg_nodes:\n",
    "                if node['boundary']=='l':\n",
    "                    bc[0] += 1\n",
    "                else:\n",
    "                    bc[1] += 1\n",
    "bc = [bc[0]%2, bc[1]%2]\n",
    "bc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11e475b7",
   "metadata": {},
   "source": [
    "And there we have it. All that remains is to paste all the above into one big function to serve as our decoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eec770f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process(game):\n",
    "    \n",
    "    d = game.d\n",
    "    dg = game.decoding_graph.graph\n",
    "    \n",
    "    error_graph = DecodingGraph(None)\n",
    "    eg = error_graph.graph\n",
    "    for n, node in enumerate(dg.nodes()):\n",
    "        if node['highlighted']:\n",
    "            eg.add_node(node)  \n",
    "\n",
    "    dg_nodes = dg.nodes()\n",
    "    # go through all pairs of nodes in eg\n",
    "    for n0, node0 in enumerate(eg.nodes()):\n",
    "        for n1, node1 in enumerate(eg.nodes()):\n",
    "            if n0<n1:\n",
    "\n",
    "                # determine which nodes these are in dg\n",
    "                dg_n0 = dg_nodes.index(node0)\n",
    "                dg_n1 = dg_nodes.index(node1)\n",
    "                # get the distance\n",
    "                dist = get_distance(dg, dg_n0, dg_n1)\n",
    "                # and add the edge\n",
    "                error_graph.graph.add_edge(n0, n1, {'distance':dist})\n",
    "\n",
    "    # get number of bulk nodes\n",
    "    bulk_node_num = len(eg.nodes())\n",
    "    \n",
    "    # get indices for boundary nodes\n",
    "    for n, node in enumerate(dg.nodes()):\n",
    "        if node['x']==-1:\n",
    "            dg_nl = n\n",
    "        if node['x']==d-1:\n",
    "            dg_nr = n\n",
    "    \n",
    "    virtuals = []\n",
    "    for n, node in enumerate(eg.nodes()):\n",
    "        # copy each bulk node and flag up the copy as being on the boundary\n",
    "        vnode = node.copy()\n",
    "        vnode['is_boundary'] = True\n",
    "        if node['x']<(d-1)/2:\n",
    "            # if the node is on the left, also note this\n",
    "            vnode['boundary'] = 'l'\n",
    "            # determine the distance to the left boundary\n",
    "            dg_n = dg_nodes.index(node)\n",
    "            dist = get_distance(dg, dg_n, dg_nl)\n",
    "            # and we place these on the left of any plot\n",
    "            error_pos.append((node['x']-(d+1)/2,-node['y']-2))\n",
    "            node_color.append('green')\n",
    "        else:\n",
    "            # similarly for the right\n",
    "            vnode['boundary'] = 'r'\n",
    "            dg_n = dg_nodes.index(node)\n",
    "            dist = get_distance(dg,dg_n, dg_nr)\n",
    "            error_pos.append((node['x']+(d+1)/2,-node['y']-2))\n",
    "            node_color.append('blue')\n",
    "        # add it to the graph\n",
    "        nv = eg.add_node(vnode)\n",
    "        # along with an edge, whose distance is of the straight line to the boundary\n",
    "        eg.add_edge(n, nv, {'distance':dist})\n",
    "        # it will be useful to keep track of which virtual nodes are where\n",
    "        virtuals.append(nv)\n",
    "    \n",
    "    for nv0 in virtuals:\n",
    "        for nv1 in virtuals:\n",
    "            if nv0<nv1:\n",
    "                eg.add_edge(nv0, nv1, {'distance':0})\n",
    "        \n",
    "    def weight_fn(edge):\n",
    "        return -int(edge['distance'])\n",
    "\n",
    "    matching = max_weight_matching(eg, max_cardinality=True, weight_fn=weight_fn) \n",
    "        \n",
    "    bc = [0,0]\n",
    "    clusters = {}\n",
    "    c = 0\n",
    "    eg_nodes = eg.nodes()\n",
    "    for pair in matching:\n",
    "        if pair[0]<bulk_node_num or pair[1]<bulk_node_num:\n",
    "            for ne in pair:\n",
    "                node = eg_nodes[ne]\n",
    "                if node in dg_nodes:\n",
    "                    nd = dg_nodes.index(eg_nodes[ne])\n",
    "                    clusters[nd] = c\n",
    "                else:\n",
    "                    if node['boundary']=='l':\n",
    "                        bc[0] += 1\n",
    "                    else:\n",
    "                        bc[1] += 1\n",
    "            c += 1\n",
    "    bc = [bc[0]%2, bc[1]%2]\n",
    "    \n",
    "    return bc, clusters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46c4834d",
   "metadata": {},
   "source": [
    "Now let's see how well it works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b283710",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = 0.05\n",
    "game = Decodoku(d=d, p=p, process=process)\n",
    "\n",
    "game.draw_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e349aa6",
   "metadata": {},
   "outputs": [],
   "source": [
    "game.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9a07d6e",
   "metadata": {},
   "source": [
    "We can also try it on a few sets of errors we made earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "081f20d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_errors = [\n",
    "    [(0,0,1,0), (d-1,0,d,0)],\n",
    "    [((d-1)/2,d,(d-1)/2+1,d)],\n",
    "    [(0,0,1,0), (d-1,0,d,0), ((d-1)/2,d,(d-1)/2+1,d)]   \n",
    "]\n",
    "\n",
    "game = Decodoku(d=d, process=process, errors=test_errors[1])\n",
    "\n",
    "game.draw_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cc405dc",
   "metadata": {},
   "source": [
    "Here's a more complex error, which we'll put on a slightly smaller graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db05a775",
   "metadata": {},
   "outputs": [],
   "source": [
    "errors = [(2,3,3,4),(5,4,7,5)]\n",
    "\n",
    "game = Decodoku(d=8, process=process, errors=errors)\n",
    "\n",
    "game.draw_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c96aa5fc",
   "metadata": {},
   "source": [
    "There are multiple sets of five errors that could have caused this. Two of these are shown in the image below. \n",
    "\n",
    "<img src=\"errors.png\" alt=\"Drawing\" style=\"width: 400px;\"/>\n",
    "\n",
    "Note that they belong to two different equivalence classes. The red set corresponds to the correct class (i.e. it corresponds to the `errors` list I used to set this puzzle up). But the decoder seems to have gone for the purple one.\n",
    "\n",
    "This is not mistake by the decoder. According to the information it has, both equivalence classes are equal, so it just arbitrarily picks one. However, we could have coded it better. Let's use $n$ to denote the total number of edges, and $p$ to represent the probability of an error on any edge. The probability of any set of 5 errors is then $p^5 (1-p)^{n-5}$.\n",
    "\n",
    "If we restrict to just these minimal error sets, the probability of the single set of errors that puts us in the purple class is\n",
    "\n",
    "$P_{\\rm purple} = p^5 (1-p)^{n-5}.$\n",
    "\n",
    "But note that there are multiple equal length paths that could take the place of the ones highlighted in red: two for the pair of nodes on the left and four for the right. So\n",
    "\n",
    "$P_{\\rm red} = 6 \\times p^5 (1-p)^{n-5}.$\n",
    "\n",
    "This shows us that the red equivalence class is actually more likely. So how could we have told our decoder about this?\n",
    "\n",
    "First let's rewrite things in a way that makes more sense for the edges of the error graph. Specifically, we want to see which terms correspond to the different pairs of nodes. In the red case, we can think of there being two independent sets of errors: one with 2 errors and 2 possibilities, the other with 3 errors and 3 possibilities. So\n",
    "\n",
    "$P_{\\rm purple} = \\left( \\frac{p}{1-p} \\right)^2 \\times \\left( \\frac{p}{1-p} \\right)^2 \\times \\left( \\frac{p}{1-p} \\right)^1 \\times (1-p)^n.$\n",
    "\n",
    "$P_{\\rm red} = 2 \\times \\left( \\frac{p}{1-p} \\right)^2 \\times 3 \\times \\left( \\frac{p}{1-p} \\right)^3 \\times (1-p)^n.$\n",
    "\n",
    "Note that the factor of $(1-p)^n$ is common to both. It contains no information about any possible set of errors. When comparing probabilities, we can therefore just ignore it.\n",
    "\n",
    "Next we will make it so that the contributions from each pair of nodes add together rather than multiply together. This is done by taking the logarithm of each side. To make things simpler, we'll replace the constantly occuring $p/(1-p)$ with\n",
    "\n",
    "$w = -\\log{\\frac{p}{1-p}}.$\n",
    "\n",
    "The minus sign is used here such that we can thing of $w$ as a positive valued 'weight'. The logarithms are then\n",
    "\n",
    "$\\log{P_{\\rm purple}} = 2 w \\, + \\, 2 w \\, + \\, 1 w \\,\\, + \\,\\, {\\rm const},$\n",
    "\n",
    "$\\log{P_{\\rm red}} = \\left(\\log{2} + 2 w \\right) + \\left(\\log{4} + 3 w \\right) \\,\\, + \\,\\,  {\\rm const}.$\n",
    "\n",
    "From this we can extract better choices for the distances in the error graph. Rather than just a simple distance, we could use a sum of edge weights that depend on probabilities, with an additional entropic correction. The comparision of matching weights then corresponds exactly to a comparison of total error probabilities.\n",
    "\n",
    "If we did this, the entropic corrections would ensure that the decoder opts for the red outcome instead, as it should. So let's update the distance function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc367eeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "def get_distance(graph, n0, n1):\n",
    "    \n",
    "    node0 = graph[n0]\n",
    "    node1 = graph[n1]\n",
    "    \n",
    "    # x distance\n",
    "    dx = abs(node0['x']-node1['x'])\n",
    "    \n",
    "    # y distance (accounting for boundary)\n",
    "    if node0['is_boundary'] or node1['is_boundary']:\n",
    "        dy = 0\n",
    "    else:\n",
    "        dy = abs(node0['y']-node1['y'])   \n",
    "    \n",
    "    # number of minimum distance paths\n",
    "    # this is dx+dy choose dx\n",
    "    num = math.comb(dx+dy, dx)\n",
    "    \n",
    "    # the final weight combines these (and needs to be an integer, and so needs a scaling factor)\n",
    "    weight = int(100*(- math.log( p/(1-p) )*(dx+dy) - math.log(num)))\n",
    "        \n",
    "    return weight"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "589ef1ee",
   "metadata": {},
   "source": [
    "Let's see what the matching does now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b103f929",
   "metadata": {},
   "outputs": [],
   "source": [
    "game = Decodoku(d=8, process=process, errors=errors)\n",
    "\n",
    "game.draw_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b321d691",
   "metadata": {},
   "source": [
    "Now it has chosen the most entropically likely error.\n",
    "\n",
    "Let's see what the threshold plot looks like!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fadee351",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "def get_logical_probs(ds,ps,process):\n",
    "    logical_probs = {}\n",
    "    for d in ds:\n",
    "        logical_probs[d] = []\n",
    "        for p in ps:\n",
    "            error_num = 0\n",
    "            sample_num = 0\n",
    "            game = Decodoku(d=d, p=p, process=process)\n",
    "            for _ in range(10**3):\n",
    "                bc, clusters = game.process(game)\n",
    "                error_num += bc!=game.boundary_errors\n",
    "                sample_num += 1\n",
    "                game.restart()\n",
    "            logical_probs[d].append(error_num/sample_num)\n",
    "    return logical_probs\n",
    "\n",
    "def threshold_plot(logical_probs):\n",
    "    plt.figure(figsize=(10,7))\n",
    "    for d in ds:\n",
    "        plt.scatter(ps,logical_probs[d],label='d='+str(d))\n",
    "    plt.xlabel('p')\n",
    "    plt.ylabel('logical prob')\n",
    "    plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b650c961",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = [7, 9, 11]\n",
    "ps = [0.01, 0.05, 0.1, 0.15, 0.2]\n",
    "\n",
    "logical_probs = get_logical_probs(ds,ps,process)\n",
    "\n",
    "threshold_plot(logical_probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91c74153",
   "metadata": {},
   "source": [
    "As a quickly run set of points for relatively small codes, it is hardly the best threshold plot ever made. Nevertheless, we can see that bigger codes allow for better correction below a threshold of around $p=0.11$. Which is exactly what we expect from the classic paper of [Dennis, Landahl, Kitaev and Preskill](https://arxiv.org/abs/quant-ph/0110143)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adfff9f1",
   "metadata": {},
   "source": [
    "## Getting weights from real data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60db37fd",
   "metadata": {},
   "source": [
    "So far we've looked at cases where every type of error occurs with equal probability $p$. But this won't be the case in real life. Instead, there will be a unique $p_j$ for each and every edge in the decoding graph. And hence a unique edge weight $w_j$.\n",
    "\n",
    "Let's see this by taking a look at real life! We are going to run some QEC $^*$ on a quantum computer $^\\dagger$!\n",
    "\n",
    "$^*$ If you count repetition codes as QEC\n",
    "\n",
    "$^\\dagger$ not a very big one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8f5ac76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import stuff\n",
    "from qiskit import IBMQ\n",
    "import qiskit.tools.jupyter\n",
    "\n",
    "# load you account credential stuff\n",
    "IBMQ.load_account()\n",
    "\n",
    "# pick one of your providers\n",
    "provider = IBMQ.providers()[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7286a8fa",
   "metadata": {},
   "source": [
    "Set up the backend object for the device you want to use. Here's the one for a publically available 7 qubit device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1c81b08",
   "metadata": {},
   "outputs": [],
   "source": [
    "backend = provider.get_backend('ibm_nairobi')\n",
    "\n",
    "backend"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb115739",
   "metadata": {},
   "source": [
    "Now we are going to set up some repetition code circuits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c947b94",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qiskit_qec.circuits import RepetitionCodeCircuit\n",
    "from qiskit_qec.decoders import DecodingGraph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cea7be4b",
   "metadata": {},
   "source": [
    "We have enough qubits for a `d=3` code (requiring a line of five qubits). We'll do one for the phase flip encoding this time. And since measurement and reset are long gates, we'll skip the resets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f301189c",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = 3\n",
    "T = 2\n",
    "code = RepetitionCodeCircuit(d, T, xbasis=True, resets=False, barriers=True)\n",
    "\n",
    "code.circuit['0'].draw(output='mpl')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eeda7c1",
   "metadata": {},
   "source": [
    "When running the circuit on the device, we'll want to choose which qubits play which roles. Let's go for the line `0-1-3-5-6`. That means the code qubits are `0`, `3` and `6`, and the link qubits are `1` and `5`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07347e2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "code_qubits = [0,3,6]\n",
    "link_qubits = [1,5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "113905d6",
   "metadata": {},
   "source": [
    "This information can provided to the transpiler to guide it when rewriting the circuit for the backend. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dfc284c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qiskit import transpile\n",
    "\n",
    "qc = transpile(code.circuit['0'], backend, initial_layout=link_qubits+code_qubits, scheduling_method='alap')\n",
    "\n",
    "qc.draw(output='mpl')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ef2a804",
   "metadata": {},
   "source": [
    "We also used the transpiler to schedule the circuit. So we can now also see a scheduled version of the circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b8dc12e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qiskit.visualization.timeline import draw\n",
    "\n",
    "draw(qc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "303a7279",
   "metadata": {},
   "source": [
    "QEC is a powerful tool, but it's also good to fix errors before they get to QEC. We can do this by adding in dynamical decoupling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcc8c40f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qiskit.circuit.library import XGate\n",
    "from qiskit.transpiler import PassManager, InstructionDurations\n",
    "from qiskit.transpiler.passes import DynamicalDecoupling\n",
    "\n",
    "dd_sequence = [XGate()]*2\n",
    "durations = InstructionDurations().from_backend(backend)\n",
    "\n",
    "pm = PassManager([DynamicalDecoupling(durations, dd_sequence, qubits=code_qubits)])\n",
    "qc = pm.run(qc)\n",
    "\n",
    "draw(qc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0d4be5c",
   "metadata": {},
   "source": [
    "Uncomment the following to run the circuit and get the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6349730c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#job = backend.run(qc)\n",
    "#counts = job.result().get_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d268e43e",
   "metadata": {},
   "source": [
    "Otherwise, here are some results I made earlier, with a code centered around qubit `2` of the 27 qubit `ibm_cairo`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4378e79",
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = {'000 00 00': 61253, '000 00 01': 430, '001 00 00': 10781, '001 00 01': 1634, '001 00 10': 23, '001 00 11': 6, '001 01 00': 8598, '001 01 01': 1362, '001 01 10': 14, '001 01 11': 1, '001 10 00': 117, '001 10 01': 17, '001 10 10': 200, '001 10 11': 31, '001 11 00': 103, '001 11 01': 22, '001 11 10': 145, '001 11 11': 26, '000 00 10': 140, '010 00 00': 389, '010 00 01': 89, '010 00 10': 37, '010 00 11': 836, '010 01 00': 203, '010 01 01': 59, '010 01 10': 99, '010 01 11': 79, '010 10 00': 115, '010 10 01': 221, '010 10 10': 28, '010 10 11': 75, '010 11 00': 596, '010 11 01': 64, '010 11 10': 50, '010 11 11': 28, '000 00 11': 4, '011 00 00': 95, '011 00 01': 24, '011 00 10': 54, '011 00 11': 101, '011 01 00': 102, '011 01 01': 29, '011 01 10': 46, '011 01 11': 113, '011 10 00': 135, '011 10 01': 48, '011 10 10': 13, '011 10 11': 6, '011 11 00': 163, '011 11 01': 55, '011 11 10': 23, '011 11 11': 19, '000 01 00': 2581, '100 00 00': 726, '100 00 01': 1, '100 00 10': 903, '100 00 11': 9, '100 01 00': 35, '100 01 01': 22, '100 01 10': 40, '100 01 11': 31, '100 10 00': 1088, '100 10 01': 10, '100 10 10': 209, '100 10 11': 5, '100 11 00': 34, '100 11 01': 31, '100 11 10': 9, '100 11 11': 12, '000 01 01': 2100, '101 00 00': 142, '101 00 01': 16, '101 00 10': 165, '101 00 11': 26, '101 01 00': 111, '101 01 01': 16, '101 01 10': 148, '101 01 11': 18, '101 10 00': 202, '101 10 01': 14, '101 10 10': 43, '101 10 11': 7, '101 11 00': 203, '101 11 01': 23, '101 11 10': 39, '101 11 11': 6, '000 01 10': 5, '110 00 00': 12, '110 00 01': 18, '110 00 10': 4, '110 00 11': 20, '110 01 00': 10, '110 01 01': 4, '110 01 10': 7, '110 01 11': 3, '110 10 00': 8, '110 10 01': 4, '110 10 10': 4, '110 10 11': 20, '110 11 00': 7, '110 11 01': 1, '110 11 10': 12, '110 11 11': 2, '000 01 11': 9, '111 00 00': 2, '111 00 01': 1, '111 00 10': 3, '111 00 11': 2, '111 01 00': 6, '111 01 01': 8, '111 01 10': 2, '111 01 11': 1, '111 10 00': 2, '111 10 01': 3, '111 10 10': 1, '111 10 11': 4, '111 11 00': 3, '111 11 10': 3, '111 11 11': 1, '000 10 00': 607, '000 10 01': 13, '000 10 10': 1118, '000 10 11': 14, '000 11 00': 30, '000 11 01': 18, '000 11 10': 41, '000 11 11': 41}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e325df38",
   "metadata": {},
   "source": [
    "Qiskit-QEC is able to analyze the syndrome information here to determine all the probabilities, $p_j$, for each edge. We use the algorithm of [Spitz, Tarasinski, Beenakker and O'Brien](https://doi.org/10.1002/qute.201800012). See [my recent paper](https://arxiv.org/abs/2207.00553) for more information.\n",
    "\n",
    "Basically, once you have the decoding graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c255797",
   "metadata": {},
   "outputs": [],
   "source": [
    "decoding_graph = DecodingGraph(code)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d10f1ac0",
   "metadata": {},
   "source": [
    "you supply the `counts` dictionary to the `get_error_probs` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1721457e",
   "metadata": {},
   "outputs": [],
   "source": [
    "error_probs = decoding_graph.get_error_probs(counts)\n",
    "error_probs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be4655c3",
   "metadata": {},
   "source": [
    "With these we can make weights that are very much tailored to the hardware on which our code is being run."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.15 ('decodoku')",
   "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.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "3631a5c11ea8272419473841f6c7a371f0144b233b78c9dd4753c6019b6bfd67"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}