hhcavity.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T18:35:53.520167Z",
     "start_time": "2022-02-09T18:35:52.934225Z"
    },
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "%gui qt5\n",
    "%matplotlib qt5\n",
    "\n",
    "from pycolleff.hhcavity import Params, HarmonicCavity\n",
    "import mathphys\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm as cm\n",
    "import matplotlib.colors as colors\n",
    "import matplotlib.gridspec as mpl_gs\n",
    "from matplotlib import rcParams, rc\n",
    "\n",
    "rcParams.update({\n",
    "    'font.size': 18, 'lines.linewidth': 2,\n",
    "    'axes.grid': True, 'grid.alpha': 0.5, 'grid.linestyle': '--', 'text.usetex': True})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A. Chao Book Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T19:26:19.824395Z",
     "start_time": "2021-08-20T19:26:19.812857Z"
    }
   },
   "outputs": [],
   "source": [
    "# A. Chao Book Example\n",
    "hcparams = Params()\n",
    "hcparams.alpha = 0.03\n",
    "hcparams.E0 = 1e9\n",
    "w0 = 9.4e6\n",
    "T0 = 2*np.pi/w0 # s\n",
    "echarge = mathphys.constants.elementary_charge\n",
    "N = 1e11\n",
    "hcparams.I0 = N*echarge/T0\n",
    "hcparams.tunes = 0.01\n",
    "hcparams.frf = 360e6\n",
    "hcparams.h = 240\n",
    "hcparams.Rs = 1e6\n",
    "hcparams.Q = 2000\n",
    "dw = 2*np.pi*(-10e3)\n",
    "wrf = 2*np.pi*hcparams.frf\n",
    "wr = wrf + dw\n",
    "hcparams.bunlen = 1e-4\n",
    "hcparams.nharm = 1\n",
    "\n",
    "hcav = HarmonicCavity()\n",
    "hcav.params = hcparams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T19:30:13.484968Z",
     "start_time": "2021-08-20T19:30:13.476931Z"
    }
   },
   "outputs": [],
   "source": [
    "hcparams.I0*1e3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T19:26:31.114856Z",
     "start_time": "2021-08-20T19:26:31.109790Z"
    }
   },
   "outputs": [],
   "source": [
    "growth1 = hcav.robinson_growth_rate(w=wrf, wr=wr, approx=False)\n",
    "growth2 = hcav.robinson_growth_rate(w=wrf, wr=wr, approx=True)\n",
    "\n",
    "1/growth1, 1/growth2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## MAX-IV Comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-11-16T14:11:24.181049Z",
     "start_time": "2021-11-16T14:11:24.177832Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "hcparams = Params()\n",
    "hcav = HarmonicCavity()\n",
    "\n",
    "# hcparams.sirius_params()\n",
    "# # hcparams.Q = 21600\n",
    "# hcparams.Q = int(2e8)\n",
    "# hcparams.Rs = 88.4 * hcparams.Q\n",
    "# hcparams.I0 = 350e-3\n",
    "hcparams.maxiv_params()\n",
    "# hcparams.Rs = 2.017e6\n",
    "hcparams.Rs = 4.2e6\n",
    "hcav.params = hcparams\n",
    "\n",
    "hcav.harmonic_phase = hcav.phih_harmonic_flat_potential\n",
    "hcav.psih_harmonic = hcav.psih_harmonic_flat_potential\n",
    "hcav.detune_angle = hcav.psih_harmonic\n",
    "\n",
    "hcav.print_flat_potential()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Equilibrium Distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-11-16T14:11:27.244450Z",
     "start_time": "2021-11-16T14:11:27.233761Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "dt = 0.9\n",
    "frf = hcav.params.frf\n",
    "c = mathphys.constants.light_speed\n",
    "npts = int(1e6)\n",
    "tau = np.linspace(-1/frf/2 * dt, 1/frf/2 * dt,  npts)\n",
    "z = tau * c\n",
    "wrf = hcav.params.wrf\n",
    "phase = wrf*z/c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-11-16T14:11:36.559981Z",
     "start_time": "2021-11-16T14:11:36.401510Z"
    },
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "vmain, vmain_pert, vharm = hcav.calc_voltages(z)\n",
    "detune = hcav.detune_passive_cavity(hcav.params.Rs)\n",
    "# hcav.detune_angle = detune\n",
    "# detune = -96.558 * np.pi/180\n",
    "hcav.detune_angle = detune\n",
    "\n",
    "# hcav.wr = 3*hcav.params.wrf + 60.36*1e3 * 2*np.pi\n",
    "# detune = hcav.detune_angle\n",
    "vharm = hcav.calc_passive_voltage(z, Rs=hcav.params.Rs, detune_phase=detune)\n",
    "fig = hcav.plot_voltages(z, vmain, vharm, vmain_pert+vharm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:11:40.232633Z",
     "start_time": "2021-07-20T19:11:40.168098Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# rho0 = hcav.calc_distribution(z, harmonic=False)\n",
    "# rhohc = hcav.calc_distribution(z, harmonic=True)\n",
    "\n",
    "rho0 = hcav.calc_distribution(z, voltage=vmain)\n",
    "rhohc = hcav.calc_distribution(z, voltage=vmain_pert + vharm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Scalar self-consistency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:11:44.218610Z",
     "start_time": "2021-07-20T19:11:43.539994Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Loop Form Factor\n",
    "rho = rhohc\n",
    "wrf = hcav.params.wrf\n",
    "fhist = []\n",
    "ghist = []\n",
    "tol = 1e-8\n",
    "for i in range(100):\n",
    "    fhist.append(hcav.form_factor)\n",
    "    vmain_pert, vharm, rhof = loop_form_factor(hcav, rho, include_phase=False, update_detune=True)\n",
    "    fnew = hcav.complex_form_factor(z, 3*wrf, rhof)\n",
    "    ghist.append(abs(fhist[-1]) - abs(fnew))\n",
    "    rho = rhof\n",
    "    print(i, hcav.form_factor, ghist[-1], hcav.detune_angle)\n",
    "    if abs(ghist[-1]) < tol:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:11:45.877888Z",
     "start_time": "2021-07-20T19:11:45.875930Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "maxiv1 = rho"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "### Full self-consistency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:11:48.698087Z",
     "start_time": "2021-07-20T19:11:46.468030Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Loop Form Factor\n",
    "rho = rhohc\n",
    "wrf = hcav.params.wrf\n",
    "fhist = []\n",
    "ghist = []\n",
    "tol = 1e-8\n",
    "for i in range(100):\n",
    "    fhist.append(hcav.form_factor)\n",
    "    vmain_pert, vharm, rhof = loop_form_factor(hcav, rho, include_phase=True, update_detune=True)\n",
    "    fnew = hcav.complex_form_factor(z, 3*wrf, rhof)\n",
    "    ghist.append(abs(fhist[-1]) - abs(fnew))\n",
    "    rho = rhof\n",
    "    print(i, hcav.form_factor, ghist[-1], hcav.detune_angle)\n",
    "    if abs(ghist[-1]) < tol:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:11:49.862100Z",
     "start_time": "2021-07-20T19:11:49.860123Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "maxiv2 = rho"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:11:52.082024Z",
     "start_time": "2021-07-20T19:11:51.283680Z"
    },
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "z0 = hcav.calc_sync_phase(z, rho0)\n",
    "z0hc = hcav.calc_sync_phase(z, rho)\n",
    "\n",
    "plt.plot((z-z0)*100, rho0*hcav.params.Ib, label='Without 3HC')\n",
    "# plt.plot((z-z0)*100, rhohc0, label='With 3HC - Flat Potential')\n",
    "plt.plot((z-z0hc)*100, rho*hcav.params.Ib, '--', label='With 3HC - Passive Cavity')\n",
    "plt.xlabel('z [cm]')\n",
    "plt.ylabel('Normalized bunch distribution [a.u.]')\n",
    "# plt.title('Cold parking +500kHz')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.5, ls='--')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:11:53.750907Z",
     "start_time": "2021-07-20T19:11:53.748131Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "maxdist = np.loadtxt('/home/murilo/Downloads/maxiv_distribution_best.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-21T13:07:50.427445Z",
     "start_time": "2021-07-21T13:07:50.283671Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "ph = maxdist[:, 0]\n",
    "val1 = maxdist[:, 1]\n",
    "val2 = maxdist[:, 2]\n",
    "\n",
    "fig = plt.figure(figsize=(8, 4))\n",
    "gs = mpl_gs.GridSpec(1, 2)\n",
    "ax1 = plt.subplot(gs[0, 0])\n",
    "ax2 = plt.subplot(gs[0, 1])\n",
    "\n",
    "ax1.plot(ph, val1 , '-', label='Paper')\n",
    "# conv1 = np.max(val1)/np.max(maxiv1)\n",
    "lambrf = 2*np.pi/(wrf/c)\n",
    "conv1 = hcav.params.I0*lambrf\n",
    "filt = np.logical_and(phase > np.min(ph), phase < np.max(ph))\n",
    "ax1.plot(phase[filt], maxiv1[filt]*conv1, '--', lw=2, label='Simul.')\n",
    "ax1.legend()\n",
    "ax1.grid(True, ls='--', alpha=0.5)\n",
    "ax1.set_xlabel('phase [rad]')\n",
    "ax1.set_ylabel('current distribution [A]')\n",
    "ax1.set_title('Scalar self-consistency')\n",
    "bun1 = hcav.calc_bunch_length(z, maxiv1)*1e3\n",
    "textstr = r'$\\sigma_z$ = {:.2f} mm'.format(bun1)\n",
    "textstr += '\\n'\n",
    "textstr += r'$I_p$ = {:.2f} A'.format(np.max(maxiv1*conv1))\n",
    "props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)\n",
    "ax1.text(0.05, 0.95, textstr, transform=ax1.transAxes, fontsize=12,\n",
    "        verticalalignment='top', bbox=props)\n",
    "\n",
    "\n",
    "ax2.plot(ph, val2, '-', label='Paper')\n",
    "# conv2 = np.max(val2)/np.max(maxiv2)\n",
    "conv2 = hcav.params.I0*lambrf\n",
    "ax2.plot(phase[filt], maxiv2[filt]*conv2, '--', lw=2, label='Simul.')\n",
    "ax2.grid(True, ls='--', alpha=0.5)\n",
    "ax2.set_xlabel('phase [rad]')\n",
    "ax2.set_ylabel('current distribution [A]')\n",
    "ax2.set_title('Full self-consistency')\n",
    "\n",
    "bun2 = hcav.calc_bunch_length(z, maxiv2)*1e3\n",
    "textstr = r'$\\sigma_z$ = {:.2f} mm'.format(bun2)\n",
    "textstr += '\\n'\n",
    "textstr += r'$I_p$ = {:.2f} A'.format(np.max(maxiv2*conv2))\n",
    "ax2.text(0.05, 0.95, textstr, transform=ax2.transAxes, fontsize=12,\n",
    "        verticalalignment='top', bbox=props)\n",
    "# ax2.legend()\n",
    "# plt.plot(max_dist_complex[:, 0], max_dist_complex[:, 1], '--')\n",
    "plt.tight_layout(True)\n",
    "plt.show()\n",
    "# plt.savefig('maxiv_comparison_paper.png', dpi=300, format='png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:12:13.641083Z",
     "start_time": "2021-07-20T19:12:13.508173Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "pot0 = hcav.integrated_potential(z, voltage=vmain)\n",
    "# pothc0 = hcav.integrated_potential(z, voltage=vmain_pert0+vharm0)\n",
    "pothc = hcav.integrated_potential(z, voltage=vmain_pert+vharm)\n",
    "\n",
    "\n",
    "norm = (hcav.params.alpha * hcav.params.espread)**2\n",
    "plt.plot(z*100, pot0/norm, label='Without 3HC')\n",
    "# plt.plot((z-z0)*100, pothc0*1e10, label='With 3HC - Flat Potential')\n",
    "plt.plot((z-z0)*100, pothc/norm, '--', label='With 3HC - Passive Cavity')\n",
    "\n",
    "\n",
    "plt.xlabel('z [cm]')\n",
    "plt.ylabel(r'Integrated RF potential $\\Phi$ [a.u.]')\n",
    "# plt.title('Cold parking +500kHz')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.5, ls='--')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "hcav.wr = 3*hcav.params.wrf + 2*np.pi*(500e3)\n",
    "detune = hcav.detune_angle\n",
    "hcav.harmonic_phase = (hcav.detune_angle + np.pi/2)/hcav.params.nharm\n",
    "vharm = hcav.calc_passive_voltage(z, Rs=hcav.params.Rs, detune_phase=detune)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "#### Bunch Length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:12:56.159873Z",
     "start_time": "2021-07-20T19:12:56.131601Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sig1 = hcav.calc_bunch_length(z, rho0)*1e3\n",
    "sig2 = hcav.calc_bunch_length(z, rho)*1e3\n",
    "sig1, sig2, sig2/sig1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:12:58.357444Z",
     "start_time": "2021-07-20T19:12:58.355197Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# detune = hcav.detune_passive_cavity(hcav.params.Rs)\n",
    "# hcav.detune_angle = detune\n",
    "detune = hcav.detune_angle\n",
    "hcav.harmonic_phase = (hcav.detune_angle + np.pi/2)/hcav.params.nharm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:12:59.088513Z",
     "start_time": "2021-07-20T19:12:58.947369Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# detune = hcav.psih_harmonic\n",
    "vpass = hcav.calc_passive_voltage(z, hcav.params.Rs, detune)\n",
    "\n",
    "\n",
    "vmain, vmain_pert, vharm = hcav.calc_voltages(z)\n",
    "fig = hcav.plot_voltages(z, vmain, vmain_pert, vpass)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:13:03.077277Z",
     "start_time": "2021-07-20T19:13:03.075337Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "angle = hcav.psih_harmonic\n",
    "wr = hcav.wr\n",
    "wrf = hcav.params.wrf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:13:03.429504Z",
     "start_time": "2021-07-20T19:13:03.427262Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "df = (wr - 3*wrf)/2/np.pi\n",
    "print(df/1e3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "### Growth Rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:13:05.129079Z",
     "start_time": "2021-07-20T19:13:05.123938Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "growth1 = hcav.robinson_growth_rate(w=3*wrf, wr=wr, approx=False)\n",
    "growth2 = hcav.robinson_growth_rate(w=3*wrf, wr=wr, approx=True)\n",
    "\n",
    "growth1, growth2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:13:07.282292Z",
     "start_time": "2021-07-20T19:13:07.219660Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "npts = int(1e6)\n",
    "w0 = wrf/hcav.params.h\n",
    "# w = np.linspace(-(wr+3*wrf), wr+3*wrf, npts)\n",
    "w = np.linspace(wr-wrf, wr + wrf, npts)\n",
    "w = np.r_[-np.flipud(w), w]\n",
    "deltaw, Zl, wp, interpol_Z, spectrum = hcav.tuneshifts_cbi(w=w, wr=wr, m=1, nbun_fill=hcav.params.h, radiation=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:13:09.607162Z",
     "start_time": "2021-07-20T19:13:09.603931Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "deltaw.imag[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:13:12.408842Z",
     "start_time": "2021-07-20T19:13:12.387595Z"
    },
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plt.plot(deltaw.imag, '.-')\n",
    "plt.xlabel('mode index')\n",
    "plt.ylabel(r'$\\tau^{-1}$ for $m=1$ [$s^{-1}$]')\n",
    "plt.title('Longitudinal growth rate for passive 3HC cold parking')\n",
    "plt.grid(True, alpha=0.5, ls='--')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T18:51:27.217469Z",
     "start_time": "2021-07-20T18:51:23.566414Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plt.semilogy(w, Zl.real/1e6, '.')\n",
    "idx = 0\n",
    "plt.semilogy(wp[idx, :], interpol_Z[idx, :].real/1e6, 'o')\n",
    "\n",
    "plt.xlabel(r'$\\omega$ [rad/s]')\n",
    "plt.ylabel(r'Re$(Z_L)$ [M$\\Omega$] ')\n",
    "plt.title('Mode {:d}'.format(idx))\n",
    "\n",
    "plt.grid(True, alpha=0.5, ls='--')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T18:51:41.165582Z",
     "start_time": "2021-07-20T18:51:36.606965Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "plt.semilogy(w, Zl.real/1e6, '-')\n",
    "\n",
    "idx1 = 0\n",
    "plt.semilogy(wp[idx1, :], interpol_Z[idx1, :].real/1e6, 'o', label=f'Mode {idx1:d}')\n",
    "idx2 = 1\n",
    "plt.semilogy(wp[idx2, :], interpol_Z[idx2, :].real/1e6, '+', label=f'Mode {idx2:d}')\n",
    "idx3 = 2\n",
    "plt.semilogy(wp[idx3, :], interpol_Z[idx3, :].real/1e6, 'd', label=f'Mode {idx3:d}')\n",
    "# idx4 = 3\n",
    "# plt.semilogy(wp[idx4, :], interpol_Z[idx4, :].real/1e6, 'x', label=f'Mode {idx4:d}')\n",
    "\n",
    "plt.legend()\n",
    "plt.xlabel(r'$\\omega$ [rad/s]')\n",
    "plt.ylabel(r'Re$(Z_L)$ [M$\\Omega$] ')\n",
    "plt.title(r'$R_s=$ {:.3f}$M\\Omega$, $Q=$ {:.1e}, Sirius'.format(\n",
    "    hcav.params.Rs/1e6, hcav.params.Q))\n",
    "\n",
    "plt.grid(True, alpha=0.5, ls='--')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sirius Studies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T18:08:15.743814Z",
     "start_time": "2022-02-09T18:08:15.738345Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "k flat potential              : +000.317 \n",
      "harmonic phase [deg]          : -002.190 \n",
      "harmonic tuning angle [deg]   : -096.571 \n",
      "detuning frequency    [kHz]   : +049.305 \n",
      "shunt impedance flat [M.ohm]  : +11.8769 \n",
      "unperturbed sync. phase [deg] : +163.122 \n",
      "perturbed sync. phase [deg]   : +160.936 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "hcparams = Params()\n",
    "hcav = HarmonicCavity()\n",
    "\n",
    "hcparams.sirius_params()\n",
    "# hcparams.Q = 21600\n",
    "# hcparams.Q = 4e8\n",
    "# hcparams.Rs = 90 * hcparams.Q\n",
    "\n",
    "# hcparams.Q = 4e8\n",
    "# hcparams.Rs = 180 * hcparams.Q\n",
    "\n",
    "# hcparams.Q = 4e8\n",
    "hcparams.Rs = 11876921.43677626\n",
    "hcparams.Q = hcparams.Rs/90\n",
    "\n",
    "# hcparams.I0 = 350e-3\n",
    "# hcparams.maxiv_params()\n",
    "# hcparams.Rs = 2.017e6\n",
    "# hcparams.Rs = 4.2e6\n",
    "hcav.params = hcparams\n",
    "\n",
    "hcav.harmonic_phase = hcav.phih_harmonic_flat_potential\n",
    "hcav.psih_harmonic = hcav.psih_harmonic_flat_potential\n",
    "hcav.detune_angle = hcav.psih_harmonic\n",
    "wr = hcav.wr\n",
    "\n",
    "hcav.print_flat_potential()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T18:07:57.635264Z",
     "start_time": "2022-02-09T18:07:57.626660Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11876921.43677626"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hcav.shunt_impedance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Equilibrium Distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T18:08:19.902699Z",
     "start_time": "2022-02-09T18:08:19.869178Z"
    }
   },
   "outputs": [],
   "source": [
    "dt = 0.5\n",
    "frf = hcav.params.frf\n",
    "c = mathphys.constants.light_speed\n",
    "npts = int(1e3)+1\n",
    "tau = np.linspace(-1/frf/2 * dt, 1/frf/2 * dt,  npts)\n",
    "z = tau * c\n",
    "wrf = hcav.params.wrf\n",
    "phase = wrf*z/c\n",
    "\n",
    "vmain, vmain_pert, vharm = hcav.calc_voltages(z)\n",
    "detune = hcav.detune_passive_cavity(hcav.params.Rs)\n",
    "hcav.detune_angle = detune\n",
    "# detune = -96.558 * np.pi/180\n",
    "# hcav.detune_angle = detune\n",
    "# hcav.wr = 3*hcav.params.wrf + 60.36*1e3 * 2*np.pi\n",
    "# detune = hcav.detune_angle\n",
    "\n",
    "Vrf = hcav.params.Vrf\n",
    "wrf = 2*np.pi*hcav.params.frf\n",
    "phis0 = hcav.params.sync_phase\n",
    "phis_pert = hcav.perturbed_sync_phase\n",
    "kh = hcav.k_harmonic_flat_potential\n",
    "nh = hcav.params.nharm\n",
    "phih = hcav.harmonic_phase\n",
    "\n",
    "phase = wrf*z/c\n",
    "vmain = Vrf*np.sin(phase + phis0)\n",
    "vmain_pert = Vrf*np.sin(phase + phis_pert)\n",
    "# vharm = Vrf*kh*np.sin(nh*phase + nh*phih)\n",
    "# vtotal = vmain_pert + vharm\n",
    "\n",
    "vharm = hcav.calc_passive_voltage(z, Rs=hcav.params.Rs, detune_phase=detune)\n",
    "vtotal = vmain_pert + vharm\n",
    "fig = hcav.plot_voltages(z, vmain, vharm, vtotal)\n",
    "\n",
    "# rho0 = hcav.calc_distribution(z, harmonic=False)\n",
    "# rhohc = hcav.calc_distribution(z, harmonic=True)\n",
    "rho0 = hcav.calc_distribution(z, voltage=vmain)\n",
    "rhohc = hcav.calc_distribution(z, voltage=vtotal)\n",
    "# fig.savefig('sirius_voltage_3hc.png', dpi=600, format='png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T17:51:13.991317Z",
     "start_time": "2022-02-09T17:51:13.961562Z"
    }
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8, 6))\n",
    "gs = mpl_gs.GridSpec(1, 1)\n",
    "ax1 = plt.subplot(gs[0, 0])\n",
    "\n",
    "ax1.plot(z*100, vmain * 1e-6, label=r'$V_{rf}$ Main', color='C0')\n",
    "ax1.plot(z*100, vharm * 1e-6, label=r'$V_{3h}$ 3HC', color='C1')\n",
    "ax1.plot(\n",
    "    z*100, vtotal * 1e-6, label=r'$V_{rf} + V_{3h}$ Total', color='C3')\n",
    "ax1.axhline(\n",
    "    y=hcav.params.U0*1e-6, color='tab:gray', ls='--', label=r'$U_0$')\n",
    "\n",
    "ax1.set_xlabel(r'$z$ [cm]')\n",
    "ax1.set_ylabel('RF voltage [MV]')\n",
    "ax1.legend(loc='upper right')\n",
    "ax1.grid(ls='--', alpha=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T18:08:36.121848Z",
     "start_time": "2022-02-09T18:08:36.058441Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "pot0 = hcav.integrated_potential(z, voltage=vmain)\n",
    "# pothc0 = hcav.integrated_potential(z, voltage=vmain_pert0+vharm0)\n",
    "pothc = hcav.integrated_potential(z, voltage=vtotal)\n",
    "\n",
    "norm = (hcav.params.alpha * hcav.params.espread)**2\n",
    "plt.plot(z*100, pot0/norm, label='Without 3HC', lw=2)\n",
    "# plt.plot((z-z0)*100, pothc0*1e10, label='With 3HC - Flat Potential')\n",
    "plt.plot(z*100, pothc/norm, '-', label='With 3HC', lw=2)\n",
    "\n",
    "plt.xlabel('z [cm]')\n",
    "plt.ylabel(r'RF potential $\\Phi$ [a.u.]')\n",
    "# plt.title('Cold parking +500kHz')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.5, ls='--')\n",
    "# plt.savefig('sirius_potential_3hc.png', dpi=600, format='png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T18:09:24.843093Z",
     "start_time": "2022-02-09T18:09:24.794398Z"
    }
   },
   "outputs": [],
   "source": [
    "lambrf = 2*np.pi/(wrf/c)\n",
    "conv = hcav.params.I0*lambrf\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(z*100, rho0*conv, label='Without 3HC', lw=2)\n",
    "# plt.plot((z-z0)*100, pothc0*1e10, label='With 3HC - Flat Potential')\n",
    "plt.plot(z*100, rhohc*conv, '-', label='With 3HC', lw=2)\n",
    "\n",
    "plt.xlabel('z [cm]')\n",
    "plt.ylabel(r'Longitudinal distribution [A]')\n",
    "# plt.title('Cold parking +500kHz')\n",
    "plt.xlim([-7.5, 7.5])\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.5, ls='--')\n",
    "# plt.savefig('sirius_distribution_3hc.png', dpi=600, format='png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T14:32:17.181461Z",
     "start_time": "2022-02-09T14:32:17.170279Z"
    }
   },
   "outputs": [],
   "source": [
    "np.trapz(rho0**2, z)/np.trapz(rhohc**2, z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Scalar self-consistency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-02T20:11:54.590477Z",
     "start_time": "2021-12-02T20:11:54.579461Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Loop Form Factor\n",
    "rho = rhohc\n",
    "wrf = hcav.params.wrf\n",
    "fhist = []\n",
    "ghist = []\n",
    "tol = 1e-12\n",
    "for i in range(100):\n",
    "    fhist.append(hcav.form_factor)\n",
    "    vmain_pert, vharm, rhof = loop_form_factor(hcav, rho, include_phase=False, update_detune=True)\n",
    "    fnew = hcav.complex_form_factor(z, 3*wrf, rhof)\n",
    "    ghist.append(abs(fhist[-1]) - abs(fnew))\n",
    "    rho = rhof\n",
    "    print(i, hcav.form_factor, ghist[-1], hcav.detune_angle)\n",
    "    if abs(ghist[-1]) < tol:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-21T21:11:01.348151Z",
     "start_time": "2021-07-21T21:11:01.346171Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sirius1 = rho\n",
    "dtune1 = hcav.wr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Complex self-consistency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T18:09:41.230609Z",
     "start_time": "2022-02-09T18:09:41.225982Z"
    }
   },
   "outputs": [],
   "source": [
    "hcav.form_factor = 1\n",
    "psi_opt = hcav.detune_passive_cavity(Rs=hcav.params.Rs)\n",
    "hcav.detune_angle = psi_opt\n",
    "dw = hcav.wr - 3*hcav.params.wrf\n",
    "df = dw/2/np.pi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### Fixed K"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-12-02T19:04:38.153679Z",
     "start_time": "2021-12-02T19:04:38.098303Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "dfreqs[np.argmin(peak_curr)]/1e3\n",
    "best_freq = 42.248525706516496*1e3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T14:34:27.133900Z",
     "start_time": "2022-02-09T14:33:20.749459Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# vari = 1e-3/100\n",
    "# detunes = np.linspace(psi_opt*(1-3*vari), psi_opt*(1+vari), 11)\n",
    "# detunes = [psi_opt, psi_opt*(1-1e-5/100)]\n",
    "peak_curr = []\n",
    "bun_lens = []\n",
    "# detunes = [4.712307380481071]\n",
    "# dfreqs = np.linspace(40, 55, 11)*1e3\n",
    "\n",
    "hcav.form_factor = 1\n",
    "psi_opt = hcav.detune_passive_cavity(Rs=hcav.params.Rs)\n",
    "hcav.detune_angle = psi_opt\n",
    "dw = hcav.wr - 3*hcav.params.wrf\n",
    "# df_opt = dw/2/np.pi\n",
    "# df_opt = 45575.687342148754\n",
    "\n",
    "rho0 = rhohc\n",
    "wrf = hcav.params.wrf\n",
    "\n",
    "lambrf = 2*np.pi/(wrf/c)\n",
    "conv = hcav.params.I0*lambrf\n",
    "# dfreqs = np.linspace(41, 45, 101)*1e3\n",
    "# dtunes = np.linspace(-5, +5, 25)*1e3\n",
    "# dfreqs = df_opt + dtunes\n",
    "# dfreqs = np.array([df_opt])\n",
    "dists = []\n",
    "errs = []\n",
    "entropys = []\n",
    "\n",
    "# dfreqs = np.array([best_freq])\n",
    "# dfreqs = np.array([df_opt])\n",
    "# nbeta = 100\n",
    "nbeta = 0\n",
    "for idx, dfreq in enumerate(dfreqs):\n",
    "    # Loop Form Factor\n",
    "    rho = rhohc\n",
    "    hcav.wr = 3*wrf + 2*np.pi*dfreq\n",
    "    rhof, err, entropy, nit = hcav.convergence_distribution(\n",
    "        z=z, rho0=rho0, niter=100, tol=1e-10, beta=nbeta, \n",
    "        method='anderson_acc', update_detune=True, include_phase=True, print_iter=False)\n",
    "    peak_curr.append(np.max(rhof)*conv)\n",
    "    bun_lens.append(hcav.calc_bunch_length(z, rhof)*1e3)\n",
    "    dists.append(rhof)\n",
    "    errs.append(err)\n",
    "    entropys.append(entropy)\n",
    "    print('{:03d}/{:03d} - err: {:.1e} - beta: {:06.2f} - niter: {:05d}'.format(idx+1, dfreqs.size, err[-1], nbeta, nit))\n",
    "    nbeta /= 1.05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T14:34:31.144829Z",
     "start_time": "2022-02-09T14:34:31.136048Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "(hcav.wr-3*wrf)/2/np.pi, df_opt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T14:34:35.607818Z",
     "start_time": "2022-02-09T14:34:32.923878Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "f = plt.figure()\n",
    "ax = plt.axes()\n",
    "jet = plt.get_cmap('jet') \n",
    "cNorm  = colors.Normalize(vmin=dtunes.min(), vmax=dtunes.max())\n",
    "scalarMap = cm.ScalarMappable(norm=cNorm, cmap=jet)\n",
    "\n",
    "# perc_diff = np.linspace(0.9, 1.1, 51)\n",
    "cNorm  = colors.Normalize(vmin=(df_opt+dtunes).min()/1e3, vmax=(df_opt+dtunes).max()/1e3)\n",
    "scalarMapdiff = cm.ScalarMappable(norm=cNorm, cmap=jet)\n",
    "ax.set_xlim([-10, 5])\n",
    "ax.set_xlabel(r'$z$ [cm]')\n",
    "ax.set_ylabel('Distribution [a.u.]')\n",
    "plt.grid(True, ls='--', alpha=0.5, color='gray')\n",
    "plt.colorbar(scalarMapdiff, label=r'Frequency detune [kHz]')\n",
    "plt.tight_layout(True)\n",
    "\n",
    "for idx, lamb in enumerate(dists):\n",
    "    clr = scalarMap.to_rgba(dtunes[idx])\n",
    "    ax.plot(-z*100, lamb, color=clr)\n",
    "    plt.pause(0.01)\n",
    "#     ax.clear()\n",
    "# textstr =r'$\\sigma_z$ = {:.2f} mm'.format(bun_lens[0])\n",
    "# textstr += '\\n'\n",
    "# textstr += r'$I_p$ = {:.2f} A'.format(peak_curr[0])\n",
    "# props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)\n",
    "# ax.text(0.05, 0.95, textstr, transform=ax.transAxes, fontsize=12,\n",
    "#         verticalalignment='top', bbox=props)\n",
    "\n",
    "# plt.title('Distribution with minimum peak current')\n",
    "plt.tight_layout(True)\n",
    "# f.savefig('dist_fp_condition.png', dpi=300)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T16:10:08.009602Z",
     "start_time": "2021-08-20T16:10:08.006744Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "bun_lens, peak_curr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed K and detune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T18:10:38.208615Z",
     "start_time": "2022-02-09T18:10:21.388396Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/facs/repos/collective_effects/pycolleff/pycolleff/hhcavity.py:440: RuntimeWarning: divide by zero encountered in log\n",
      "  sint = rho*_np.log(rho)\n",
      "/home/facs/repos/collective_effects/pycolleff/pycolleff/hhcavity.py:440: RuntimeWarning: invalid value encountered in multiply\n",
      "  sint = rho*_np.log(rho)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.95\n",
      "001/001 - err: 9.7e-11 - beta: 001.00 - niter: 00442\n"
     ]
    }
   ],
   "source": [
    "# vari = 1e-3/100\n",
    "# detunes = np.linspace(psi_opt*(1-3*vari), psi_opt*(1+vari), 11)\n",
    "# detunes = [psi_opt, psi_opt*(1-1e-5/100)]\n",
    "peak_curr = []\n",
    "bun_lens = []\n",
    "# detunes = [4.712307380481071]\n",
    "# dfreqs = np.linspace(40, 55, 11)*1e3\n",
    "\n",
    "hcav.form_factor = 1\n",
    "psi_opt = hcav.detune_passive_cavity(Rs=hcav.params.Rs)\n",
    "hcav.detune_angle = psi_opt\n",
    "dw = hcav.wr - 3*hcav.params.wrf\n",
    "df = dw/2/np.pi\n",
    "\n",
    "dfreqs = np.array([df])\n",
    "wrs = 3*hcav.params.wrf + 2*np.pi*dfreqs\n",
    "rho0 = rhohc\n",
    "wrf = hcav.params.wrf\n",
    "\n",
    "lambrf = 2*np.pi/(wrf/c)\n",
    "conv = hcav.params.I0*lambrf\n",
    "factor = 1\n",
    "# perc_diff = np.linspace(0.95, 1.05, 51*factor)\n",
    "# perc_diff = np.array([1, 1.01, 1.05])\n",
    "# khs = hcav.k_harmonic_flat_potential * perc_diff\n",
    "khs = np.array([hcav.k_harmonic_flat_potential])\n",
    "dists = []\n",
    "errs = []\n",
    "entropys = []\n",
    "\n",
    "nbeta = 1\n",
    "# beta_min = 1\n",
    "# beta_max = 10\n",
    "# nbeta = np.linspace(beta_min, beta_max, khs.size)\n",
    "for idx, k in enumerate(khs):\n",
    "    # Loop Form Factor\n",
    "    rho = rhohc\n",
    "    rhof, err, entropy, nit = hcav.convergence_distribution(\n",
    "        z=z, rho0=rho0, niter=1000, tol=1e-10, beta=nbeta, \n",
    "        method='anderson_acc', kh=k, update_detune=True, include_phase=True, print_iter=False)\n",
    "    peak_curr.append(np.max(rhof)*conv)\n",
    "    bun_lens.append(hcav.calc_bunch_length(z, rhof)*1e3)\n",
    "    dists.append(rhof)\n",
    "    errs.append(err)\n",
    "    entropys.append(entropy)\n",
    "#     nbeta *= 1.05/factor\n",
    "    print(perc_diff[idx])\n",
    "    print('{:03d}/{:03d} - err: {:.1e} - beta: {:06.2f} - niter: {:05d}'.format(idx+1, khs.size, err[-1], nbeta, nit))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T18:10:42.238847Z",
     "start_time": "2022-02-09T18:10:42.235416Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "45942.96523777304\n"
     ]
    }
   ],
   "source": [
    "dw = hcav.wr - 3*hcav.params.wrf\n",
    "df = dw/2/np.pi\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T18:10:57.885300Z",
     "start_time": "2022-02-09T18:10:57.592440Z"
    }
   },
   "outputs": [],
   "source": [
    "f = plt.figure()\n",
    "ax = plt.axes()\n",
    "jet = plt.get_cmap('jet') \n",
    "cNorm  = colors.Normalize(vmin=khs[0], vmax=khs[-1])\n",
    "scalarMap = cm.ScalarMappable(norm=cNorm, cmap=jet)\n",
    "\n",
    "# perc_diff = np.linspace(0.9, 1.1, 51)\n",
    "diff = (1-perc_diff)*100\n",
    "cNorm  = colors.Normalize(vmin=diff[0], vmax=diff[-1])\n",
    "scalarMapdiff = cm.ScalarMappable(norm=cNorm, cmap=jet)\n",
    "# ax.set_xlim([-10, 5])\n",
    "ax.set_xlabel(r'$z$ [cm]')\n",
    "ax.set_ylabel('Distribution [a.u.]')\n",
    "plt.grid(True, ls='--', alpha=0.5, color='gray')\n",
    "plt.colorbar(scalarMapdiff, label=r'Diff. to FP voltage [\\%]')\n",
    "plt.tight_layout(True)\n",
    "\n",
    "for idx, lamb in enumerate(dists):\n",
    "    clr = scalarMap.to_rgba(khs[idx])\n",
    "    ax.plot(z*100, lamb, color=clr)\n",
    "    plt.pause(0.2)\n",
    "#     ax.clear()\n",
    "\n",
    "\n",
    "# f.savefig('scan_hvoltage_dist.png', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T17:27:54.631380Z",
     "start_time": "2022-02-09T17:27:54.628645Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(180.0, 4.0)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hcav.params.Rs/hcav.params.Q, hcav.params.Q/1e8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T16:45:44.204286Z",
     "start_time": "2022-02-09T16:45:44.126397Z"
    }
   },
   "outputs": [],
   "source": [
    "f = plt.figure()\n",
    "jet = plt.get_cmap('jet') \n",
    "cNorm  = colors.Normalize(vmin=khs[0], vmax=khs[-1])\n",
    "scalarMap = cm.ScalarMappable(norm=cNorm, cmap=jet)\n",
    "for idx, err in enumerate(errs):\n",
    "    clr = scalarMap.to_rgba(khs[idx])\n",
    "    plt.plot(err, color=clr)\n",
    "plt.yscale('log')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T16:46:08.029885Z",
     "start_time": "2022-02-09T16:46:07.928520Z"
    }
   },
   "outputs": [],
   "source": [
    "f = plt.figure()\n",
    "jet = plt.get_cmap('jet') \n",
    "cNorm  = colors.Normalize(vmin=khs[0], vmax=khs[-1])\n",
    "scalarMap = cm.ScalarMappable(norm=cNorm, cmap=jet)\n",
    "off = 0\n",
    "for idx, etpy in enumerate(entropys):\n",
    "    clr = scalarMap.to_rgba(khs[idx])\n",
    "    plt.plot(etpy+off, '-', color=clr)\n",
    "    off += 0.5\n",
    "plt.yscale('log')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-02-09T16:46:16.763277Z",
     "start_time": "2022-02-09T16:46:16.714426Z"
    }
   },
   "outputs": [],
   "source": [
    "f = plt.figure()\n",
    "plt.plot(khs, peak_curr, '-', label='Peak Current')\n",
    "plt.plot(khs, bun_lens, '-', label='Bunch Length')\n",
    "plt.vlines(hcav.k_harmonic_flat_potential, ymin=5, ymax=25, ls='-')\n",
    "plt.vlines(khs[np.argmin(peak_curr)], ymin=5, ymax=25, ls='--')\n",
    "plt.grid(True, ls='--', alpha=0.5, color='gray')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-19T14:42:42.484620Z",
     "start_time": "2021-08-19T14:42:42.474463Z"
    }
   },
   "outputs": [],
   "source": [
    "khs[np.argmin(peak_curr)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-18T20:25:46.333958Z",
     "start_time": "2021-08-18T20:25:46.297888Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.plot(dfreqs/1e3, peak_curr, 'o-', label='Peak Current')\n",
    "plt.plot(dfreqs/1e3, bun_lens, 'o-', label='Bunch Length')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T18:57:51.335507Z",
     "start_time": "2021-08-20T18:57:50.948498Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax1 = plt.subplots()\n",
    "\n",
    "ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis\n",
    "\n",
    "ax1.plot(dfreqs/1e3, bun_lens, '-', color='C0')\n",
    "ax2.plot(dfreqs/1e3, peak_curr, '-', color='C1')\n",
    "\n",
    "ax1.vlines(df_opt/1e3, ymin=np.min(bun_lens), ymax=np.max(bun_lens), label=r'Detune FP - $\\Delta f$ = {:.3f}kHz'.format(df_opt/1e3))\n",
    "# ax1.vlines(dfreqs[np.argmax(bun_lens)]/1e3, ymin=np.min(bun_lens), ymax=np.max(peak_curr), ls='--', label=r'Max. $\\sigma_z$ - $\\Delta f$ = {:.3f}kHz'.format(dfreqs[np.argmax(bun_lens)]/1e3))\n",
    "ax1.vlines(dfreqs[np.argmin(peak_curr)]/1e3, ymin=np.min(bun_lens), ymax=np.max(bun_lens), ls=':', label=r'Min. $I_p$ - $\\Delta f$ = {:.3f}kHz'.format(dfreqs[np.argmin(peak_curr)]/1e3))\n",
    "ax1.set_xlabel('$\\Delta f$ [kHz]')\n",
    "ax1.set_ylabel('Bunch Length [mm]', color='C0')\n",
    "ax2.set_ylabel('Peak Current [A]', color='C1')\n",
    "\n",
    "ax1.grid(True, ls='--', alpha=0.5)\n",
    "ax2.grid(True, ls='--', alpha=0.5)\n",
    "ax1.legend(loc='upper center')\n",
    "\n",
    "ax1.tick_params(axis='y', labelcolor='C0')\n",
    "ax2.tick_params(axis='y', labelcolor='C1')\n",
    "fig.tight_layout()\n",
    "# fig.savefig('detune_scan_bun_peak_fixed_k.png', dpi=600, format='png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T18:54:57.704846Z",
     "start_time": "2021-08-20T18:54:57.602954Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax1 = plt.subplots()\n",
    "\n",
    "ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis\n",
    "\n",
    "ax1.plot(khs, bun_lens, '-', color='C0')\n",
    "ax2.plot(khs, peak_curr, '-', color='C1')\n",
    "\n",
    "k_opt = hcav.k_harmonic_flat_potential\n",
    "ax1.vlines(k_opt, ymin=np.min(bun_lens), ymax=np.max(bun_lens), label=r'FP - $k$ = {:.3f}'.format(k_opt))\n",
    "# ax1.vlines(dfreqs[np.argmax(bun_lens)]/1e3, ymin=np.min(bun_lens), ymax=np.max(peak_curr), ls='--', label=r'Max. $\\sigma_z$ - $\\Delta f$ = {:.3f}kHz'.format(dfreqs[np.argmax(bun_lens)]/1e3))\n",
    "ax1.vlines(khs[np.argmin(peak_curr)], ymin=np.min(bun_lens), ymax=np.max(bun_lens), ls=':', label=r'Min. $I_p$ - $k$ = {:.3f}'.format(khs[np.argmin(peak_curr)]))\n",
    "ax1.set_xlabel('$k$ harmonic')\n",
    "ax1.set_ylabel('Bunch Length [mm]', color='C0')\n",
    "ax2.set_ylabel('Peak Current [A]', color='C1')\n",
    "\n",
    "ax1.grid(True, ls='--', alpha=0.5)\n",
    "ax2.grid(True, ls='--', alpha=0.5)\n",
    "ax1.legend(loc='upper center')\n",
    "\n",
    "ax1.tick_params(axis='y', labelcolor='C0')\n",
    "ax2.tick_params(axis='y', labelcolor='C1')\n",
    "fig.tight_layout()\n",
    "# fig.savefig('detune_scan_bun_peak_fixed_k.png', dpi=600, format='png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-11-16T14:24:48.900960Z",
     "start_time": "2021-11-16T14:24:48.891677Z"
    }
   },
   "outputs": [],
   "source": [
    "pot0 = hcav.integrated_potential(z, voltage=vmain)\n",
    "# pothc0 = hcav.integrated_potential(z, voltage=vmain_pert0+vharm0)\n",
    "vmain_pert, vharm, rho = hcav.loop_form_factor(\n",
    "                z, rho, include_phase=True, update_detune=False)\n",
    "pothc = hcav.integrated_potential(z, voltage=vmain_pert+vharm)\n",
    "\n",
    "norm = (hcav.params.alpha * hcav.params.espread)**2\n",
    "plt.plot(z*100, pot0/norm, label='Without 3HC')\n",
    "# plt.plot((z-z0)*100, pothc0*1e10, label='With 3HC - Flat Potential')\n",
    "plt.plot(z*100, pothc/norm, '--', label='With 3HC - Passive Cavity')\n",
    "\n",
    "plt.xlabel('z [cm]')\n",
    "plt.ylabel(r'Integrated RF potential $\\Phi$ [a.u.]')\n",
    "# plt.title('Cold parking +500kHz')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.5, ls='--')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-20T19:39:55.703543Z",
     "start_time": "2021-07-20T19:39:55.701954Z"
    },
    "heading_collapsed": true
   },
   "source": [
    "### Bunch-Length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-22T00:04:49.636990Z",
     "start_time": "2021-07-22T00:04:49.611911Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sig1 = hcav.calc_bunch_length(z, rho0)*1e3\n",
    "sig2 = hcav.calc_bunch_length(z, rho)*1e3\n",
    "sig1, sig2, sig2/sig1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-21T23:46:22.834573Z",
     "start_time": "2021-07-21T23:46:22.687313Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# detune = hcav.detune_passive_cavity(hcav.params.Rs)\n",
    "# hcav.detune_angle = detune\n",
    "detune = hcav.detune_angle\n",
    "hcav.harmonic_phase = (hcav.detune_angle + np.pi/2)/hcav.params.nharm\n",
    "\n",
    "# detune = hcav.psih_harmonic\n",
    "vpass = hcav.calc_passive_voltage(z, hcav.params.Rs, detune)\n",
    "vmain, vmain_pert, vharm = hcav.calc_voltages(z)\n",
    "fig = hcav.plot_voltages(z, vmain, vmain_pert, vpass)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-22T00:04:59.134255Z",
     "start_time": "2021-07-22T00:04:59.131840Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# angle = hcav.psih_harmonic\n",
    "angle = hcav.detune_angle\n",
    "wr = hcav.wr\n",
    "wrf = hcav.params.wrf\n",
    "\n",
    "df = (wr - 3*wrf)/2/np.pi\n",
    "print(df/1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-22T00:05:00.559580Z",
     "start_time": "2021-07-22T00:05:00.555889Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "angle*180/np.pi-360"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Growth Rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T19:25:45.745783Z",
     "start_time": "2021-08-20T19:25:45.734532Z"
    }
   },
   "outputs": [],
   "source": [
    "# wr = 3*wrf + dfreqs[np.argmin(peak_curr)]*2*np.pi\n",
    "# wr = 3*wrf + df_opt*2*np.pi\n",
    "growth1 = hcav.robinson_growth_rate(w=3*wrf, wr=wr, approx=False)\n",
    "growth2 = hcav.robinson_growth_rate(w=3*wrf, wr=wr, approx=True)\n",
    "\n",
    "growth1, growth2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T18:59:05.530287Z",
     "start_time": "2021-08-20T18:59:05.522265Z"
    }
   },
   "outputs": [],
   "source": [
    "1/growth1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T18:59:13.106632Z",
     "start_time": "2021-08-20T18:59:13.021518Z"
    }
   },
   "outputs": [],
   "source": [
    "npts = int(1e6)\n",
    "w0 = wrf/hcav.params.h\n",
    "# w = np.linspace(-(wr+3*wrf), wr+3*wrf, npts)\n",
    "w = np.linspace(wr-wrf, wr + wrf, npts)\n",
    "w = np.r_[-np.flipud(w), w]\n",
    "deltaw, Zl, wp, interpol_Z, spectrum = hcav.tuneshifts_cbi(w=w, wr=wr, m=1, nbun_fill=hcav.params.h, radiation=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T18:59:16.405822Z",
     "start_time": "2021-08-20T18:59:16.374966Z"
    }
   },
   "outputs": [],
   "source": [
    "deltaw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-08-20T18:59:18.883699Z",
     "start_time": "2021-08-20T18:59:18.842065Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.plot(deltaw.imag, '.-')\n",
    "plt.xlabel('mode index')\n",
    "plt.ylabel(r'$\\tau^{-1}$ for $m=1$ [$s^{-1}$]')\n",
    "plt.title('Longitudinal growth rate for maximum bunch length')\n",
    "plt.grid(True, alpha=0.5, ls='--')\n",
    "# plt.savefig('growth_rate_max_bun_len.png', dpi=600, format='png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-21T23:51:21.460905Z",
     "start_time": "2021-07-21T23:51:21.392143Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.semilogy(w, Zl.real/1e6, '-')\n",
    "idx = 0\n",
    "plt.semilogy(wp[idx, :], interpol_Z[idx, :].real/1e6, 'o')\n",
    "\n",
    "plt.xlabel(r'$\\omega$ [rad/s]')\n",
    "plt.ylabel(r'Re$(Z_L)$ [M$\\Omega$] ')\n",
    "plt.title('Mode {:d}'.format(idx))\n",
    "\n",
    "plt.grid(True, alpha=0.5, ls='--')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-21T23:48:32.246924Z",
     "start_time": "2021-07-21T23:48:32.175210Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.semilogy(w, Zl.real/1e6, '-')\n",
    "\n",
    "idx1 = 0\n",
    "plt.semilogy(wp[idx1, :], interpol_Z[idx1, :].real/1e6, 'o', label=f'Mode {idx1:d}')\n",
    "idx2 = 1\n",
    "plt.semilogy(wp[idx2, :], interpol_Z[idx2, :].real/1e6, '+', label=f'Mode {idx2:d}')\n",
    "idx3 = 2\n",
    "plt.semilogy(wp[idx3, :], interpol_Z[idx3, :].real/1e6, 'd', label=f'Mode {idx3:d}')\n",
    "# idx4 = 3\n",
    "# plt.semilogy(wp[idx4, :], interpol_Z[idx4, :].real/1e6, 'x', label=f'Mode {idx4:d}')\n",
    "\n",
    "plt.legend()\n",
    "plt.xlabel(r'$\\omega$ [rad/s]')\n",
    "plt.ylabel(r'Re$(Z_L)$ [M$\\Omega$] ')\n",
    "plt.title(r'$R_s=$ {:.3f}$M\\Omega$, $Q=$ {:.1e}, Sirius'.format(\n",
    "    hcav.params.Rs/1e6, hcav.params.Q))\n",
    "\n",
    "plt.grid(True, alpha=0.5, ls='--')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Initialization Cell",
  "hide_input": false,
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}