{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "N = 10100\n",
      "ion has left the boundary!\n",
      "n = 1\n",
      "simulation time =  9.000900090009e-09\n",
      "number of positiions =  2\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35952 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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",
    "    \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 = 100\n",
    "N_points_per_period = 101\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,1)\n",
    "\n",
    "# quadrupole\n",
    "quad = infiniteLengthQuadrupole(0.0101, 0.01)\n",
    "\n",
    "#quad.check()\n",
    "\n",
    "positions = []\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(2.4764,17.335,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",
    "        positions.append(newPos)\n",
    "    \n",
    "    if 1:\n",
    "        positions.append(newPos)\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",
    "\n",
    "if 0:\n",
    "    print(positions[:10])\n",
    "\n",
    "if 1:\n",
    "    ax = plt.axes(projection='3d')\n",
    "    #ax.scatter3D(positions[:][0], positions[:][1], positions[:][2], c=positions[:][2], cmap='Greens')\n",
    "    ax.scatter3D(quad.forces[:][0], quad.forces[:][1], quad.forces[:][2], cmap='Greens')\n",
    "    ax.set_xlabel('x')\n",
    "    ax.set_ylabel('y')\n",
    "    ax.set_zlabel('z')\n",
    "    plt.show()\n",
    "\n",
    "print('done!')\n",
    "# dummy check\n",
    "if 0:\n",
    "    ion_1 = Ion(np.array([[0.0000005],[0],[0]]), 1000 * sc.e, 1,0)\n",
    "    ion_2 = Ion(np.array([[-0.0000005],[0],[0]]), 1000 * sc.e, 1,0)\n",
    "    print(quad.coulombForceOnQ1(ion_1.charge,ion_2.charge,ion_1.position,ion_2.position))\n",
    "                 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[array([[1.e-100],\n",
      "       [1.e-100],\n",
      "       [0.e+000]]), array([[1.00000000e-100],\n",
      "       [1.00000000e-100],\n",
      "       [2.27272727e-007]])]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34092 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f8f553ca198>]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12572 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(positions[:2])\n",
    "\n",
    "\n",
    "# https://jakevdp.github.io/PythonDataScienceHandbook/04.12-three-dimensional-plotting.html\n",
    "ax = plt.axes(projection='3d')\n",
    "#ax.scatter3D(positions[:][0], positions[:][1], positions[:][2], c=positions[:][2], cmap='Greens')\n",
    "ax.scatter3D(positions[:][0], positions[:][2], positions[:][1], cmap='Greens')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('z')\n",
    "ax.set_zlabel('y')\n",
    "plt.show()\n",
    "\n",
    "plt.plot(positions[:][0],positions[:][1])\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
