{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Routing with t|ket>\n",
    "\n",
    "We wrap tket's compilation unit framework to keep track of qubit mappings and work with generic devices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cirq\n",
    "import recirq\n",
    "import networkx as nx\n",
    "from cirq.contrib.svg import SVGCircuit\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytket.predicates import CompilationUnit, ConnectivityPredicate\n",
    "from pytket.passes import SequencePass, RoutingPass, DecomposeSwapsToCXs\n",
    "from pytket.routing import GraphPlacement"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example circuit\n",
    "We'll route a 3-regular circuit to Sycamore23. To try to clear up some of the confusion about which indices are which, we'll construct the initial circuit with `LineQubits` 10 through 19 which should be thought of as \"logical indices\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.problem_circuits import get_generic_qaoa_circuit\n",
    "from recirq.qaoa.gates_and_compilation import compile_problem_unitary_to_arbitrary_zz, \\\n",
    "    compile_driver_unitary_to_rx\n",
    "\n",
    "problem_graph = nx.random_regular_graph(d=3, n=10)\n",
    "nx.set_edge_attributes(problem_graph, values=1, name='weight')\n",
    "circuit_qubits = cirq.LineQubit.range(10, 20)\n",
    "gammas = np.random.randn(2)\n",
    "betas = np.random.randn(2)\n",
    "circuit = get_generic_qaoa_circuit(\n",
    "    problem_graph=problem_graph,\n",
    "    qubits=circuit_qubits,\n",
    "    gammas=gammas,\n",
    "    betas=betas)\n",
    "circuit = compile_problem_unitary_to_arbitrary_zz(circuit)\n",
    "circuit = compile_driver_unitary_to_rx(circuit)\n",
    "SVGCircuit(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We need to \"route\" this circuit\n",
    "Let's look at the \"connectivity graph\" of the circuit vs. that of the device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cirq.contrib.routing as ccr\n",
    "uncompiled_c_graph = ccr.get_circuit_connectivity(circuit)\n",
    "nx.draw_networkx(uncompiled_c_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cirq.google as cg\n",
    "dev_graph = ccr.xmon_device_to_graph(cg.Sycamore23)\n",
    "nx.draw_networkx(dev_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# alias for the device. If this notebook were wrapped\n",
    "# in a function, `circuit` and `device` would be the arguments\n",
    "device = cg.Sycamore23"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mapping to device *indices*\n",
    "We'll keep a set of secret indices that number device qubits contiguously from zero instead of `(row, col)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index_to_qubit = sorted(device.qubit_set())\n",
    "qubit_to_index = {q: i for i, q in enumerate(index_to_qubit)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert to pytket `Device`\n",
    "The provided function doesn't work with `SerializableDevice`. We use existing functionality to turn Devices into graphs to provide a more robust solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytket\n",
    "from pytket.circuit import Node\n",
    "\n",
    "def _qubit_index_edges():\n",
    "    dev_graph = ccr.xmon_device_to_graph(device)\n",
    "    for n1, n2 in dev_graph.edges:\n",
    "        #yield Node('q', n1.row, n1.col), Node('q', n2.row, n2.col)\n",
    "        yield (qubit_to_index[n1], qubit_to_index[n2])\n",
    "\n",
    "def _device_to_tket_device():\n",
    "    arc = pytket.routing.Architecture(\n",
    "        list(_qubit_index_edges())\n",
    "    )\n",
    "    return pytket.device.Device({}, {}, arc)\n",
    "\n",
    "tk_circuit = pytket.cirq.cirq_to_tk(circuit)\n",
    "tk_device = _device_to_tket_device()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### tket understands LineQubit and uses our strange indexing convention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tk_circuit.qubits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### but our device uses our secret indices\n",
    "\n",
    "There seems to be a bug if you use their built-in support for two-index qubits (nodes): `Existing register q cannot support id: q[6, 1]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tk_device.coupling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Placement and Routing pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytket.predicates import CompilationUnit, ConnectivityPredicate\n",
    "from pytket.passes import SequencePass, RoutingPass, DecomposeSwapsToCXs, PlacementPass\n",
    "from pytket.routing import GraphPlacement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unit = CompilationUnit(tk_circuit, [ConnectivityPredicate(tk_device)])\n",
    "passes = SequencePass([\n",
    "    PlacementPass(GraphPlacement(tk_device)),\n",
    "    RoutingPass(tk_device)])\n",
    "passes.apply(unit)\n",
    "valid = unit.check_all_predicates()\n",
    "assert valid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The initial mapping\n",
    "This maps from logical LineQubits to secret device indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unit.initial_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bookkept initial mapping\n",
    "We \"decode\" our tket conventions back into Cirq idioms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tk_to_i(tk):\n",
    "    i = tk.index\n",
    "    assert len(i) == 1, i\n",
    "    return i[0]\n",
    "\n",
    "initial_map = {cirq.LineQubit(tk_to_i(n1)): index_to_qubit[tk_to_i(n2)] for n1, n2 in unit.initial_map.items()}\n",
    "initial_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The final mapping\n",
    "This maps from logical LineQubits to final secret device indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unit.final_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_map = {cirq.LineQubit(tk_to_i(n1)): index_to_qubit[tk_to_i(n2)]\n",
    "             for n1, n2 in unit.final_map.items()}\n",
    "final_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The compilation unit applies the mapping\n",
    "So our circuit qubits use secret device indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unit.circuit.qubits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### So we map the circuit to Grid Qubits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "routed_circuit = pytket.cirq.tk_to_cirq(unit.circuit)\n",
    "routed_circuit = routed_circuit.transform_qubits(lambda q: index_to_qubit[q.x])\n",
    "SVGCircuit(routed_circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now it's nice and compiled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "routed_c_graph = ccr.get_circuit_connectivity(routed_circuit)\n",
    "nx.draw_networkx(routed_c_graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check that circuits are equivalent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _, op, _ in routed_circuit.findall_operations_with_gate_type(cirq.TwoQubitGate):\n",
    "    a, b = op.qubits\n",
    "    assert a.is_adjacent(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cirq.contrib.acquaintance as cca\n",
    "def permute_gate(qubits, permutation):\n",
    "    return cca.LinearPermutationGate(\n",
    "        num_qubits=len(qubits),\n",
    "        permutation={i: permutation[i] for i in range(len(permutation))}\n",
    "    ).on(*qubits)\n",
    "\n",
    "final_to_initial_map = {final_map[cq]: initial_map[cq]\n",
    "                              for cq in circuit_qubits}\n",
    "initial_qubits = [initial_map[cq] for cq in circuit_qubits]\n",
    "final_permutation = [initial_qubits.index(final_to_initial_map[q])\n",
    "                     for q in initial_qubits]\n",
    "rcircuit_with_perm = routed_circuit.copy()\n",
    "rcircuit_with_perm.append(permute_gate(initial_qubits, final_permutation))\n",
    "expected = circuit.unitary(qubit_order=cirq.QubitOrder.explicit(circuit_qubits))\n",
    "actual = rcircuit_with_perm.unitary(qubit_order=cirq.QubitOrder.explicit(initial_qubits))\n",
    "cirq.testing.assert_allclose_up_to_global_phase(expected, actual, atol=1e-8)"
   ]
  }
 ],
 "metadata": {
  "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}