{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Arbitrary waveform generator\n",
    "\n",
    "This notebook demonstrates how to create a simple signal chain consisting of an arbitrary waveform generator (AWG), which represents the most minimalistic control stack. The AWG is an indispensable element, enabling precise control over the shape and amplitude of the desired output pulse. Such control is essential for performing high-fidelity gates on qubits.\n",
    "\n",
    "The `qruise-toolset` provides several pre-defined source signals, such as a Gaussian, a local oscillator, and a Heaviside (step) function. In this example, we'll construct a Gaussian source signal.\n",
    "\n",
    "The tutorial consists of the following:\n",
    "\n",
    "1. Defining parameters\n",
    "2. Constructing the signal chain\n",
    "3. Visualising the Gaussian AWG signal\n",
    "\n",
    "---\n",
    "\n",
    "**Tip:** Make sure that your environment is correctly set to handle `float64` precision by setting `JAX_ENABLE_X64=True` or add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "\n",
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "to your scripts' preamble.\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Defining parameters\n",
    "\n",
    "We start by defining the time parameters for our pulse. These include the start and end times of the pulse ($t_0$ and $t_\\text{final}$, respectively), and the grid size, which specifies the number of evenly spaced points used to discretise the time interval. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "\n",
    "t0 = 0.0  # start time of pulse (s)\n",
    "tfinal = 20e-9  # end (final) time of pulse (s)\n",
    "grid_size = int(1e3)  # simulation grid size (number of time points)\n",
    "t_span = jnp.linspace(t0, tfinal, grid_size)  # defines time span array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need to define the parameters of the Gaussian function, $G(t)$, which in the `qruise-toolset` takes the generic form:\n",
    "\n",
    "$$ G(t; a, \\sigma, \\mu) = \\frac{a}{\\sigma\\sqrt{2\\pi}} \\mathrm{exp}\\left[-\\frac{(t - \\mu)^2}{2\\sigma^2}\\right].$$\n",
    "\n",
    "This is characterised by three parameters:\n",
    "\n",
    "- $a$, a scalar that adjusts the amplitude of the pulse\n",
    "\n",
    "- $\\sigma$, the pulse variance\n",
    "\n",
    "- $\\mu$, the time at which it reaches its maximum amplitude\n",
    "\n",
    "For this example, we'll set $a = \\sigma \\sqrt{2\\pi}$ and $\\mu=\\frac{t_\\text{final}}{2}$, so that the Gaussian function reaches a peak value of 1.0 in the centre of the simulation time window. We arbitrarily define $\\sigma = \\frac{1}{4} t_\\text{final}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sigma = 0.25 * tfinal  # pulse variance\n",
    "mu = tfinal / 2.0  # mu, set to centre of simulation time window\n",
    "amp = sigma * jnp.sqrt(2.0 * jnp.pi)  # a, set so peak value = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Constructing the signal chain\n",
    "\n",
    "As mentioned above, the signal chain will consist of a single component. We start by storing the parameters of the Gaussian function in a dictionary (`gauss_params`), and then we define an instance of the `Gaussian()` class. This alone defines our signal chain which consists of only an AWG component that outputs a Gaussian pulse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qruise.toolset import Gaussian\n",
    "\n",
    "# store Gaussian parameters in dictionary\n",
    "gauss_params = {\n",
    "    \"amp\": (amp, True),\n",
    "    \"mu\": (mu, True),\n",
    "    \"sigma\": (sigma, True),\n",
    "}\n",
    "\n",
    "gauss = Gaussian(\"awg_gauss\", gauss_params)  # define instance of Gaussian class"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "**Note:** The argument `True` in the `gauss_params` dictionary indicates that the parameter is modifiable (e.g., for optimisation), but this is not relevant for this tutorial.\n",
    "\n",
    "---\n",
    "\n",
    "Now we need to evaluate the signal chain over the time span we defined earlier (`t_span`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluate signal over defined time span\n",
    "awg_result = gauss(t_span, gauss.params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Visualising the Gaussian AWG signal\n",
    "\n",
    "Let's plot our AWG signal to see if it looks as we'd expect. For this we'll use the `PlotUtil` module from the `qruise-toolset`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qruise.toolset import PlotUtil\n",
    "\n",
    "canvas = PlotUtil(x_axis_label=\"t [s]\", y_axis_label=\"Amplitude [a.u.]\", notebook=True)\n",
    "canvas.plot(t_span, awg_result, labels=[\"Gaussian AWG\"])\n",
    "canvas.show_canvas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nice! Our signal looks as we'd expect, and now you know how to create a signal chain containing an AWG."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "qruise-simple",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}