{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1000001",
   "metadata": {},
   "source": [
    "# Displacement Calibration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000002",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "This guide assumes you have configured a `DeviceSetup` and created `BosonicQubit` objects with assigned parameters. Please see [our tutorials](https://docs.zhinst.com/labone_q_user_manual/applications_library/tutorials/index.html) if you need to create your setup and qubits for the first time.\n",
    "\n",
    "You can run this notebook on real hardware in the lab. If you don't have the hardware at your disposal, you can also run the notebook \"as is\" using an emulated session (see below).\n",
    "\n",
    "Before running the displacement calibration experiment, the following calibrations must have been performed:\n",
    "\n",
    "- **Memory spectroscopy**: the memory cavity resonance frequency `memory_resonance_frequency` must be known (see the [Memory Spectroscopy how-to guide](01_memory_spectroscopy.ipynb)).\n",
    "- **Transmon spectroscopy**: the steady-state ancilla transmon resonance frequency `transmon_resonance_frequency_at_n0` must be known.\n",
    "- **Transmon Rabi calibration**: the steady-state transmon $\\pi$-pulse amplitude `selective_transmon_X180_amplitudes[0]` must be calibrated.\n",
    "- **Dispersive shift**: the dispersive shift `chi` between the memory cavity and the ancilla transmon must be measured.\n",
    "\n",
    "If you are just getting started with the LabOne Q Applications Library, please don't hesitate to reach out to us at info@zhinst.com."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000003",
   "metadata": {},
   "source": [
    "## Background\n",
    "\n",
    "In this how-to guide, you will calibrate the **displacement operation** on the memory cavity of a bosonic qubit using the `displacement_calibration` experiment workflow in the LabOne Q Applications Library.\n",
    "\n",
    "### The displacement operation\n",
    "\n",
    "Applying a resonant coherent drive to the memory cavity at frequency $f_\\mathrm{memory}$ displaces the cavity state from the vacuum $|0\\rangle$ to a coherent state $|\\beta\\rangle$, where $\\beta$ is a complex number whose magnitude $|\\beta|$ characterises the displacement in phase space. The mean photon number of the resulting coherent state is:\n",
    "\n",
    "$$\\bar{n} = |\\beta|^2$$\n",
    "\n",
    "The photon number distribution of a coherent state is Poisson:\n",
    "\n",
    "$$P(n) = \\frac{e^{-\\bar{n}}\\,\\bar{n}^n}{n!}$$\n",
    "\n",
    "### Calibration goal\n",
    "\n",
    "The goal of this experiment is to find the hardware drive amplitude $A$ that produces a displacement of $|\\beta| = 1$. This conversion factor,\n",
    "\n",
    "$$A = \\texttt{displacement\\_amp\\_per\\_unit\\_beta} \\times |\\beta|$$\n",
    "\n",
    "is stored in the qubit parameter `displacement_amp_per_unit_beta`. Once calibrated, the `displacement` quantum operation can accept a target $|\\beta|$ (or mean photon number $\\bar{n}$) and compute the required hardware amplitude automatically.\n",
    "\n",
    "### Measuring the photon number distribution via number-resolved spectroscopy\n",
    "\n",
    "The photon number distribution is measured indirectly through the ancilla transmon using the dispersive interaction. Each photon in the memory shifts the transmon resonance frequency by the dispersive shift $\\chi$:\n",
    "\n",
    "$$f_\\mathrm{transmon}(n) = f_\\mathrm{transmon}(0) + n \\cdot \\chi$$\n",
    "\n",
    "The experiment performs a **2D sweep** over displacement amplitude $A$ and transmon drive frequency $f$:\n",
    "\n",
    "1. **Displacement pulse** — A drive tone is applied to the memory cavity at $f_\\mathrm{memory}$ with swept amplitude $A$, preparing the cavity in a coherent state $|\\beta(A)\\rangle$.\n",
    "2. **Transmon $\\pi$ pulse** — A $\\pi$ pulse is applied at swept RF frequency $f$, calibrated for the transmon when the memory is in $|0\\rangle$.\n",
    "3. **Readout** — The ancilla transmon is read out dispersively.\n",
    "\n",
    "When the drive frequency $f$ coincides with $f_\\mathrm{transmon}(n) = f_\\mathrm{transmon}(0) + n \\cdot \\chi$ for some photon number $n$, the $\\pi$ pulse selectively inverts the transmon conditioned on the memory containing $n$ photons, producing a peak in the readout signal. The heights of the peaks at $n = 0, 1, 2, \\ldots$ are proportional to the photon number probabilities $P(n)$.\n",
    "\n",
    "For each amplitude $A$, the extracted photon number distribution is fitted to a Poisson distribution to extract the mean photon number $\\bar{n}(A)$. A linear fit of $A$ vs. $\\sqrt{\\bar{n}}$ yields `displacement_amp_per_unit_beta`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000004",
   "metadata": {},
   "source": [
    "## Imports\n",
    "\n",
    "You'll start by importing `laboneq.simple`, as well as the `BosonicQubit` and `BosonicQubitOperations` classes from the bosonic qubits contribution."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000005",
   "metadata": {},
   "source": [
    "import copy\n",
    "\n",
    "import numpy as np\n",
    "from laboneq.simple import *\n",
    "\n",
    "from laboneq_applications.contrib.qpu_types.bosonic_qubits import (\n",
    "    BosonicQubit,\n",
    "    BosonicQubitOperations,\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000006",
   "metadata": {},
   "source": [
    "## Define your experimental setup\n",
    "\n",
    "Let's define our experimental setup. We will need:\n",
    "\n",
    "- a [DeviceSetup](https://docs.zhinst.com/labone_q_user_manual/core/functionality_and_concepts/00_device_setup/concepts/index.html)\n",
    "- one or more [BosonicQubits](https://docs.zhinst.com/labone_q_user_manual/applications_library/reference/contrib/index.html)\n",
    "- a set of `BosonicQubitOperations`\n",
    "- a [QPU](https://docs.zhinst.com/labone_q_user_manual/core/reference/dsl/quantum.html#laboneq.dsl.quantum.qpu.QPU)\n",
    "\n",
    "Here we will be brief and mainly provide the code to obtain these objects. For full details, see the tutorials on [defining your experimental setup](../../../tutorials/sources/getting_started.ipynb) and [logbooks and data saving](../../../tutorials/sources/logbooks.ipynb).\n",
    "\n",
    "We will use **2 bosonic qubits** in this guide. Change this number to match your setup."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000007",
   "metadata": {},
   "source": [
    "### Instrument serial numbers"
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000008",
   "metadata": {},
   "source": [
    "# Replace with the serial number of your SHFQC instrument\n",
    "SHFQC_SERIAL = \"DEV12073\"\n",
    "\n",
    "DATASERVER_HOST = \"localhost\"\n",
    "DATASERVER_PORT = \"8004\"\n",
    "\n",
    "number_of_qubits = 2"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000009",
   "metadata": {},
   "source": [
    "### DeviceSetup\n",
    "\n",
    "The displacement calibration experiment requires signal lines for the memory cavity drive, the ancilla transmon drive, and qubit readout. Here we use a single **SHFQC** instrument.\n",
    "\n",
    "Each bosonic qubit uses:\n",
    "- One **QA channel** pair (OUTPUT + INPUT) for readout\n",
    "- One **SG channel** for the memory cavity drive (`drive_memory`)\n",
    "- One **SG channel** for all number-selective transmon drives (`drive_transmon_at_n0` … `drive_transmon_at_n{max_photon_number}`)\n",
    "- Optionally, one **SG channel** in LF mode for the SWAP drive (`swap_drive`)\n",
    "\n",
    "If you already have a `DeviceSetup`, you can reuse it and skip this cell. Just make sure the signal line names match those expected by `BosonicQubit`."
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Maximum photon number — determines how many selective transmon signal lines are created.\n",
    "MAX_PHOTON_NUMBER = 5\n",
    "\n",
    "device_setup = DeviceSetup(uid=\"my_bosonic_setup\")\n",
    "device_setup.add_dataserver(\n",
    "    uid=\"my_bosonic_setup\",\n",
    "    host=DATASERVER_HOST,\n",
    "    port=DATASERVER_PORT,\n",
    ")\n",
    "\n",
    "device_setup.add_instruments(\n",
    "    SHFQC(\n",
    "        uid=\"device_shfqc\",\n",
    "        address=SHFQC_SERIAL,\n",
    "        device_options=\"SHFQC/LRT/PLUS/QC6CH/RTR\",\n",
    "    ),\n",
    ")\n",
    "\n",
    "# ── Signal lines for qubit q0 ─────────────────────────────────────────\n",
    "device_setup.add_connections(\n",
    "    \"device_shfqc\",\n",
    "    # Readout (QA channel 0)\n",
    "    create_connection(to_signal=\"q0/measure\", ports=\"QACHANNELS/0/OUTPUT\"),\n",
    "    create_connection(to_signal=\"q0/acquire\", ports=\"QACHANNELS/0/INPUT\"),\n",
    "    # Memory cavity drive (SG channel 0)\n",
    "    create_connection(to_signal=\"q0/drive_memory\", ports=\"SGCHANNELS/0/OUTPUT\"),\n",
    "    # SWAP drive — optional, used for Fock-state preparation (SG channel 4)\n",
    "    create_connection(to_signal=\"q0/swap_drive\", ports=\"SGCHANNELS/4/OUTPUT\"),\n",
    ")\n",
    "# Photon-number-selective transmon drive lines for q0 (all share SG channel 2)\n",
    "for n in range(MAX_PHOTON_NUMBER + 1):\n",
    "    device_setup.add_connections(\n",
    "        \"device_shfqc\",\n",
    "        create_connection(\n",
    "            to_signal=f\"q0/drive_transmon_at_n{n}\",\n",
    "            ports=\"SGCHANNELS/2/OUTPUT\",\n",
    "        ),\n",
    "    )\n",
    "\n",
    "# ── Signal lines for qubit q1 (optional) ─────────────────────────────\n",
    "if number_of_qubits == 2:\n",
    "    device_setup.add_connections(\n",
    "        \"device_shfqc\",\n",
    "        create_connection(to_signal=\"q1/measure\", ports=\"QACHANNELS/0/OUTPUT\"),\n",
    "        create_connection(to_signal=\"q1/acquire\", ports=\"QACHANNELS/0/INPUT\"),\n",
    "        create_connection(to_signal=\"q1/drive_memory\", ports=\"SGCHANNELS/1/OUTPUT\"),\n",
    "        create_connection(to_signal=\"q1/swap_drive\", ports=\"SGCHANNELS/5/OUTPUT\"),\n",
    "    )\n",
    "    for n in range(MAX_PHOTON_NUMBER + 1):\n",
    "        device_setup.add_connections(\n",
    "            \"device_shfqc\",\n",
    "            create_connection(\n",
    "                to_signal=f\"q1/drive_transmon_at_n{n}\",\n",
    "                ports=\"SGCHANNELS/3/OUTPUT\",\n",
    "            ),\n",
    "        )"
   ],
   "id": "c274381f0b32fb77"
  },
  {
   "cell_type": "markdown",
   "id": "a1000011",
   "metadata": {},
   "source": [
    "### BosonicQubits\n",
    "\n",
    "We will create `BosonicQubit` objects from the logical signal groups in our `DeviceSetup` and assign qubit parameters. The qubit UID will match the name of the logical signal group (`q0`, `q1`, …).\n",
    "\n",
    "Adjust the parameter values to reflect the properties of your device. The parameters relevant to displacement calibration are highlighted in the comments."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000012",
   "metadata": {},
   "source": [
    "# Qubit parameters — adjust these to match your device.\n",
    "# Parameters marked (*) are directly relevant to the displacement calibration experiment.\n",
    "QUBIT_PARAMS = {\n",
    "    # ── Memory cavity ────────────────────────────────────────────────\n",
    "    \"memory_lo_frequency\": 5.8e9,  # LO for memory cavity drive\n",
    "    \"memory_resonance_frequency\": 6.0e9,  # (*) Memory resonance frequency (set from memory spectroscopy)\n",
    "    \"memory_drive_amplitude\": 0.5,\n",
    "    \"memory_drive_length\": 0.1e-6,  # (*) Duration of the displacement pulse\n",
    "    \"memory_drive_pulse\": {\"function\": \"gaussian\", \"sigma\": 0.25},  # (*) Pulse shape\n",
    "    \"memory_spectroscopy_amplitude\": 0.5,\n",
    "    \"memory_spectroscopy_length\": 4e-6,\n",
    "    \"memory_drive_range\": 10,\n",
    "    # ── Ancilla transmon ─────────────────────────────────────────────\n",
    "    \"transmon_lo_frequency\": 4.8e9,\n",
    "    \"transmon_resonance_frequency_at_n0\": 5.0e9,  # (*) Transmon freq when memory in |0>\n",
    "    \"transmon_drive_range\": 10,\n",
    "    # ── Selective transmon Rx drives (*) ─────────────────────────────\n",
    "    # These are used for the pi pulse that probes the photon number state.\n",
    "    \"max_photon_number\": MAX_PHOTON_NUMBER,\n",
    "    \"selective_transmon_Rx_length\": 0.08e-6,\n",
    "    \"selective_transmon_Rx_pulse\": {\"function\": \"gaussian\", \"sigma\": 0.25},\n",
    "    \"selective_transmon_X180_amplitudes\": [0.8]\n",
    "    * (MAX_PHOTON_NUMBER + 1),  # (*) pi amplitudes\n",
    "    \"selective_transmon_X90_amplitudes\": [0.4] * (MAX_PHOTON_NUMBER + 1),\n",
    "    # ── SWAP drive ────────────────────────────────────────────────────\n",
    "    # This setup drives the SWAP line from an SG channel, which is put into\n",
    "    # LF mode with its local oscillator at 0 Hz. Leave this out for a SWAP\n",
    "    # line on a baseband channel, such as an HDAWG output.\n",
    "    \"swap_drive_on_sg_channel\": True,\n",
    "    # ── Readout ───────────────────────────────────────────────────────\n",
    "    \"readout_lo_frequency\": 6.8e9,\n",
    "    \"readout_resonator_frequency\": 7.0e9,\n",
    "    \"readout_amplitude\": 0.5,\n",
    "    \"readout_length\": 1e-6,\n",
    "    \"readout_pulse\": {\"function\": \"const\"},\n",
    "    \"readout_integration_length\": 1e-6,\n",
    "    \"readout_integration_delay\": 20e-9,\n",
    "    # ── Reset ─────────────────────────────────────────────────────────\n",
    "    \"reset_delay_length\": 1e-6,\n",
    "    # ── Memory-transmon coupling (*) ──────────────────────────────────\n",
    "    \"chi\": -1e6,  # (*) Dispersive shift in Hz\n",
    "    # ── Displacement calibration result (*) ───────────────────────────\n",
    "    \"displacement_amp_per_unit_beta\": None,  # (*) To be calibrated — output of this experiment\n",
    "}"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "a1000013",
   "metadata": {},
   "source": [
    "q0 = BosonicQubit.from_logical_signal_group(\n",
    "    uid=\"q0\",\n",
    "    lsg=device_setup.logical_signal_groups[\"q0\"],\n",
    "    parameters=copy.deepcopy(QUBIT_PARAMS),\n",
    ")\n",
    "qubits = [q0]\n",
    "\n",
    "if number_of_qubits == 2:\n",
    "    q1 = BosonicQubit.from_logical_signal_group(\n",
    "        uid=\"q1\",\n",
    "        lsg=device_setup.logical_signal_groups[\"q1\"],\n",
    "        parameters=copy.deepcopy(QUBIT_PARAMS),\n",
    "    )\n",
    "    qubits.append(q1)\n",
    "\n",
    "for q in qubits:\n",
    "    print(\"-------------\")\n",
    "    print(\"Qubit UID:\", q.uid)\n",
    "    print(\"Qubit logical signals:\")\n",
    "    for sig, lsg in q.signals.items():\n",
    "        print(f\"  {sig:<30} ('{lsg}')\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000014",
   "metadata": {},
   "source": [
    "### Quantum Operations\n",
    "\n",
    "Create the set of `BosonicQubitOperations`:"
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000015",
   "metadata": {},
   "source": [
    "qops = BosonicQubitOperations()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000016",
   "metadata": {},
   "source": [
    "### QPU\n",
    "\n",
    "Create the `QPU` object from the qubits and the quantum operations:"
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000017",
   "metadata": {},
   "source": [
    "from laboneq.dsl.quantum import QPU\n",
    "\n",
    "qpu = QPU(qubits, quantum_operations=qops)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000018",
   "metadata": {},
   "source": [
    "## Connect to Session"
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000019",
   "metadata": {},
   "source": [
    "session = Session(device_setup)\n",
    "session.connect(do_emulation=True)  # set do_emulation=False when at a real setup"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000020",
   "metadata": {},
   "source": [
    "## Create a `FolderStore` for Saving Data\n",
    "\n",
    "The experiment `Workflows` can automatically save the inputs and outputs of all their tasks to the folder path we specify when instantiating the `FolderStore`. Here, we choose the current working directory."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000021",
   "metadata": {},
   "source": [
    "from pathlib import Path\n",
    "\n",
    "folder_store = workflow.logbook.FolderStore(Path.cwd() / \"data\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "43a7c46b",
   "metadata": {},
   "source": [
    "We disable saving in this guide. To enable it, simply run `folder_store.activate()`."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000022",
   "metadata": {},
   "source": [
    "folder_store.deactivate()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000023",
   "metadata": {},
   "source": [
    "## Optional: Configure the LoggingStore\n",
    "\n",
    "You can also activate/deactivate the `LoggingStore`, which controls the display of `Workflow` logging information in the notebook. See the [tutorial on Recording Experiment Workflow Results](../../../tutorials/sources/logbooks.ipynb) for details.\n",
    "\n",
    "Displaying the `Workflow` logging information is activated by default, but here we deactivate it to shorten the outputs, which are not very meaningful in emulation mode.\n",
    "\n",
    "**We recommend that you do not deactivate the Workflow logging in practice.**"
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000024",
   "metadata": {},
   "source": [
    "from laboneq.workflow.logbook import LoggingStore\n",
    "\n",
    "logging_store = LoggingStore()\n",
    "logging_store.deactivate()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000025",
   "metadata": {},
   "source": [
    "## Running the Experiment Workflow\n",
    "\n",
    "You'll now instantiate the experiment workflow and run it. For more details on what experiment workflows are and what tasks they execute, see the [Experiment Workflows tutorial](../../../tutorials/sources/experiment_workflows.ipynb).\n",
    "\n",
    "Start by importing the `displacement_calibration` experiment workflow from the bosonic qubits contribution, as well as `plot_simulation` for inspecting the pulse sequence."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000026",
   "metadata": {},
   "source": [
    "from laboneq.contrib.example_helpers.plotting.plot_helpers import plot_simulation\n",
    "\n",
    "from laboneq_applications.contrib.experiments.bosonic_qubits import (\n",
    "    displacement_calibration,\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000027",
   "metadata": {},
   "source": [
    "### Configure the experiment options\n",
    "\n",
    "Let's create the options object for the displacement calibration experiment and inspect it using `workflow.show_fields`:"
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000028",
   "metadata": {},
   "source": [
    "options = displacement_calibration.experiment_workflow.options()\n",
    "workflow.show_fields(options)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000029",
   "metadata": {},
   "source": [
    "Notice that, unless we change it:\n",
    "\n",
    "- The experiment is run in `AcquisitionType.INTEGRATION` and `AveragingMode.CYCLIC`, using 1024 averages (`count`).\n",
    "- The analysis workflow runs automatically after the experiment (`do_analysis=True`). It extracts the photon number distributions, fits them to Poisson distributions, and returns the calibrated `displacement_amp_per_unit_beta` for each qubit.\n",
    "- The figures produced by the analysis are automatically closed (`close_figures=True`)\n",
    "- The qubit parameters will not be updated (`update=False`).\n",
    "\n",
    "The `create_experiment` sub-options expose additional pulse overrides for advanced use:\n",
    "- `memory_drive_length`: override for the displacement pulse duration (defaults to `qubit.parameters.memory_drive_length`).\n",
    "- `memory_drive_pulse`: pulse shape override for the displacement pulse.\n",
    "- `transmon_pi_pulse`: pulse shape override for the transmon $\\pi$ pulse.\n",
    "\n",
    "Here, let's reduce the number of averages for a quicker run in emulation mode, and disable closing the figures produced by the analysis so we see them in the cell output. Note that the fits attempted by the analysis in emulation mode are not representative, because we do not acquire data from a real experiment."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000030",
   "metadata": {},
   "source": [
    "options.count(1)\n",
    "options.close_figures(False)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000031",
   "metadata": {},
   "source": [
    "### Define the displacement amplitude and frequency sweeps\n",
    "\n",
    "The displacement calibration requires two sweep axes:\n",
    "\n",
    "1. **Amplitude sweep** — the hardware drive amplitude of the displacement pulse on the memory cavity. Start from 0 and increase to a value where the mean photon number reaches a few photons (typically $\\bar{n} \\approx 3$–$5$).\n",
    "2. **Frequency sweep** — the RF frequency of the transmon $\\pi$ pulse. The sweep must cover the photon-number-resolved spectroscopy peaks at $f_\\mathrm{transmon}(n) = f_\\mathrm{transmon}(0) + n \\cdot \\chi$ for $n = 0, 1, \\ldots, N_\\mathrm{max}$, with sufficient margin on either side.\n",
    "\n",
    "Both arrays must have the **same number of points** for all qubits when running in parallel (since the hardware sweeps them in lock-step). However, each qubit can have an independently chosen sweep range."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000032",
   "metadata": {},
   "source": "qubit = qpu.quantum_elements[0]\n\nf0 = qubit.parameters.transmon_resonance_frequency_at_n0\nchi = qubit.parameters.chi\n\n# Few sweep points keep the emulated run short; use denser sweeps on hardware\n# (see the parallel example below).\n\n# Displacement amplitude sweep\namplitudes_q0 = np.linspace(0, 0.8, 2)\n\n# Transmon drive frequency sweep (RF frequencies in Hz).\n# Cover the photon-number peaks at f(n) = f0 + n * chi for n = 0 ... MAX_PHOTON_NUMBER + 1,\n# with a margin of |chi| on each side. This assumes chi < 0, so the peaks lie below f0.\n# The analysis needs more frequency points than peaks in the sweep; a step of |chi| / 2\n# puts a point on every peak.\nfrequencies_q0 = np.linspace(\n    f0 + (MAX_PHOTON_NUMBER + 1) * chi - abs(chi),\n    f0 + abs(chi),\n    2 * (MAX_PHOTON_NUMBER + 3) + 1,\n)\n\nprint(\n    f\"Amplitude sweep: {amplitudes_q0[0]:.3f} to {amplitudes_q0[-1]:.3f} ({len(amplitudes_q0)} points)\"\n)\nprint(\n    f\"Frequency sweep: {frequencies_q0[0] / 1e9:.4f} to {frequencies_q0[-1] / 1e9:.4f} GHz ({len(frequencies_q0)} points)\"\n)\nprint(\n    f\"Expected peaks at (GHz): {[f'{(f0 + n * chi) / 1e9:.4f}' for n in range(MAX_PHOTON_NUMBER + 1)]}\"\n)",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "a1000033",
   "metadata": {},
   "source": [
    "# Run on a single qubit\n",
    "exp_workflow = displacement_calibration.experiment_workflow(\n",
    "    session=session,\n",
    "    qpu=qpu,\n",
    "    qubits=[qubit.uid],\n",
    "    amplitudes=[amplitudes_q0],\n",
    "    frequencies=[frequencies_q0],\n",
    "    options=options,\n",
    ")\n",
    "\n",
    "workflow_results = exp_workflow.run()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000034",
   "metadata": {},
   "source": [
    "### Inspect the Tasks That Were Run\n",
    "\n",
    "The displacement calibration workflow executes four steps in sequence: `create_experiment`, `compile_experiment`, `run_experiment`, and (when `do_analysis=True`) the `analysis_workflow`, which runs the `fit_data`, `extract_qubit_parameters` and plotting tasks. With `update=True`, an `update_qpu` task then stores the new parameter values in the `qpu`."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000035",
   "metadata": {},
   "source": [
    "for t in workflow_results.tasks:\n",
    "    print(t)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000036",
   "metadata": {},
   "source": [
    "### Inspect the Output Simulation\n",
    "\n",
    "You can inspect the compiled experiment and plot the simulated output to verify that the pulse sequence looks as expected. The displacement calibration sequence consists of a short displacement pulse on the memory cavity drive line, followed by a transmon $\\pi$ pulse and a readout pulse.\n",
    "\n",
    "The total duration of one repetition is approximately:\n",
    "\n",
    "$$T \\approx t_\\mathrm{displacement} + t_{\\pi} + t_\\mathrm{readout} + t_\\mathrm{reset}$$\n",
    "\n",
    "Note that the full experiment sweeps over `len(amplitudes) × len(frequencies)` repetitions. We plot just the first repetition here to inspect the pulse shapes."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000037",
   "metadata": {},
   "source": [
    "q = qpu.quantum_elements[0]\n",
    "total_time = (\n",
    "    q.parameters.memory_drive_length\n",
    "    + q.parameters.selective_transmon_Rx_length\n",
    "    + q.parameters.readout_length\n",
    "    + q.parameters.reset_delay_length\n",
    ")\n",
    "\n",
    "compiled_experiment = workflow_results.tasks[\"compile_experiment\"].output\n",
    "plot_simulation(\n",
    "    compiled_experiment,\n",
    "    length=total_time * 1.05,\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000038",
   "metadata": {},
   "source": [
    "### Inspect the Source Code of the Pulse-Sequence Creation Task\n",
    "\n",
    "You can inspect the source code of the `create_experiment` task to see exactly how the pulse sequence is constructed using quantum operations. To learn more about quantum operations, see the [Quantum Operations tutorial](../../../tutorials/sources/quantum_operations.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000039",
   "metadata": {},
   "source": [
    "displacement_calibration.create_experiment.src"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000040",
   "metadata": {},
   "source": [
    "## Analyse the Results\n",
    "\n",
    "Because `do_analysis=True` by default, the analysis workflow ran automatically after the experiment. It performed the following steps for each qubit:\n",
    "\n",
    "1. Reshaped the raw 1D result data into a 2D array of shape `(n_amplitudes, n_frequencies)`.\n",
    "2. For each amplitude slice, extracted the heights of the photon-number peaks at $f_\\mathrm{transmon}(n) = f_\\mathrm{transmon}(0) + n \\cdot \\chi$ for $n = 0, 1, \\ldots$\n",
    "3. Fitted the peak heights to a Poisson distribution to extract the mean photon number $\\bar{n}$.\n",
    "4. Fitted $A$ vs. $\\sqrt{\\bar{n}}$ with a line through the origin to extract `displacement_amp_per_unit_beta`.\n",
    "\n",
    "Three plots are generated automatically:\n",
    "- **2D displacement map** — the transmon spectrum magnitude as a function of displacement amplitude and transmon drive frequency, with dashed lines indicating the expected photon-number peak positions.\n",
    "- **Photon number distributions** — the extracted photon number distributions $P(n)$ with Poisson fits for a selection of amplitude values.\n",
    "- **$\\bar{n}$ vs. amplitude** — the fitted and directly computed mean photon numbers as a function of displacement amplitude, showing the $\\bar{n} \\propto A^2$ relationship."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1000041",
   "metadata": {},
   "source": [
    "### Access the Analysis Outputs"
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000042",
   "metadata": {},
   "source": [
    "# Access the fit results of the analysis workflow.\n",
    "analysis_result = workflow_results.tasks[\"analysis_workflow\"].tasks[\"fit_data\"].output\n",
    "\n",
    "q_uid = qpu.quantum_elements[0].uid\n",
    "q_result = analysis_result[q_uid]\n",
    "\n",
    "# Keys available in the per-qubit result dict:\n",
    "print(\"Available result keys:\", list(q_result.keys()))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000045",
   "metadata": {},
   "source": [
    "### Extract the Calibration Constant\n",
    "\n",
    "The analysis fits the relationship $A = \\texttt{displacement\\_amp\\_per\\_unit\\_beta} \\times |\\beta|$ where $|\\beta| = \\sqrt{\\bar{n}}$. The fitted slope is stored in the result dict:"
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000046",
   "metadata": {},
   "source": [
    "disp_cal = q_result[\"displacement_amp_per_unit_beta\"]\n",
    "print(f\"displacement_amp_per_unit_beta for {q_uid}: {disp_cal}\")\n",
    "\n",
    "if np.isfinite(disp_cal):\n",
    "    print(\"\\nExample conversions:\")\n",
    "    print(f\"  |beta| = 1.0  (n_bar = 1.0) -> amplitude = {disp_cal * 1.0:.4f}\")\n",
    "    print(f\"  |beta| = 2.0  (n_bar = 4.0) -> amplitude = {disp_cal * 2.0:.4f}\")\n",
    "    print(f\"  |beta| = sqrt(5) (n_bar = 5) -> amplitude = {disp_cal * 5**0.5:.4f}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000047",
   "metadata": {},
   "source": [
    "### Update the Qubit Parameters\n",
    "\n",
    "The analysis workflow returns the extracted qubit parameters together with their previous values. Set `options.update(True)` to store the new values in the `qpu` automatically when the workflow runs. Here, we update the `qpu` manually instead. Qubits without a valid fit have no new values and keep their previous parameters.\n",
    "\n",
    "Once `displacement_amp_per_unit_beta` is set, the `displacement` quantum operation can accept `beta` or `n_bar` directly as arguments, and will compute the required hardware amplitude automatically."
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000048",
   "metadata": {},
   "source": [
    "from pprint import pprint\n",
    "\n",
    "qubit_parameters = workflow_results.tasks[\"analysis_workflow\"].output\n",
    "pprint(qubit_parameters)\n",
    "\n",
    "qpu.update(qubit_parameters[\"new_parameter_values\"])"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000049",
   "metadata": {},
   "source": [
    "## Parallel Displacement Calibration\n",
    "\n",
    "The `displacement_calibration` workflow natively supports running multiple qubits in parallel. When you pass a list of qubit UIDs together with corresponding lists of amplitude and frequency arrays, the pulse sequences for all qubits are compiled into a single experiment and executed simultaneously on the hardware — halving experiment time compared to sequential single-qubit runs.\n",
    "\n",
    "Each qubit gets its own independently chosen amplitude and frequency sweep, so qubits with different transmon frequencies or dispersive shifts can still be characterised in one shot. The only constraint is that all amplitude arrays must have the **same number of points**, and all frequency arrays must have the **same number of points**, because the hardware sweeps them in lock-step.\n",
    "\n",
    "### Run the two-qubit experiment\n",
    "\n",
    "Pass both qubit UIDs and their sweep arrays to the workflow:"
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000050",
   "metadata": {},
   "source": [
    "# Store the calibrated values in the qpu automatically.\n",
    "options.update(True)\n",
    "\n",
    "q0 = qpu.quantum_elements[0]\n",
    "q1 = qpu.quantum_elements[1]\n",
    "\n",
    "# Each qubit gets its own sweep, independently centred on its transmon frequency.\n",
    "# As above, the frequency ranges assume chi < 0.\n",
    "f0_q0 = q0.parameters.transmon_resonance_frequency_at_n0\n",
    "chi_q0 = q0.parameters.chi\n",
    "f0_q1 = q1.parameters.transmon_resonance_frequency_at_n0\n",
    "chi_q1 = q1.parameters.chi\n",
    "\n",
    "amplitudes_q0 = np.linspace(0, 0.8, 21)\n",
    "amplitudes_q1 = np.linspace(0, 0.8, 21)\n",
    "\n",
    "frequencies_q0 = np.linspace(\n",
    "    f0_q0 + (MAX_PHOTON_NUMBER + 1) * chi_q0 - abs(chi_q0),\n",
    "    f0_q0 + abs(chi_q0),\n",
    "    501,\n",
    ")\n",
    "frequencies_q1 = np.linspace(\n",
    "    f0_q1 + (MAX_PHOTON_NUMBER + 1) * chi_q1 - abs(chi_q1),\n",
    "    f0_q1 + abs(chi_q1),\n",
    "    501,\n",
    ")\n",
    "\n",
    "exp_workflow_2q = displacement_calibration.experiment_workflow(\n",
    "    session=session,\n",
    "    qpu=qpu,\n",
    "    qubits=[q0.uid, q1.uid],\n",
    "    amplitudes=[amplitudes_q0, amplitudes_q1],\n",
    "    frequencies=[frequencies_q0, frequencies_q1],\n",
    "    options=options,\n",
    ")\n",
    "\n",
    "workflow_results_2q = exp_workflow_2q.run()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000051",
   "metadata": {},
   "source": [
    "### Check the Updated Qubit Parameters\n",
    "\n",
    "With `update=True`, the workflow stored the new parameter values in the `qpu`. Qubits without a valid fit keep their previous parameters:"
   ]
  },
  {
   "cell_type": "code",
   "id": "a1000052",
   "metadata": {},
   "source": [
    "qubit_parameters_2q = workflow_results_2q.tasks[\"analysis_workflow\"].output\n",
    "pprint(qubit_parameters_2q[\"new_parameter_values\"])\n",
    "\n",
    "for q in [q0, q1]:\n",
    "    print(\n",
    "        f\"{q.uid}: displacement_amp_per_unit_beta = {q.parameters.displacement_amp_per_unit_beta}\"\n",
    "    )"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a1000053",
   "metadata": {},
   "source": [
    "Great! You have now calibrated the displacement amplitude for all qubits. The `displacement_amp_per_unit_beta` parameter is stored in the qubit objects and will be used by all subsequent experiments that perform displacement operations on the memory cavity — including state preparation for SWAP calibration, SNAP gate calibration, and Wigner tomography.\n",
    "\n",
    "The next step in the bosonic qubit tune-up sequence is the [**Selective Transmon Rx calibration**](03_selective_transmon_rx_calibration.ipynb), which calibrates the number-selective transmon $\\pi$ and $\\pi/2$ pulse amplitudes."
   ]
  }
 ],
 "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.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
