{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Allocating Qubits on QPU Devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use Braket SDK Cost Tracking to estimate the cost to run this example\n",
    "from braket.tracking import Tracker\n",
    "t = Tracker().start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates how you can specify explicitly which qubits to use when you run a quantum circuit on QPU devices from Rigetti.\n",
    "\n",
    "When you submit a circuit for execution on a QPU, Amazon Braket performs a series of compilation steps: it maps the _abstract qubits_ in your circuit to _physical qubits_ in the device; it synthesizes gates into the native gate set of the device; it optimizes the circuit to reduce the number of gates; and finally, it translates the gates into executable pulses.\n",
    "\n",
    "This section shows how the first step, called qubit allocation, works for the Rigetti Aspen-M-3 device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# general imports\n",
    "from braket.aws import AwsDevice\n",
    "from braket.circuits import Circuit\n",
    "import numpy as np\n",
    "import random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Automatic qubit allocation\n",
    "\n",
    "Qubit allocation for Rigetti devices on Amazon Braket utilizes [the Quil Compilers](https://pyquil-docs.rigetti.com/en/latest/compiler.html#the-quil-compiler)'s _rewiring_ strategies. By default, when you submit a circuit on Amazon Braket to a Rigetti device, the circuit is rewired according to the [PARTIAL](https://pyquil-docs.rigetti.com/en/latest/compiler.html#partial) rewiring strategy. Specifically, the compiler starts with an empty mapping from logical to physical qubits. Taking into account the latest calibration data of the device, the compiler fills in the mapping with the goal, sequentially, to maximize the overall fidelity of the circuit.\n",
    "\n",
    "The example that follows shows how to create a GHZ state on qubits that are not physically connected. After the task is completed, you can obtain a list of the actual gates executed on the device, by viewing the result metadata.\n",
    "\n",
    "First, instantiate the Rigetti Aspen-M-3 device and retrieve its connectivity graph, which shows the qubits that are directly connected on the chip."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "the connectivity of Aspen-M-3 is: {'0': ['1', '7'], '1': ['0', '16', '2'], '10': ['11', '113', '17'], '100': ['101', '107'], '101': ['100', '102', '116'], '102': ['101', '103', '115'], '103': ['102', '104'], '104': ['103', '105', '7'], '105': ['104', '106'], '106': ['105', '107'], '107': ['100', '106'], '11': ['10', '12', '26'], '110': ['111', '117'], '111': ['110', '112', '126'], '112': ['111', '113'], '113': ['10', '112', '114'], '114': ['113', '115', '17'], '115': ['102', '114', '116'], '116': ['101', '115', '117'], '117': ['110', '116'], '12': ['11', '13', '25'], '120': ['121', '127'], '121': ['120', '122', '136'], '122': ['121', '123', '135'], '123': ['122', '124', '20'], '124': ['123', '125', '27'], '125': ['124', '126'], '126': ['111', '125', '127'], '127': ['120', '126'], '13': ['12', '14'], '130': ['131', '137'], '131': ['130', '132', '146'], '132': ['131', '133', '145'], '133': ['132', '134', '30'], '134': ['133', '135', '37'], '135': ['122', '134', '136'], '136': ['121', '135', '137'], '137': ['130', '136'], '14': ['13', '15'], '140': ['141', '147'], '141': ['140', '142'], '142': ['141', '143'], '143': ['142', '144', '40'], '144': ['143', '145', '47'], '145': ['132', '144', '146'], '146': ['131', '145', '147'], '147': ['140', '146'], '15': ['14', '16', '2'], '16': ['1', '15', '17'], '17': ['10', '16', '114'], '2': ['1', '15', '3'], '20': ['123', '21', '27'], '21': ['20', '22', '36'], '22': ['21', '23', '35'], '23': ['22', '24'], '24': ['23', '25'], '25': ['12', '24', '26'], '26': ['11', '25', '27'], '27': ['20', '26', '124'], '3': ['2', '4'], '30': ['133', '31', '37'], '31': ['30', '32', '46'], '32': ['31', '33', '45'], '33': ['32', '34'], '34': ['33', '35'], '35': ['22', '34', '36'], '36': ['21', '35', '37'], '37': ['30', '36', '134'], '4': ['3', '5'], '40': ['143', '41', '47'], '41': ['40', '42'], '42': ['41', '43'], '43': ['42', '44'], '44': ['43', '45'], '45': ['32', '44', '46'], '46': ['31', '45', '47'], '47': ['40', '46', '144'], '5': ['4', '6'], '6': ['5', '7'], '7': ['0', '6', '104']}\n"
     ]
    }
   ],
   "source": [
    "device = AwsDevice(\"arn:aws:braket:us-west-1::device/qpu/rigetti/Aspen-M-3\")\n",
    "\n",
    "connectivity_graph = device.properties.paradigm.connectivity.connectivityGraph\n",
    "print(f\"the connectivity of {device.name} is: {connectivity_graph}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next, create a GHZ circuit with three qubits 0, 2, 4, and run it on the Aspen-M-3 device. Notice that none of these qubits are connected on the Aspen-M-3 connectivity graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "T  : |0|1|2|\n",
      "            \n",
      "q0 : -H-C-C-\n",
      "        | | \n",
      "q2 : ---X-|-\n",
      "          | \n",
      "q4 : -----X-\n",
      "\n",
      "T  : |0|1|2|\n"
     ]
    }
   ],
   "source": [
    "# create a GHZ state with non-neighboring qubits\n",
    "circuit = Circuit()\n",
    "circuit.h(0).cnot(0,2).cnot(0,4)\n",
    "print(circuit)\n",
    "\n",
    "rigetti_rewiring = device.run(circuit, shots=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Status of task: CREATED\n"
     ]
    }
   ],
   "source": [
    "print(\"Status of task:\", rigetti_rewiring.state())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To verify the final qubit allocation, retrieve the compiled program that was executed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Measurement counts: Counter({'001': 4, '111': 3, '000': 2, '101': 1})\n",
      "The compiled circuit is:\n",
      " DECLARE ro BIT[3]\n",
      "PRAGMA INITIAL_REWIRING \"PARTIAL\"\n",
      "RESET\n",
      "RZ(-pi/2) 12\n",
      "RX(-pi/2) 12\n",
      "RZ(pi) 13\n",
      "XY(pi) 12 13\n",
      "RZ(pi/2) 12\n",
      "RX(pi/2) 12\n",
      "RZ(-pi/2) 12\n",
      "XY(pi) 12 13\n",
      "RZ(pi/2) 25\n",
      "RX(-pi/2) 25\n",
      "CZ 25 12\n",
      "RZ(pi) 12\n",
      "RX(-pi/2) 13\n",
      "RX(pi/2) 25\n",
      "RZ(-pi/2) 25\n",
      "MEASURE 25 ro[2]\n",
      "MEASURE 13 ro[1]\n",
      "MEASURE 12 ro[0]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "result = rigetti_rewiring.result()\n",
    "counts = result.measurement_counts\n",
    "print(\"Measurement counts:\", counts)\n",
    "print(\"The compiled circuit is:\\n\", result.additional_metadata.rigettiMetadata.compiledProgram)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Notice that the PARTIAL rewiring was applied. The qubits 0, 2, 4 in the original circuit were mapped to three other qubits in the Rigetti device, and the gates were compiled into native gates."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User-defined qubit allocation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In Amazon Braket, you can choose to prescribe a qubit mapping manually, and prevent further rewiring for Rigetti devices. To enable manual mapping, set `disable_qubit_rewiring=True` when submitting the task to run.\n",
    "\n",
    "If all the gates in the circuit satisfy the topological constraints of the device, Amazon Braket maps abstract qubit $i$ in the circuit to the physical qubit $i$ in the device, and maps qubit pair $(i, j)$ to the connection $(i, j)$ in the device. On the other hand, Amazon Braket raises an exception if a specified qubit or qubit pair do not exist in the device connectivity graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "T  : |   0    |1|\n",
      "                 \n",
      "q0 : -Rz(1.57)-C-\n",
      "               | \n",
      "q1 : ----------X-\n",
      "                 \n",
      "q3 : -X----------\n",
      "\n",
      "T  : |   0    |1|\n"
     ]
    }
   ],
   "source": [
    "# create a random state with neighboring qubits\n",
    "q1=random.choice(list(connectivity_graph))\n",
    "q2=int(connectivity_graph[q1][0])\n",
    "q1=int(q1)\n",
    "\n",
    "circuit = Circuit()\n",
    "circuit.rz(0,np.pi/2).cnot(q1,q2).x(7)\n",
    "print(circuit)\n",
    "rigetti_task = device.run(circuit, shots=10, disable_qubit_rewiring=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Status of task: COMPLETED\n"
     ]
    }
   ],
   "source": [
    "print(\"Status of task:\", rigetti_task.state())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Measurement counts: Counter({'011': 6, '001': 2, '101': 1, '000': 1})\n",
      "The compiled circuit is:\n",
      " DECLARE ro BIT[3]\n",
      "PRAGMA INITIAL_REWIRING \"NAIVE\"\n",
      "RESET\n",
      "RZ(pi) 1\n",
      "RX(-pi/2) 1\n",
      "XY(pi) 0 1\n",
      "RZ(-pi/2) 0\n",
      "RX(pi/2) 0\n",
      "RZ(pi/2) 0\n",
      "XY(pi) 0 1\n",
      "RX(pi) 3\n",
      "MEASURE 3 ro[2]\n",
      "MEASURE 1 ro[1]\n",
      "MEASURE 0 ro[0]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "result = rigetti_task.result()\n",
    "counts = result.measurement_counts\n",
    "print(\"Measurement counts:\", counts)\n",
    "print(\"The compiled circuit is:\\n\", result.additional_metadata.rigettiMetadata.compiledProgram)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The qubits in the original circuit followed a one-to-one mapping to the physical qubits in the device. Other compilation steps, such as gate synthesis and circuit optimization, are still performed. These steps allow the circuit to run successfully and improve the overall fidelity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using the qubits with the highest two-qubit gate fidelity\n",
    "\n",
    "Additionally, the device properties include calibration data, which you can use to find the qubits and qubit pairs with the highest fidelities for particular gates.\n",
    "\n",
    "The following function finds the qubit pair that has the highest two-qubit fidelity of an input gate, which can be any of the gates native to the Rigetti device, namely CPHASE, XY or CZ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def find_qubit_pair(gate):\n",
    "    \"Function to find the qubit pair that has the highest gate fidelity of a particular gate\"\n",
    "    \n",
    "    # check whether the input gate is a string\n",
    "    if not isinstance(gate, str):\n",
    "        raise ValueError('The input gate must be a string type.') \n",
    "    \n",
    "    # check whether the input gate is a native gate\n",
    "    gate_list = ['CPHASE', 'CZ', 'XY']\n",
    "    if gate not in gate_list:\n",
    "        raise ValueError('The input gate must be either CPHASE, CZ or XY.')\n",
    "    \n",
    "    # load all calibration data from device.properties\n",
    "    calibration_2Q = device.properties.provider.specs['2Q']\n",
    "    highest_fidelity = 0\n",
    "\n",
    "    # iterate through all calibration data to find the highest fidelity\n",
    "    for pair in calibration_2Q.keys():\n",
    "        \n",
    "        # if the particular gate type is supported by the qubit pair\n",
    "        if ('f'+ gate) in calibration_2Q[pair].keys(): \n",
    "            \n",
    "            if calibration_2Q[pair]['f'+ gate] > highest_fidelity:\n",
    "                \n",
    "                # update the highest_fidelity and the best_pair\n",
    "                highest_fidelity = calibration_2Q[pair]['f'+ gate]\n",
    "                best_pair = pair\n",
    "\n",
    "    # generate the two qubits as integers       \n",
    "    q1 = best_pair[0]\n",
    "    i = 1\n",
    "    while best_pair[i] is not '-':\n",
    "        q1 += best_pair[i]\n",
    "        i += 1\n",
    "\n",
    "    q1 = int(q1)\n",
    "    q2 = int(best_pair[i+1:])\n",
    "        \n",
    "    return q1, q2, highest_fidelity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The example in the following code applies a native two-qubit gate on the qubit pair that has the highest fidelity of that gate. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The highest fidelity for CZ gate is: 0.9794230141816626\n",
      "And the corresponding qubit pair is: qubit 21 and qubit 36\n",
      "T   : |0|\n",
      "         \n",
      "q21 : -C-\n",
      "       | \n",
      "q36 : -Z-\n",
      "\n",
      "T   : |0|\n"
     ]
    }
   ],
   "source": [
    "# the gate must be either 'CZ', 'CPHASE' or 'XY'\n",
    "gate = 'CZ'\n",
    "# find the qubit pair with the highest gate fidelity\n",
    "q1, q2, highest_fidelity = find_qubit_pair(gate)\n",
    "print('The highest fidelity for '+gate+' gate is:', highest_fidelity)\n",
    "print(f'And the corresponding qubit pair is: qubit {q1} and qubit {q2}')\n",
    "\n",
    "# create a circuit with the gate applied to the discovered qubit pair.\n",
    "# note that CPHASE in Rigetti corresponds to cphaseshift in Braket\n",
    "circuit = Circuit()\n",
    "circuit.cz(q1,q2)\n",
    "print(circuit)\n",
    "rigetti_task = device.run(circuit, shots=1000, disable_qubit_rewiring=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Measurement counts: Counter({'00': 910, '01': 55, '10': 33, '11': 2})\n",
      "The compiled circuit is:\n",
      " DECLARE ro BIT[2]\n",
      "PRAGMA INITIAL_REWIRING \"NAIVE\"\n",
      "RESET\n",
      "CZ 21 36\n",
      "MEASURE 21 ro[1]\n",
      "MEASURE 36 ro[0]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"Status of task:\", rigetti_task.state())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The qubits in the original circuit followed a one-to-one mapping to the physical qubits in the device. Since only native gates were used, the actual gates executed are the same as the gates in the original circuit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b> The IonQ device does not support manual allocation. For circuits submitted to the IonQ device, qubits are allocated automatically.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task Summary\n",
      "{'arn:aws:braket:us-west-1::device/qpu/rigetti/Aspen-M-3': {'shots': 1020, 'tasks': {'QUEUED': 1, 'CREATED': 2}}}\n",
      "Note: Charges shown are estimates based on your Amazon Braket simulator and quantum processing unit (QPU) task usage. Estimated charges shown may differ from your actual charges. Estimated charges do not factor in any discounts or credits, and you may experience additional charges based on your use of other services such as Amazon Elastic Compute Cloud (Amazon EC2).\n",
      "Estimated cost to run this example: 1.26 USD\n"
     ]
    }
   ],
   "source": [
    "print(\"Task Summary\")\n",
    "print(t.quantum_tasks_statistics())\n",
    "print('Note: Charges shown are estimates based on your Amazon Braket simulator and quantum processing unit (QPU) task usage. Estimated charges shown may differ from your actual charges. Estimated charges do not factor in any discounts or credits, and you may experience additional charges based on your use of other services such as Amazon Elastic Compute Cloud (Amazon EC2).')\n",
    "print(f\"Estimated cost to run this example: {t.qpu_tasks_cost() + t.simulator_tasks_cost():.2f} USD\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('venv': venv)",
   "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.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "590fab68195cf107911461461f81d5c472d3d6127f579badfcfad30f03e5cab2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}