{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tests the thin fixed order multipole implementations (`ThinQuadrupole`, `ThinSkewQuadrupole`, `ThinSextupole` and `ThinOctupole`) against the multipole kick formula and times the results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# general imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import division, print_function\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from scipy.constants import m_p, c, e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# sets the PyHEADTAIL directory etc.\n",
    "try:\n",
    "    from settings import *\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyHEADTAIL imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyHEADTAIL v1.10.5.271\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import PyHEADTAIL.multipoles.multipoles as multip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class B(object):\n",
    "    def __init__(self, x, y, xp, yp):\n",
    "        self.x = np.array(x, dtype=float)\n",
    "        self.y = np.array(y, dtype=float)\n",
    "        self.xp = np.array(xp, dtype=float)\n",
    "        self.yp = np.array(yp, dtype=float)\n",
    "\n",
    "def make_beam(n_particles):\n",
    "    np.random.seed(42)\n",
    "    return B(np.random.random(n_particles), \n",
    "             np.random.random(n_particles),\n",
    "             np.random.random(n_particles),\n",
    "             np.random.random(n_particles))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's go"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Normal Quadrupole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "quadrupole = multip.ThinQuadrupole(np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "beam = make_beam(int(1e6))\n",
    "xp0, yp0 = beam.xp.copy(), beam.yp.copy()\n",
    "quadrupole.track(beam)\n",
    "res_quad_xp = beam.xp - xp0\n",
    "res_quad_yp = beam.yp - yp0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "multipole = multip.ThinMultipole([0, np.pi])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xp kicks correct: True\n",
      "yp kicks correct: True\n"
     ]
    }
   ],
   "source": [
    "beam = make_beam(int(1e6))\n",
    "xp0, yp0 = beam.xp.copy(), beam.yp.copy()\n",
    "multipole.track(beam)\n",
    "print (\"xp kicks correct:\", np.allclose(beam.xp - xp0, res_quad_xp))\n",
    "print (\"yp kicks correct:\", np.allclose(beam.yp - yp0, res_quad_yp))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### timing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10 loops, best of 3: 77.2 ms per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit make_beam(int(1e6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10 loops, best of 3: 82.6 ms per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit quadrupole.track(make_beam(int(1e6)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10 loops, best of 3: 97 ms per loop\n"
     ]
    }
   ],
   "source": [
    "%timeit multipole.track(make_beam(int(1e6)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10 loops, best of 3: 134 ms per loop\n"
     ]
    }
   ],
   "source": [
    "# using ztaylor\n",
    "proper_ctaylor = multipole.ctaylor\n",
    "multipole.ctaylor = multipole.ztaylor\n",
    "\n",
    "%timeit multipole.track(make_beam(int(1e6)))\n",
    "\n",
    "multipole.ctaylor = proper_ctaylor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Skew Quadrupole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "squadrupole = multip.ThinSkewQuadrupole(np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "beam = make_beam(int(1e6))\n",
    "xp0, yp0 = beam.xp.copy(), beam.yp.copy()\n",
    "squadrupole.track(beam)\n",
    "res_squad_xp = beam.xp - xp0\n",
    "res_squad_yp = beam.yp - yp0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "multipole = multip.ThinMultipole([0, 0], [0, np.pi])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xp kicks correct: True\n",
      "yp kicks correct: True\n"
     ]
    }
   ],
   "source": [
    "beam = make_beam(int(1e6))\n",
    "xp0, yp0 = beam.xp.copy(), beam.yp.copy()\n",
    "multipole.track(beam)\n",
    "print (\"xp kicks correct:\", np.allclose(beam.xp - xp0, res_squad_xp))\n",
    "print (\"yp kicks correct:\", np.allclose(beam.yp - yp0, res_squad_yp))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sextupole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sextupole = multip.ThinSextupole(np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "beam = make_beam(int(1e6))\n",
    "xp0, yp0 = beam.xp.copy(), beam.yp.copy()\n",
    "sextupole.track(beam)\n",
    "res_sext_xp = beam.xp - xp0\n",
    "res_sext_yp = beam.yp - yp0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "multipole = multip.ThinMultipole([0, 0, np.pi])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xp kicks correct: True\n",
      "yp kicks correct: True\n"
     ]
    }
   ],
   "source": [
    "beam = make_beam(int(1e6))\n",
    "xp0, yp0 = beam.xp.copy(), beam.yp.copy()\n",
    "multipole.track(beam)\n",
    "print (\"xp kicks correct:\", np.allclose(beam.xp - xp0, res_sext_xp))\n",
    "print (\"yp kicks correct:\", np.allclose(beam.yp - yp0, res_sext_yp))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Octupole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "octupole = multip.ThinOctupole(np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "beam = make_beam(int(1e6))\n",
    "xp0, yp0 = beam.xp.copy(), beam.yp.copy()\n",
    "octupole.track(beam)\n",
    "res_octu_xp = beam.xp - xp0\n",
    "res_octu_yp = beam.yp - yp0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "multipole = multip.ThinMultipole([0, 0, 0, np.pi])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xp kicks correct: True\n",
      "yp kicks correct: True\n"
     ]
    }
   ],
   "source": [
    "beam = make_beam(int(1e6))\n",
    "xp0, yp0 = beam.xp.copy(), beam.yp.copy()\n",
    "multipole.track(beam)\n",
    "print (\"xp kicks correct:\", np.allclose(beam.xp - xp0, res_octu_xp))\n",
    "print (\"yp kicks correct:\", np.allclose(beam.yp - yp0, res_octu_yp))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}