1-qubitEvolution.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import qutip\n",
    "from scipy import signal\n",
    "from scipy.constants import hbar\n",
    "from scipy.linalg import expm\n",
    "from scipy.optimize import curve_fit\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "# Qubit evolution in the lab frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "Define the Pauli matrices\n",
    "\n",
    "$\\sigma_z = \\begin{pmatrix} 1 & 0 \\\\ 0 & -1 \\end{pmatrix}$ $\\sigma_x = \\begin{pmatrix} 0 & 1 \\\\ 1 & 0 \\end{pmatrix}$ $\\sigma_y = \\begin{pmatrix} 0 & -i \\\\ i & 0 \\end{pmatrix}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "sigmaZ = np.array([[1,0],\n",
    "                   [0,-1]])\n",
    "sigmaX = np.array([[0,1],\n",
    "                   [1,0]])\n",
    "sigmaY = np.array([[0,-1j],\n",
    "                   [1j,0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "Define the Hamiltonian $\\mathcal{H} = \\frac{\\hbar}{2} \\omega \\sigma_z$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "omega = 2 * np.pi\n",
    "\n",
    "H = hbar/2 * omega * sigmaZ\n",
    "print(H)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "Define an initial state $| \\psi_0 \\rangle = c_0 \\begin{pmatrix} 1 \\\\ 0 \\end{pmatrix} + c_1 \\begin{pmatrix} 0 \\\\ 1 \\end{pmatrix}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# psi0 = np.array([[1],[0]])\n",
    "psi0 = np.array([[1],[1]])*1/np.sqrt(2)\n",
    "print(psi0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "Evolve the state in time using $| \\psi_t \\rangle = U(t) | \\psi_0 \\rangle$ with $U(t) = e^{\\frac{-i \\mathcal{H} t}{\\hbar}}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "time = np.linspace(0,0.95,96)\n",
    "psiT = []\n",
    "\n",
    "for t in time:\n",
    "    psiT.append(expm(-1j*H*t/hbar) @ psi0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "Save the expectation values of x, y and z (e.g. $\\langle \\psi_t | \\sigma_z | \\psi_t \\rangle$)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ev = lambda psi,oper: (psi.conj().T @ oper @ psi).real[0,0]\n",
    "\n",
    "x = [ev(p, sigmaX) for p in psiT]\n",
    "y = [ev(p, sigmaY) for p in psiT]\n",
    "z = [ev(p, sigmaZ) for p in psiT]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "Plot the evolution of the state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "-"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def qubitPlots(time,x,y,z):\n",
    "    bsPlot = qutip.Bloch()\n",
    "    bsPlot.add_points([x, y, z])\n",
    "    bsPlot.render()\n",
    "    bsPlot.fig.set_size_inches([25,5])\n",
    "\n",
    "    ax1 = bsPlot.fig.add_subplot(133)\n",
    "    ax1.plot(time,x,label=\"x\")\n",
    "    ax1.plot(time,y,label=\"y\")\n",
    "    ax1.plot(time,z,label=\"z\")\n",
    "    ax1.legend()\n",
    "    ax1.set_xlabel(\"Time\")\n",
    "    ax1.set_ylabel(\"Expectation value\")\n",
    "    ax1.set_ylim((-1,1))\n",
    "    bsPlot.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "qubitPlots(time,x,y,z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Let's add a drive term $\\mathcal{H} = \\frac{\\hbar}{2} \\left( \\omega \\sigma_z + \\Omega \\cos \\left( \\omega t \\right) \\sigma_x \\right)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "-"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "omegaStatic = 20 * np.pi\n",
    "omegaDrive = 2 * np.pi\n",
    "\n",
    "H = lambda t: hbar / 2 * (omegaStatic * sigmaZ + omegaDrive * np.cos(omegaStatic * t) * sigmaX)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "With a time-dependent Hamiltonian, we need to evolve step-by-step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "-"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "psi0 = np.array([[1],[0]])\n",
    "\n",
    "time = np.linspace(0,1,1001)\n",
    "dTime = time[1] - time[0]\n",
    "psiT = []\n",
    "\n",
    "for t in time:\n",
    "    if t == 0:\n",
    "        psiT.append(psi0)\n",
    "    else:\n",
    "        psiT.append(expm(-1j*H(t)*dTime/hbar) @ psiT[-1])\n",
    "    \n",
    "x = [ev(p, sigmaX) for p in psiT]\n",
    "y = [ev(p, sigmaY) for p in psiT]\n",
    "z = [ev(p, sigmaZ) for p in psiT]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "qubitPlots(time,x,y,z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "# Qubit evolution in the rotating frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "Define the Hamiltonian $\\mathcal{H} = \\frac{\\hbar}{2} ( \\Delta \\sigma_z + \\Omega_x \\sigma_x + \\Omega_y \\sigma_y )$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "-"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "delta = 0 # 4 * np.pi\n",
    "omegaX = 2 * np.pi\n",
    "omegaY = 0 * np.pi\n",
    "\n",
    "H = hbar/2 * ( delta*sigmaZ + omegaX*sigmaX + omegaY*sigmaY )\n",
    "print(H)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "Define an initial state and evolve in time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "-"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "psi0 = np.array([[1],[0]])\n",
    "\n",
    "time = np.linspace(0,0.95,96)\n",
    "dTime = time[1] - time[0]\n",
    "psiT = []\n",
    "\n",
    "for t in time:\n",
    "    if t == 0:\n",
    "        psiT.append(psi0)\n",
    "    else:\n",
    "        psiT.append(expm(-1j*H*dTime/hbar) @ psiT[-1])\n",
    "    \n",
    "x = [ev(p, sigmaX) for p in psiT]\n",
    "y = [ev(p, sigmaY) for p in psiT]\n",
    "z = [ev(p, sigmaZ) for p in psiT]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "qubitPlots(time,x,y,z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "Let's add a detuning axis now"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "-"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "psi0 = np.array([[1],[0]])\n",
    "\n",
    "time = np.linspace(0,3,301)\n",
    "dTime = time[1] - time[0]\n",
    "\n",
    "detunings = np.linspace(-4*np.pi, 4*np.pi, 51)\n",
    "psiT = []\n",
    "z = np.zeros([len(detunings),len(time)])\n",
    "\n",
    "for iDet, delta in enumerate(detunings):\n",
    "    H = hbar/2 * ( delta*sigmaZ + omegaX*sigmaX + omegaY*sigmaY )\n",
    "    for iTime, t in enumerate(time):\n",
    "        if t == 0:\n",
    "            psiT.append(psi0)\n",
    "        else:\n",
    "            psiT.append(expm(-1j*H*dTime/hbar) @ psiT[-1])\n",
    "            z[iDet, iTime] = ev(psiT[-1], sigmaZ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.pcolormesh(time, detunings, z, shading=\"nearest\")\n",
    "plt.xlabel(\"Time\")\n",
    "plt.ylabel(\"Detuning\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Energy-level diagram"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another way to understand the Chevron pattern, is by looking at the energy-level diagram of the qubit Hamiltonian as a function of detuning. We do that by computing the eigenvalues and eigenstates of the Hamiltonian:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "omegaX = 0.1 * np.pi\n",
    "\n",
    "eig_vals = []\n",
    "eig_states = []\n",
    "for iDet, delta in enumerate(detunings):\n",
    "    H = hbar/2 * ( delta*sigmaZ + omegaX*sigmaX + omegaY*sigmaY )\n",
    "    ev, es = np.linalg.eigh(H)\n",
    "    eig_vals.append(ev)\n",
    "    eig_states.append(es)\n",
    "\n",
    "plt.plot(detunings, eig_vals, 'b')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eig_states[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Qubit decoherence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "## Dephasing noise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Noise that modulates the qubit splitting (i.e. along the quantization axis). In most cases, this noise is pink. Let's create an array sampled from a spectrum $S_N = \\frac{S_0}{\\omega}$, and add it to the detuning term in the Hamiltonian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "N = 1001\n",
    "fs = 1000\n",
    "wNoise = np.random.normal(0,50,N)\n",
    "b,a = signal.butter(1, [0.01, 0.5], btype=\"band\", fs=fs)\n",
    "pNoise = signal.lfilter(b,a,wNoise)\n",
    "\n",
    "H_noise = lambda n: hbar/2 * (n * sigmaZ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Start with a superposition state and evolve our noisy Hamiltonian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "psi0 = np.array([[1],[1]])*1/np.sqrt(2)\n",
    "\n",
    "time = np.linspace(0,0.95,N)\n",
    "dTime = 1 / fs\n",
    "psiT = []\n",
    "\n",
    "for t,n in zip(time,pNoise):\n",
    "    if t == 0:\n",
    "        psiT.append(psi0)\n",
    "    else:\n",
    "        psiT.append(expm(-1j*H_noise(n)*dTime/hbar) @ psiT[-1])\n",
    "    \n",
    "x = [ev(p, sigmaX) for p in psiT]\n",
    "y = [ev(p, sigmaY) for p in psiT]\n",
    "z = [ev(p, sigmaZ) for p in psiT]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "qubitPlots(time,x,y,z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Characterising and decoupling dephasing noise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Ramsey fringes - $T_2^*$\n",
    "##### $\\frac{\\pi}{2} - \\tau - \\frac{\\pi}{2} - \\textrm{M}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "In this experiment we apply a $\\pi / 2$ gate to prepare the spin on the x-y plane of the Bloch sphere, wait some time $\\tau$ and apply another $\\pi / 2$ gate to attempt to bring the spin back to the quantization axis.\n",
    "\n",
    "When we perform this experiment on a qubit, we need to repeat the pulse sequence many times to build up statistics on the single-shot measurements. We will capture that by adding a repetition loop to the simulation.\n",
    "\n",
    "In the simulation of this pulse sequence, we assume that we can apply a perfect $\\pi / 2$ gate, and apply the unitary directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "X_2 = (sigmaX + sigmaZ) / np.sqrt(2) # This defines the unitary of a pi/2 gate\n",
    "psi0 = X_2 @ np.array([[1],[0]]) # First pi/2 pulse on the initialised state\n",
    "\n",
    "# We can play with two terms on the Hamiltonian, a static detuning and a time-dependent noise on the qubit splitting\n",
    "H_det = hbar/2 * (2*np.pi * sigmaZ)\n",
    "H_noise = lambda n: hbar/2 * (n * sigmaZ)\n",
    "\n",
    "# Define the sampling frequency, the Ramsey wait times, the number of measurements to perform and noise amplitude\n",
    "fs = 1000\n",
    "taus = np.linspace(0.01, 10, int(10/0.01))\n",
    "reps = 100\n",
    "noise_amp = 30\n",
    "\n",
    "# Generate an array of pink noise\n",
    "dTime = 1 / fs\n",
    "tauPoints = int(taus[0] / dTime)\n",
    "N = int(taus[-1] / dTime) * reps\n",
    "wNoise = np.random.normal(0,noise_amp,N)\n",
    "b,a = signal.butter(1, [0.01, 0.5], btype=\"band\", fs=fs)\n",
    "pNoise = signal.lfilter(b,a,wNoise)\n",
    "\n",
    "# Start simulation\n",
    "zTau = np.zeros((reps, len(taus)))\n",
    "for r in tqdm(range(reps)):\n",
    "    psiT = [psi0] # Initialise the state at the start of each measurement\n",
    "    for i, tau in enumerate(taus):\n",
    "        for j in range(tauPoints):\n",
    "            nIdx = r * len(taus) + i * tauPoints + j\n",
    "            psiT.append(expm(-1j*(H_noise(pNoise[nIdx]) + H_det)*dTime/hbar) @ psiT[-1])\n",
    "        # Apply the second pi/2 gate and measure along z after each tau interval\n",
    "        psiRamsey = X_2 @ psiT[-1]\n",
    "        zTau[r,i] = ev(psiRamsey, sigmaZ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "Now let's visualise the data. Let's first plot the bare data to give you a chance to try to think about what it means..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.plot(taus, np.mean(zTau, axis=0), 'o')\n",
    "plt.ylabel(\"Average qubit measurement\")\n",
    "plt.xlabel(r\"Wait time $\\tau$\")\n",
    "plt.ylim([-1,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "Notice that there are two characteristic features in the data. A single _frequency_ and an exponential _decay_.\n",
    "\n",
    "The frequency corresponds to the detuning frequency (i.e. the difference between the qubit frequency and the frequency of the rf source). \n",
    "In traditional NMR measurements which may be performed on samples with several spin species (different molecules), the spectrum of the Ramsey signal can be used to identify which molecules are present.\n",
    "In qubits, it can be used for a coarse calibration of the qubit frequency.\n",
    "\n",
    "The decay time is known as $T_2^*$ (also sometimes known as _free induction decay_) and it provides an initial characterisation of the environmental noise that couples to the energy splitting of the qubit, a.k.a. _dephasing noise_."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "To extract these two important parameters, we can fit our data to a function of the form:\n",
    "\n",
    "$\\langle \\sigma_Z \\rangle \\left( \\tau \\right) = A \\cos \\left( \\omega \\tau \\right) e^{-\\tau / T_2^*}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define the fitting function\n",
    "def t2_star_decay(time, amplitude, omega, decay):\n",
    "    return amplitude * np.cos(omega * time) * np.exp(-time / decay)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "#Fit the data using the function previously defined\n",
    "popt,_ = curve_fit(t2_star_decay, taus, np.mean(zTau, axis=0))\n",
    "\n",
    "# Plot data and function with extracted parameters\n",
    "plt.plot(taus, np.mean(zTau, axis=0), 'o')\n",
    "plt.plot(taus, t2_star_decay(taus, *popt), label=rf\"$\\omega$ = {popt[1]:.2f}\"+\"\\n\"+rf\"$T_2^*$ = {popt[2]:.2f}\")\n",
    "plt.ylabel(\"Average qubit measurement\")\n",
    "plt.xlabel(r\"Wait time $\\tau$\")\n",
    "plt.ylim([-1,1])\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Hahn echo - $T_2$\n",
    "##### $\\frac{\\pi}{2} - {\\scriptsize \\tau / 2} - \\pi - {\\scriptsize \\tau / 2} - \\frac{\\pi}{2} - \\textrm{M}$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Over short time scales, the pink dephasing noise looks like a constant detuning, i.e., a constant drift in the phase of the qubit. A constant drift can be _refocussed_ by adding a $\\pi$ gate in the middle of our wait time.\n",
    "Try to think about what this does to the qubit on the Bloch sphere before running the simulation..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "X = sigmaX # This defines the unitary of a pi gate\n",
    "psi0 = X_2 @ np.array([[1],[0]]) # First pi/2 pulse on the initialised state\n",
    "\n",
    "# We can now set the detuning to 0, since it will get refocussed anyway\n",
    "H_det = 0 #hbar/2 * (2*np.pi * sigmaZ)\n",
    "H_noise = lambda n: hbar/2 * (n * sigmaZ)\n",
    "\n",
    "# Define the sampling frequency, the Ramsey wait times, the number of measurements to perform and noise amplitude\n",
    "fs = 1000\n",
    "taus = np.linspace(0.01, 10, int(10/0.01))\n",
    "reps = 100\n",
    "noise_amp = 3000\n",
    "\n",
    "# Generate an array of pink noise\n",
    "dTime = 1 / fs\n",
    "tauPoints = int(taus[0] / dTime)\n",
    "N = int(taus[-1] / dTime) * reps\n",
    "wNoise = np.random.normal(0,noise_amp,N)\n",
    "b,a = signal.butter(1, [0.01, 0.5], btype=\"band\", fs=fs)\n",
    "pNoise = signal.lfilter(b,a,wNoise)\n",
    "\n",
    "# Start simulation\n",
    "zTau = np.zeros((reps, len(taus)))\n",
    "for r in tqdm(range(reps)):\n",
    "    psiT = [psi0] # Initialise the state at the start of each measurement\n",
    "    for i, tau in enumerate(taus):\n",
    "        for j in range(tauPoints):\n",
    "            nIdx = r * len(taus) + i * tauPoints + j\n",
    "            if j == int(tauPoints / 2): # At the half-way point, add an X-gate after the evolution\n",
    "                psiT.append(X @ expm(-1j*(H_noise(pNoise[nIdx]) + H_det)*dTime/hbar) @ psiT[-1])\n",
    "            else:\n",
    "                psiT.append(expm(-1j*(H_noise(pNoise[nIdx]) + H_det)*dTime/hbar) @ psiT[-1])\n",
    "        # Apply the second pi/2 gate and measure along z after each tau interval\n",
    "        psiRamsey = X_2 @ psiT[-1]\n",
    "        zTau[r,i] = ev(psiRamsey, sigmaZ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Now let's plot the data from this pulse sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(taus, np.mean(zTau, axis=0), 'o')\n",
    "plt.ylabel(\"Average qubit measurement\")\n",
    "plt.xlabel(r\"Wait time $\\tau$\")\n",
    "plt.ylim([-1,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Note that we have increased the noise amplitude 100x to observe the decay on the same timescale. In a real experiment, the decay would just be observed over a longer timescale, but that would take too long to simulate. For example, for single-spin qubits in 28Si, $T_2^* \\approx 120 \\mu s$, while $T_2 \\approx 1.2 ms$.\n",
    "\n",
    "This decay or _dephasing time_ is typically referred to as $T_2$, and can be extracted by fitting the Hahn echo data to a function of the form:\n",
    "\n",
    "$\\langle \\sigma_Z \\rangle \\left( \\tau \\right) = A e^{(-\\tau / T_2)} + B$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the fitting function\n",
    "def t2_decay(time, amplitude, offset, decay):#, alpha):\n",
    "    return amplitude * np.exp((-time / decay)**1) + offset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "#Fit the data using the function previously defined\n",
    "popt,_ = curve_fit(t2_decay, taus, np.mean(zTau, axis=0))\n",
    "\n",
    "# Plot data and function with extracted parameters\n",
    "plt.plot(taus, np.mean(zTau, axis=0), 'o')\n",
    "plt.plot(taus, t2_decay(taus, *popt), label=rf\"$T_2$ = {popt[2]:.2f}\")\n",
    "plt.ylabel(\"Average qubit measurement\")\n",
    "plt.xlabel(r\"Wait time $\\tau$\")\n",
    "plt.ylim([-1,1])\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}