{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 133: EPU Flux - Multi-particle\n",
    "\n",
    "In this example we calculate the multi-particle flux near the peak on a rectangular surface 30 [m] downstream from the center of the EPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "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": {
    "collapsed": true
   },
   "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": {
    "collapsed": true
   },
   "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(\n",
    "    type='electron',\n",
    "    name='beam_0',\n",
    "    energy_GeV=3,\n",
    "    x0=[0, 0, -1],\n",
    "    current=0.500,\n",
    "    sigma_energy_GeV=0.001*3,\n",
    "    beta=[1.5, 0.8],\n",
    "    emittance=[0.9e-9, 0.008e-9]\n",
    ")\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": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Run the particle trajectory calculation for the ideal particle\n",
    "# First must load the ideal particle initial conditions\n",
    "osr.set_new_particle(particle='ideal')\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": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Calculate spectrum zoom for ideal particle.  In reality you may want to run the multi-particle\n",
    "# spectrum calculation, then pick from there what energy we want to see the spectrum at\n",
    "osr.set_new_particle(particle='ideal')\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": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Calculate the multi-particle Flux.  We'll pick a point around 154 [eV]\n",
    "# Although 5 particles are used here, you shuold use a lot more.\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",
    "    nparticles=5\n",
    ")\n",
    "\n",
    "plot_flux(flux)"
   ]
  },
  {
   "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
}