{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9f4fc0e8",
   "metadata": {},
   "source": [
    "# QCoDeS example with Rigol DSA800\n",
    "\n",
    "The Rigol DSA800 series are USB-controlled spectrum analysers.\n",
    "This notebook shows how to connect to the instrument, acquire and save a trace with QCoDeS, use a marker, and, on `-TG` models, enable the tracking generator.\n",
    "\n",
    "## Connect to the instrument\n",
    "1. Make sure you have QCoDeS set up (see the [QCoDeS website](https://microsoft.github.io/Qcodes/index.html) or my notebook [14 minutes to QCoDeS](https://github.com/lairdgrouplancaster/14-minutes-to-QCoDeS/blob/main/14_minutes_to_QCoDeS.ipynb)).\n",
    "2. Connect the analyser over USB and find its VISA resource string.\n",
    "3. Replace the example address below with your instrument resource and check that you get a connect message.\n",
    "4. A good test configuration is to connect the input of your DSA800 to a signal generator at 400 MHz and -10 dBm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2395400",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "from qcodes.instrument import Instrument\n",
    "\n",
    "from qcodes_contrib_drivers.drivers.Rigol.Rigol_DSA800 import RigolDSA800\n",
    "\n",
    "Instrument.close_all()\n",
    "\n",
    "address = \"USB0::0x1AB1::0x0960::DSA8A191400227::INSTR\"  # Replace with your VISA resource\n",
    "sa = RigolDSA800(\"sa\", address) # Should print something like \"Connected to Rigol DSA800 series spectrum analyzer with ID: DSA8A191400227\"\n",
    "sa.get_idn() # Should print something like \"{'vendor': 'Rigol Technologies', 'model': 'DSA815', 'serial': 'DSA8A191400227', 'firmware': '00.01.18.00.02'}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6707a3ad",
   "metadata": {},
   "source": [
    "## Initialise QCoDeS control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4eccba7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qcodes.dataset import (\n",
    "    Measurement,\n",
    "    do0d,\n",
    "    do1d,\n",
    "    initialise_or_create_database_at,\n",
    "    load_or_create_experiment,\n",
    "    plot_by_id,\n",
    ")\n",
    "\n",
    "initialise_or_create_database_at(\"./test_rigol_dsa800.db\")\n",
    "load_or_create_experiment(\n",
    "    experiment_name=\"testing_rigol_dsa800\",\n",
    "    sample_name=\"signal_source\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "412fc0fb",
   "metadata": {},
   "source": [
    "## Set up the analyser and save a trace\n",
    "\n",
    "This example configures a simple sweep around a signal near 400 MHz, runs a single sweep, and stores the trace together with the matching frequency axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35179e2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up instrument parameters for a spectrum analyzer measurement:\n",
    "\n",
    "# Frequency range and points:\n",
    "sa.frequency_center(400e6)\n",
    "sa.frequency_span(50e6)\n",
    "sa.sweep_points(1001)\n",
    "# Display settings:\n",
    "sa.reference_level(0)\n",
    "sa.video_bandwidth_auto(\"ON\")\n",
    "# Detector setup:\n",
    "sa.detector(\"RMS\")\n",
    "sa.resolution_bandwidth(100e3)\n",
    "# Other settings:\n",
    "sa.trace1.mode(\"WRIT\") # Clear Write mode: each new trace overwrites the previous one\n",
    "\n",
    "# Acquire trace data and frequency axis:\n",
    "trace_dbm = sa.acquire_trace(trace_index=1)\n",
    "freq_hz = sa.trace_axis()\n",
    "\n",
    "# Set the instrument back to continuous sweeping\n",
    "sa.run_continuous()\n",
    "\n",
    "# Save the data to the database and plot it:\n",
    "meas = Measurement()\n",
    "meas.register_parameter(sa.trace_axis)\n",
    "meas.register_parameter(sa.trace1.data, setpoints=(sa.trace_axis,))\n",
    "\n",
    "with meas.run() as datasaver:\n",
    "    datasaver.add_result(\n",
    "        (sa.trace_axis, freq_hz),\n",
    "        (sa.trace1.data, trace_dbm),\n",
    "    )\n",
    "    run_id = datasaver.run_id\n",
    "\n",
    "plot_by_id(run_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "337f80f7",
   "metadata": {},
   "source": [
    "## Make a measurement using do0d()\n",
    "Use the workhorse `do0d(...)` command to read traces.\n",
    "\n",
    "There are two ways to do this. The usual one is to grab whatever is on the screen:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d5dbd7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "do0d(sa.trace1.data, do_plot=True)  # This reads the trace exactly as it is on the screen."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0072cf4",
   "metadata": {},
   "source": [
    "\n",
    "If you want to guarantee that all the data is fresh, you can do this:\n",
    "- Put the analyser in HOLD mode\n",
    "- Estimate the scan time\n",
    "- Trigger one sweep\n",
    "-  Call `do0d(...)` after that sweep has finished.\n",
    "\n",
    "Like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4713c0e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sa.resolution_bandwidth(3e3)   # Deliberately slow down the scan for demonstration purposes.\n",
    "\n",
    "sa.hold()\n",
    "\n",
    "sweep_time_s = sa.sweep_time()\n",
    "sweep_count = sa.sweep_count()\n",
    "if sweep_time_s is None or sweep_count is None:\n",
    "    raise RuntimeError(\"The instrument did not return sweep timing information.\")\n",
    "\n",
    "estimated_scan_time_s = float(sweep_time_s) * int(sweep_count)\n",
    "previous_timeout_s = sa.timeout()\n",
    "if previous_timeout_s is None:\n",
    "    scan_timeout_s = estimated_scan_time_s + 1.0\n",
    "else:\n",
    "    scan_timeout_s = max(float(previous_timeout_s), estimated_scan_time_s + 1.0)\n",
    "\n",
    "try:\n",
    "    sa.timeout(scan_timeout_s)\n",
    "    sa.trigger_sweep()\n",
    "    do0d(sa.trace1.data, do_plot=True)\n",
    "finally:\n",
    "    sa.timeout(previous_timeout_s)      # Restore the original timeout value to avoid affecting other operations.\n",
    "\n",
    "sa.run_continuous() # Set the instrument back to continuous sweeping when you have finished.\n",
    "sa.resolution_bandwidth(100e3)  # Restore the resolution bandwidth to a more typical value."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e72e269e",
   "metadata": {},
   "source": [
    "## Make a measurement using do1d()\n",
    "\n",
    "To acquire a 2D scan, we use `do1d()`. (This is how QCoDeS works.)\n",
    "\n",
    "In this example, we acquire a trace versus a dummy sweep parameter and frequency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8ac3fe6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from qcodes.parameters import ManualParameter\n",
    "\n",
    "dummy_step = ManualParameter(\"dummy_step\", label=\"Dummy step\", unit=\"arb\", initial_value=0)\n",
    "\n",
    "sa.run_continuous()\n",
    "\n",
    "do1d_dataset, _, _ = do1d(\n",
    "    dummy_step,\n",
    "    0,\n",
    "    9,\n",
    "    10,\n",
    "    0.2,\n",
    "    sa.trace1.data,\n",
    "    do_plot=True,\n",
    "    show_progress=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e7eea1c",
   "metadata": {},
   "source": [
    "## Use a marker on the latest trace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec4c9cba",
   "metadata": {},
   "outputs": [],
   "source": [
    "sa.marker1.enabled(\"ON\")\n",
    "sa.marker1.peak_search()\n",
    "\n",
    "peak_frequency_hz = sa.marker1.x()\n",
    "peak_amplitude = sa.marker1.y()\n",
    "if peak_frequency_hz is None or peak_amplitude is None:\n",
    "    raise RuntimeError(\"The marker did not return a valid readout.\")\n",
    "\n",
    "print(f\"Peak frequency: {peak_frequency_hz/1e6:.3f} MHz\")\n",
    "print(f\"Peak amplitude: {peak_amplitude:.2f} {sa.power_unit()}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0408858",
   "metadata": {},
   "source": [
    "## Optional: use the tracking generator\n",
    "\n",
    "The following cells are intended for models with a tracking generator.\n",
    "\n",
    "Set the tracking generator up and turn it on:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d52dd21",
   "metadata": {},
   "outputs": [],
   "source": [
    "sa.tracking_generator.enable()\n",
    "sa.tracking_generator.fixed_power_amplitude(-20)\n",
    "sa.tracking_generator.power_mode(\"FIX\")\n",
    "print(\"Tracking generator enabled.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f9c90c1",
   "metadata": {},
   "source": [
    "Turn the tracking generator off:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51f662e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sa.tracking_generator.disable()\n",
    "print(\"Tracking generator disabled.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d5c1106",
   "metadata": {},
   "source": [
    "## Close the connection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a46ac96",
   "metadata": {},
   "outputs": [],
   "source": [
    "sa.close()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".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.11.9"
  },
  "nbsphinx": {
   "execute": "never"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}