{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hardware Grid Circuits\n",
    "\n",
    "The \"hardware grid\" problem is defined by a Hamiltonian whose topology matches the hardware graph natively. This permits a simple compilation (\"routing\") with circuit depth per p-step going like $O(1)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cirq\n",
    "import recirq\n",
    "\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "from cirq.contrib.svg import SVGCircuit, circuit_to_svg\n",
    "\n",
    "from recirq.qaoa.classical_angle_optimization import OptimizationResult\n",
    "from recirq.qaoa.problems import get_all_hardware_grid_problems\n",
    "\n",
    "# theme colors\n",
    "QBLUE = '#1967d2'\n",
    "QRED = '#ea4335ff'\n",
    "QGOLD = '#fbbc05ff'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we'll generate a 3x3 grid with aribitrarily chosen (fake!) beta, gamma parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fake_device_graph = nx.grid_2d_graph(3, 3)\n",
    "fake_device_graph = nx.relabel_nodes(\n",
    "    fake_device_graph, mapping={(r, c): cirq.GridQubit(r, c)\n",
    "                                for r, c in fake_device_graph.nodes})\n",
    "\n",
    "problems = get_all_hardware_grid_problems(fake_device_graph, central_qubit=cirq.GridQubit(1, 1),\n",
    "                                          n_instances=10, rs=np.random.RandomState(52))\n",
    "n_qubits = 9\n",
    "instance_i = 0\n",
    "problem = problems[n_qubits, instance_i]\n",
    "\n",
    "optimum = OptimizationResult(p=1, f_val=None, gammas=[0.123], betas=[0.456], min_c=None, max_c=None)\n",
    "nx.draw_networkx(problem.graph, \n",
    "                 pos={i: problem.coordinates[i] for i in range(problem.graph.number_of_nodes())},\n",
    "                 node_color=QBLUE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If, however, you've been following along, we can load in the results of `HardwareGridProblemGenerationTask`s for which we've actually pre-computed the optimal angles. TODO: enable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "from recirq.qaoa.experiments.problem_generation_tasks import HardwareGridProblemGenerationTask\n",
    "from recirq.qaoa.experiments.angle_precomputation_tasks import AnglePrecomputationTask\n",
    "\n",
    "gen_task = HardwareGridProblemGenerationTask(\n",
    "    dataset_id = '2020-03-19',\n",
    "    device_name = 'Sycamore23',\n",
    "    instance_i = 0,\n",
    "    n_qubits = 5,\n",
    ")\n",
    "\n",
    "pre_task = AnglePrecomputationTask(\n",
    "    dataset_id = '2020-03-23',\n",
    "    generation_task = gen_task,\n",
    "    p = 1,\n",
    ")\n",
    "print(gen_task)\n",
    "print(pre_task)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "from recirq.qaoa.experiments.problem_generation_tasks import DEFAULT_BASE_DIR as PGEN_BASE_DIR\n",
    "from recirq.qaoa.experiments.angle_precomputation_tasks import DEFAULT_BASE_DIR as APRE_BASE_DIR\n",
    "\n",
    "gen_data = recirq.load(gen_task, base_dir=PGEN_BASE_DIR)\n",
    "pre_data = recirq.load(pre_task, base_dir=APRE_BASE_DIR)\n",
    "problem = gen_data['problem']\n",
    "optimum = pre_data['optimum']\n",
    "print(optimum)\n",
    "nx.draw_networkx(problem.graph, \n",
    "                 pos={i: problem.coordinates[i] for i in range(problem.graph.number_of_nodes())},\n",
    "                 node_color=QBLUE\n",
    "                )\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ansatz\n",
    "\n",
    "As always, the circuit ansatz involves $|+\\rangle$ initialization followed by alternating applications of the problem and driver unitaries. We first construct a highly abstracted circuit with these multi-qubit operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.gates_and_compilation import ProblemUnitary, DriverUnitary\n",
    "qubits = cirq.LineQubit.range(problem.graph.number_of_nodes())\n",
    "\n",
    "circuit = cirq.Circuit(\n",
    "    cirq.H.on_each(qubits),\n",
    "    ProblemUnitary(problem.graph, gamma=optimum.gammas[0]).on(*qubits),\n",
    "    DriverUnitary(len(qubits), beta=optimum.betas[0]).on(*qubits)\n",
    ")\n",
    "SVGCircuit(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Harware topology\n",
    "\n",
    "We can enact the problem unitary with four entangling layers per p-step. \n",
    "\n",
    " 1. Horizontal links from even columns\n",
    " 2. Horizontal links from odd columns\n",
    " 3. Vertical links from even rows\n",
    " 4. Vertical links from odd rows\n",
    " \n",
    "To help the algorithm, we must specify `coordinates` to the compilation routine. This maps from bit indices $\\in \\{0, 1, \\dots n\\}$ to `(row, column)` coordinates so the compilation routine can categorize the various links into the above four categories. This is a little roundabout since we'll be mapping to `GridQubit`s, but I'm trying to emphasize the distinction between the problem (which is not related to quantum computing) and the implementation (which is).\n",
    " \n",
    "As always, the driver unitary is nothing more than single-qubit X rotations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.gates_and_compilation import compile_problem_unitary_to_hardware_graph, \\\n",
    "                                              compile_driver_unitary_to_rx\n",
    "circuit = compile_problem_unitary_to_hardware_graph(circuit, problem.coordinates)\n",
    "circuit = compile_driver_unitary_to_rx(circuit)\n",
    "SVGCircuit(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Aside: Compiling $e^{i \\gamma w_{ij} Z_i Z_j}$\n",
    "\n",
    "For this problem, we need to express the `ZZ` interaction as three rounds of `SYC` gates. We take a brief aside to look at this compilation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "zz = cirq.Circuit(cirq.ZZ(*qubits[:2])**(2*0.345/np.pi))\n",
    "SVGCircuit(zz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.gates_and_compilation import compile_to_syc\n",
    "zz = compile_to_syc(zz)\n",
    "SVGCircuit(zz)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function `zz_as_syc` is included for convenience"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.gates_and_compilation import zz_as_syc\n",
    "zz = zz_as_syc(0.345, *qubits[:2])\n",
    "SVGCircuit(zz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cirq.testing.assert_allclose_up_to_global_phase(\n",
    "    cirq.Circuit(cirq.ZZ(*qubits[:2])**(2*0.345/np.pi)).unitary(),\n",
    "    zz_as_syc(0.345, *qubits[:2]).unitary(),\n",
    "    atol=1e-8\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cirq.testing.assert_allclose_up_to_global_phase(\n",
    "    compile_to_syc(cirq.Circuit(cirq.ZZ(*qubits[:2])**(2*0.345/np.pi))).unitary(),\n",
    "    zz_as_syc(0.345, *qubits[:2]).unitary(),\n",
    "    atol=1e-8\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We make sure all our gates are \"well-structured\"\n",
    "This means each layer is composed of homogeneous operations which are native to the device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.circuit_structure import validate_well_structured\n",
    "_, stats = validate_well_structured(zz)\n",
    "stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compiling to native operations\n",
    "\n",
    "We use the above compilation of `ZZ` to compile our circuit to native operations. Because our compilation produces well-structured gates and our starting circuit was structured, the resulting circuit is well-structured."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.gates_and_compilation import compile_to_syc\n",
    "circuit = compile_to_syc(circuit)\n",
    "SVGCircuit(circuit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, stats = validate_well_structured(circuit)\n",
    "stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Append Measurement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mcircuit = circuit + cirq.measure(*qubits, key='z')\n",
    "SVGCircuit(mcircuit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, stats = validate_well_structured(mcircuit)\n",
    "stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile out Z's\n",
    "Z gates commute through SYC so we can remove them. This step is not necessary: the quantum operating system will track the virtual Zs if we don't remove them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.gates_and_compilation import compile_out_virtual_z\n",
    "mcircuit = compile_out_virtual_z(mcircuit)\n",
    "SVGCircuit(mcircuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile out Negligible gates\n",
    "We've left several `PhX^0` to keep our circuits structured. As the very last compilation step, we can drop these."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.gates_and_compilation import compile_to_non_negligible\n",
    "mcircuit = compile_to_non_negligible(mcircuit)\n",
    "SVGCircuit(mcircuit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, stats = validate_well_structured(mcircuit)\n",
    "stats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Place on device\n",
    "\n",
    " - Our problem has integer nodes because it should be specified independently of a quantum implementation\n",
    " - Our circuit has LineQubit qubits to emphasize the fact that we can place this circuit in multiple locations on a device\n",
    " - Our `coordinates` list was used only as a helper for the compilation\n",
    " \n",
    "We now place the compiled circuit onto a compatible part of the device. Here, we use networkx's subgraph isomorphism routine to find all the possibilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cirq.contrib.routing import xmon_device_to_graph\n",
    "device_graph = xmon_device_to_graph(recirq.get_device_obj_by_name('Sycamore23'))\n",
    "nx.draw_networkx(device_graph, pos={q: (q.row, q.col) for q in device_graph.nodes}, node_color=QRED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "from cirq.contrib.routing import xmon_device_to_graph\n",
    "device_graph = xmon_device_to_graph(recirq.get_device_obj_by_name('Sycamore23'))\n",
    "matcher = nx.algorithms.isomorphism.GraphMatcher(device_graph, problem.graph)\n",
    "\n",
    "# There's a \"rotational\" freedom which we remove here:\n",
    "each_set_of_qubits_only_one_subgraph = {}\n",
    "for q_to_i in matcher.subgraph_isomorphisms_iter():\n",
    "    each_set_of_qubits_only_one_subgraph[frozenset(q_to_i.keys())] = q_to_i\n",
    "\n",
    "for q_to_i in each_set_of_qubits_only_one_subgraph.values():\n",
    "    nx.draw_networkx(device_graph, pos={q: (q.row, q.col) for q in device_graph.nodes},\n",
    "                     node_color=[QRED if q in q_to_i else QBLUE for q in device_graph.nodes])\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i_to_q = {i: q for q, i in q_to_i.items()}\n",
    "# Since our nodes are contiguous integers starting from 0, we can flatten into a list\n",
    "device_qubits = [i_to_q[i] for i in range(len(i_to_q))]\n",
    "del i_to_q\n",
    "\n",
    "def _mapq(q):\n",
    "    return device_qubits[q.x]\n",
    "\n",
    "mcircuit = mcircuit.transform_qubits(_mapq)\n",
    "SVGCircuit(mcircuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Circuit Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.problem_circuits import get_generic_qaoa_circuit\n",
    "circuit = get_generic_qaoa_circuit(\n",
    "    problem_graph=problem.graph, \n",
    "    qubits=qubits, \n",
    "    gammas=[0.123], \n",
    "    betas=[0.456],\n",
    ")\n",
    "SVGCircuit(circuit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.problem_circuits import get_routed_hardware_grid_circuit\n",
    "circuit = get_routed_hardware_grid_circuit(\n",
    "    problem_graph=problem.graph,\n",
    "    qubits=qubits,\n",
    "    coordinates=problem.coordinates,\n",
    "    gammas=[0.123],\n",
    "    betas=[0.456],\n",
    ")\n",
    "SVGCircuit(circuit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recirq.qaoa.problem_circuits import get_compiled_hardware_grid_circuit\n",
    "circuit, qubits = get_compiled_hardware_grid_circuit(\n",
    "    problem=problem,\n",
    "    qubits=device_qubits,\n",
    "    gammas=[0.123],\n",
    "    betas=[0.456],\n",
    ")\n",
    "SVGCircuit(circuit)"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}