notebooks/4. Phase Circuit Optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Phase Circuit Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Changing working directory from `./notebooks/` to `./`, in order to import the Python packages defined in the repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "source": [
    "import os\n",
    "DEBUG = True\n",
    "try:\n",
    "    print(\"Original working directory: %s\"%str(original_wd)) # type: ignore\n",
    "    \"\"\"\n",
    "        You only get here if you---by mistake or otherwise---are re-running this cell, \n",
    "        in which case the working should not be changed again.\n",
    "    \"\"\"\n",
    "except NameError:\n",
    "    original_wd = os.getcwd()\n",
    "    os.chdir('../')\n",
    "print(\"Current working directory: %s\"%str(os.getcwd()))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "General purpose imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import set_matplotlib_formats\n",
    "%matplotlib inline\n",
    "set_matplotlib_formats('svg')\n",
    "import numpy as np\n",
    "# import qiskit"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Some utility code to display HTML elements and images/figures side-by-side:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "source": [
    "from io import BytesIO\n",
    "from typing import Union\n",
    "from IPython.display import Image, HTML # type: ignore\n",
    "def figure_to_image(fig: plt.Figure, fmt: str = \"png\") -> Image:\n",
    "    \"\"\" Converts a Matplotlib figure to a PNG IPython Image. \"\"\"\n",
    "    buffer = BytesIO()\n",
    "    fig.savefig(buffer, format=fmt)\n",
    "    buffer.seek(0)\n",
    "    image_png = buffer.getvalue()\n",
    "    buffer.close()\n",
    "    return Image(image_png, format=fmt, embed=True)\n",
    "def side_by_side(*elements: Union[str, Image]) -> HTML:\n",
    "    \"\"\"\n",
    "        Returns an HTML Div element with the given elements\n",
    "        displayed side by side. Accepts raw HTML code or\n",
    "        IPython Image objects.\n",
    "    \"\"\"\n",
    "    html = f\"<div style='display:flex; align-items: center;'>\"\n",
    "    for el in elements:\n",
    "        if isinstance(el, str):\n",
    "            html += el\n",
    "        elif isinstance(el, Image):\n",
    "            html += f\"<img src='data:image/png;base64,{el._repr_png_()}'/>\"\n",
    "    html += f\"</div>\"\n",
    "    return HTML(html)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Phase Circuit Optimization\n",
    "\n",
    "The optimized version of a given circuit of mixed ZX phase gadgets takes the form of a **Phase block**, consisting of mixed ZX phase gadgets, conjugated by a **CX block**, consisting of CX gates. If the original circuit is repeated a given number `circuit_rep` of times, then so is the phase block part of the optimized circuit:\n",
    "\n",
    "<img src=\"phase_circuit_opt.svg\" width=60%/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Optimization is done through the `OptimizedPhaseCircuit` class. An optimizer is instantiated from:\n",
    "\n",
    "- the original phase circuit to be optimized;\n",
    "- a topology, constraining the CX circuit used for the optimization;\n",
    "- a number of layers to use for the CX block;\n",
    "- the number `circuit_rep` of times that the original circuit is to be repeated (default: 1);\n",
    "- optionally, a fixed seed `rng_seed` for the random number generator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "source": [
    "from pauliopt.phase import OptimizedPhaseCircuit"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "For our running example, we use a 3x3 grid qubit topology:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "source": [
    "from pauliopt.topologies import Topology\n",
    "topology = Topology.grid(3, 3)\n",
    "topology.draw(figsize=(3, 3))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We construct a small random phase circuit on the same qubits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "source": [
    "from pauliopt.phase import PhaseCircuit, CXCircuit\n",
    "orig_circuit = PhaseCircuit.random(topology.num_qubits, 6, rng_seed=0, parametric=\"t\")\n",
    "orig_circuit"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "To optimize it, we instantiate an optimizer with a single-layer CX circuit.\n",
    "The optional parameter `rng_seed` can be used to pass a seed to the RNG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "source": [
    "opt = OptimizedPhaseCircuit(orig_circuit, topology, cx_block=1, rng_seed=0)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The topology and number of qubits are made available through suitably named properties of the optimizer. The original circuit is not saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "source": [
    "print(f\"{opt.num_qubits = }\", end=\"\\n\\n\")\n",
    "print(f\"{opt.topology = }\", end=\"\\n\\n\")"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Phase and CX Blocks\n",
    "\n",
    "Circuit optimization is progressive. The optimizer makes the current phase block and CX block available (in readonly view) through suitably named properties `phase_block` and `cx_block`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "At the beginning, the phase block is equal to the original circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "source": [
    "opt.phase_block.to_svg(scale=0.8)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "At the beginning, the CX block is empty (given number of layers, no CX gates)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "source": [
    "opt.cx_block.draw(figsize=(3, 3))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Random CX Flips\n",
    "\n",
    "Optimization is performed by simulated annealing. Random CX gates are flipped in the CX block for a given number of iterations, in an attempt to reduce the complexity of the overall optimized circuit (CX block adjoint + phase block repeated `circuit_rep` times + CX block). The method `OptimizedPhaseCircuit.random_flip_cx()` is used to perform such a random flip: it modifies the internal phase circuit and CX circuit in place, returning information about the layer index and CX gate that was just flipped in case it has to be undone.\n",
    "\n",
    "The method `OptimizedPhaseCircuit.flip_cx(layer_idx, ctrl, trgt)` is used to perform a specific CX gate flip, while the method `OptimizedPhaseCircuit.is_cx_flippable(layer_idx, ctrl, trgt)` can be used to check if a specific flip can be performed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Here we see an example with a single random flip:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "source": [
    "opt = OptimizedPhaseCircuit(orig_circuit, topology, cx_block=1, rng_seed=0)\n",
    "print(\"=== Phase block and CX block before random flip ===\")\n",
    "display(opt.phase_block.to_svg(scale=0.6))\n",
    "opt.cx_block.draw(figsize=(3, 3))\n",
    "layer_idx, gate = opt.random_flip_cx()\n",
    "print(f\"\\n\\n=== Phase block and CX block after flipping gate {gate} in layer {layer_idx} ===\")\n",
    "display(opt.phase_block.to_svg(scale=0.6))\n",
    "opt.cx_block.draw(figsize=(3, 3))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Here we see another example, at the beginning and after 8 random flips:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "source": [
    "opt = OptimizedPhaseCircuit(orig_circuit, topology, cx_block=1, rng_seed=0)\n",
    "print(\"=== Phase block and CX block before random flip ===\")\n",
    "display(opt.phase_block.to_svg(scale=0.6))\n",
    "opt.cx_block.draw(figsize=(3, 3))\n",
    "for _ in range(3):\n",
    "    opt.random_flip_cx()\n",
    "print(f\"\\n\\n=== Phase block and CX block after 3 random flips ===\")\n",
    "display(opt.phase_block.to_svg(scale=0.6))\n",
    "opt.cx_block.draw(figsize=(3, 3))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Annealing - Prelude\n",
    "\n",
    "Randomly flipping CX gates and hoping for the best is not much good as an optimisation strategy: simulated annealing is used instead as a global optimisation technique.\n",
    "\n",
    "The method `OptimizedPhaseCircuit.anneal(num_iters, schedule)` allows a given number of iterations of simulated annealing to be performed on the circuit, for a given temperature schedule (default: linear temperature schedule with initial temperature 1.0 and final temperature 0.1).\n",
    "The method also accepts an optional keyword argument `loggers`, for detailed logging."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Temperature Schedule\n",
    "\n",
    "A **temperature schedule** is any function which fits the `TempSchedule` protocol below: given the index `it` of the current iteration an the total number `num_iters` of iterations, it returns the temperature (as a `float`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "source": [
    "from typing import Protocol\n",
    "class TempSchedule(Protocol):\n",
    "    \"\"\"\n",
    "        Protocol for a temperature schedule.\n",
    "        The temperature is a number computed from the iteration number `it`\n",
    "        (starting from 0) and the total number of iterations `num_iter`.\n",
    "    \"\"\"\n",
    "\n",
    "    def __call__(self, it: int, num_iters: int) -> float:\n",
    "        ..."
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The library offers 4 utility standard temperature schedules: linear, geometric, reciprocal and logarithmic, each given initial and final temperatures. The schedules are taken from [this paper](https://link.springer.com/article/10.1007/BF00143921)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "source": [
    "from pauliopt.utils import (linear_temp_schedule, geometric_temp_schedule, \n",
    "                            reciprocal_temp_schedule, log_temp_schedule)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Here is a plot of the four temperature schedules over 1000 iterations, for initial temperature 1 and final temperature 0.1. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "source": [
    "num_iters = 1000\n",
    "plt.figure(figsize=(10,5))\n",
    "for temp_schedule_maker in [linear_temp_schedule, geometric_temp_schedule, \n",
    "                            reciprocal_temp_schedule, log_temp_schedule]:\n",
    "    temp_schedule = temp_schedule_maker(t_init=1, t_final=1e-1)\n",
    "    temp = [temp_schedule(it, num_iters) for it in range(num_iters)]\n",
    "    plt.plot(range(num_iters), temp, label=temp_schedule_maker.__name__)\n",
    "plt.legend()\n",
    "plt.show()"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Energy Function\n",
    "\n",
    "The energy function used for annealing is the number of nearest-neighbour CX gates used to implement the optimized circuit.\n",
    "For the CX blocks, this is simply the number of CX gates, because all gates are already nearest-neighbour by construction.\n",
    "For the phase blocks, the count is preformed by adding up the CX count for a topologically-aware implementation of phase gadgets based on minimum spanning trees (MST).\n",
    "The CX counts for the original circuit and optimized circuit (both with the desired number `circuit_rep` of repetitions) are exposed by the properties `OptimizedPhaseCircuit.init_cx_count` (for the original circuit) and `OptimizedPhaseCircuit.cx_count` (for the current circuit).\n",
    "\n",
    "The gadgets themselves have a `PhaseGadget.cx_count(topology)` method to compute the gadget cost for this implementation, and a `PhaseGadget.print_impl_info(topology)` method which prints detailed information about the MST and the CX count for individual branches."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "To understand how the cost of a single phase gadget is computed, we consider the following example on a 3x3 grid topology."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "source": [
    "test_circ = PhaseCircuit.random(topology.num_qubits, 1, rng_seed=0, parametric=\"t\")\n",
    "test_circ"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The phase gadget above spans qubits `{0, 1, 5, 7, 8}`. One possible minimal spanning tree implementation involves two mirror ladders of the following CX gates (with a phase gate on qubit 7 in between):\n",
    "\n",
    "- 1 CX gate between 0 and 1 (dist: 1)\n",
    "- 3 CX gate between 1 and 5 (dist: 2)\n",
    "- 1 CX gate between 5 and 8 (dist: 1)\n",
    "- 1 CX gate between 8 and 7 (dist: 1)\n",
    "\n",
    "Each ladder has 6 CX gates, so the total CX count for the phase gadget is 12."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "source": [
    "test_gadget = test_circ.gadgets[0]\n",
    "test_gadget.print_impl_info(topology)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Let's now consider a second example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "source": [
    "test_circ = PhaseCircuit.random(topology.num_qubits, 1, rng_seed=1, parametric=\"t\")\n",
    "test_circ"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "The phase gadget above spans qubits `{0, 1, 2, 4, 5, 6, 8}`. One possible minimal spanning tree implementation involves two mirror ladders of the following CX gates (with a phase gate on qubit 7 in between):\n",
    "\n",
    "- 1 CX gate between 0 and 1 (dist: 1)\n",
    "- 1 CX gate between 1 and 2 (dist: 1)\n",
    "- 1 CX gate between 1 and 4 (dist: 1)\n",
    "- 1 CX gate between 2 and 5 (dist: 1)\n",
    "- 1 CX gate between 5 and 8 (dist: 1)\n",
    "- 3 CX gates between 0 and 6 (dist: 2)\n",
    "\n",
    "Each ladder has 8 CX gates, so the total CX count for the phase gadget is 16."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "source": [
    "test_gadget = test_circ.gadgets[0]\n",
    "test_gadget.print_impl_info(topology)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Now let's look at an example of cost reduction in a circuit with several phase gadgets.\n",
    "Consider the example from before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": false
   },
   "source": [
    "opt = OptimizedPhaseCircuit(orig_circuit, topology, 1, rng_seed=1)\n",
    "print(f\"CX count before flip: {opt.init_cx_count}\")\n",
    "print(f\"  - Gadget CX counts: {[g.cx_count(topology) for g in opt.phase_block.gadgets]}\")\n",
    "print(f\"  - CX count for the CX blocks: {2*opt.cx_block.num_gates}\")\n",
    "display(opt.phase_block.to_svg(scale=0.8))\n",
    "opt.cx_block.draw(figsize=(2, 2))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "If we conjugate the circuit by a CX gate with control 0 and target 1, the cost changes as follows:\n",
    "\n",
    "- we remove 4 CX gates from the phase block (3rd and 5th phase gadgets lose a leg);\n",
    "- we add 2 CX gates to the CX blocks;\n",
    "\n",
    "Overall, the cost is reduced by 2. It would be reduced by `4*circuit_rep-2` in the general case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "source": [
    "opt.flip_cx(0, 1, 0)\n",
    "print(f\"CX count after flip: {opt.cx_count}\")\n",
    "print(f\"  - Gadget CX counts: {[g.cx_count(topology) for g in opt.phase_block.gadgets]}\")\n",
    "print(f\"  - CX count for the CX blocks: {2*opt.cx_block.num_gates}\")\n",
    "display(opt.phase_block.to_svg(scale=0.8))\n",
    "opt.cx_block.draw(figsize=(2, 2))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "If we further conjugate the circuit by a CX gate with control 2 and target 5, the cost changes as follows:\n",
    "\n",
    "- we remove 4 CX gates from the phase block (3rd and 5th phase gadget lose a leg);\n",
    "- we add 2 CX gates to the phase block (4th phase gadget gains a leg);\n",
    "- we add 2 CX gates to the CX blocks;\n",
    "\n",
    "Overall, the cost is reduced by 0. It would be reduced by `2*circuit_rep-2` in the general case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "source": [
    "opt.flip_cx(0, 2, 5)\n",
    "print(f\"CX count after second flip: {opt.cx_count}\")\n",
    "print(f\"  - Gadget CX counts: {[g.cx_count(topology) for g in opt.phase_block.gadgets]}\")\n",
    "print(f\"  - CX count for the CX blocks: {2*opt.cx_block.num_gates}\")\n",
    "display(opt.phase_block.to_svg(scale=0.8))\n",
    "opt.cx_block.draw(figsize=(2, 2))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Annealing\n",
    "\n",
    "The method `OptimizedPhaseCircuit.anneal(num_iters, schedule)` allows a given number of iterations of simulated annealing to be performed on the circuit, for a given temperature schedule.\n",
    "\n",
    "The method also accepts an optional keyword argument `loggers`, for detailed logging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "source": [
    "from math import log10, ceil\n",
    "def log_iter(it, prev_cost, new_cost, accepted, flip, t, num_iters):\n",
    "    if new_cost < prev_cost:\n",
    "        print(f\"Iter #{it:>0{ceil(log10(num_iters-1))}}, new cost: {new_cost}\")\n",
    "loggers = {\n",
    "    \"log_start\": lambda cost, num_iters: print(f\"Init cost: {cost}\"),\n",
    "    \"log_iter\": log_iter,\n",
    "    \"log_end\": lambda cost, num_iters: print(f\"Final cost: {cost}\"),\n",
    "}"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "As our first example, we anneal for 100 iterations on a geometric temperature schedule, from an initial temperature of 1 to a final temperature of 1e-5.\n",
    "The circuit consists of 6 phase gadgets on a 3x3 qubit grid.\n",
    "The cost function is based on the circuit not being repeated (`circuit_rep=1`).\n",
    "We use a single-layer CX block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "source": [
    "topology = Topology.grid(3, 3)\n",
    "circuit = PhaseCircuit.random(topology.num_qubits, 6, rng_seed=0, parametric=\"t\")\n",
    "opt = OptimizedPhaseCircuit(circuit, topology, cx_block=1, circuit_rep=1, rng_seed=1)\n",
    "opt.anneal(100, schedule=(\"geometric\", 1, 1e-5), loggers=loggers)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "After annealing, the cost has been reduced by 16%: the improvements in gadget CX counts are somewhat offset by the CX count of the CX blocks. In the limit of a large number of repetitions of the original circuit, the cost would be reduced by up to 28%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "source": [
    "print(f\"CX count before annealing: {opt.init_cx_count}\")\n",
    "print(f\"  - Gadget CX counts: {[g.cx_count(topology) for g in circuit.gadgets]}\")\n",
    "print()\n",
    "change = (opt.cx_count-opt.init_cx_count)/opt.init_cx_count\n",
    "max_change = (opt.cx_count-opt.init_cx_count-2*opt.cx_block.num_gates)/opt.init_cx_count\n",
    "print(f\"Cost after annealing: {opt.cx_count} \"\n",
    "      f\"({opt.circuit_rep} reps: {change:.1%}, many reps limit: {max_change:.1%})\")\n",
    "print(f\"  - Gadget CX counts: {[g.cx_count(topology) for g in opt.phase_block.gadgets]}\")\n",
    "print(f\"  - CX blocks CX count: {2*opt.cx_block.num_gates}\")\n",
    "display(opt.phase_block.to_svg(scale=0.8))\n",
    "opt.cx_block.draw(figsize=(3, 3))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "As our second example, we anneal for 1000 iterations on the same schedule. Circuit and number of repetitions as before. However, this time we use a 4-layer CX block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "source": [
    "topology = Topology.grid(3, 3)\n",
    "circuit = PhaseCircuit.random(topology.num_qubits, 6, rng_seed=0, parametric=\"t\")\n",
    "opt = OptimizedPhaseCircuit(circuit, topology, cx_block=4, circuit_rep=1, rng_seed=1)\n",
    "opt.anneal(1000, schedule=(\"geometric\", 1, 1e-5), loggers=loggers)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "After annealing, the cost has been reduced by 36%. In the limit of a large number of repetitions of the original circuit, the cost would be reduced by up to 60%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "source": [
    "print(f\"CX count before annealing: {opt.init_cx_count}\")\n",
    "print(f\"  - Gadget CX counts: {[g.cx_count(topology) for g in circuit.gadgets]}\")\n",
    "display(circuit.to_svg(scale=0.8))\n",
    "change = (opt.cx_count-opt.init_cx_count)/opt.init_cx_count\n",
    "max_change = (opt.cx_count-opt.init_cx_count-2*opt.cx_block.num_gates)/opt.init_cx_count\n",
    "print(f\"Cost after annealing: {opt.cx_count} \"\n",
    "      f\"({opt.circuit_rep} reps: {change:.1%}, many reps limit: {max_change:.1%})\")\n",
    "print(f\"  - Gadget CX counts: {[g.cx_count(topology) for g in opt.phase_block.gadgets]}\")\n",
    "print(f\"  - CX blocks CX count: {2*opt.cx_block.num_gates}\")\n",
    "display(opt.to_svg(scale=0.8))\n",
    "opt.cx_block.draw(figsize=(3, 3))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "As an example, we look at the effect on 3 repetitions of the circuit, where annealing achieves a 48% reduction in CX count:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "source": [
    "circuit_3 = circuit >> circuit >> circuit\n",
    "opt_3 = OptimizedPhaseCircuit(circuit, topology, cx_block=3, circuit_rep=3, rng_seed=4)\n",
    "print(f\"CX count before annealing: {opt_3.init_cx_count}\")\n",
    "display(circuit_3.to_svg(scale=0.5))\n",
    "opt_3.anneal(1000, schedule=(\"geometric\", 1, 1e-5))\n",
    "change_3 = (opt_3.cx_count-opt_3.init_cx_count)/opt_3.init_cx_count\n",
    "print(f\"CX count after annealing: {opt_3.cx_count} ({opt_3.circuit_rep} reps: {change_3:.1%})\")\n",
    "display(opt_3.to_svg(hscale=0.42, vscale=0.5))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "As our final example, we anneal for 1000 iterations on a linear temperature schedule, from an initial temperature of 1 to a final temperature of 1e-5.\n",
    "The circuit consists of 50 phase gadgets on a 5x5 qubit grid, with a number of legs between 1 and 3.\n",
    "The cost function is based on the circuit being repeated 5 times (`circuit_rep=5`).\n",
    "We use a 3-layer CX block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "source": [
    "topology = Topology.grid(5, 5)\n",
    "circuit = PhaseCircuit.random(topology.num_qubits, 50, min_legs=1, max_legs=3, rng_seed=0, parametric=\"t\")\n",
    "opt = OptimizedPhaseCircuit(circuit, topology, cx_block=3, circuit_rep=5, rng_seed=1)\n",
    "opt.anneal(1000, schedule=(\"linear\", 1, 1e-5), loggers=loggers)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "After annealing, the cost has been reduced by 23.6% for 5 repetitions of the original circuit, close the 24.8% saving that could be obtained in the limit of a large number of repetitions of the original circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "source": [
    "print(f\"CX count before annealing: {opt.init_cx_count}\")\n",
    "print(f\"  - Gadget CX counts: {[g.cx_count(topology) for g in circuit.gadgets]}\")\n",
    "print()\n",
    "change = (opt.cx_count-opt.init_cx_count)/opt.init_cx_count\n",
    "max_change = (opt.cx_count-opt.init_cx_count-2*opt.cx_block.num_gates)/opt.init_cx_count\n",
    "print(f\"Cost after annealing: {opt.cx_count} \"\n",
    "      f\"({opt.circuit_rep} reps: {change:.1%}, many reps limit: {max_change:.1%})\")\n",
    "print(f\"  - Gadget CX counts: {[g.cx_count(topology) for g in opt.phase_block.gadgets]}\")\n",
    "print(f\"  - CX blocks CX count: {2*opt.cx_block.num_gates}\")\n",
    "opt.cx_block.draw(figsize=(3, 3))"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We can use the loggers to perform some basic performance benchmarking of our annealing algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "source": [
    "from time import perf_counter\n",
    "iter_timestamps = []\n",
    "def log_start_time(cost, num_iters):\n",
    "    del iter_timestamps[:]\n",
    "    iter_timestamps.append(perf_counter())\n",
    "def log_iter_time(it, prev_cost, new_cost, accepted, flip, t, num_iters):\n",
    "    iter_timestamps.append(perf_counter())\n",
    "def log_end_time(cost, num_iters):\n",
    "    a = np.array(iter_timestamps)\n",
    "    iter_times = 1_000_000*(a[1:]-a[:-1]) # in μs\n",
    "    avg_iter_time = np.mean(iter_times)\n",
    "    iter_time_std = np.std(iter_times)\n",
    "    print(f\"Average iter time: {avg_iter_time:.0f} μs \"\n",
    "          f\"({avg_iter_time/circuit.num_gadgets:.0f} μs/gadget)\")\n",
    "    print(f\"Iter time std dev: {iter_time_std:.0f} μs\")\n",
    "time_loggers = {\n",
    "    \"log_start\": log_start_time,\n",
    "    \"log_iter\": log_iter_time,\n",
    "    \"log_end\": log_end_time,\n",
    "}"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We first consider a small 9-qubit circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "source": [
    "topology = Topology.grid(3, 3)\n",
    "circuit = PhaseCircuit.random(topology.num_qubits, 50, min_legs=1, max_legs=3, rng_seed=0)\n",
    "opt = OptimizedPhaseCircuit(circuit, topology, cx_block=4, circuit_rep=1, rng_seed=1)\n",
    "opt.anneal(100, schedule=(\"geometric\", 1, 1e-5), loggers=time_loggers)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We then consider a larger 36-qubit circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "scrolled": true
   },
   "source": [
    "topology = Topology.grid(5, 5)\n",
    "circuit = PhaseCircuit.random(topology.num_qubits, 50, min_legs=1, max_legs=3, rng_seed=0)\n",
    "opt = OptimizedPhaseCircuit(circuit, topology, cx_block=4, circuit_rep=1, rng_seed=1)\n",
    "opt.anneal(100, schedule=(\"linear\", 1, 1e-5), loggers=time_loggers)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "We then consider a large 81-qubit circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "source": [
    "topology = Topology.grid(9, 9)\n",
    "circuit = PhaseCircuit.random(topology.num_qubits, 50, min_legs=1, max_legs=3, rng_seed=0)\n",
    "opt = OptimizedPhaseCircuit(circuit, topology, cx_block=4, circuit_rep=1, rng_seed=0)\n",
    "opt.anneal(100, schedule=(\"linear\", 1, 1e-5), loggers=time_loggers)"
   ],
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Finally, we consider a very large 144-qubit circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "scrolled": true
   },
   "source": [
    "topology = Topology.grid(12, 12)\n",
    "circuit = PhaseCircuit.random(topology.num_qubits, 50, min_legs=1, max_legs=3, rng_seed=0)\n",
    "opt = OptimizedPhaseCircuit(circuit, topology, cx_block=4, circuit_rep=1, rng_seed=0)\n",
    "opt.anneal(10, schedule=(\"linear\", 1, 1e-5), loggers=time_loggers)"
   ],
   "outputs": []
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "kernelspec": {
   "display_name": "Python 3",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}