{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$$\\overrightarrow{E} = \\left[\\begin{matrix}- \\frac{x \\left(U + V \\cos{\\left (\\omega t \\right )}\\right)}{r_{0}^{2}}\\\\\\frac{y \\left(U + V \\cos{\\left (\\omega t \\right )}\\right)}{r_{0}^{2}}\\\\0\\end{matrix}\\right]$$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$$\\overrightarrow{F} = \\left[\\begin{matrix}- \\frac{q x \\left(U + V \\cos{\\left (\\omega t \\right )}\\right)}{r_{0}^{2}}\\\\\\frac{q y \\left(U + V \\cos{\\left (\\omega t \\right )}\\right)}{r_{0}^{2}}\\\\0\\end{matrix}\\right]$$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$$\\overrightarrow{a} = \\left[\\begin{matrix}- \\frac{q x \\left(U + V \\cos{\\left (\\omega t \\right )}\\right)}{m r_{0}^{2}}\\\\\\frac{q y \\left(U + V \\cos{\\left (\\omega t \\right )}\\right)}{m r_{0}^{2}}\\\\0\\end{matrix}\\right]$$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from sympy import *\n",
    "# https://docs.sympy.org/latest/tutorials/intro-tutorial/printing.html\n",
    "\n",
    "# https://stackoverflow.com/questions/40806591/how-do-i-display-a-full-expression-in-sympy\n",
    "from IPython.display import display, Math\n",
    "\n",
    "U, V, omega, t, x, y, r_0, q, m, v_v, v_p = symbols('U V omega t x y r_0 q m v_v v_p')\n",
    "init_printing()\n",
    "\n",
    "v_ex = Matrix([1,0,0]) \n",
    "v_ey = Matrix([0,1,0])\n",
    "v_ez = Matrix([0,0,1])\n",
    "\n",
    "# theory prerequisite\n",
    "\n",
    "# RF + DC offset of the quadrupole\n",
    "phi_0 = U + V * cos(omega * t)\n",
    "\n",
    "# potential inside a quadrupole at a given point in space\n",
    "phi  = phi_0 * (x**2 - y**2) / (2 * r_0**2)\n",
    "\n",
    "# In the first step the electric field is derived from the potiential distribution\n",
    "# with the formula E = -\\nabla * phi.\n",
    "\n",
    "# calculating the partial derivateives for the nabla operation(gradient)\n",
    "dx = diff(phi, x)\n",
    "dy = diff(phi, y)\n",
    "\n",
    "v_E =  -(dx * v_ex + dy* v_ey)\n",
    "\n",
    "display(Math(latex('\\overrightarrow{E} = ') + latex(v_E)))\n",
    "\n",
    "# in the seccond step the lorenz force is determined with the formulafor considering B = 0!\n",
    "# F = E * q\n",
    "\n",
    "v_F = v_E * q\n",
    "\n",
    "display(Math('\\overrightarrow{F} = '+latex(v_F)))\n",
    "\n",
    "# With the force the acceleration can be determined with respect to the ions mass\n",
    "# a = F / m\n",
    "\n",
    "v_a = v_F / m\n",
    "\n",
    "display(Math('\\overrightarrow{a} = '+latex(v_a)))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(755.0120770093375, 10570.169078130726)\n",
      "N = 1100000\n",
      "ion has left the boundary!\n",
      "n = 1210\n",
      "simulation time =  9.999999999999773e-05\n",
      "number of positiions =  1\n",
      "[-2.0391744630349912e-13, 3.56317754718436e-15, -9.979659008986673e-17, 3.0235407792913644e-18, -7.139985036292582e-20, 7.19867582572237e-22, 4.321977872069379e-24, 8.443545107936075e-26, 2.2137865823636512e-27, 5.297458959531464e-29]\n",
      "done!\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import scipy.constants as sc\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "class Ion:\n",
    "    def __init__(self,position, charge, mass, v_z):\n",
    "        self.position = position\n",
    "        self.charge = charge\n",
    "        self.m = mass\n",
    "        self.v_z = v_z\n",
    "        self.F = np.array([[0],[0],[0]])\n",
    "\n",
    "class infiniteLengthQuadrupole:\n",
    "    def __init__(self,R,r_0):\n",
    "        # rod radius\n",
    "        self.R = R\n",
    "        # rod placement radius\n",
    "        self.r_0 = r_0\n",
    "        # rod positions\n",
    "        self.pPole = np.array([[[ (R + r_0)/np.sqrt(2)],\n",
    "                                [(R + r_0)/np.sqrt(2)],\n",
    "                                [0]],\n",
    "                                [[-(R + r_0)/np.sqrt(2)],\n",
    "                                 [(R + r_0)/np.sqrt(2)],\n",
    "                                 [0]],\n",
    "                                [[-(R + r_0)/np.sqrt(2)],\n",
    "                                 [-(R + r_0)/np.sqrt(2)],\n",
    "                                 [0]],\n",
    "                                [[(R + r_0)/np.sqrt(2)],\n",
    "                                 [-(R + r_0)/np.sqrt(2)],\n",
    "                                 [0]]])\n",
    "        # rod pseudo charges\n",
    "        self.rodsPseudoQ = np.zeros(4)\n",
    "        # K factor for simplification ofcalculating the equations\n",
    "        self.K = 1 / (4 * np.pi * sc.epsilon_0)\n",
    "        \n",
    "        self.forces= []\n",
    "        \n",
    "    def check(self):\n",
    "        print('pole positions')\n",
    "        print(self.pPole)\n",
    "        print('pseudo charges')\n",
    "        print(self.rodsPseudoQ)\n",
    "        \n",
    "    def phi_0(self,U, V, f, t):\n",
    "        return U + V * np.sin(2 * np.pi * f * t)\n",
    "    \n",
    "    def coulombForceOnQ1(self,Q1,Q2,r1,r2):\n",
    "        a = self.K * Q1 * Q2\n",
    "        r12 = r1 - r2\n",
    "        mag_r12 = np.linalg.norm(r12)\n",
    "        # calculate force\n",
    "        return a * r12 / (mag_r12**3)\n",
    "    \n",
    "    def getUV(self, mass_amu, f, a, q):\n",
    "        U = a * 1.66054e-27 * mass_amu * (2 * np.pi * f)**2 *self.r_0**2 /  (8 * sc.e)\n",
    "        V = q * 1.66054e-27 * mass_amu * (2 * np.pi * f)**2 *self.r_0**2 /  (4 * sc.e)\n",
    "        return U, V\n",
    "        \n",
    "    \n",
    "    def calcNewIonPos_1(self,U,V,f,ion,t_s,t):\n",
    "        # check if ion has left the r_0 boundary\n",
    "        if np.linalg.norm(ion.position) >= self.r_0:\n",
    "            return None\n",
    "        \n",
    "        phi_0 = U + V * np.sin(2 * np.pi * f * t)\n",
    "        a_x = -(float(ion.position[0]) * phi_0) / (self.r_0**2) * ion.charge / ion.m\n",
    "        a_y = (float(ion.position[1]) * phi_0) / (self.r_0**2) * ion.charge / ion.m\n",
    "                \n",
    "        newPosition = 0.5 * t_s**2 * np.array([[a_x],[a_y],[0]])\n",
    "        newPosition[2] = float(ion.position[2])\n",
    "        \n",
    "        # velocity component in z\n",
    "        newPosition[2] =  newPosition[2] + ion.v_z * t_s\n",
    "        \n",
    "        # ion position update\n",
    "        ion.position = newPosition\n",
    "\n",
    "        #print(newPosition)\n",
    "        \n",
    "        return newPosition\n",
    "        \n",
    "    def calcNewIonPos(self,U,V,f,ion,t_s,t):\n",
    "        \n",
    "        # check if ion has left the r_0 boundary\n",
    "        if np.linalg.norm(ion.position) >= self.r_0:\n",
    "            return None\n",
    "        \n",
    "        signLUT = [1,-1,1,-1]\n",
    "        index = 0;\n",
    "        \n",
    "        # reset force\n",
    "        F_c = np.array([[0],[0],[0]])\n",
    "        \n",
    "        for Q in self.rodsPseudoQ:\n",
    "\n",
    "            \n",
    "            # calculate the pseudo charge points of the quadrupole\n",
    "            Q = signLUT[index] * self.phi_0(U,V,f,t)/2 * self.K * self.R\n",
    "            \n",
    "            # match pseudo charge with ion in z direction\n",
    "            self.pPole[index][2] = ion.position[2]\n",
    "            \n",
    "            #print('ion position', ion.position)\n",
    "            #print(self.pPole[index])\n",
    "            \n",
    "            # calculate force on ion\n",
    "            F_q = self.coulombForceOnQ1(ion.charge,\n",
    "                                        Q,\n",
    "                                        self.pPole[index],\n",
    "                                        ion.position)\n",
    "            #print(F_q)\n",
    "            \n",
    "            F_c = F_c + F_q\n",
    "            self.forces.append(F_c)\n",
    "            \n",
    "            # update index\n",
    "            index = index + 1\n",
    "            \n",
    "        ion.F = F_c\n",
    "        # http://lampx.tugraz.at/~hadley/physikm/problems/rvaF/index.en.php\n",
    "        newPosition = 0.5 * t_s**2 * 1 / ion.m * ion.F + ion.position\n",
    "        \n",
    "        # velocity component in z\n",
    "        newPosition[2] =  newPosition[2] + ion.v_z * t_s\n",
    "\n",
    "        #print(newPosition)\n",
    "        \n",
    "        return newPosition\n",
    "    \n",
    "N_periodes = 100000\n",
    "N_points_per_period = 11\n",
    "f = 1.1e6    \n",
    "t_step = 1 / f / N_points_per_period\n",
    "t_sim = 0\n",
    "N = N_periodes * N_points_per_period\n",
    "\n",
    "# coffein atom 524.50002164 dalton (8.709527e-25 kg)\n",
    "\n",
    "\n",
    "#ion = Ion(np.array([[1e-10],[1e-10],[0]]), sc.e, 1.661e-27,100)\n",
    "ion = Ion(np.array([[1e-10],[1e-10],[0]]), sc.e, 1220 * 1.66054e-27,100)\n",
    "\n",
    "# quadrupole\n",
    "quad = infiniteLengthQuadrupole(0.0101, 0.01)\n",
    "\n",
    "#quad.check()\n",
    "\n",
    "positions = []\n",
    "\n",
    "p_x = []\n",
    "p_y = []\n",
    "p_z = []\n",
    "\n",
    "print(quad.getUV(1220,f,0.1,0.7))\n",
    "\n",
    "\n",
    "positions.append(ion.position)\n",
    "print('N =',N)\n",
    "for n in range(0,N-1,1):\n",
    "    \n",
    "    newPos = quad.calcNewIonPos_1(755.0120770093375, 10570.169078130726,f,ion,t_step,t_sim)\n",
    "    #newPos = quad.calcNewIonPos_1(100,100,f,ion,t_step,t_sim)\n",
    "    \n",
    "    if(newPos is None):\n",
    "        print('ion has left the boundary!')\n",
    "        break\n",
    "        \n",
    "    \n",
    "    ion.position = newPos\n",
    "    \n",
    "    if n % 100 is 0 and 0:\n",
    "        print(newPos)\n",
    "        #positions.append(newPos)\n",
    "    \n",
    "    if 0:\n",
    "        positions.append(newPos)\n",
    "    \n",
    "    if 1:\n",
    "        p_x.append(float(newPos[0]))\n",
    "        p_y.append(float(newPos[1]))\n",
    "        p_z.append(float(newPos[2]))\n",
    "        \n",
    "    \n",
    "    t_sim += t_step\n",
    "\n",
    "print('n =',n)\n",
    "print('simulation time = ', t_sim)\n",
    "print('number of positiions = ', len(positions))\n",
    "\n",
    "if 1:\n",
    "    print(p_x[:10])\n",
    "\n",
    "\n",
    "if 0:\n",
    "    print(positions[:10])\n",
    "\n",
    "\n",
    "\n",
    "print('done!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33876 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 5976 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "if 0:\n",
    "    n_start = 500\n",
    "    n_end = 3000\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111, projection = '3d')\n",
    "    ax.scatter(p_x[n_start:n_end],p_y[n_start:n_end],p_z[n_start:n_end],c = 'r')\n",
    "    ax.set_xlabel('x')\n",
    "    ax.set_ylabel('y')\n",
    "    ax.set_zlabel('z')\n",
    "    plt.show()\n",
    "\n",
    "if 1:\n",
    "\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111, projection = '3d')\n",
    "    ax.scatter(p_x,p_y,p_z,c = 'r')\n",
    "    ax.set_xlabel('x')\n",
    "    ax.set_ylabel('y')\n",
    "    ax.set_zlabel('z')\n",
    "    plt.show()\n",
    "    \n",
    "if 1:\n",
    "    plt.plot(p_x, p_y, 'o', color='black');\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'list' object has no attribute 'shape'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-64-1f7302c8ffde>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mmpl_toolkits\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmplot3d\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mAxes3D\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;31m#print(positions[1:2](0))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpositions\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0mfig\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mplt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfigure\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'list' object has no attribute 'shape'"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "#print(positions[1:2](0))\n",
    "print(np.array(positions.shape))\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.scatter(positions[:][0], positions[:][1], positions[:][2], zdir='z', c= 'red')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "ax.set_zlabel('z')\n",
    "plt.savefig(\"demo.png\")"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
