{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Undulator Power Density on a CAD Generated Object\n",
    "\n",
    "This is a simple example of an undulator power density calculated on a CAD object.  Currently OSCARS supports the STL (binary) format.  This is easily exported from your typical CAD software.\n",
    "\n",
    "This example folder contains one simple cad object file which is used below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OSCARS v2.1.1 - Open Source Code for Advanced Radiation Simulation\n",
      "Brookhaven National Laboratory, Upton NY, USA\n",
      "http://oscars.bnl.gov\n",
      "oscars@bnl.gov\n"
     ]
    }
   ],
   "source": [
    "# Plots inline for notebook\n",
    "%matplotlib inline\n",
    "\n",
    "# Import the OSCARS SR module\n",
    "import oscars.sr\n",
    "\n",
    "# Import the 3D and parametric surfaces utilities\n",
    "from oscars.plots3d_mpl import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU is not available: Setting gpu global setting to 0.\n"
     ]
    }
   ],
   "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": [
    "## Undulator Field\n",
    "\n",
    "Create a simple undulator field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Clear all existing fields and create an undulator field\n",
    "osr.clear_bfields()\n",
    "osr.add_bfield_undulator(bfield=[0, 1, 0], period=[0, 0, 0.050], nperiods=31)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Particle Beam and Trajectory\n",
    "\n",
    "Define a particle beam, in this case a 3 [GeV] electron beam.  You must also define the start and stop times for the calculation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Define simple electron beam\n",
    "osr.set_particle_beam(energy_GeV=3, x0=[0, 0, -1], current=0.5)\n",
    "\n",
    "# Define the start and stop times for the calculation\n",
    "osr.set_ctstartstop(0, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate Power Density on CAD STL Surface\n",
    "\n",
    "Will import the STL file and calculate the power density.  The return object is different from a normal power density calculation and allows for visualization of the object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "module 'matplotlib.cm' has no attribute 'default'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-5-1fec842d4187>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      9\u001b[0m )\n\u001b[1;32m     10\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 11\u001b[0;31m \u001b[0mplot_power_density_stl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpd\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcolorbar\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/venv/py3.6/lib/python3.6/site-packages/oscars-2.1.1-py3.6-macosx-10.6-intel.egg/oscars/plots3d_mpl.py\u001b[0m in \u001b[0;36mplot_power_density_stl\u001b[0;34m(P, title, elev, azim, alpha, bbox_inches, transparent, ofile, colorbar)\u001b[0m\n\u001b[1;32m    456\u001b[0m     \u001b[0mpmin\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mp\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mp\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mP\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    457\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 458\u001b[0;31m     \u001b[0mm\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mScalarMappable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdefault\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    459\u001b[0m \u001b[0;31m#m.set_array([p[1]/pmax for p in P])\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    460\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAttributeError\u001b[0m: module 'matplotlib.cm' has no attribute 'default'"
     ]
    }
   ],
   "source": [
    "# Clculate power density on CAD surface from STL file.\n",
    "# Here 'scale' is used since the input file units are in mm\n",
    "# while OSCARS expects m (there is no STL standard so you must\n",
    "# keep track of this yourself)\n",
    "pd = osr.calculate_power_density_stl(\n",
    "    ifile='sphere.stl',\n",
    "    scale=0.001,\n",
    "    translation=[0, 0, 5]\n",
    ")\n",
    "\n",
    "plot_power_density_stl(pd, colorbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_power_density_stl(pd, colorbar=True)"
   ]
  },
  {
   "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
}