{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 005: Canted Striaght with 2 Undulators\n",
    "\n",
    "In this example the electron trajectory in a mock-straight section with two undulators and center canting magnet is calculated in two ways.  Firstly it is calculated with initial conditions given at z=0 (in the middle), then for the initial conditions given before the first undulator.\n",
    "\n"
   ]
  },
  {
   "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 and Canting Magnet Fields\n",
    "\n",
    "Here we create the undulator fields and Bending Magnet fields.\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",
    "For the canting magnets we will use the gaussian field.  *bfield* represents maximum magnetic field [$B_x, B_y, B_z$].  *sigma* is the 1-sigma width and can be given an any direction (if zero, that coordinate is ignored).\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.\n",
    "\n",
    "The undulators and canting magnets are placed in space using the *translate* option."
   ]
  },
  {
   "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",
    "dist_between_centers = 0.049*15\n",
    "\n",
    "osr.add_bfield_gaussian(bfield=[0, +0.08, 0], sigma=[0, 0, 0.05], translation=[0, 0, -1.8])\n",
    "osr.add_bfield_undulator(bfield=[0, 1, 0], period=[0, 0, 0.049], nperiods=21, translation=[0, 0, -dist_between_centers])\n",
    "osr.add_bfield_gaussian(bfield=[0, -0.16, 0], sigma=[0, 0, 0.05], translation=[0, 0, 0])\n",
    "osr.add_bfield_undulator(bfield=[0, 1, 0], period=[0, 0, 0.049], nperiods=21, translation=[0, 0, +dist_between_centers])\n",
    "osr.add_bfield_gaussian(bfield=[0, +0.08, 0], sigma=[0, 0, 0.05], translation=[0, 0, 1.8])\n",
    "\n",
    "# Just to check the field that we added seems visually correct\n",
    "plot_bfield(osr, -2, 2)"
   ]
  },
  {
   "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",
    "    * x0=[0, 0, 0]\n",
    "    * d0=[0, 0, 1]\n",
    "\n",
    "For perfectly symmetric data sometimes it is convenient to set the initial conditions in the middle of your scenario.  Often with real data it is more practical to set it before entry of the first device of interest.  Both of these are shown below.\n",
    "\n",
    "One must specify ctstartstop, and they are notebly different for the two examples of beam setups below.\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(energy_GeV=3, current=0.500)\n",
    "\n",
    "# Set the start and stop times for the calculation\n",
    "osr.set_ctstartstop(-3, 3)"
   ]
  },
  {
   "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\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": [
    "Now set the initial conditions before entry of the first canting magnet (often this is more practical)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Setup beam similar to NSLSII with different starting position from above\n",
    "# (this makes more sense for some scenarios)\n",
    "osr.clear_particle_beams()\n",
    "osr.set_particle_beam(x0=[0, 0, -3], energy_GeV=3, current=0.500)\n",
    "\n",
    "# Set the start and stop times for the calculation\n",
    "osr.set_ctstartstop(0, 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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": "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
}