{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "This notebook calculates and plots: \n",
    "* $E_z$ field distribution at time_start and time_end;\n",
    "* total EM energy evolution in time;\n",
    "* NCI growth rate, $Im( \\omega )/\\omega_{p,r}$, calculated between time_start and time_end (within the linear growth of the total energy).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from scipy.constants import c\n",
    "import numpy as np\n",
    "import scipy.constants as scc\n",
    "import yt ; yt.funcs.mylog.setLevel(50)\n",
    "import glob\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot NCI growth rate\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "path_wx   = 'path to diags folder'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "file_list_warpx = glob.glob(path_wx + 'diag1?????')\n",
    "iterations_warpx = [ int(file_name[len(file_name)-5:]) for file_name in file_list_warpx ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def calculate_parameters(path):\n",
    "    iteration=200\n",
    "    dsx = yt.load( path + 'diag1%05d/' %iteration )\n",
    "    dxx = dsx.domain_width/dsx.domain_dimensions\n",
    "    dx=dxx[0];\n",
    "    dx = 1.*dx.ndarray_view()\n",
    "\n",
    "    dz=dxx[1];\n",
    "    \n",
    "    dz = 1.*dz.ndarray_view()\n",
    "    cell_volume_x = np.prod(dxx)\n",
    "\n",
    "    ds1 = yt.load(path+'/diag100100/')\n",
    "    ds2 = yt.load(path+'/diag100200/')\n",
    "    \n",
    "    cur_t1 = ds1.current_time\n",
    "    cur_t2 = ds2.current_time \n",
    "    cur_t2.to_ndarray\n",
    "    dt = (cur_t2-cur_t1)/100\n",
    "    dt = 1.*dt.ndarray_view();\n",
    "    return dx, dz, dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dx, dz, dt = calculate_parameters(path_wx)\n",
    "print(dx,dz,dt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_fourier_transform_wx( path, fieldcomp,  \n",
    "                          iteration, plot=False, remove_last=True ):\n",
    "    \"\"\"\n",
    "    Calculate the Fourier transform of the field at a given iteration\n",
    "    \"\"\"\n",
    "    \n",
    "    ds = yt.load(path + '/diag1%05d/' %iteration )\n",
    "\n",
    "    grid = ds.index.grids[0]\n",
    "    F = grid[fieldcomp]\n",
    "    F = F.ndarray_view()\n",
    "\n",
    "    \n",
    "    if remove_last:\n",
    "        F = F[:-1,:-1]\n",
    "    F = F[:,:,0]\n",
    "\n",
    "    kxmax = np.pi/dx\n",
    "    kzmax = np.pi/dz\n",
    "    Nx = F.shape[0]\n",
    "    Nz = F.shape[1]\n",
    "    spectralF = np.fft.fftshift( np.fft.fft2(F) )[int(Nx/2):, int(Nz/2):]\n",
    "\n",
    "    if plot:\n",
    "        plt.imshow( np.log(abs(spectralF)), origin='lower',\n",
    "                    extent=[0, kxmax, 0, kzmax] )\n",
    "        plt.colorbar()\n",
    "    \n",
    "    return( spectralF, kxmax, kzmax )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def growth_rate_between_wx( path, iteration1, iteration2, \n",
    "                        remove_last=False, threshold=-13 ):\n",
    "    \"\"\"\n",
    "    Calculate the difference in spectral amplitude between two iterations, \n",
    "    return the growth rate\n",
    "    \n",
    "    \"\"\"\n",
    "    spec1, kxmax, kzmax = \\\n",
    "            get_fourier_transform_wx( path, 'Ez', iteration=iteration1, remove_last=remove_last )\n",
    "    spec1 = np.where( abs(spec1) > np.exp(threshold), spec1, np.exp(threshold) )\n",
    "    \n",
    "\n",
    "    spec2, kxmax, kzmax = \\\n",
    "            get_fourier_transform_wx( path, 'Ez', iteration=iteration2, remove_last=remove_last )\n",
    "    \n",
    "    spec2 = np.where( abs(spec2) > np.exp(threshold), spec2, np.exp(threshold) )\n",
    "    diff_growth = np.log( abs(spec2) ) - np.log( abs(spec1) )\n",
    "\n",
    "    diff_time = (iteration2-iteration1)*dt;\n",
    "    growth_rate = diff_growth/diff_time/c;\n",
    "\n",
    "    return( growth_rate, [0, kxmax, 0, kzmax] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def energy( ts ):\n",
    "    Ex= ts.index.grids[0]['boxlib', 'Ex'].squeeze().v\n",
    "    Ey= ts.index.grids[0]['boxlib', 'Ey'].squeeze().v\n",
    "    Ez= ts.index.grids[0]['boxlib', 'Ez'].squeeze().v\n",
    "    \n",
    "    Bx= ts.index.grids[0]['boxlib', 'Ex'].squeeze().v\n",
    "    By= ts.index.grids[0]['boxlib', 'Ey'].squeeze().v\n",
    "    Bz= ts.index.grids[0]['boxlib', 'Ez'].squeeze().v\n",
    "\n",
    "    energyE = scc.epsilon_0*np.sum(Ex**2+Ey**2+Ez**2)\n",
    "    energyB= np.sum(Bx**2+By**2+Bz**2)/scc.mu_0 \n",
    "    energy = energyE  + energyB\n",
    "    return energy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_list = []\n",
    "for iter in iterations_warpx:\n",
    "    path = path_wx + '/diag1%05d/' %iter\n",
    "    ds = yt.load( path) \n",
    "    energy_list.append( energy(ds) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iteration_start = 1700\n",
    "iteration_end= 1900\n",
    "iter_delta = iterations_warpx[2] - iterations_warpx[1]\n",
    "\n",
    "\n",
    "ds_start = yt.load(path_wx + 'diag1%05d/' %iteration_start)\n",
    "Ez_start= ds.index.grids[0]['boxlib', 'Ez'].squeeze().v\n",
    "\n",
    "ds_end = yt.load(path_wx + 'diag1%05d/' %iteration_end)\n",
    "Ez_end= ds_end.index.grids[0]['boxlib', 'Ez'].squeeze().v\n",
    "\n",
    "gr_wx, extent = growth_rate_between_wx(path_wx, iteration_start, iteration_end )\n",
    "\n",
    "\n",
    "fs = 13\n",
    "vmax = 0.05\n",
    "vmin=-vmax\n",
    "cmap_special = 'bwr'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots( ncols=2, figsize=(10,4.) )\n",
    "\n",
    "fs=14\n",
    "cmap = \"viridis\"\n",
    "aspect='auto'\n",
    "\n",
    "img1 = ax1.imshow(Ez_start,aspect=aspect, cmap=cmap, extent=extent)\n",
    "ax1.set_title('$t_{step}=$%i' % iteration_start, size=fs)\n",
    "ax1.set_xlabel(' $k_{p,r} z$ ',size=fs)\n",
    "ax1.set_ylabel(' $k_{p,r} x $ ',size=fs)\n",
    "fig.colorbar(img1, ax=ax1, label = '$E_z, [V/m]$')\n",
    "\n",
    "img2 = ax2.imshow(Ez_end,aspect=aspect, cmap=cmap, extent=extent)\n",
    "ax2.set_title('$t_{step}=$%i' % iteration_end,size=fs)\n",
    "ax2.set_xlabel(' $k_{p,r} z$ ',size=fs)\n",
    "ax2.set_ylabel(' $k_{p,r} x $ ',size=fs)\n",
    "fig.colorbar(img2, ax=ax2, label = '$E_z, [V/m]$')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots( ncols=2,nrows=1, figsize=(13,5.) )\n",
    "\n",
    "fs=14\n",
    "\n",
    "img1 = ax1.semilogy(iterations_warpx,energy_list)\n",
    "ax1.semilogy(iteration_start,energy_list[iteration_start/iter_delta],'ro')\n",
    "ax1.semilogy(iteration_end,energy_list[iteration_end/iter_delta],'ro')\n",
    "ax1.grid()\n",
    "ax1.legend()\n",
    "ax1.set_xlabel('time step',size=fs)\n",
    "ax1.set_ylabel('Total EM energy',size=fs)\n",
    "\n",
    "img2 = ax2.imshow( gr_wx, origin='lower' ,cmap='bwr', vmax=0.05, vmin=-vmax, interpolation='nearest', extent=[0, 1, 0, 1] )\n",
    "ax2.set_title('NCI growth rate',size=fs)\n",
    "ax2.set_xlabel('$k_{p,r} z$ ',size=fs)\n",
    "ax2.set_ylabel('$k_{p,r} x $ ',size=fs)\n",
    "fig.colorbar(img2, ax=ax2, label = '$Im(\\omega)/\\omega_{p,r}$')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}