{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 033: EPU Flux\n",
    "\n",
    "In this example we calculate the flux near the peak on a rectangular surface 30 [m] downstream from the center of the EPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# This has nothing to do with OSCARS, but it puts the matplotlib plots inline in the notebook\n",
    "%matplotlib inline\n",
    "\n",
    "# Import the OSCARS SR module\n",
    "import oscars.sr\n",
    "\n",
    "# Import basic plot utilities (matplotlib).  You don't need these to run OSCARS, but it's used here for basic plots\n",
    "from oscars.plots_mpl import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new OSCARS object.  Default to 8 threads and always use the GPU if available\n",
    "osr = oscars.sr.sr(nthreads=8, gpu=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the undulator field\n",
    "\n",
    "Here we create an EPU by adding two ideal undulators, one with a vertical field and one with a horizontal field, out of phase.\n",
    "\n",
    "For the undulator, *bfield* represents maximum magnetic field [$B_x, B_y, B_z$].  The *period* is also in vector form which allows you to orient the axis of the undulator in any arbitrary direction.  The number of periods is given by *nperiods*.  This is the number of FULL periods.  A terminating field of 1 period length is added to each side in addition to *nperiods*.\n",
    "\n",
    "Typically clear_magnetic_fields() is called before adding a field in notebooks only to save time when making changes and rerunning sections of the code so it is not strictly necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Phase difference between fields in [rad]\n",
    "phase = osr.pi()/2.\n",
    "\n",
    "# Clear any existing fields (just good habit in notebook style) and add an undulator field\n",
    "osr.clear_bfields()\n",
    "osr.add_bfield_undulator(bfield=[0, 0.7, 0], period=[0, 0, 0.049], nperiods=21, phase=-phase/2.)\n",
    "osr.add_bfield_undulator(bfield=[0.7, 0, 0], period=[0, 0, 0.049], nperiods=21, phase=+phase/2.)\n",
    "\n",
    "# Just to check the field that we added seems visually correct\n",
    "plot_bfield(osr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add a particle beam\n",
    "\n",
    "Here we add a particle beam making use of some of the defaults, namely:\n",
    "    * type='electron'\n",
    "    * t0=0\n",
    "    * d0=[0, 0, 1]\n",
    "\n",
    "One must specify ctstartstop.  This is the start and stop time of the calculation.  In this example we will start the calculation at t=0 and go to t=2 (given in units of ct) since the beam is relativistic.  In this example you can specify the start time as less than 0 which is useful if you want to propogate the particle backwars in time.  This is useful for instance if you have a bending magnet before the undulator that you wish to include.\n",
    "\n",
    "clear_particle_beams() is called, again for convenience, but it is not necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Setup beam similar to NSLSII\n",
    "osr.clear_particle_beams()\n",
    "osr.set_particle_beam(x0=[0, 0, -1], energy_GeV=3, current=0.500)\n",
    "\n",
    "# Set the start and stop times for the calculation\n",
    "osr.set_ctstartstop(0, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate Trajectory\n",
    "\n",
    "Now we calculate the trajectory and plot it.  It is enough to call calculate_trajectory().  If you are doing other calculations (flux, spectra, power density) it is not necesary to call this since it is called internally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the particle trajectory calculation\n",
    "trajectory = osr.calculate_trajectory()\n",
    "\n",
    "# Plot the trajectory position and velocity\n",
    "plot_trajectory_position(trajectory)\n",
    "plot_trajectory_velocity(trajectory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate Spectrum\n",
    "\n",
    "Calculate the spectrum in a given range with the given number of points equally spaced (not a requirement)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate spectrum zoom\n",
    "spectrum = osr.calculate_spectrum(obs=[0, 0, 30], energy_range_eV=[140, 170], npoints=200)\n",
    "plot_spectrum(spectrum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate Flux"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate Flux\n",
    "flux = osr.calculate_flux_rectangle(\n",
    "    plane='XY',\n",
    "    energy_eV=154,\n",
    "    width=[0.01, 0.01],\n",
    "    npoints=[101, 101],\n",
    "    translation=[0, 0, 30]\n",
    ")\n",
    "\n",
    "plot_flux(flux)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Polarization\n",
    "\n",
    "It may be of interested to calculate polarization components for such a device.  Below is shown the circular-right and circular left components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate flux, circular right and left, short versions are: 'cr', 'cl'\n",
    "flux = osr.calculate_flux_rectangle(\n",
    "    plane='XY',\n",
    "    energy_eV=154,\n",
    "    width=[0.01, 0.01],\n",
    "    npoints=[101, 101],\n",
    "    translation=[0, 0, 30],\n",
    "    polarization='circular-right'\n",
    ")\n",
    "\n",
    "plot_flux(flux, title='circular-right Polarization')\n",
    "\n",
    "flux = osr.calculate_flux_rectangle(\n",
    "    plane='XY',\n",
    "    energy_eV=154,\n",
    "    width=[0.01, 0.01],\n",
    "    npoints=[101, 101],\n",
    "    translation=[0, 0, 30],\n",
    "    polarization='circular-left'\n",
    ")\n",
    "\n",
    "plot_flux(flux, title='circular-left Polarization', xlabel='X [m]', ylabel='Y [m]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}