{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cc2a6468",
   "metadata": {},
   "source": [
    "# NV Center Qubit\n",
    "\n",
    "## Setup\n",
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b78c1fda",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from typing import TYPE_CHECKING, Callable\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from qcodes.instrument.parameter import ManualParameter\n",
    "\n",
    "if TYPE_CHECKING:\n",
    "    from qblox_instruments.qcodes_drivers.module import Module\n",
    "\n",
    "from qcodes.instrument import find_or_create_instrument\n",
    "\n",
    "import quantify_core.visualization.pyqt_plotmon as pqm\n",
    "from qblox_instruments import Cluster, ClusterType\n",
    "from quantify_core.data import handling as dh\n",
    "from quantify_core.measurement.control import MeasurementControl\n",
    "from quantify_core.visualization.instrument_monitor import InstrumentMonitor\n",
    "from quantify_scheduler.backends import SerialCompiler\n",
    "from quantify_scheduler.device_under_test.quantum_device import QuantumDevice\n",
    "from quantify_scheduler.gettables import ScheduleGettable\n",
    "from quantify_scheduler.instrument_coordinator import InstrumentCoordinator\n",
    "from quantify_scheduler.instrument_coordinator.components.generic import (\n",
    "    GenericInstrumentCoordinatorComponent,\n",
    ")\n",
    "from quantify_scheduler.instrument_coordinator.components.qblox import ClusterComponent\n",
    "from quantify_scheduler.schedules.schedule import Schedule"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55c9ee46",
   "metadata": {},
   "source": [
    "### Scan For Clusters\n",
    "\n",
    "We scan for the available devices connected via ethernet using the Plug & Play functionality of the Qblox Instruments package (see [Plug & Play](https://qblox-qblox-instruments.readthedocs-hosted.com/en/main/api_reference/tools.html#api-pnp) for more info)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86a0887c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!qblox-pnp list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b90acb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_ip = None\n",
    "cluster_name = \"cluster0\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d066833e",
   "metadata": {},
   "source": [
    "### Connect to Cluster\n",
    "\n",
    "We now make a connection with the Cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0ed0c29",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "cluster = find_or_create_instrument(\n",
    "    Cluster,\n",
    "    recreate=True,\n",
    "    name=cluster_name,\n",
    "    identifier=cluster_ip,\n",
    "    dummy_cfg=(\n",
    "        {\n",
    "            2: ClusterType.CLUSTER_QCM,\n",
    "            4: ClusterType.CLUSTER_QRM,\n",
    "            6: ClusterType.CLUSTER_QCM_RF,\n",
    "            8: ClusterType.CLUSTER_QRM_RF,\n",
    "        }\n",
    "        if cluster_ip is None\n",
    "        else None\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "825bd0d2",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "#### Get connected modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "054dd916",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_connected_modules(cluster: Cluster, filter_fn: Callable | None = None) -> dict[int, Module]:\n",
    "    def checked_filter_fn(mod: ClusterType) -> bool:\n",
    "        if filter_fn is not None:\n",
    "            return filter_fn(mod)\n",
    "        return True\n",
    "\n",
    "    return {\n",
    "        mod.slot_idx: mod for mod in cluster.modules if mod.present() and checked_filter_fn(mod)\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f3304a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# QRM baseband modules\n",
    "readout_modules = get_connected_modules(cluster, lambda mod: mod.is_qrm_type and not mod.is_rf_type)\n",
    "readout_modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74a035cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "readout_module = readout_modules[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6467eed4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# QCM baseband modules\n",
    "control_modules = get_connected_modules(\n",
    "    cluster, lambda mod: not mod.is_qrm_type and not mod.is_rf_type\n",
    ")\n",
    "control_modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4a1e9c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "control_module = control_modules[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "131f248b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# QCM RF modules\n",
    "rf_control_modules = get_connected_modules(\n",
    "    cluster, lambda mod: not mod.is_qrm_type and mod.is_rf_type\n",
    ")\n",
    "rf_control_modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a09a462",
   "metadata": {},
   "outputs": [],
   "source": [
    "rf_control_module = rf_control_modules[6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7548880f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from quantify_scheduler.helpers.mock_instruments import MockLocalOscillator\n",
    "\n",
    "lo_readout_laser = find_or_create_instrument(\n",
    "    MockLocalOscillator, recreate=True, name=\"lo_readout_laser\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d927fe4c",
   "metadata": {},
   "source": [
    "### Connect to the Cluster in Quantify\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21265830",
   "metadata": {},
   "outputs": [],
   "source": [
    "dh.set_datadir(dh.default_datadir())\n",
    "\n",
    "meas_ctrl = find_or_create_instrument(MeasurementControl, recreate=True, name=\"meas_ctrl\")\n",
    "# Create the live plotting isntrument which handles the graphical interface\n",
    "# Two windows will be created, the main will feature 1D plots and any 2D plots will go\n",
    "# to the secondary\n",
    "\n",
    "plotmon = find_or_create_instrument(pqm.PlotMonitor_pyqt, recreate=True, name=\"plotmon\")\n",
    "# Connect the live plotting monitor to the measurement control\n",
    "meas_ctrl.instr_plotmon(plotmon.name)\n",
    "\n",
    "# # The instrument monitor will give an overview of all parameters of all instruments\n",
    "insmon = find_or_create_instrument(InstrumentMonitor, recreate=True, name=\"Instrument_Monitor\")\n",
    "\n",
    "ic = find_or_create_instrument(InstrumentCoordinator, recreate=True, name=\"ic\")\n",
    "ic.add_component(ClusterComponent(cluster))\n",
    "ic.add_component(GenericInstrumentCoordinatorComponent(lo_readout_laser))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afb8779c",
   "metadata": {},
   "source": [
    "### Quantum Device Configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11ab3bbb",
   "metadata": {},
   "source": [
    "To generate the device configuration on the fly, `quantify_scheduler` provides the `QuantumDevice` and `BasicElectronicNVElement` classes, with the latter being specifically for NV-Center experiments. `QuantumDevice` can contain multiple `BasicElectronicNVElement`, but in this case we will add a single element to it, `qe0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9c26835",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "from quantify_scheduler.device_under_test.nv_element import BasicElectronicNVElement\n",
    "\n",
    "nv_center = find_or_create_instrument(QuantumDevice, recreate=True, name=\"nv_center\")\n",
    "\n",
    "nv_center.instr_measurement_control(meas_ctrl.name)\n",
    "nv_center.instr_instrument_coordinator(ic.name)\n",
    "nv_center.cfg_sched_repetitions(1)\n",
    "\n",
    "qe0 = find_or_create_instrument(BasicElectronicNVElement, recreate=True, name=\"qe0\")\n",
    "nv_center.add_element(qe0)\n",
    "\n",
    "qe0.clock_freqs.ge0.set(470.4e12)  # 637 nm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7187ea9b",
   "metadata": {},
   "source": [
    "### Hardware Configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f1d270b",
   "metadata": {},
   "source": [
    "We define the hardware configuration for the experiment. Three different Cluster modules are used:\n",
    "- A readout module which receives triggers from the single photon detector as input. The triggers are counted if their amplitude passes a certain threshold `ttl_acq_threshold` defined in the hardware configuration below. This module is also used to drive the AOM connected to the readout laser through Output 1, at 200 MHz;\n",
    "- A control module which drives the AOMs for the spinpump and green lasers, at 200 MHz;\n",
    "- An RF control module which is used for spectroscopy operations - the RF current is injected into the stripline which generates an alternating field at the NV center."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a614be3b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "[![image1](figures/NVCenterSetup.svg)](figures/NVCenterSetup.svg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33f07bac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_hardware_config(ttl_acq_threshold: float) -> dict:\n",
    "    hardware_cfg = {\n",
    "        \"backend\": \"quantify_scheduler.backends.qblox_backend.hardware_compile\",\n",
    "        f\"{cluster.name}\": {\n",
    "            \"ref\": \"internal\",\n",
    "            \"sequence_to_file\": \"true\",\n",
    "            \"instrument_type\": \"Cluster\",\n",
    "            f\"{readout_module.name}\": {\n",
    "                \"instrument_type\": \"QRM\",\n",
    "                \"real_input_0\": {\n",
    "                    \"portclock_configs\": [\n",
    "                        {\n",
    "                            \"port\": \"qe0:optical_readout\",\n",
    "                            \"clock\": \"qe0.ge0\",\n",
    "                            \"interm_freq\": 0,\n",
    "                            \"ttl_acq_threshold\": ttl_acq_threshold,\n",
    "                        },\n",
    "                    ],\n",
    "                },\n",
    "                \"real_output_0\": {\n",
    "                    \"lo_name\": \"lo_readout_laser\",\n",
    "                    \"mix_lo\": False,\n",
    "                    \"portclock_configs\": [\n",
    "                        {\"port\": \"qe0:optical_control\", \"clock\": \"qe0.ge0\", \"interm_freq\": 200e6},\n",
    "                    ],\n",
    "                },\n",
    "            },\n",
    "            f\"{control_module.name}\": {\n",
    "                \"instrument_type\": \"QCM\",\n",
    "                \"real_output_0\": {\n",
    "                    \"lo_name\": \"lo_spinpump_laser\",\n",
    "                    \"mix_lo\": False,\n",
    "                    \"portclock_configs\": [\n",
    "                        {\"port\": \"qe0:optical_control\", \"clock\": \"qe0.ge1\", \"interm_freq\": 200e6}\n",
    "                    ],\n",
    "                },\n",
    "                \"real_output_1\": {\n",
    "                    \"lo_name\": \"lo_green_laser\",\n",
    "                    \"mix_lo\": False,\n",
    "                    \"portclock_configs\": [\n",
    "                        {\n",
    "                            \"port\": \"qe0:optical_control\",\n",
    "                            \"clock\": \"qe0.ionization\",\n",
    "                            \"interm_freq\": 200e6,\n",
    "                        }\n",
    "                    ],\n",
    "                },\n",
    "            },\n",
    "            f\"{rf_control_module.name}\": {\n",
    "                \"instrument_type\": \"QCM_RF\",\n",
    "                \"complex_output_0\": {\n",
    "                    \"lo_freq\": None,\n",
    "                    \"dc_mixer_offset_I\": 0.0,\n",
    "                    \"dc_mixer_offset_Q\": 0.0,\n",
    "                    \"portclock_configs\": [\n",
    "                        {\n",
    "                            \"port\": \"qe0:mw\",\n",
    "                            \"clock\": \"qe0.spec\",\n",
    "                            \"interm_freq\": 200e6,\n",
    "                            \"mixer_amp_ratio\": 0.9999,\n",
    "                            \"mixer_phase_error_deg\": -4.2,\n",
    "                        }\n",
    "                    ],\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "        \"lo_readout_laser\": {\"instrument_type\": \"LocalOscillator\", \"frequency\": None, \"power\": 1},\n",
    "        \"lo_spinpump_laser\": {\"instrument_type\": \"LocalOscillator\", \"frequency\": None, \"power\": 1},\n",
    "        \"lo_green_laser\": {\"instrument_type\": \"LocalOscillator\", \"frequency\": None, \"power\": 1},\n",
    "    }\n",
    "    return hardware_cfg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dcfd409",
   "metadata": {},
   "outputs": [],
   "source": [
    "hardware_cfg = set_hardware_config(ttl_acq_threshold=0.2)\n",
    "nv_center.hardware_config.set(hardware_cfg)\n",
    "\n",
    "readout_module.scope_acq_sequencer_select(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83783669",
   "metadata": {},
   "source": [
    "## Trace Acquisition\n",
    "\n",
    "We start by running a trace acquisition to establish the threshold voltage at which we count a photon detection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0003418",
   "metadata": {},
   "outputs": [],
   "source": [
    "from quantify_scheduler.enums import BinMode\n",
    "from quantify_scheduler.operations.gate_library import Measure\n",
    "\n",
    "qe0.measure.pulse_amplitude(0.5)\n",
    "qe0.measure.pulse_duration(16e-6)\n",
    "qe0.measure.acq_duration(16e-6)\n",
    "qe0.measure.acq_delay(0)\n",
    "\n",
    "\n",
    "def trace_acquisition_sched(qubit: str, repetitions: int = 1) -> Schedule:\n",
    "    schedule = Schedule(\"Trace Acquisition\")\n",
    "    schedule.add(Measure(\"qe0\", acq_protocol=\"Trace\"))\n",
    "    return schedule\n",
    "\n",
    "\n",
    "trace_schedule = trace_acquisition_sched(\"qe0\")\n",
    "compiler = SerialCompiler(\"compiler\")\n",
    "compiled_trace_sched = compiler.compile(trace_schedule, nv_center.generate_compilation_config())\n",
    "\n",
    "compiled_trace_sched.plot_pulse_diagram(plot_backend=\"plotly\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6b5a160",
   "metadata": {},
   "outputs": [],
   "source": [
    "ic.prepare(compiled_trace_sched)\n",
    "ic.start()\n",
    "data = ic.retrieve_acquisition()\n",
    "ic.stop()\n",
    "trace = data[0].real[0]\n",
    "time = np.arange(0, 16000, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc11e20d",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1)\n",
    "ax.plot(time / 1000, trace, color=\"#00839F\")\n",
    "ax.set_xlabel(r\"Time ($\\mu$s)\")\n",
    "ax.set_ylabel(\"Relative amplitude\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6dd95b4",
   "metadata": {},
   "source": [
    "[![image2](figures/Trace_Acquisition.png)](figures/Trace_Acquisition.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "691a244a",
   "metadata": {},
   "source": [
    "## TTL Acquisition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8b244fd",
   "metadata": {},
   "source": [
    "We can now set a correct threshold for the Trigger Count acquisition to ensure all the photon arrivals are counted a single time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "366982e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "hardware_cfg = set_hardware_config(ttl_acq_threshold=0.3)\n",
    "nv_center.hardware_config.set(hardware_cfg)\n",
    "\n",
    "qe0.measure.pulse_duration(20e-6)\n",
    "qe0.measure.acq_duration(20e-6)\n",
    "\n",
    "\n",
    "def ttl_acquisition_sched(qubit: str, repetitions: int = 1):\n",
    "    schedule = Schedule(\"TTL Acquisition\")\n",
    "    schedule.add(Measure(\"qe0\", acq_protocol=\"TriggerCount\"))\n",
    "    schedule.repetitions = repetitions\n",
    "    return schedule\n",
    "\n",
    "\n",
    "schedule = ttl_acquisition_sched(\"qe0\")\n",
    "compiled_sched = compiler.compile(schedule, nv_center.generate_compilation_config())\n",
    "\n",
    "compiled_sched.plot_pulse_diagram(plot_backend=\"plotly\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19172d0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "sched_kwargs = {\n",
    "    \"qubit\": \"qe0\",\n",
    "}\n",
    "\n",
    "ttl_acquisition_gettable = ScheduleGettable(\n",
    "    quantum_device=nv_center,\n",
    "    schedule_function=ttl_acquisition_sched,\n",
    "    schedule_kwargs=sched_kwargs,\n",
    "    batched=True,\n",
    "    data_labels=[\"Trigger Count\"],\n",
    ")\n",
    "ttl_acquisition_gettable.unit = [\"\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4796dac",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = ttl_acquisition_gettable.get()\n",
    "print(\"Photons counted: \" + str(result[0][0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f988180b",
   "metadata": {},
   "outputs": [],
   "source": [
    "rep_par = ManualParameter(\n",
    "    name=\"repetitions\",\n",
    "    unit=\"\",\n",
    "    label=\"Repetitions\",\n",
    ")\n",
    "rep_par.batched = True\n",
    "repetitions = 200\n",
    "meas_ctrl.settables(rep_par)\n",
    "meas_ctrl.setpoints(np.asarray(range(repetitions)))\n",
    "meas_ctrl.gettables(ttl_acquisition_gettable)\n",
    "dataset = meas_ctrl.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9bd138d",
   "metadata": {},
   "outputs": [],
   "source": [
    "photon_count = dataset.y0\n",
    "\n",
    "plt.hist(photon_count, bins=range(1, 10), align=\"left\", rwidth=0.9, color=\"#00839F\")\n",
    "plt.title(r\"Photon counts for 200 acquisitions of 20 $\\mu$s\")\n",
    "plt.xlabel(\"Trigger Count\")\n",
    "plt.ylabel(\"# of Occurences\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecfc61cf",
   "metadata": {},
   "source": [
    "[![image3](figures/TTLAcquisitionBins.png)](figures/TTLAcquisitionBins.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b4eaca2",
   "metadata": {},
   "source": [
    "## Quantify - NV Center Native Library\n",
    "\n",
    "We define a schedule consisting of a Charge Reset operation which sets the NV to its negative charge state NV$^-$, followed by a Charge and Resonance Count operation which runs the ionization and the spin pump lasers with a photon count, on the element `qe0`. These operations are defined in `quantify_scheduler.operations.nv_native_library`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0484e019",
   "metadata": {},
   "outputs": [],
   "source": [
    "from quantify_scheduler.operations.nv_native_library import ChargeReset, CRCount\n",
    "\n",
    "qe0.clock_freqs.ionization.set(564e12)  # 532 nm\n",
    "qe0.clock_freqs.ge1.set(470.4e12 - 5e9)  # slightly detuned\n",
    "qe0.clock_freqs.f01.set(3.592e9)\n",
    "qe0.clock_freqs.spec.set(2.895e9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59c46b70",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define ChargeReset parameters\n",
    "green_voltage = 0.697  # 30 uW\n",
    "\n",
    "qe0.charge_reset.amplitude(green_voltage)\n",
    "qe0.charge_reset.duration(40e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b3917de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define CRCount parameters\n",
    "readout_voltage = 0.5\n",
    "spinpump_voltage = 0.15\n",
    "\n",
    "qe0.cr_count.readout_pulse_amplitude(readout_voltage)\n",
    "qe0.cr_count.spinpump_pulse_amplitude(spinpump_voltage)\n",
    "qe0.cr_count.acq_delay(5e-6)\n",
    "qe0.cr_count.acq_duration(40e-6)\n",
    "qe0.cr_count.readout_pulse_duration(45e-6)\n",
    "qe0.cr_count.spinpump_pulse_duration(45e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f568c117",
   "metadata": {},
   "source": [
    "For this schedule, three lasers are used, as well as the other modules in the cluster. Since these are not present in the hardware configuration defined above, a new hardware configuration listing all the hardware is necessary. We load an example json hardware configuration file present in `quantify_scheduler`, which contains all the instruments required. Additionally, these are added to the instrument coordinator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a1fe580",
   "metadata": {},
   "outputs": [],
   "source": [
    "lo_spinpump_laser = find_or_create_instrument(\n",
    "    MockLocalOscillator, recreate=True, name=\"lo_spinpump_laser\"\n",
    ")\n",
    "lo_green_laser = find_or_create_instrument(\n",
    "    MockLocalOscillator, recreate=True, name=\"lo_green_laser\"\n",
    ")\n",
    "\n",
    "ic.add_component(GenericInstrumentCoordinatorComponent(lo_spinpump_laser))\n",
    "ic.add_component(GenericInstrumentCoordinatorComponent(lo_green_laser))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "341595cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Schedule\n",
    "def nv_schedule(qubit: str, repetitions: int = 1):\n",
    "    sched = Schedule(\"NV Center Schedule\")\n",
    "    sched.add(ChargeReset(qubit))\n",
    "    sched.add(CRCount(qubit, acq_index=0, acq_protocol=\"TriggerCount\", bin_mode=BinMode.APPEND))\n",
    "    return sched\n",
    "\n",
    "\n",
    "compiled_nv_sched = compiler.compile(nv_schedule(\"qe0\"), nv_center.generate_compilation_config())\n",
    "compiled_nv_sched.plot_pulse_diagram(plot_backend=\"plotly\")\n",
    "compiled_nv_sched.plot_circuit_diagram()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50a0314d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sched_kwargs = {\n",
    "    \"qubit\": \"qe0\",\n",
    "}\n",
    "\n",
    "nv_gettable = ScheduleGettable(\n",
    "    quantum_device=nv_center,\n",
    "    schedule_function=nv_schedule,\n",
    "    schedule_kwargs=sched_kwargs,\n",
    "    batched=True,\n",
    "    data_labels=[\"Trigger Count\"],\n",
    ")\n",
    "nv_gettable.unit = [\"\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b9259e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = nv_gettable.get()\n",
    "print(\"Photons counted: \" + str(result[0][0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16233c85",
   "metadata": {},
   "source": [
    "## Dark ESR Schedule\n",
    "\n",
    "We turn now to the dark electron spin resonance (ESR) experiment.\n",
    "A predefined schedule has been defined in the Quantify Scheduler library, which uses the numerically controled oscillator frequency sweeping capabilities (see NCO Control tutorial for more information) to perform the spectroscopy operation in a defined frequency range. To ensure the NV center is in the correct state, after the Reset operation, a Charge Resonance check is performed before and after each Measure operation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "061135b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reset Operation\n",
    "spinpump_reset_voltage = 0.502  # 0.5 uW\n",
    "qe0.reset.amplitude(spinpump_reset_voltage)\n",
    "qe0.reset.duration(3e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9d7be8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Spectroscopy Operation\n",
    "qe0.spectroscopy_operation.amplitude(0.25)\n",
    "qe0.spectroscopy_operation.duration(3e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b0db056",
   "metadata": {},
   "outputs": [],
   "source": [
    "from quantify_scheduler.schedules.spectroscopy_schedules import nv_dark_esr_sched_nco\n",
    "\n",
    "spec_frequencies = np.asarray(np.linspace(2.898e9, 3.098e9, 20))\n",
    "\n",
    "compiler = SerialCompiler(\"compiler\")\n",
    "compiled_nv_sched = compiler.compile(\n",
    "    schedule=nv_dark_esr_sched_nco(\"qe0\", spec_clock=\"qe0.spec\", spec_frequencies=spec_frequencies),\n",
    "    config=nv_center.generate_compilation_config(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10a9c674",
   "metadata": {},
   "outputs": [],
   "source": [
    "schedule_kwargs = {\n",
    "    \"qubit\": \"qe0\",\n",
    "    \"spec_clock\": \"qe0.spec\",\n",
    "    \"spec_frequencies\": spec_frequencies,\n",
    "}\n",
    "\n",
    "dark_esr_gettable = ScheduleGettable(\n",
    "    quantum_device=nv_center,\n",
    "    schedule_function=nv_dark_esr_sched_nco,\n",
    "    schedule_kwargs=schedule_kwargs,\n",
    "    batched=True,\n",
    "    data_labels=[\"Trigger Count\"],\n",
    ")\n",
    "dark_esr_gettable.unit = [\"\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ebbdbbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "acq_par = ManualParameter(\n",
    "    name=\"acq_index\",\n",
    "    unit=\"\",\n",
    "    label=\"Acquisition Index\",\n",
    ")\n",
    "acq_par.batched = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b20a5525",
   "metadata": {},
   "outputs": [],
   "source": [
    "meas_ctrl.settables([acq_par, qe0.clock_freqs.spec])\n",
    "meas_ctrl.setpoints_grid([np.asarray([0, 1, 2]), spec_frequencies])\n",
    "meas_ctrl.gettables(dark_esr_gettable)\n",
    "dataset = meas_ctrl.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a99ecf5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "rep_par = ManualParameter(\n",
    "    name=\"repetitions\",\n",
    "    unit=\"\",\n",
    "    label=\"Repetitions\",\n",
    ")\n",
    "rep_par.batched = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7a59853",
   "metadata": {},
   "outputs": [],
   "source": [
    "repetitions = 10\n",
    "\n",
    "meas_ctrl.settables([acq_par, rep_par, qe0.clock_freqs.spec])\n",
    "meas_ctrl.setpoints_grid([np.asarray([0, 1, 2]), np.asarray(range(repetitions)), spec_frequencies])\n",
    "meas_ctrl.gettables(dark_esr_gettable)\n",
    "dataset = meas_ctrl.run()"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}