{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 142: Undulator Power Density Multi-Particle\n",
    "\n",
    "In this example multi-particle power density calculation from an undulator for a non-zero emittance electron beam."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": 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 the undulator field.  Here *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*.  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": [
    "# 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, 1, 0], period=[0, 0, 0.049], nperiods=31)\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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Setup beam similar to NSLSII\n",
    "osr.set_particle_beam(\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 on the ideal case\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate Power Density\n",
    "\n",
    "Calcuate the power density on a rectangular surface downstream 30 [m] from the center of the device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Calcuate flux at 30 [m] for 152 eV\n",
    "pd_multi = osr.calculate_power_density_rectangle(\n",
    "    plane='XY',\n",
    "    width=[0.05, 0.05],\n",
    "    npoints=[101, 101],\n",
    "    translation=[0, 0, 30],\n",
    "    nparticles=5\n",
    ")\n",
    "\n",
    "plot_power_density(pd_multi, ofile='UndulatorPowerDensity_MultiParticle.png')"
   ]
  },
  {
   "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
}