{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compress the spectrum of an exponential decaying pulse with a cavity and an electro-optical modulator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'parmap'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-1-57ef86982410>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mscipy\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mconstants\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mspeed_of_light\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[1;31m# import tqdm\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 3\u001b[1;33m \u001b[1;32mimport\u001b[0m \u001b[0mparmap\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      4\u001b[0m \u001b[1;32mfrom\u001b[0m \u001b[0mlmfit\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmodels\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mLorentzianModel\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mConstantModel\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mmatplotlib\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mpyplot\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mplt\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'parmap'"
     ]
    }
   ],
   "source": [
    "from scipy.constants import speed_of_light\n",
    "import tqdm\n",
    "import parmap\n",
    "from lmfit.models import LorentzianModel, ConstantModel\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "%matplotlib notebook\n",
    "\n",
    "\n",
    "def exponential_decay(t, tau):\n",
    "    'Electric field envelope for an exponential decay'\n",
    "    return 1/np.sqrt(tau) * np.exp(-t / (2 * tau)) * (t > 0)\n",
    "\n",
    "# def Adrian_FP_cavity_transferfunction(d_freq, R1, R2, cav_length):\n",
    "# # From Adrians paper need to check this function\n",
    "#     domega = 2 * np.pi * d_freq\n",
    "#     finess = np.pi * (R1*R2)**(1/4) / (1 - np.sqrt(R1*R2))\n",
    "#     k = np.pi * speed_of_light / (2 * L * finess)\n",
    "#     k_l = k * np.sqrt(1-R1) / (np.sqrt(1-R1) + np.sqrt(1-R2))\n",
    "#     return 1 - (2 * k_l) / (1j * domega + k)\n",
    "\n",
    "\n",
    "def finesse(R1, R2):\n",
    "    return np.pi * (R1 * R2)**.25 / (1 - np.sqrt(R1 * R2))\n",
    "\n",
    "\n",
    "def FSR(R1, R2, cavity_lw):\n",
    "    return finesse(R1, R2) * cavity_lw\n",
    "\n",
    "\n",
    "def FP_cavity_transferfunction(d_freq, R1, R2, cavity_lw):\n",
    "    d_freq = d_freq * (2 * np.pi)\n",
    "    fsr = FSR(R1, R2, cavity_lw)\n",
    "    return (np.sqrt(R1) - np.sqrt(R2) * np.exp(-1j*d_freq / fsr)) / (1 - np.sqrt(R1*R2) * np.exp(-1j*d_freq / fsr))\n",
    "\n",
    "\n",
    "def spectral_compression(pulse, freq_vec, R1, R2, cavity_lw, cavity_detuning):\n",
    "    freq_vec = np.fft.fftshift(freq_vec)\n",
    "    pulse_spectr = np.fft.fft(pulse, norm='ortho')\n",
    "    # cavity transfer function\n",
    "    transf_func = FP_cavity_transferfunction(\n",
    "        (freq_vec - cavity_detuning), R1, R2, cavity_lw)\n",
    "    # pulse spectrum after cavity\n",
    "    pulse_spectr_after_cavity = transf_func * pulse_spectr\n",
    "    # ifft back to time domain\n",
    "    pulse_after_cavity = np.fft.ifft(pulse_spectr_after_cavity, norm='ortho')\n",
    "    # remove the phase to compresses the spectrum\n",
    "    lensed_pulse_spectrum = np.fft.fft(\n",
    "        np.abs(pulse_after_cavity), norm='ortho')\n",
    "    return np.fft.fftshift(lensed_pulse_spectrum) , pulse_after_cavity#, np.angle(pulse_after_cavity)\n",
    "\n",
    "\n",
    "def FWHM(X, Y):\n",
    "    half_max = max(Y) / 2.\n",
    "    d = np.sign(half_max - np.array(Y[0:-1])) - \\\n",
    "        np.sign(half_max - np.array(Y[1:]))\n",
    "    left_idx = np.where(d > 0)[0]\n",
    "    right_idx = np.where(d < 0)[-1]\n",
    "    # return the difference (full width half max)\n",
    "    return X[right_idx] - X[left_idx]\n",
    "\n",
    "\n",
    "def bandwidth(R1, R2, L):\n",
    "    L_prime = L\n",
    "    return speed_of_light / (2*L_prime) * (1 - np.sqrt(R1*R2)) / (np.pi * (R1*R2)**0.25)\n",
    "\n",
    "\n",
    "# def find_fraction_freq(freq_vec, spectrum, fraction):\n",
    "#     power_spectrum = np.abs(spectrum)**2\n",
    "#     i_max = np.argmax(power_spectrum)\n",
    "#     norm_factor = np.sum(power_spectrum)\n",
    "#     running_sum = 0\n",
    "#     for counter, i in enumerate(power_spectrum[i_max:]/norm_factor):\n",
    "#         running_sum += i\n",
    "#         if running_sum >= fraction/2:\n",
    "#             break\n",
    "#     return freq_vec[i_max+counter]*2, i_max, i_max + counter\n",
    "\n",
    "\n",
    "def find_fraction_freq(freq_vec, spectrum, fraction):\n",
    "    power_spectrum = np.abs(spectrum)**2\n",
    "    i_max = np.argmax(power_spectrum)\n",
    "    norm_factor = np.sum(power_spectrum)\n",
    "    running_sum = 0\n",
    "    counter_outside = 0\n",
    "    for counter, i in enumerate(power_spectrum[i_max:]/norm_factor):\n",
    "        running_sum += i\n",
    "        if running_sum >= fraction/2:\n",
    "            counter_outside = counter\n",
    "            break\n",
    "    return freq_vec[i_max+counter_outside]*2, i_max, i_max + counter_outside"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define photon and cavity parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "inputHidden": false,
    "outputHidden": false
   },
   "outputs": [],
   "source": [
    "# photon decay time\n",
    "tau = 6e-9\n",
    "photon_bandwidth = 1 / (tau*2*np.pi)\n",
    "# cavity parameters\n",
    "R1 = 0.97\n",
    "R2 = 0.9998\n",
    "cavity_bandwidth = 7e6\n",
    "cavity_detuning = 0e6  # bandwidth(R1,R2, cav_length)/2\n",
    "\n",
    "print(f'Photon bandwidth {photon_bandwidth/1e6 :.2f} MHz')\n",
    "print(f'Cavity bandwidth {cavity_bandwidth/1e6 :.2f} MHz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "inputHidden": false,
    "outputHidden": false
   },
   "outputs": [],
   "source": [
    "time_vec, dt = np.linspace(0, 1e-3, int(2e6), retstep=True)\n",
    "# time_vec, dt = np.linspace(0, 1e-5, int(1e6), retstep=True)\n",
    "pulse = exponential_decay(time_vec, tau)\n",
    "freq_vec = np.fft.fftshift(np.fft.fftfreq(pulse.size, dt))\n",
    "init_spectrum = np.fft.fftshift(np.fft.fft(pulse, norm='ortho'))\n",
    "new_spectrum, new_pulse = spectral_compression(\n",
    "    pulse, freq_vec, R1, R2, cavity_bandwidth, cavity_detuning)\n",
    "\n",
    "\n",
    "df_init, i, di = find_fraction_freq(freq_vec, init_spectrum, .5)\n",
    "df_new, _, _ = find_fraction_freq(freq_vec, new_spectrum, .5)\n",
    "print(df_init/1e6, df_new/1e6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "inputHidden": false,
    "outputHidden": false
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(time_vec/1e-9, np.abs(pulse)**2)\n",
    "plt.plot(time_vec/1e-9, np.abs(new_pulse)**2)\n",
    "plt.xlim(-20, 200)\n",
    "\n",
    "# plot new spectra\n",
    "plt.figure(figsize=[8.6, 5.6])\n",
    "\n",
    "plt.plot(freq_vec/1e6, np.abs(init_spectrum)**2)\n",
    "plt.plot(freq_vec/1e6, np.abs(new_spectrum)**2)\n",
    "plt.xlim(-100, 100)\n",
    "\n",
    "\n",
    "mask = (freq_vec > -df_init/2) & (freq_vec < df_init/2)\n",
    "mask_new = (freq_vec > -1*df_new/2) & (freq_vec < df_new/2)\n",
    "\n",
    "plt.fill_between(freq_vec[mask]/1e6, np.abs(init_spectrum[mask])**2, alpha=0.7)\n",
    "plt.fill_between(freq_vec[mask_new]/1e6,\n",
    "                 np.abs(new_spectrum[mask_new])**2, alpha=0.8)\n",
    "plt.xlabel('Detuning (MHz)')\n",
    "plt.ylabel('Power spectrum density')\n",
    "plt.ylim(0);\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Scan over cavity detunings\n",
    "Cavity bandwidth and photon bandwidth kept constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "inputHidden": false,
    "outputHidden": false
   },
   "outputs": [],
   "source": [
    "detunings = np.arange(-150, 150, 1) * 1e6\n",
    "cavity_bandwidth = 7e6\n",
    "\n",
    "def objf(det):\n",
    "    spec_tmp, _ = spectral_compression(\n",
    "        pulse, freq_vec, R1, R2, cavity_bandwidth, det)\n",
    "    freq, _, _ = find_fraction_freq(freq_vec, spec_tmp, .5)\n",
    "    return freq\n",
    "\n",
    "\n",
    "freqs = parmap.map(objf, detunings, pm_pbar=True)\n",
    "freqs = np.array(freqs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "inputHidden": false,
    "outputHidden": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=[8.6, 5.6])\n",
    "init_bw = find_fraction_freq(freq_vec, init_spectrum, .5)[0]\n",
    "plt.plot(detunings/1e6, freqs/1e6, '.-')\n",
    "plt.axhline(init_bw/1e6, color='r')\n",
    "plt.axhline(6.12, color='g')\n",
    "plt.text(-160, 6.2, '6 MHz Atom linewidth')\n",
    "plt.xlabel('Cavity detunings (MHz)')\n",
    "plt.ylabel('Spectral width at 50 percent total power (MHz)')\n",
    "plt.ylim(0);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Scan over dispersion cavity bandwidths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def objf_bw(detuning, bandwidth):\n",
    "    spec_tmp, _ = spectral_compression(\n",
    "        pulse, freq_vec, R1, R2, bandwidth, detuning)\n",
    "    dfreq, _, _ = find_fraction_freq(freq_vec, spec_tmp, .5)\n",
    "    return dfreq\n",
    "\n",
    "bandwidths = np.array([0.1, 0.25, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2])\n",
    "detunings = np.arange(-150, 150, 1) * 1e6\n",
    "freqs_bw = []\n",
    "for i in bandwidths:\n",
    "    freqs_bw.append(np.array(parmap.map(objf_bw, detunings,\n",
    "                              photon_bandwidth * i, pm_pbar=True)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=[8.6, 5.6])\n",
    "for (i, factor) in zip(freqs_bw, bandwidths):\n",
    "    plt.plot(detunings/1e6, i/1e6, label=f'{factor} $\\Gamma_p$')\n",
    "    \n",
    "plt.legend()\n",
    "plt.axhline(init_bw/1e6, color='r')\n",
    "plt.axhline(6.12, color='g')\n",
    "plt.xlabel('Dispersion cavity detuning (MHz)')\n",
    "plt.ylabel('Compressed power spectrum width at 50% tot.power (MHz)')\n",
    "plt.legend(title='Dispersion cavity bandwidth')\n",
    "plt.ylim(0, 30)\n",
    "plt.text(-160, 6.4, '6.12 MHz Atom linewidth')\n",
    "plt.text(-160, 27, 'inital photon 50% bandwidth')\n",
    "plt.tight_layout()\n",
    "# plt.savefig('power_spectrum_width_over_detunings_different_cavbw.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernel_info": {
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.0"
  },
  "nteract": {
   "version": "0.12.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}