{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Particle Beams\n",
    "\n",
    "This example describes all of the major capabilities of particle beams in OSCARS.  There are pre-defined particle types (default is electron, but proton, pion, etc are included) and the user is free to describe any other particle type they wish.  Non-zero beam widths can be defined using the emittance, beta functions, and lattice reference (point at which the beta functions are defined).  One can also define multiple particle beams in one OSCARS object and give weights to each.\n",
    "\n",
    "Patricle beams can as well be defined, like other OSCARS objects, in any position and direction.  Although examples are often given where $+\\hat z$ is the beam direction (mostly to align with accelerator nomenclature), there is no preferred direction in OSCARS.\n",
    "\n",
    "The minimum definition for a particle beam must include calling the two functions:\n",
    "* set_particle_beam()\n",
    "* set_ctstartstop(tstart, tstop)\n",
    "\n",
    "where tstart and tstop are the start and stop times for particle propoation.  In the above set_particle_beam() will use all of its defaults (since no arguments are given) which is essentially an electron at rest at the origin.  The defaults for set_particle_beam() are:\n",
    "* type='electron'\n",
    "* name  (auto-assigned)\n",
    "* energy_GeV=0\n",
    "* x0=[0, 0, 0]\n",
    "* d0=[0, 0, 1]\n",
    "* t0=0\n",
    "* current=1.602176462e-19  (single electron charge)\n",
    "* emittance=[0, 0]\n",
    "\n",
    "The elements which you may want to define for a beam are:\n",
    "* type  (one of the predefined types, or 'custom')\n",
    "* name  (any unique name you like)\n",
    "* energy_GeV  (the energy in GeV)\n",
    "* x0  (position where initial conditions are given)\n",
    "* d0  (direction of beam at initial position)\n",
    "* current  (the current in [A])\n",
    "\n",
    "For a non-zero emittance beam you must additionally define the following (the beam shape is taken to be gaussian, though other distributions are possible):\n",
    "* sigma_energy_GeV  (width in energy spread in [GeV])\n",
    "* emittance  (the beam emittance in [m.rad] as [h, v])\n",
    "* beta  (the beta function in [m] as [h, v])\n",
    "* lattice_reference  (the reference point at which beta is defined)\n",
    "* horizontal_direction  (the horizontal direction of the beam)\n",
    "\n",
    "For a particle beam you can additionally define the following:\n",
    "* rotations  (rotate the beam about x, then y, then z given as [$\\theta_x, \\theta_y, \\theta_z$])\n",
    "* translation  (translate the beam in x, y, z given as [x, y, z])\n",
    "* t0  (time at the position x0 defined above, default is 0)\n",
    "* weight  (for randome particle generation with multiple beams, default is 1)\n",
    "\n",
    "Custom:  If you set type='custom' you will need to also define the following:\n",
    "* mass (in [kg])\n",
    "* charge (in [C])\n",
    "\n",
    "The two functions for adding a particle beam are:\n",
    "* set_particle_beam() - Removes all other particle beams\n",
    "* add_particle_beam() - The same as set but does not remove previous beams\n",
    "\n",
    "To see what beam(s) are setup in OSCARS you may call:\n",
    "* print_particle_beams()\n",
    "\n",
    "or to see everything we recommend:\n",
    "* print_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# matplotlib plots inline\n",
    "%matplotlib inline\n",
    "\n",
    "# Import the OSCARS SR module\n",
    "import oscars.sr\n",
    "\n",
    "# Import OSCARS plots (matplotlib)\n",
    "from oscars.plots_mpl import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an OSCARS SR object\n",
    "osr = oscars.sr.sr()\n",
    "\n",
    "# For these examples we will make use of a simple undulator field\n",
    "osr.add_bfield_undulator(bfield=[0, 1, 0], period=[0, 0, 0.100], nperiods=5)\n",
    "\n",
    "plot_bfield(osr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic  beams /  Particles\n",
    "\n",
    "These are basic beams with zero emittance, one electron (NSLSII energy) and one proton (LHC energy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Electron\n",
    "\n",
    "First example is showing full definition, second using the defaults for less typing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a basic electron beam with zero emittance\n",
    "osr.set_particle_beam(type='electron', name='beam_0', energy_GeV=3, x0=[0, 0, -1], d0=[0, 0, 1], current=0.500)\n",
    "\n",
    "# You MUST set the start and stop time for the calculation\n",
    "osr.set_ctstartstop(0, 2)\n",
    "\n",
    "# Plot trajectory of beam\n",
    "plot_trajectory_position(osr.calculate_trajectory())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a basic electron beam with zero emittance\n",
    "osr.set_particle_beam(energy_GeV=3, x0=[0, 0, -1], current=0.500)\n",
    "\n",
    "# You MUST set the start and stop time for the calculation\n",
    "osr.set_ctstartstop(0, 2)\n",
    "\n",
    "# Plot trajectory of beam\n",
    "plot_trajectory_position(osr.calculate_trajectory())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Proton beam\n",
    "\n",
    "This example shows how to define a proton beam at 13 TeV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a basic proton beam with zero emittance\n",
    "osr.set_particle_beam(type='proton', energy_GeV=13000, x0=[0, 0, -1], current=0.500)\n",
    "\n",
    "# You MUST set the start and stop time for the calculation\n",
    "osr.set_ctstartstop(0, 2)\n",
    "\n",
    "# Plot trajectory of beam\n",
    "plot_trajectory_position(osr.calculate_trajectory())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rotating Beams\n",
    "\n",
    "Beams can easily be rotated using the *rotations* variable.  This essentialyl rotates the initial conditions given.  It rotates the initial position *x0* in space and rotates the direction vector *d0*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a basic proton beam and rotate it 45 degrees\n",
    "osr.set_particle_beam(type='proton',\n",
    "                      energy_GeV=13000,\n",
    "                      x0=[0, 0, -1],\n",
    "                      current=0.500,\n",
    "                      rotations=[osr.pi()/4., 0, 0]\n",
    "                     )\n",
    "\n",
    "# You MUST set the start and stop time for the calculation\n",
    "osr.set_ctstartstop(0, 2)\n",
    "\n",
    "# Plot trajectory of beam\n",
    "plot_trajectory_position(osr.calculate_trajectory())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple beams\n",
    "\n",
    "This examples defined multiple beams in the same object, then plots the trajectory for two random particles selected from these beams.\n",
    "\n",
    "Instead of using set_particle_beam we will use add_particle_beam (which is really the same thing, but set_particle_beam issues a clear_particle_beams() internally)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add basic electron beams with zero emittance\n",
    "osr.clear_particle_beams()\n",
    "osr.add_particle_beam(type='electron', name='beam_0', energy_GeV=3,  x0=[0, 0, -1], current=0.500)\n",
    "osr.add_particle_beam(type='positron', name='beam_1', energy_GeV=3,  x0=[0, 0, -1], current=0.500)\n",
    "osr.add_particle_beam(type='proton',   name='beam_2', energy_GeV=13, x0=[0, 0, -1], current=0.500)\n",
    "osr.add_particle_beam(type='muon',     name='beam_3', energy_GeV=2,  x0=[0, 0, -1], current=0.500)\n",
    "osr.add_particle_beam(type='pi+',      name='beam_4', energy_GeV=1,  x0=[0, 0, -1], current=0.500)\n",
    "\n",
    "# You MUST set the start and stop time for the calculation\n",
    "osr.set_ctstartstop(0, 2)\n",
    "\n",
    "# Plot trajectory of beam\n",
    "plot_trajectory_position(osr.calculate_trajectory())\n",
    "osr.set_new_particle()\n",
    "plot_trajectory_position(osr.calculate_trajectory())\n",
    "osr.set_new_particle()\n",
    "plot_trajectory_position(osr.calculate_trajectory())\n",
    "\n",
    "# Print particle beam info for checking\n",
    "osr.print_particle_beams()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Beam with non-zero emittance\n",
    "\n",
    "Below is an example of a beam with non-zero memittance.  We use the emittance and the beta function.  You must give the reference point for the beta function definition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a basic electron beam with zero emittance\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",
    "# You MUST set the start and stop time for the calculation\n",
    "osr.set_ctstartstop(0, 2)\n",
    "\n",
    "# Plot trajectory of beam\n",
    "osr.set_new_particle()\n",
    "plot_trajectory_position(osr.calculate_trajectory())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's plot the distribution of particles at the initial position 'position'\n",
    "energy = []\n",
    "x = []\n",
    "y = []\n",
    "bx = []\n",
    "by = []\n",
    "for i in range(5000):\n",
    "    osr.set_new_particle()\n",
    "    energy.append(osr.get_particle_e0())\n",
    "    position = osr.get_particle_x0()\n",
    "    beta = osr.get_particle_beta0()\n",
    "\n",
    "    x.append(position[0])\n",
    "    y.append(position[1])\n",
    "\n",
    "    bx.append(beta[0])\n",
    "    by.append(beta[1])\n",
    "    \n",
    "plt.figure(1, figsize=[15, 3])\n",
    "\n",
    "plt.subplot(131)\n",
    "plt.scatter(x, y)\n",
    "plt.xlim(-0.0005, 0.0005)\n",
    "plt.ylim(-0.00005, 0.00005)\n",
    "plt.xlabel('x0 [m]')\n",
    "plt.ylabel('y0 [m]')\n",
    "plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))\n",
    "\n",
    "plt.subplot(132)\n",
    "plt.scatter(bx, by)\n",
    "plt.xlim(-0.0002, 0.0002)\n",
    "plt.ylim(-0.00002, 0.00002)\n",
    "plt.xlabel('x0 [m]')\n",
    "plt.ylabel('y0 [m]')\n",
    "plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))\n",
    "\n",
    "plt.subplot(133)\n",
    "plt.hist(energy)\n",
    "plt.xlabel('Energy [GeV]')\n",
    "plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))\n",
    "plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple beams with non-zero emittance\n",
    "\n",
    "A simple example of 2 beams with different beam parameters in the same object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add basic electron beams with zero emittance\n",
    "osr.clear_particle_beams()\n",
    "osr.add_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",
    "    weight=0.7\n",
    ")\n",
    "\n",
    "osr.add_particle_beam(\n",
    "    energy_GeV=3.02,\n",
    "    x0=[0, 0.000035, -1],\n",
    "    d0=[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",
    "    weight=0.3\n",
    ")\n",
    "\n",
    "\n",
    "# You MUST set the start and stop time for the calculation\n",
    "osr.set_ctstartstop(0, 2)\n",
    "\n",
    "# And plot the results\n",
    "energy = []\n",
    "x = []\n",
    "y = []\n",
    "for i in range(5000):\n",
    "    osr.set_new_particle()\n",
    "    energy.append(osr.get_particle_e0())\n",
    "    position = osr.get_particle_x0()\n",
    "\n",
    "    x.append(position[0])\n",
    "    y.append(position[1])\n",
    "    \n",
    "plt.figure(1, figsize=[10, 3])\n",
    "plt.subplot(121)\n",
    "plt.scatter(x, y)\n",
    "plt.xlim(-0.0005, 0.0005)\n",
    "plt.ylim(-0.00005, 0.00005)\n",
    "plt.xlabel('x0 [m]')\n",
    "plt.ylabel('y0 [m]')\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.hist(energy)\n",
    "plt.xlabel('Energy [GeV]')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Particles\n",
    "\n",
    "This is an example of a custom particle in OSCARS.  All that additionally needs definition are the charge and mass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a custom beam with zero emittance.\n",
    "# This is for a particle with twice the chargs as the electron, but the same mass\n",
    "osr.set_particle_beam(\n",
    "    type='custom',\n",
    "    charge=-2*osr.qe(),\n",
    "    mass=osr.me(),\n",
    "    energy_GeV=3,\n",
    "    x0=[0, 0, -1],\n",
    "    current=0.500\n",
    ")\n",
    "\n",
    "# You MUST set the start and stop time for the calculation\n",
    "osr.set_ctstartstop(0, 2)\n",
    "\n",
    "# Plot trajectory of beam\n",
    "plot_trajectory_position(osr.calculate_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
}