{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 006: Imported Field Trajectory\n",
    "\n",
    "In this example the trajectory is calculated for a field coming from an imported data file.  We first creat the data file by adding some undulator field, some noise, and exporting the field to a file.  Several formats are possible (for more information see the All About Magnetic Fields tutorial."
   ]
  },
  {
   "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 a Magnetic field and Export it to a file\n",
    "\n",
    "We are only doing this so you don't have to download many data files."
   ]
  },
  {
   "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=41)\n",
    "\n",
    "# Add some random imperfections\n",
    "for i in range(10):\n",
    "    osr.add_bfield_gaussian(\n",
    "        bfield=[0, 0.001 * (osr.rand() - 0.5), 0],\n",
    "        sigma=[0, 0, 0.1*(osr.rand())],\n",
    "        translation=[0, 0, 2 * (osr.rand() - 0.5)]\n",
    "    )\n",
    "\n",
    "\n",
    "# Export the field\n",
    "osr.write_bfield(ofile='EX006.dat', oformat='OSCARS', zlim=[-3, 3], nz=50000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import the magnetic field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Clear all magnetic fields!\n",
    "osr.clear_bfields()\n",
    "\n",
    "# Import the field from the file created above\n",
    "osr.add_bfield_file(ifile='EX006.dat', iformat='OSCARS')\n",
    "\n",
    "# Plot imported field\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",
    "    * d0=[0, 0, 1]\n",
    "\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=6 (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 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": "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": "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
}