{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 007: Importing Trajectories\n",
    "\n",
    "It is possible to import an externally calculated trajectory to oscars.sr.  The trajectory may be input in the OSCARS trajectory format directly from python or it can be input from a text file with user specified input format.  For inputting from a file the minimum required is time and position in at least 1 spatial dimension.\n",
    "\n",
    "Defining a beam is required before importing so that OSCARS can associate the trajectory with a beam (charge required for calculations).  Note that ctstart and ctstop are not needed.  Calculations will be based on the input trajectry time.\n",
    "\n",
    "This example will serve as a guide for inputting trajectories of various types.\n",
    "\n",
    "If only the position is given the first and second derivatives are calculated from the derivative coefficients of a cubic spline.  Higher order derivatives are less accurate and it is up to the user to dertermine if this is sufficient for the purpose."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Standard setup and importing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inline plots\n",
    "%matplotlib inline\n",
    "\n",
    "# Import oscars modules and plotting tools\n",
    "import oscars.sr\n",
    "from oscars.plots_mpl import *\n",
    "\n",
    "# Other python imports\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an oscars.sr object for calculations (use the GPU if you can!)\n",
    "osr = oscars.sr.sr(nthreads=10, gpu=1)\n",
    "\n",
    "# Define a beam before we import any trajectories\n",
    "# For imported trajectories ctstartstop is not needed, but is if you use internal propogation\n",
    "osr.set_particle_beam(type='electron', current=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Trajectory as a Python List\n",
    "\n",
    "The trajectory must be in the OSCARS trajectory format:\n",
    "* [[time, [x, y, z], [bx, by, bz], [bpx, bpy, bpz]], ...]\n",
    "where b refers to the relativistic beta (velicoty divided by speed of light) and bp is beta' or acceleration divided by the speed of light."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importation of a python list trajectory.  This is just a straight trajectory.\n",
    "beta = 0.9999999854933347\n",
    "my_trajectory = [[t, [0, 0, t*beta*osr.c()], [0, 0, beta], [0, 0, 0]] for t in np.linspace(0, 1e-7, 100)]\n",
    "osr.set_trajectory(trajectory=my_trajectory)\n",
    "\n",
    "# Plot the trajectory\n",
    "plot_trajectory_position(osr.get_trajectory())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# It is possible to let OSCARS calculate the beta and beta' terms, just leave them out\n",
    "# of the trajectory\n",
    "beta = 0.9999999854933347\n",
    "my_trajectory = [[t, [1e-6 * np.sin(2.*osr.pi()*t*beta*osr.c()/0.050), 0, t*beta*osr.c()]] for t in np.linspace(0, 1e-8, 10000)]\n",
    "osr.set_trajectory(trajectory=my_trajectory)\n",
    "\n",
    "# Plot the trajectory\n",
    "plot_trajectory_position(osr.get_trajectory())\n",
    "plot_trajectory_velocity(osr.get_trajectory())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# As an example we calculate the power density\n",
    "pd = osr.calculate_power_density_rectangle(plane='XY', width=[0.01, 0.01], npoints=[51, 51], translation=[0, 0, 30])\n",
    "plot_power_density(pd)\n",
    "print('Total Power:', osr.calculate_total_power(), '[W]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing Data: Text File\n",
    "\n",
    "Data can be imported from a text file in column format.  These must be spece delimited columns.  You can specify the file format with the 'iformat' parameter.  The default format for this is 'T X Y Z BX BY BZ BPX BPY BPZ'.  At a minimum T and one of (X, Y, Z) must be specified.  '*' in iformat will ignore a column.\n",
    "\n",
    "Below we first create some trajectory in the default format, then read it in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Create an EPU trajectory.  Here we need to define a beam with ctstartstop\n",
    "osr.set_particle_beam(beam='NSLSII', x0=[0, 0, -1], ctstartstop=[0, 2])\n",
    "osr.clear_bfields()\n",
    "osr.add_bfield_undulator(bfield=[0, 1, 0], period=[0, 0, 0.042], nperiods=31, phase=-0.25*osr.pi())\n",
    "osr.add_bfield_undulator(bfield=[1, 0, 0], period=[0, 0, 0.042], nperiods=31, phase=+0.25*osr.pi())\n",
    "\n",
    "# Calculate trajectory and export it to a file\n",
    "t = osr.calculate_trajectory(ofile='Example_007_Trajectory1.txt')\n",
    "plot_trajectory_position(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the trajectory has been created we can read it in using the 'ifile' argument"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Read trajectory in default format\n",
    "osr.set_trajectory(ifile='Example_007_Trajectory1.txt')\n",
    "\n",
    "# Get and plot the trajectory\n",
    "t = osr.get_trajectory()\n",
    "plot_trajectory_position(t)\n",
    "plot_trajectory_velocity(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# As an example we can now calculate the power density\n",
    "pd = osr.calculate_power_density_rectangle(plane='XY', width=[0.1, 0.1], npoints=[51, 51], translation=[0, 0, 30])\n",
    "plot_power_density(pd)\n",
    "print('Total Power:', osr.calculate_total_power(), '[W]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we only ask for the time and space components, OSCARS will calculate the beta and beta' terms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read trajectory in default format\n",
    "osr.set_trajectory(ifile='Example_007_Trajectory1.txt', iformat='T X Y Z * * * * * *')\n",
    "\n",
    "# Get and plot the trajectory\n",
    "t = osr.get_trajectory()\n",
    "plot_trajectory_position(t)\n",
    "plot_trajectory_velocity(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# As an example we can now calculate the power density\n",
    "pd = osr.calculate_power_density_rectangle(plane='XY', width=[0.1, 0.1], npoints=[51, 51], translation=[0, 0, 30])\n",
    "plot_power_density(pd)\n",
    "print('Total Power:', osr.calculate_total_power(), '[W]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ignore Columns\n",
    "\n",
    "One can ignore columns in the input format using '*'.  This is given as an example of how to ignore a column of data only and it is not recommended to ignore non-zero trajectory components since they are used in SR calculations.  Note that in the example below beta and beta' will be technically incorrect.\n",
    "\n",
    "***This feature is meant only for ignoring extra data columns.  Do not ignore trajectory data.***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read trajectory in default format\n",
    "osr.set_trajectory(ifile='Example_007_Trajectory1.txt', iformat='T X * Z')\n",
    "\n",
    "# Get and plot the trajectory\n",
    "t = osr.get_trajectory()\n",
    "plot_trajectory_position(t)\n",
    "plot_trajectory_velocity(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing Data: Binary File\n",
    "\n",
    "Data can be imported from a binary file.  Here the iformat argument is only needed for non-oscars binary files.  It follows the same convention as the text file iformat.  The default format for this is 'T X Y Z BX BY BZ BPX BPY BPZ'.  For an OSCARS generated binary file one should NOT specify iformat because the format is described within the binary file itself.  In binary format each element should be a 64-bit double.\n",
    "\n",
    "Below we first create some trajectory in the default format, then read it in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an EPU trajectory and save it to a file\n",
    "osr.set_particle_beam(beam='NSLSII', x0=[0, 0, -1], ctstartstop=[0, 2])\n",
    "osr.clear_bfields()\n",
    "osr.add_bfield_undulator(bfield=[0, 1, 0], period=[0, 0, 0.042], nperiods=31, phase=-0.25*osr.pi())\n",
    "osr.add_bfield_undulator(bfield=[1, 0, 0], period=[0, 0, 0.042], nperiods=31, phase=+0.25*osr.pi())\n",
    "\n",
    "t = osr.calculate_trajectory(bofile='Example_007_Trajectory1.dat')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read trajectory in default format\n",
    "osr.set_trajectory(bifile='Example_007_Trajectory1.dat')\n",
    "\n",
    "# Get and plot the trajectory\n",
    "t = osr.get_trajectory()\n",
    "plot_trajectory_position(t)\n",
    "plot_trajectory_velocity(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# As an example we can now calculate the power density\n",
    "pd = osr.calculate_power_density_rectangle(plane='XY', width=[0.1, 0.1], npoints=[51, 51], translation=[0, 0, 30])\n",
    "plot_power_density(pd)\n",
    "print('Total Power:', osr.calculate_total_power(), '[W]')"
   ]
  }
 ],
 "metadata": {
  "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": 2
}