examples/6_Customize_gate_calibrations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dda432a3",
   "metadata": {},
   "source": [
    "# Customize gate calibrations\n",
    "\n",
    "The pulse implementation of a gate is called \"gate calibration\". Gaining access to gate calibrations empowers you to experiment with quantum gate designs and execute quantum programs with customized, improved gate calibrations. This notebook shows you how to access the native gate calibrations provided by quantum hardware providers, and shows you how to define custom gate calibrations for your quantum programs.\n",
    "\n",
    "We start with basic imports, and initialize the Rigetti Ankaa-2 device which is the targeted device for this example notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "30fd0b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "# general imports\n",
    "import numpy as np\n",
    "\n",
    "# AutoQASM imports\n",
    "import autoqasm as aq\n",
    "from autoqasm import pulse\n",
    "from autoqasm.instructions import measure, rx, rz\n",
    "\n",
    "# AWS imports: Import Braket SDK modules\n",
    "from braket.aws import AwsDevice\n",
    "from braket.circuits import FreeParameter, Gate, QubitSet\n",
    "from braket.devices import Devices\n",
    "from braket.pulse import DragGaussianWaveform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1dda4463-2dbb-40e1-a90c-66f1278ce5cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = AwsDevice(Devices.Rigetti.Ankaa2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9fdfd5e",
   "metadata": {},
   "source": [
    "## Gate calibrations from hardware providers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d08537b",
   "metadata": {},
   "source": [
    "To customize a gate calibration, you often need to be intimately familiar with how quantum gates are implemented with pulses. It is often a good place to start by viewing the native gate calibrations supplied by hardware providers. To do so, we first retrieve the information from the `gate_calibrations` attribute of a device. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f1e5a16f",
   "metadata": {},
   "outputs": [],
   "source": [
    "provider_calibrations = device.gate_calibrations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "959c5fe8",
   "metadata": {},
   "source": [
    "Even for the same gate, different physical qubits and angles usually have a different pulse implementation. Because of this, a gate calibration may be associated with specific physical qubits and angles. In the code snippet below, we retrieve the gate calibration of the Rx gate for physical qubit `$0` and `angle=pi/2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3248ed04",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OPENQASM 3.0;\n",
      "cal {\n",
      "    waveform wf_drag_gaussian_153 = drag_gaussian(40.0ns, 8.49321800288ns, -2.618211817882948e-10, 0.19930451095449803, false);\n",
      "    barrier $0;\n",
      "    play(Transmon_0_charge_tx, wf_drag_gaussian_153);\n",
      "    barrier $0;\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "pulse_sequence_rx_pi_2_q0 = provider_calibrations.pulse_sequences[(Gate.Rx(np.pi / 2), QubitSet(0))]\n",
    "\n",
    "print(pulse_sequence_rx_pi_2_q0.to_ir())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5579f6e",
   "metadata": {},
   "source": [
    "Some gate calibrations are defined not with fixed values of angles, but rather with variables. In this case, these gate calibrations can be retrieved with angles of `FreeParameter` type. For example, the code snippet below retrieves the gate calibration of the Rz gate on the qubit `$1` with a variable angle `theta`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "81bb487f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OPENQASM 3.0;\n",
      "cal {\n",
      "    input float theta;\n",
      "    shift_phase(Transmon_1_charge_tx, -1.0 * theta);\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "theta = FreeParameter(\"theta\")\n",
    "pulse_sequence_rz_theta_q0 = provider_calibrations.pulse_sequences[(Gate.Rz(theta), QubitSet(1))]\n",
    "\n",
    "print(pulse_sequence_rz_theta_q0.to_ir())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d4d9641",
   "metadata": {},
   "source": [
    "## Compose custom gate calibrations with AutoQASM\n",
    "In this section, we show you how to customize a gate calibration and use it in a quantum program. To create a custom gate calibration, you need to \n",
    "1. Find the definition of the gate in AutoQASM.\n",
    "2. Decide the qubits and angles you want to define the calibration.\n",
    "3. Compose a pulse program that defines the calibration.\n",
    "\n",
    "As an example, let's define the gate calibration of `rx` on qubit `$1` and at angle `pi/2`. We first inspect the definition of the `rx` gate. You can find the definition of gate instructions in the [gate module](https://github.com/amazon-braket/autoqasm/blob/main/src/autoqasm/instructions/gates.py) of AutoQASM. You can also use the Python built-in `help` function to retrieve the definition. A gate calibration for the `rx` gate must fully specify the inputs of the gate, `target` and `theta`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "63e15583",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on function rx in module autoqasm.instructions.gates:\n",
      "\n",
      "rx(target: Union[int, str, braket.registers.qubit.Qubit, oqpy.classical_types._ClassicalVar, oqpy.base.OQPyExpression, oqpy.quantum_types.Qubit], theta: Union[float, braket.parametric.free_parameter_expression.FreeParameterExpression, oqpy.classical_types._ClassicalVar], **kwargs) -> None\n",
      "    X-axis rotation gate.\n",
      "    \n",
      "    Args:\n",
      "        target (QubitIdentifierType): Target qubit.\n",
      "        theta (GateParameterType): Rotation angle in radians.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(rx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1189626",
   "metadata": {},
   "source": [
    "To specify fixed values for the parameters `target` and `theta`, we set the values through keyword arguments of the decorator. The pulse implementation for the gate calibration is in the body of the `my_rx_cal` function decorated by `@aq.gate_calibration`. In the example in the next cell, we want to experiment with offsetting the frequency of the pulse by 100 Hz away from the hardware provider's implementation. The gate calibration `my_rx_cal` recreates the hardware provider's implementation but with an offset in the frequency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "64ebd3a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "wf = DragGaussianWaveform(\n",
    "    24e-9, 2.547965400864e-9, 2.370235498840002e-10, 0.293350447987059, False, \"wf_drag_gaussian_1\"\n",
    ")\n",
    "q0_control_frame = device.frames[\"Transmon_0_charge_tx\"]\n",
    "\n",
    "\n",
    "@aq.gate_calibration(implements=rx, target=\"$0\", theta=np.pi / 2)\n",
    "def my_rx_cal():\n",
    "    pulse.barrier(\"$0\")\n",
    "    pulse.shift_frequency(q0_control_frame, -321047.14178613486 - 100)\n",
    "    pulse.play(q0_control_frame, wf)\n",
    "    pulse.shift_frequency(q0_control_frame, 321047.14178613486 + 100)\n",
    "    pulse.barrier(\"$0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9300f0f8",
   "metadata": {},
   "source": [
    "Next, we will demonstrate how to attach `my_rx_cal` to the main quantum program `my_program`. The program is defined in the code block below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bae7731a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OPENQASM 3.0;\n",
      "rx(1.5707963267948966) $0;\n",
      "rz(0.123) $1;\n",
      "bit __bit_0__;\n",
      "__bit_0__ = measure $0;\n"
     ]
    }
   ],
   "source": [
    "@aq.main\n",
    "def my_program():\n",
    "    rx(\"$0\", np.pi / 2)\n",
    "    rz(\"$1\", 0.123)\n",
    "    measure(\"$0\")\n",
    "\n",
    "\n",
    "print(my_program.build().to_ir())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7579516",
   "metadata": {},
   "source": [
    "To attach gate calibrations to the program, call the `with_calibrations` method on the built program to create a new program with `my_rx_cal` attached, leaving the original intact. This allows you to experiment with different gate calibrations on the same program without the need to redefine the main program every time. \n",
    "\n",
    "In the printed OpenQASM script of the program, the gate definition for the `rx` gate becomes a `defcal` block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6b7cc158",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OPENQASM 3.0;\n",
      "cal {\n",
      "    waveform wf_drag_gaussian_1 = drag_gaussian(24.0ns, 2.547965400864ns, 2.370235498840002e-10, 0.293350447987059, false);\n",
      "}\n",
      "defcal rx(1.5707963267948966) $0 {\n",
      "    barrier $0;\n",
      "    shift_frequency(Transmon_0_charge_tx, -321147.14178613486);\n",
      "    play(Transmon_0_charge_tx, wf_drag_gaussian_1);\n",
      "    shift_frequency(Transmon_0_charge_tx, 321147.14178613486);\n",
      "    barrier $0;\n",
      "}\n",
      "rx(1.5707963267948966) $0;\n",
      "rz(0.123) $1;\n",
      "bit __bit_0__;\n",
      "__bit_0__ = measure $0;\n"
     ]
    }
   ],
   "source": [
    "custom_program = my_program.build().with_calibrations(my_rx_cal)\n",
    "print(custom_program.to_ir())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f038b8c8",
   "metadata": {},
   "source": [
    "You can also define a gate calibration with variable parameters. Variable parameters are used in the body of the decorated function. The variable parameters must be arguments of the decorated Python function, instead of being arguments to the decorator `@aq.gate_calibration`. The variable parameters must have a type hint of either `aq.Qubit` for qubits or `float` for angles. Let's define another gate calibration with a variable parameter. This time, we define a calibration for the `rz` gate on qubit `$1` and a variable angle `theta`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "49d9155b",
   "metadata": {},
   "outputs": [],
   "source": [
    "q1_control_frame = device.frames[\"Transmon_1_charge_tx\"]\n",
    "\n",
    "\n",
    "@aq.gate_calibration(implements=rz, target=\"$1\")\n",
    "def my_rz_cal(theta: float):\n",
    "    pulse.barrier(\"$1\")\n",
    "    pulse.shift_frequency(q1_control_frame, -1.0 * theta)\n",
    "    pulse.barrier(\"$1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6690a735",
   "metadata": {},
   "source": [
    "We then attach this gate calibration, `my_rz_cal`, to the main program we composed previously, together with the other gate calibration, `my_rx_cal`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a40c95c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OPENQASM 3.0;\n",
      "cal {\n",
      "    waveform wf_drag_gaussian_1 = drag_gaussian(24.0ns, 2.547965400864ns, 2.370235498840002e-10, 0.293350447987059, false);\n",
      "}\n",
      "defcal rx(1.5707963267948966) $0 {\n",
      "    barrier $0;\n",
      "    shift_frequency(Transmon_0_charge_tx, -321147.14178613486);\n",
      "    play(Transmon_0_charge_tx, wf_drag_gaussian_1);\n",
      "    shift_frequency(Transmon_0_charge_tx, 321147.14178613486);\n",
      "    barrier $0;\n",
      "}\n",
      "defcal rz(angle[32] theta) $1 {\n",
      "    barrier $1;\n",
      "    shift_frequency(Transmon_1_charge_tx, -1.0 * theta);\n",
      "    barrier $1;\n",
      "}\n",
      "rx(1.5707963267948966) $0;\n",
      "rz(0.123) $1;\n",
      "bit __bit_0__;\n",
      "__bit_0__ = measure $0;\n"
     ]
    }
   ],
   "source": [
    "custom_program = my_program.build().with_calibrations([my_rx_cal, my_rz_cal])\n",
    "print(custom_program.to_ir())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10d4cc3a",
   "metadata": {},
   "source": [
    "In the printed OpenQASM script of the program, the gate definitions for the `rx` gate and the `rz` become `defcal` blocks of the program. The variable parameter `angle` for the `rz` gate is captured."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cbc7366",
   "metadata": {},
   "source": [
    "## Summary\n",
    "This example notebook shows you how to retrieve and view the gate calibrations supplied by hardware providers. With AutoQASM, you can customize gate calibrations to explore different pulse implementations of gates. You can define gate calibrations against fixed values of gate parameters, as well as using parametric definitions for variable parameters. Multiple gate calibrations can be attached to a program to create a new program. This makes it easy to experiment with different gate calibrations on a single program."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}