{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Introduction to openPMD-viewer laser-plasma tools\n",
    "\n",
    "In addition to the general methods `get_field` and `get_particle`, openPMD-viewer can also include a set of domain-specific tools. For instance, this notebook describes a set of methods that are useful when analyzing simulations of **laser-plasma acceleration**. \n",
    "\n",
    "If you are not interested in laser-plasma simulations, you can skip this notebook as these methods are only an add-on to openPMD-viewer. \n",
    "\n",
    "You can run this notebook locally by downloading it from [this link](https://github.com/openPMD/openPMD-viewer/blob/dev/docs/source/tutorials/5_Laser-plasma_tools.ipynb).\n",
    "\n",
    "## (optional) Preparing this notebook to run it locally\n",
    "\n",
    "If you choose to run this notebook on your local machine, you will need to download the openPMD data files which will then be analysed. To do so, execute the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "def download_if_absent( dataset_name ):\n",
    "    \"Function that downloads and decompress a chosen dataset\"\n",
    "    if os.path.exists( dataset_name ) is False:\n",
    "        import wget, tarfile\n",
    "        tar_name = \"%s.tar.gz\" %dataset_name\n",
    "        url = \"https://github.com/openPMD/openPMD-example-datasets/raw/draft/%s\" %tar_name\n",
    "        wget.download(url, tar_name)\n",
    "        with tarfile.open( tar_name ) as tar_file:\n",
    "            tar_file.extractall()\n",
    "        os.remove( tar_name )\n",
    "\n",
    "download_if_absent( 'example-2d' )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, we choose here to incorporate the plots inside the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The LpaDiagnostics class\n",
    "\n",
    "To use the laser-plasma acceleration (LPA) tools:\n",
    "- Load the class `LpaDiagnostics` from the module `openpmd_viewer.addons`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openpmd_viewer.addons import LpaDiagnostics"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Create an `LpaDiagnostics` instance in the same way, as you would do for an `OpenPMDTimeSeries`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d = LpaDiagnostics('./example-2d/hdf5/')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `LpaDiagnostics` class inherits from the `OpenPMDTimeSeries` class, and therefore it also has the methods `get_field`, `get_particle` and `slider`. For instance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ts_2d.slider()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Diagnostic methods\n",
    "\n",
    "In addition to the methods `get_particle`, `get_field` and `slider`, the `LpaDiagnotics` class has an extra set of specialized methods.\n",
    "\n",
    "These diagnostic methods are currently only implemented for API-like usage. In particular, either the time or iteration have to be specified when calling a method, by setting the `t` or `iteration` parameter. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Particle diagnostics\n",
    "\n",
    "For methods involving particle data, the desired particle species also needs to be specified with the `species` parameter. Optionally, these methods can be applied to only a subset of the particles, by using the `select` parameter.\n",
    "\n",
    "For instance, `select={'uz' : [-1, 2]}` will select only the particles which have a longitudinal normalized momentum between `-1` and `2`. \n",
    "\n",
    "In the following the available particle diagnostic methods will be explained.\n",
    "\n",
    "For more information a method's documentation can be called with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_mean_gamma?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Charge\n",
    "`get_charge` calculates the charge of the given particle selection in Coulomb."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_charge(iteration=300, species='electrons')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the evolution of the charge (or of any of the quantities below) can be easily obtained with `ts.iterate`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.iterate( ts_2d.get_charge, species='electrons' )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Mean gamma\n",
    "To calculate the mean gamma value and standard deviation of the selected particles `get_mean_gamma` can be used. In the example below, only the particles with $u_z > 0.05$ are selected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_mean_gamma(iteration=300, species='electrons', select={'uz' : [0.05, None]})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Divergence\n",
    "This method calculates the divergence of the selected particles, using $\\langle \\arctan{u_{x/y}/u_z} \\rangle$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_divergence(iteration=300, species='electrons')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Emittance\n",
    "This method can be used to calculate the normalized emittance in the x and y plane for selected particles by evaluating\n",
    "$$\\epsilon_{n,rms}=\\sqrt{\\langle x^2 \\rangle \\langle u_x^2 \\rangle - \\langle x u_x \\rangle^2}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_emittance(iteration=300, species='electrons')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Current\n",
    "`get_current` can be used to calculate the instantaneous current along the z_axis generated by the selected particles. When setting `plot=True` the resulting current profile is directly plotted. Otherwise an array with the data is returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_current(iteration=300, species='electrons', plot=True);"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Laser diagnostic\n",
    "The laser diagnostic methods require the user to specify the plane of laser polarisation by setting the argument `pol=` to either `'x'` or `'y'`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Laser envelope\n",
    "\n",
    "This method calculates the envelope of a given laser field. This can be done for a 1D slice of the field or for an entire 2D plane. The resulting data is returned in form of a 1D or 2D array, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_laser_envelope(iteration=300, pol='y');"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Spectrum\n",
    "This function helps to easily calculate (and plot) the spectrum of a given laser field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_spectrum(iteration=300, pol='y', plot=True);"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Spectrogram\n",
    "This method does a time-frequency analysis of the laser, by applying the FROG method.\n",
    "\n",
    "Mathematically:\n",
    "        $$ s(\\omega, \\tau) = \\left| \\int_{-\\infty}^{\\infty} E(t) |E(t-\\tau)|^2\n",
    "            \\exp( -i\\omega t) dt \\right|^2 $$\n",
    "            \n",
    "(Additional matplotlib arguments for the plotting option can directly be passed to the function, e.g `cmap='coolwarm'`\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_spectrogram(iteration=300, pol='y', plot=True, cmap='YlGnBu_r');"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Main frequency\n",
    "To calculate the main frequency (i.e. maximum of the spectrum) call. This returns the frequency in $rad.s^{-1}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_main_frequency(iteration=300, pol='y')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Laser $a_0$\n",
    "A method to calculate the laser strength a0 as given by $$a0 = E_{max} e / (m_e c \\omega)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_a0(iteration=300, pol='y')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Laser waist\n",
    "\n",
    "Calculate the waist of a (gaussian) laser pulse (i.e. $\\sqrt{2} \\sigma_r$, where $\\sigma_r$ is the transverse RMS of the field)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_laser_waist(iteration=300, pol='y')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pulse length\n",
    "Calculate the length of a (gaussian) laser pulse. Here 'length' means the 'longitudinal waist' (i.e $\\sqrt{2} \\sigma_z$ where $\\sigma_z$ is the longitudinal RMS of the field).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_2d.get_ctau(iteration=300, pol='y')"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.5"
  },
  "widgets": {
   "state": {
    "7f2b0c4c82444cb881753f9785f81e6c": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "8c3f18cf657d48afaffabcf9950a8d20": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "b0a0574bc9c34b10a7c2c58f30a7ab78": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    },
    "dcbbddf44d6742b7bf3ccfe452fa56aa": {
     "views": [
      {
       "cell_index": 9
      }
     ]
    }
   },
   "version": "1.2.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}