{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a4246751",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Calculator import *\n",
    "from Potentials import *\n",
    "from Helpers import *\n",
    "from Plotter import *\n",
    "# %matplotlib notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c68468e4",
   "metadata": {},
   "source": [
    "## Plot ideal trap potential resulting for given parameters only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "38c770ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Karthik\\anaconda3\\envs\\Py39VEnv\\lib\\site-packages\\scipy\\optimize\\_minpack_py.py:1010: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 71808 characters]",
      "text/plain": [
       "<Figure size 900x700 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Power = 40*u.W\n",
    "Wavelength = 1.064*u.um\n",
    "w_x, w_z = 30*u.um, 30*u.um # Beam Waists in the x and y directions\n",
    "\n",
    "#Power = 11*u.W\n",
    "#w_x, w_z = 67*u.um, 67*u.um # Beam Waists in the x and y directions\n",
    "\n",
    "options = {\n",
    "    'axis': 2, # axis referenced to the beam along which you want the dipole trap potential\n",
    "    'extent': 1e2, # range of spatial coordinates in one direction to calculate trap potential over\n",
    "    'crossed': False,\n",
    "    'delta': 70, # angle between arms in degrees\n",
    "    'modulation': True,\n",
    "    'aspect_ratio': 4, # required aspect ratio of modulated arm\n",
    "    'gravity': True,\n",
    "    'tilt_gravity': False,\n",
    "    'theta': 0.75, # gravity tilt angle in degrees\n",
    "    'tilt_axis': [1, 0, 0], # lab space coordinates are rotated about x-axis in reference frame of beam\n",
    "    'astigmatism': False,\n",
    "    'disp_foci': 2.5*u.mm, #0.9 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um, # difference in position of the foci along the propagation direction (Astigmatism)\n",
    "    'extract_trap_frequencies': False\n",
    "}\n",
    "\n",
    "ComputedPotentials = [] \n",
    "Params = []  \n",
    "\n",
    "Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, Power, options)\n",
    "ComputedPotentials.append(IdealTrappingPotential)\n",
    "ComputedPotentials.append(TrappingPotential)\n",
    "Params.append([TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies])\n",
    "\n",
    "cpots = np.asarray(ComputedPotentials)\n",
    "plotPotential(Positions, cpots, options, Params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc9809de",
   "metadata": {},
   "source": [
    "## Plot harmonic fit for trap potential resulting for given parameters only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0f3e80f7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 89520 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "v, dv, popt, pcov = extractTrapFrequency(Positions, TrappingPotential, options['axis'])\n",
    "plotHarmonicFit(Positions, TrappingPotential, TrapDepthsInKelvin, options['axis'], popt, pcov)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37b40607",
   "metadata": {},
   "source": [
    "## Plot trap potential resulting for given parameters (with one parameter being a list of values and the potential being computed for each of these values) only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8504f99f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Karthik\\anaconda3\\envs\\Py39VEnv\\lib\\site-packages\\scipy\\optimize\\_minpack_py.py:1010: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 127368 characters]",
      "text/plain": [
       "<Figure size 900x700 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Potentials = [] \n",
    "Params = []  \n",
    "Power = [10, 30, 40]*u.W # Single Beam Power\n",
    "for p in Power: \n",
    "    Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, p, options)\n",
    "    Potentials.append(IdealTrappingPotential)\n",
    "    Potentials.append(TrappingPotential)\n",
    "    Params.append([TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies])\n",
    "\n",
    "cpots = np.asarray(Potentials)\n",
    "plotPotential(Positions, cpots, options, Params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "951010c6",
   "metadata": {},
   "source": [
    "## Plot transverse intensity profile and trap potential resulting for given parameters only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f3e4afd9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 118324 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Power = 40*u.W\n",
    "Wavelength = 1.064*u.um\n",
    "w_x, w_z = 30*u.um, 30*u.um # Beam Waists in the x and y directions\n",
    "\n",
    "options = {\n",
    "    'axis': 2, # axis referenced to the beam along which you want the dipole trap potential\n",
    "    'extent': 60, # range of spatial coordinates in one direction to calculate trap potential over\n",
    "    'crossed': False,\n",
    "    'delta': 70, # angle between arms in degrees\n",
    "    'modulation': True,\n",
    "    'modulation_function': 'arccos',\n",
    "    'modulation_amplitude': 2.16,\n",
    "    'aspect_ratio': 4, # required aspect ratio of modulated arm\n",
    "    'gravity': True,\n",
    "    'tilt_gravity': False,\n",
    "    'theta': 0.75, # gravity tilt angle in degrees\n",
    "    'tilt_axis': [1, 0, 0], # lab space coordinates are rotated about x-axis in reference frame of beam\n",
    "    'astigmatism': False,\n",
    "    'disp_foci': 2.5*u.mm, #0.9 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um, # difference in position of the foci along the propagation direction (Astigmatism)\n",
    "    'extract_trap_frequencies': False\n",
    "}\n",
    "\n",
    "positions, waists, I, U, p = computeIntensityProfileAndPotentials(Power, [w_x, w_z], Wavelength, options)\n",
    "plotIntensityProfileAndPotentials(positions, waists, I, U)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db0df307",
   "metadata": {},
   "source": [
    "## Plot gaussian fit for trap potential resulting from modulation for given parameters only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7afa7d82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 82168 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 73640 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "x_Positions = positions[0].value\n",
    "z_Positions = positions[1].value\n",
    "x_Potential = U[:, np.where(z_Positions==0)[0][0]].value\n",
    "z_Potential = U[np.where(x_Positions==0)[0][0], :].value\n",
    "poptx, pcovx = p[0], p[1]\n",
    "poptz, pcovz = p[2], p[3]\n",
    "plotGaussianFit(x_Positions, x_Potential, poptx, pcovx)\n",
    "plotGaussianFit(z_Positions, z_Potential, poptz, pcovz)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e5b8123",
   "metadata": {},
   "source": [
    "## Calculate relevant parameters for evaporative cooling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "95ab43bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Particle Density = 4.63E+13 1 / cm3\n",
      "Elastic Collision Rate = 2813.75 1 / s\n",
      "PSD = 9.00E-05 \n",
      "v_x = 2588.17 Hz\n",
      "v_y = 20.66 Hz\n",
      "v_z = 2588.17 Hz\n",
      "a_s = 111.31 \n"
     ]
    }
   ],
   "source": [
    "Power = 40*u.W\n",
    "Wavelength = 1.064*u.um\n",
    "w_x, w_z = 30*u.um, 30*u.um # Beam Waists in the x and y directions\n",
    "\n",
    "AtomNumber = 1.00 * 1e7\n",
    "BField = 2.5 * u.G\n",
    "\n",
    "modulation_depth = 0.0\n",
    "Temperature = convert_modulation_depth_to_temperature(modulation_depth)[0] * u.uK\n",
    "n = particleDensity(w_x, w_z, Power, N = AtomNumber, T = Temperature).decompose().to(u.cm**(-3))\n",
    "Gamma_elastic = calculateElasticCollisionRate(w_x, w_z, Power, N = AtomNumber, T = Temperature, B = BField)\n",
    "PSD = calculatePSD(w_x, w_z, Power, N = AtomNumber, T = Temperature).decompose()\n",
    "\n",
    "print('Particle Density = %.2E ' % (n.value) + str(n.unit))\n",
    "print('Elastic Collision Rate = %.2f ' % (Gamma_elastic.value) + str(Gamma_elastic.unit))\n",
    "print('PSD = %.2E ' % (PSD.value))\n",
    "\n",
    "v_x = calculateTrapFrequency(w_x, w_z, Power, dir = 'x')\n",
    "v_y = calculateTrapFrequency(w_x, w_z, Power, dir = 'y')\n",
    "v_z = calculateTrapFrequency(w_x, w_z, Power, dir = 'z')\n",
    "\n",
    "print('v_x = %.2f ' %(v_x.value) + str(v_x.unit))\n",
    "print('v_y = %.2f ' %(v_y.value) + str(v_y.unit))\n",
    "print('v_z = %.2f ' %(v_z.value) + str(v_z.unit))\n",
    "\n",
    "print('a_s = %.2f ' %(scatteringLength(BField)[0] / ac.a0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff252fbe",
   "metadata": {},
   "source": [
    "## Plot alphas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dd7fc03d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 34768 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plotAlphas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c09cb260",
   "metadata": {},
   "source": [
    "## Plot Temperatures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5c79840e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 44544 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plotTemperatures(w_x, w_z, plot_against_mod_depth = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "063879ee",
   "metadata": {},
   "source": [
    "##  Calculate and Plot calculated trap frequencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0ddff726",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 66848 characters]",
      "text/plain": [
       "<Figure size 800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "AtomNumber = 1.00 * 1e7\n",
    "BField = 1.4 * u.G\n",
    "Wavelength = 1.064*u.um\n",
    "w_x, w_z = 30*u.um, 30*u.um # Beam Waists in the x and y directions\n",
    "modulation_depth = np.arange(0, 1.0, 0.08)\n",
    "\n",
    "w_xs = w_x * convert_modulation_depth_to_alpha(modulation_depth)[0]\n",
    "new_aspect_ratio  = w_xs / w_z\n",
    "\n",
    "v_x = np.zeros(len(modulation_depth))\n",
    "v_y = np.zeros(len(modulation_depth))\n",
    "v_z = np.zeros(len(modulation_depth))\n",
    "\n",
    "for i in range(len(modulation_depth)):\n",
    "    v_x[i] = calculateTrapFrequency(w_xs[i], w_z, Power, dir = 'x').value\n",
    "    v_y[i] = calculateTrapFrequency(w_xs[i], w_z, Power, dir = 'y').value\n",
    "    v_z[i] = calculateTrapFrequency(w_xs[i], w_z, Power, dir = 'z').value\n",
    "\n",
    "plotTrapFrequencies(v_x, v_y, v_z, modulation_depth, new_aspect_ratio, plot_against_mod_depth = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76ff8301",
   "metadata": {},
   "source": [
    "## Plot measured trap frequencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7a85ec41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39232 characters]",
      "text/plain": [
       "<Figure size 800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "modulation_depth_radial = [0, 0.5, 0.3, 0.7, 0.9, 0.8, 1.0, 0.6, 0.4, 0.2, 0.1]\n",
    "fx = [3.135, 0.28, 0.690, 0.152, 0.102, 0.127, 0.099, 0.205, 0.404, 1.441, 2.813]\n",
    "dfx = [0.016, 0.006, 0.005, 0.006, 0.003, 0.002, 0.002,0.002, 0.003, 0.006, 0.024]\n",
    "fz = [2.746, 1.278, 1.719, 1.058, 0.923, 0.994, 0.911, 1.157, 1.446, 2.191, 2.643]\n",
    "dfz = [0.014, 0.007, 0.009, 0.007, 0.005, 0.004, 0.004, 0.005, 0.007, 0.009, 0.033]\n",
    "\n",
    "modulation_depth_axial = [1, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1]\n",
    "fy = [3.08, 3.13, 3.27, 3.46, 3.61, 3.82, 3.51, 3.15, 3.11, 3.02]\n",
    "dfy = [0.03, 0.04, 0.04, 0.05, 0.07, 0.06, 0.11, 0.07, 0.1, 1.31]\n",
    "\n",
    "plotMeasuredTrapFrequencies(fx, dfx, fy, dfy, fz, dfz,  modulation_depth_radial, modulation_depth_axial, w_x, w_z, plot_against_mod_depth = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a4843d2",
   "metadata": {},
   "source": [
    "## Plot ratio of measured to calculated trap frequencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "58cf3f64",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30068 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "modulation_depth = [0.5, 0.3, 0.7, 0.9, 0.8, 1.0, 0.6, 0.4, 0.2, 0.1]\n",
    "w_xs = w_x * convert_modulation_depth_to_alpha(modulation_depth)[0]\n",
    "\n",
    "v_x = np.zeros(len(modulation_depth))\n",
    "v_y = np.zeros(len(modulation_depth))\n",
    "v_z = np.zeros(len(modulation_depth))\n",
    "\n",
    "for i in range(len(modulation_depth)):\n",
    "    v_x[i] = calculateTrapFrequency(w_xs[i], w_z, Power, dir = 'x').value / 1e3\n",
    "    v_y[i] = calculateTrapFrequency(w_xs[i], w_z, Power, dir = 'y').value\n",
    "    v_z[i] = calculateTrapFrequency(w_xs[i], w_z, Power, dir = 'z').value / 1e3\n",
    "\n",
    "fx = [0.28, 0.690, 0.152, 0.102, 0.127, 0.099, 0.205, 0.404, 1.441, 2.813]\n",
    "dfx = [0.006, 0.005, 0.006, 0.003, 0.002, 0.002,0.002, 0.003, 0.006, 0.024]\n",
    "fy = [3.08, 3.13, 3.27, 3.46, 3.61, 3.82, 3.51, 3.15, 3.11, 3.02]\n",
    "dfy = [0.03, 0.04, 0.04, 0.05, 0.07, 0.06, 0.11, 0.07, 0.1, 1.31]\n",
    "fz = [1.278, 1.719, 1.058, 0.923, 0.994, 0.911, 1.157, 1.446, 2.191, 2.643]\n",
    "dfz = [0.007, 0.009, 0.007, 0.005, 0.004, 0.004, 0.005, 0.007, 0.009, 0.033]\n",
    "\n",
    "plotRatioOfTrapFrequencies(fx, fy, fz, dfx, dfy, dfz, v_x, v_y, v_z, modulation_depth, w_x, w_z, plot_against_mod_depth = True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44e92099",
   "metadata": {},
   "source": [
    "## Plot Feshbach Resonances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d15205ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 62716 characters]",
      "text/plain": [
       "<Figure size 900x700 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plotScatteringLengths(B_range = [0, 3.6])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a2e113f",
   "metadata": {},
   "source": [
    "##  Calculate and Plot Collision Rates and PSD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "86e9ba21",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 67752 characters]",
      "text/plain": [
       "<Figure size 800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "AtomNumber = 1.00 * 1e7\n",
    "BField = 1.4 * u.G\n",
    "Power = 40*u.W\n",
    "Wavelength = 1.064*u.um\n",
    "w_x, w_z = 30*u.um, 30*u.um # Beam Waists in the x and y directions\n",
    "modulation_depth = np.arange(0, 1.0, 0.08)\n",
    "\n",
    "w_xs = w_x * convert_modulation_depth_to_alpha(modulation_depth)[0]\n",
    "new_aspect_ratio  = w_xs / w_z\n",
    "Temperatures = convert_modulation_depth_to_temperature(modulation_depth)[0] * u.uK\n",
    "\n",
    "# n = np.zeros(len(modulation_depth))\n",
    "Gamma_elastic = np.zeros(len(modulation_depth))\n",
    "PSD = np.zeros(len(modulation_depth))\n",
    "\n",
    "for i in range(len(modulation_depth)):\n",
    "    # n[i] = particleDensity(w_xs[i], w_z, Power, N = AtomNumber, T = Temperatures[i]).decompose().to(u.cm**(-3))\n",
    "    Gamma_elastic[i] = calculateElasticCollisionRate(w_xs[i], w_z, Power, N = AtomNumber, T = Temperatures[i], B = BField).value\n",
    "    PSD[i] = calculatePSD(w_xs[i], w_z, Power, N = AtomNumber, T = Temperatures[i]).decompose().value\n",
    "\n",
    "\n",
    "plotCollisionRatesAndPSD(Gamma_elastic, PSD, modulation_depth, new_aspect_ratio, plot_against_mod_depth = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74d353c9",
   "metadata": {},
   "source": [
    "## Plot Collision Rates and PSD from only measured trap frequencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6c81d9da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61560 characters]",
      "text/plain": [
       "<Figure size 800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "AtomNumber = 1.00 * 1e7\n",
    "BField = 1.4 * u.G\n",
    "Wavelength = 1.064*u.um\n",
    "w_x, w_z = 30*u.um, 30*u.um # Beam Waists in the x and y directions\n",
    "\n",
    "fin_mod_depth = [0.5, 0.3, 0.7, 0.9, 0.8, 1.0, 0.6, 0.4, 0.2, 0.1]\n",
    "v_x = [0.28, 0.690, 0.152, 0.102, 0.127, 0.099, 0.205, 0.404, 1.441, 2.813] * u.kHz\n",
    "dv_x = [0.006, 0.005, 0.006, 0.003, 0.002, 0.002,0.002, 0.003, 0.006, 0.024] * u.kHz\n",
    "v_z = [1.278, 1.719, 1.058, 0.923, 0.994, 0.911, 1.157, 1.446, 2.191, 2.643] * u.kHz\n",
    "dv_z = [0.007, 0.009, 0.007, 0.005, 0.004, 0.004, 0.005, 0.007, 0.009, 0.033] * u.kHz\n",
    "sorted_modulation_depth, sorted_v_x = zip(*sorted(zip(fin_mod_depth, v_x)))\n",
    "sorted_modulation_depth, sorted_dv_x = zip(*sorted(zip(fin_mod_depth, dv_x)))\n",
    "sorted_modulation_depth, sorted_v_z = zip(*sorted(zip(fin_mod_depth, v_z)))\n",
    "sorted_modulation_depth, sorted_dv_z = zip(*sorted(zip(fin_mod_depth, dv_z)))\n",
    "\n",
    "fin_mod_depth = [1, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1] \n",
    "v_y = [3.08, 3.13, 3.27, 3.46, 3.61, 3.82, 3.51, 3.15, 3.11, 3.02] * u.Hz\n",
    "dv_y = [0.03, 0.04, 0.04, 0.05, 0.07, 0.06, 0.11, 0.07, 0.1, 1.31] * u.Hz\n",
    "sorted_modulation_depth, sorted_v_y = zip(*sorted(zip(fin_mod_depth, v_y)))\n",
    "sorted_modulation_depth, sorted_dv_y = zip(*sorted(zip(fin_mod_depth, dv_y)))\n",
    "\n",
    "fin_mod_depth = [1.0, 0.8, 0.6, 0.4, 0.2, 0.9, 0.7, 0.5, 0.3, 0.1]\n",
    "T_x = [22.1, 27.9, 31.7, 42.2, 145.8, 27.9, 33.8, 42.4, 61.9, 136.1] * u.uK\n",
    "dT_x = [1.7, 2.6, 2.4, 3.7, 1.1, 2.2, 3.2, 1.7, 2.2, 1.2] * u.uK\n",
    "T_y = [13.13, 14.75, 18.44, 26.31, 52.55, 13.54, 16.11, 21.15, 35.81, 85.8] * u.uK\n",
    "dT_y = [0.05, 0.05, 0.07, 0.16, 0.28, 0.04, 0.07, 0.10, 0.21, 0.8] * u.uK\n",
    "\n",
    "sorted_modulation_depth, sorted_T_x = zip(*sorted(zip(fin_mod_depth, T_x)))\n",
    "sorted_modulation_depth, sorted_dT_x = zip(*sorted(zip(fin_mod_depth, dT_x)))\n",
    "sorted_modulation_depth, sorted_T_y = zip(*sorted(zip(fin_mod_depth, T_y)))\n",
    "sorted_modulation_depth, sorted_dT_y = zip(*sorted(zip(fin_mod_depth, dT_y)))\n",
    "\n",
    "modulation_depth = sorted_modulation_depth\n",
    "\n",
    "pd, dpd, T, dT, new_aspect_ratio = calculateParticleDensityFromMeasurements(sorted_v_x, sorted_dv_x, sorted_v_y, sorted_dv_y, sorted_v_z, sorted_dv_z, w_x, w_z, sorted_T_x, sorted_T_y, sorted_dT_x, sorted_dT_y, sorted_modulation_depth, AtomNumber, m = 164*u.u)\n",
    "\n",
    "Gamma_elastic = [(pd[i] * scatteringCrossSection(BField) * meanThermalVelocity(T[i]) / (2 * np.sqrt(2))).decompose()  for i in range(len(pd))]\n",
    "Gamma_elastic_values = [(Gamma_elastic[i]).value  for i in range(len(Gamma_elastic))]\n",
    "dGamma_elastic = [(Gamma_elastic[i] * ((dpd[i]/pd[i]) + (dT[i]/(2*T[i])))).decompose() for i in range(len(Gamma_elastic))]\n",
    "dGamma_elastic_values = [(dGamma_elastic[i]).value for i in range(len(dGamma_elastic))]\n",
    "\n",
    "PSD = [((pd[i] * thermaldeBroglieWavelength(T[i])**3).decompose()).value for i in range(len(pd))]\n",
    "dPSD = [((PSD[i] * ((dpd[i]/pd[i]) - (1.5 * dT[i]/T[i]))).decompose()).value for i in range(len(Gamma_elastic))]\n",
    "\n",
    "fig, ax1 = plt.subplots(figsize=(8, 6))\n",
    "ax2 = ax1.twinx()\n",
    "ax1.errorbar(modulation_depth, Gamma_elastic_values, yerr = dGamma_elastic_values, fmt = 'ob', markersize=5, capsize=5)\n",
    "ax2.errorbar(modulation_depth, PSD, yerr = dPSD, fmt = '-^r', markersize=5, capsize=5)\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1e'))\n",
    "ax1.set_xlabel('Modulation depth', fontsize= 12, fontweight='bold')\n",
    "ax1.set_ylabel('Elastic Collision Rate (' + str(Gamma_elastic[0].unit) + ')', fontsize= 12, fontweight='bold')\n",
    "ax1.tick_params(axis=\"y\", labelcolor='b')\n",
    "ax2.set_ylabel('Phase Space Density', fontsize= 12, fontweight='bold')\n",
    "ax2.tick_params(axis=\"y\", labelcolor='r')\n",
    "plt.tight_layout()\n",
    "plt.grid(visible=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9329ee37",
   "metadata": {},
   "source": [
    "## Investigate deviation in alpha"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7b8191f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Karthik\\anaconda3\\envs\\Py39VEnv\\lib\\site-packages\\scipy\\optimize\\_minpack_py.py:1010: OptimizeWarning: Covariance of the parameters could not be estimated\n",
      "  warnings.warn('Covariance of the parameters could not be estimated',\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 33788 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 36308 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Power = 40*u.W\n",
    "Wavelength = 1.064*u.um\n",
    "w_x, w_z = 30*u.um, 30*u.um\n",
    "\n",
    "options = {\n",
    "    'axis': 0, # axis referenced to the beam along which you want the dipole trap potential\n",
    "    'extent': 3e2, # range of spatial coordinates in one direction to calculate trap potential over\n",
    "    'crossed': False,\n",
    "    'delta': 70, # angle between arms in degrees\n",
    "    'modulation': False,\n",
    "    'aspect_ratio': 5, # required aspect ratio of modulated arm\n",
    "    'gravity': True,\n",
    "    'tilt_gravity': False,\n",
    "    'theta': 0.75, # gravity tilt angle in degrees\n",
    "    'tilt_axis': [1, 0, 0], # lab space coordinates are rotated about x-axis in reference frame of beam\n",
    "    'astigmatism': True,\n",
    "    'foci_disp_single': 2.5*u.mm, #0.9 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um, # difference in position of the foci along the propagation direction\n",
    "    'foci_disp_crossed': [2.5*u.mm, 1.5*u.mm], # astigmatism of each of the two beams in the cODT\n",
    "    'foci_shift': [0.0*u.mm, 0.0*u.mm],\n",
    "    'beam_disp': [[0.0, 0.0, 0.0]*u.mm, [0, 0, 0.03]*u.mm],\n",
    "    'extract_trap_frequencies': False # Flag to extract trap frequencies by fitting the computed potential\n",
    "}\n",
    "\n",
    "modulation_depth = np.arange(0, 1.1, 0.1)\n",
    "Alphas, fin_mod_dep, meas_alpha_x, meas_alpha_z, dalpha_x, dalpha_z = convert_modulation_depth_to_alpha(modulation_depth) \n",
    "meas_alpha_deviation = [(g - h) for g, h in zip(meas_alpha_x, meas_alpha_z)]\n",
    "sorted_fin_mod_dep, sorted_meas_alpha_deviation = zip(*sorted(zip(fin_mod_dep, meas_alpha_deviation)))\n",
    "avg_alpha = [(g + h) / 2 for g, h in zip(meas_alpha_x, meas_alpha_z)]\n",
    "sorted_fin_mod_dep, new_aspect_ratio = zip(*sorted(zip(fin_mod_dep, (w_x * avg_alpha) / w_z)))\n",
    "\n",
    "current_ar = w_x/w_z\n",
    "aspect_ratio = np.arange(current_ar, 10*current_ar, 0.8)\n",
    "w_x = w_x * (aspect_ratio / current_ar)\n",
    "\n",
    "v_x = np.zeros(len(w_x))\n",
    "#v_y = np.zeros(len(w_x))\n",
    "v_z = np.zeros(len(w_x))\n",
    "\n",
    "for i in range(len(w_x)):\n",
    "    \n",
    "    options['axis'] = 0\n",
    "    ExtractedTrapFrequencies = computeTrapPotential(w_x[i], w_z, Power, options)[5]\n",
    "    tmp = ExtractedTrapFrequencies[0][0]\n",
    "    v_x[i] = tmp if not np.isinf(tmp) else np.nan\n",
    "    \n",
    "    # options['axis'] = 1\n",
    "    # ExtractedTrapFrequencies = computeTrapPotential(w_x[i], w_z, Power, options)[5]\n",
    "    # tmp = ExtractedTrapFrequencies[1][0]\n",
    "    # v_y[i] = tmp if not np.isinf(tmp) else np.nan\n",
    "    \n",
    "    options['axis'] = 2\n",
    "    ExtractedTrapFrequencies = computeTrapPotential(w_x[i], w_z, Power, options)[5]\n",
    "    tmp = ExtractedTrapFrequencies[0][0]\n",
    "    v_z[i] = tmp if not np.isinf(tmp) else np.nan\n",
    "\n",
    "    #v_x[i] = calculateTrapFrequency(w_x[i], w_z, Power, dir = 'x').value\n",
    "    #v_y[i] = calculateTrapFrequency(w_x[i], w_z, Power, dir = 'y').value\n",
    "    #v_z[i] = calculateTrapFrequency(w_x[i], w_z, Power, dir = 'z').value\n",
    "\n",
    "alpha_x = [(v_x[0]/v)**(2/3) for v in v_x]\n",
    "alpha_z = [(v_z[0]/v)**2 for v in v_z]\n",
    "\n",
    "calc_alpha_deviation = [(g - h) for g, h in zip(alpha_x, alpha_z)]\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(aspect_ratio, alpha_x, '-o', label = 'From horz TF')\n",
    "plt.plot(aspect_ratio, alpha_z, '-^', label = 'From vert TF')\n",
    "plt.xlabel('Aspect Ratio', fontsize= 12, fontweight='bold')\n",
    "plt.ylabel('$\\\\alpha$', fontsize= 12, fontweight='bold')\n",
    "plt.tight_layout()\n",
    "plt.grid(visible=1)\n",
    "plt.legend(prop={'size': 12, 'weight': 'bold'})\n",
    "plt.show()\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(aspect_ratio, calc_alpha_deviation, '--ob', label = 'Extracted')\n",
    "plt.plot(new_aspect_ratio, sorted_meas_alpha_deviation, '-or', label = 'Measured')\n",
    "plt.xlabel('Aspect Ratio', fontsize= 12, fontweight='bold')\n",
    "plt.ylabel('$\\\\Delta \\\\alpha$', fontsize= 12, fontweight='bold')\n",
    "plt.ylim([-0.5, 1])\n",
    "plt.tight_layout()\n",
    "plt.grid(visible=1)\n",
    "plt.legend(prop={'size': 12, 'weight': 'bold'})\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29a1cf87",
   "metadata": {},
   "source": [
    "## Quick calculation of PSD and elastic collision rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3950f1b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Particle Density = 1.87E+12 1 / cm3\n",
      "Elastic Collision Rate = 37.24 1 / s\n",
      "PSD = 2.50E-05 \n"
     ]
    }
   ],
   "source": [
    "Power = 9.5*u.W\n",
    "Wavelength = 1.064*u.um\n",
    "w_x, w_z = 50*u.um, 45*u.um # Beam Waists in the x and y directions\n",
    "\n",
    "NCount = 11000\n",
    "AtomNumber = calculateAtomNumber(NCount, pixel_size = 3.45 * u.um, magnification = 0.5, eta = 0.5)\n",
    "#AtomNumber = calculateAtomNumber(NCount, pixel_size = 5.86 * u.um, magnification = 0.5, eta = 0.5)\n",
    "\n",
    "BField = 3.75 * u.G\n",
    "modulation_depth = 0.00\n",
    "\n",
    "w_x = w_x * convert_modulation_depth_to_alpha(modulation_depth)[0]\n",
    "new_aspect_ratio  = w_x / w_z\n",
    "Temperature = 33 * u.uK\n",
    "\n",
    "n = particleDensity(w_x, w_z, Power, N = AtomNumber, T = Temperature).decompose().to(u.cm**(-3))\n",
    "Gamma_elastic = calculateElasticCollisionRate(w_x, w_z, Power, N = AtomNumber, T = Temperature, B = BField)\n",
    "PSD = calculatePSD(w_x, w_z, Power, N = AtomNumber, T = Temperature).decompose()\n",
    "\n",
    "print('Particle Density = %.2E ' % (n.value) + str(n.unit))\n",
    "print('Elastic Collision Rate = %.2f ' % (Gamma_elastic.value) + str(Gamma_elastic.unit))\n",
    "print('PSD = %.2E ' % (PSD.value))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d975586b",
   "metadata": {},
   "source": [
    "## Plot measured PSDs and collision rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b11b6aae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 65112 characters]",
      "text/plain": [
       "<Figure size 800x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Gamma_elastic = [96.01, 82.53, 102.95, 143.5, 110.92, 471.77, 908.46]\n",
    "PSD = [8.95e-05, 4.19e-04, 2.94e-04, 2.17e-04, 8.98e-04, 5.52e-04, 6.10e-04]\n",
    "\n",
    "fig, ax1 = plt.subplots(figsize=(8, 6))\n",
    "ax2 = ax1.twinx()\n",
    "\n",
    "ax1.plot(Gamma_elastic, '-ob')\n",
    "ax2.plot(PSD, '-*r')\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1e'))\n",
    "\n",
    "ax1.set_xlabel('Optimization Steps', fontsize= 12, fontweight='bold')\n",
    "ax1.set_ylabel('Elastic Collision Rate', fontsize= 12, fontweight='bold')\n",
    "ax1.tick_params(axis=\"y\", labelcolor='b')\n",
    "ax2.set_ylabel('Phase Space Density', fontsize= 12, fontweight='bold')\n",
    "ax2.tick_params(axis=\"y\", labelcolor='r')\n",
    "plt.tight_layout()\n",
    "plt.grid(visible=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da1a3346",
   "metadata": {},
   "source": [
    "## Plot crossed beam trap potential resulting for given parameters only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f17a4d01",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x11ebf064970>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 52720 characters]",
      "text/plain": [
       "<Figure size 900x700 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Powers = [1, 11] * u.W\n",
    "Wavelength = 1.064*u.um\n",
    "w_x = [27, 67]*u.um # Beam Waists in the x direction\n",
    "w_z = [31, 67]*u.um # Beam Waists in the y direction\n",
    "\n",
    "options = {\n",
    "    'axis':2, # axis referenced to the beam along which you want the dipole trap potential\n",
    "    'extent': 3e2, # range of spatial coordinates in one direction to calculate trap potential over\n",
    "    'crossed': True, # Flag for either a crossed beam configuration or a single focussed beam\n",
    "    'delta': 70, # angle between arms in degrees\n",
    "    'modulation': False, # Flag for spatial modulation of a single beam\n",
    "    'aspect_ratio': 5, # required aspect ratio of modulated arm\n",
    "    'gravity': False, # Flag for adding levitation/gravitation potential\n",
    "    'tilt_gravity': False, # Flag for tilt of the beam wrt to gravity axis\n",
    "    'theta': 0.75, # gravity tilt angle in degrees\n",
    "    'tilt_axis': [1, 0, 0], # lab space coordinates are rotated about x-axis in reference frame of beam\n",
    "    'astigmatism': False, # Flag for astigmatism of beam\n",
    "    'foci_disp_single': 2.5*u.mm, #0.9 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um, # difference in position of the foci along the propagation direction\n",
    "    'foci_disp_crossed': [2.5*u.mm, 2.8*u.mm], # astigmatism of each of the two beams in the cODT\n",
    "    'foci_shift': [0.0*u.mm, 0.0*u.mm],\n",
    "    'beam_disp': [[0.0, 0.0, 0.05]*u.mm, [0.0, 0.0, 0.0]*u.mm], # shift arm 1 along x, z; shift arm 2 along y, z\n",
    "    'extract_trap_frequencies': True # Flag to extract trap frequencies by fitting the computed potential\n",
    "}\n",
    "\n",
    "Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, Powers, options)\n",
    "\n",
    "EffectiveTrapDepthInKelvin = TrapDepthsInKelvin[1]\n",
    "v  = ExtractedTrapFrequencies[1][0]\n",
    "dv = ExtractedTrapFrequencies[1][1]\n",
    "\n",
    "plt.figure(figsize=(9, 7))\n",
    "plt.plot(Positions[options['axis']], TrappingPotential[options['axis']], label = 'Effective Trap Depth = ' + str(round(EffectiveTrapDepthInKelvin.value, 2)) + ' ' + str(EffectiveTrapDepthInKelvin.unit) + '; ' + generate_label(v, dv)) \n",
    "#plt.plot(Positions[options['axis']], TrappingPotential[options['axis']]) \n",
    "axis = options['axis']\n",
    "if axis == 0:\n",
    "    dir = 'X - Horizontal'\n",
    "elif axis == 1:\n",
    "    dir = 'Y - Propagation'\n",
    "else:\n",
    "    dir = 'Z - Vertical'\n",
    "    \n",
    "plt.ylim(top = 0)\n",
    "plt.xlabel(dir + ' Direction (um)', fontsize= 12, fontweight='bold')\n",
    "plt.ylabel('Trap Potential (uK)', fontsize= 12, fontweight='bold')\n",
    "plt.tight_layout()\n",
    "plt.grid(visible=1)\n",
    "plt.legend(loc=1, prop={'size': 12, 'weight': 'bold'})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a28b820",
   "metadata": {},
   "source": [
    "## Calculate trap frequencies in a crossed beam trap for given parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "09149423",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "v_x = 480.60 Hz\n",
      "v_y = 255.73 Hz\n",
      "v_z = 492.63 Hz\n"
     ]
    }
   ],
   "source": [
    "Powers = [1, 11] * u.W\n",
    "w_x = [27, 67]*u.um # Beam Waists in the x direction\n",
    "w_z = [31, 67]*u.um # Beam Waists in the y direction\n",
    "\n",
    "options = {\n",
    "    'axis': 1, # axis referenced to the beam along which you want the dipole trap potential\n",
    "    'extent': 1e3, # range of spatial coordinates in one direction to calculate trap potential over\n",
    "    'crossed': True, # Flag for either a crossed beam configuration or a single focussed beam\n",
    "    'delta': 70, # angle between arms in degrees\n",
    "    'modulation': False, # Flag for spatial modulation of a single beam\n",
    "    'aspect_ratio': 5, # required aspect ratio of modulated arm\n",
    "    'gravity': False, # Flag for adding levitation/gravitation potential\n",
    "    'tilt_gravity': False, # Flag for tilt of the beam wrt to gravity axis\n",
    "    'theta': 0.75, # gravity tilt angle in degrees\n",
    "    'tilt_axis': [1, 0, 0], # lab space coordinates are rotated about x-axis in reference frame of beam\n",
    "    'astigmatism': False, # Flag for astigmatism of beam\n",
    "    'foci_disp_single': 2.5*u.mm, #0.9 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um, # difference in position of the foci along the propagation direction\n",
    "    'foci_disp_crossed': [2.5*u.mm, 1.5*u.mm], # astigmatism of each of the two beams in the cODT\n",
    "    'foci_shift': [0.01*u.mm, 0.20*u.mm],\n",
    "    'beam_disp': [[0, 0, 0.0]*u.mm, [0, 0, 0.0]*u.mm],\n",
    "    'extract_trap_frequencies': False # Flag to extract trap frequencies by fitting the computed potential\n",
    "}\n",
    "\n",
    "v_x = calculateCrossedBeamTrapFrequency(options['delta'], [w_x, w_z], Powers, dir = 'x')\n",
    "v_y = calculateCrossedBeamTrapFrequency(options['delta'], [w_x, w_z], Powers, dir = 'y')\n",
    "v_z = calculateCrossedBeamTrapFrequency(options['delta'], [w_x, w_z], Powers, dir = 'z')\n",
    "\n",
    "print('v_x = %.2f ' %(v_x.value) + str(v_x.unit))\n",
    "print('v_y = %.2f ' %(v_y.value) + str(v_y.unit))\n",
    "print('v_z = %.2f ' %(v_z.value) + str(v_z.unit))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cf7eb64",
   "metadata": {},
   "source": [
    "## Calculate trap frequencies, trap depths in a crossed beam trap for different shifts of Arm 2 along Y and Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc585c24",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "Powers = [1, 11] * u.W\n",
    "Wavelength = 1.064*u.um\n",
    "w_x = [27, 67]*u.um # Beam Waists in the x direction\n",
    "w_z = [31, 67]*u.um # Beam Waists in the y direction\n",
    "\n",
    "options = {\n",
    "    'axis':1, # axis referenced to the beam along which you want the dipole trap potential\n",
    "    'extent': 2e3, # range of spatial coordinates in one direction to calculate trap potential over\n",
    "    'crossed': True, # Flag for either a crossed beam configuration or a single focussed beam\n",
    "    'delta': 70, # angle between arms in degrees\n",
    "    'modulation': False, # Flag for spatial modulation of a single beam\n",
    "    'aspect_ratio': 5, # required aspect ratio of modulated arm\n",
    "    'gravity': False, # Flag for adding levitation/gravitation potential\n",
    "    'tilt_gravity': False, # Flag for tilt of the beam wrt to gravity axis\n",
    "    'theta': 0.75, # gravity tilt angle in degrees\n",
    "    'tilt_axis': [1, 0, 0], # lab space coordinates are rotated about x-axis in reference frame of beam\n",
    "    'astigmatism': False, # Flag for astigmatism of beam\n",
    "    'foci_disp_single': 2.5*u.mm, #0.9 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um, # difference in position of the foci along the propagation direction\n",
    "    'foci_disp_crossed': [2.5*u.mm, 2.8*u.mm], # astigmatism of each of the two beams in the cODT\n",
    "    'foci_shift': [0.0*u.mm, 0.0*u.mm],\n",
    "    'beam_disp': [[0.0, 0.0, 0.0]*u.mm, [0.0, 0.0, 0.0]*u.mm],\n",
    "    'extract_trap_frequencies': True # Flag to extract trap frequencies by fitting the computed potential\n",
    "}\n",
    "\n",
    "disp_range = 100\n",
    "yvals = np.arange(-disp_range, disp_range, 5) \n",
    "zvals = np.arange(-disp_range, disp_range, 5)\n",
    "\n",
    "TrapDepths = np.zeros((len(yvals), len(zvals))) * u.uK\n",
    "TrapFrequencies = np.zeros((len(yvals), len(zvals)))\n",
    "\n",
    "idx = 0\n",
    "if options['extract_trap_frequencies']:\n",
    "    idx = 1\n",
    "            \n",
    "for i in range(len(yvals)):\n",
    "    options['beam_disp'][1][1] = yvals[i] * u.um\n",
    "    for j in range(len(zvals)):\n",
    "        options['beam_disp'][1][2] = zvals[j] * u.um\n",
    "        Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, Powers, options)\n",
    "        TrapDepths[i][j]      = TrapDepthsInKelvin[1]        \n",
    "        try:\n",
    "            if np.isnan(ExtractedTrapFrequencies[idx][0].value):\n",
    "                TrapFrequencies[i][j] = np.nan\n",
    "            else:\n",
    "                TrapFrequencies[i][j] = ExtractedTrapFrequencies[idx][0]    \n",
    "        except:\n",
    "            TrapFrequencies[i][j] = ExtractedTrapFrequencies[idx][0]                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e060601f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "axis = options['axis']\n",
    "if axis == 0:\n",
    "    dir = 'X'\n",
    "elif axis == 1:\n",
    "    dir = 'Y'\n",
    "else:\n",
    "    dir = 'Z'\n",
    "plt.title('Trap Frequency along ' + dir)\n",
    "c = ax.matshow(np.asarray(TrapFrequencies), origin='lower', cmap = 'coolwarm', aspect = 1)\n",
    "x_pos = np.arange(len(zvals))\n",
    "plt.xticks(x_pos, zvals)\n",
    "#plt.xticks(np.arange(min(zvals.value), max(zvals.value), 30))\n",
    "ax.xaxis.set_ticks_position('bottom')\n",
    "y_pos = np.arange(len(yvals))\n",
    "plt.yticks(y_pos, yvals)\n",
    "#plt.yticks(np.arange(min(yvals.value), max(yvals.value), 30))\n",
    "plt.xlabel('Z - shift (um)', fontsize= 12, fontweight='bold')\n",
    "plt.ylabel('Y - shift (um)', fontsize= 12, fontweight='bold')\n",
    "plt.setp(ax.get_xticklabels(), rotation=45, horizontalalignment='right')\n",
    "fig.colorbar(c, ax = ax)\n",
    "plt.tight_layout()\n",
    "plt.grid(visible=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fe6e6de",
   "metadata": {},
   "source": [
    "## Calculate trap frequencies, trap depths in a crossed beam trap for different shifts of Arm 1 along X and Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "faf6d746",
   "metadata": {},
   "outputs": [],
   "source": [
    "Powers = [1, 11] * u.W\n",
    "Wavelength = 1.064*u.um\n",
    "w_x = [27, 67]*u.um # Beam Waists in the x direction\n",
    "w_z = [31, 67]*u.um # Beam Waists in the y direction\n",
    "\n",
    "options = {\n",
    "    'axis':2, # axis referenced to the beam along which you want the dipole trap potential\n",
    "    'extent': 3e3, # range of spatial coordinates in one direction to calculate trap potential over\n",
    "    'crossed': True, # Flag for either a crossed beam configuration or a single focussed beam\n",
    "    'delta': 70, # angle between arms in degrees\n",
    "    'modulation': False, # Flag for spatial modulation of a single beam\n",
    "    'aspect_ratio': 5, # required aspect ratio of modulated arm\n",
    "    'gravity': False, # Flag for adding levitation/gravitation potential\n",
    "    'tilt_gravity': False, # Flag for tilt of the beam wrt to gravity axis\n",
    "    'theta': 0.75, # gravity tilt angle in degrees\n",
    "    'tilt_axis': [1, 0, 0], # lab space coordinates are rotated about x-axis in reference frame of beam\n",
    "    'astigmatism': False, # Flag for astigmatism of beam\n",
    "    'foci_disp_single': 2.5*u.mm, #0.9 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um, # difference in position of the foci along the propagation direction\n",
    "    'foci_disp_crossed': [2.5*u.mm, 2.8*u.mm], # astigmatism of each of the two beams in the cODT\n",
    "    'foci_shift': [0.0*u.mm, 0.0*u.mm],\n",
    "    'beam_disp': [[0.0, 0.0, 0.0]*u.mm, [0.0, 0.0, 0.0]*u.mm],\n",
    "    'extract_trap_frequencies': True # Flag to extract trap frequencies by fitting the computed potential\n",
    "}\n",
    "\n",
    "disp_range = 100\n",
    "xvals = np.arange(-disp_range, disp_range, 5) \n",
    "zvals = np.arange(-disp_range, disp_range, 5)\n",
    "\n",
    "TrapDepths = np.zeros((len(xvals), len(zvals))) * u.uK\n",
    "TrapFrequencies = np.zeros((len(xvals), len(zvals)))\n",
    "\n",
    "idx = 0\n",
    "if options['extract_trap_frequencies']:\n",
    "    idx = 1\n",
    "            \n",
    "for i in range(len(xvals)):\n",
    "    options['beam_disp'][0][0] = xvals[i] * u.um\n",
    "    for j in range(len(zvals)):\n",
    "        options['beam_disp'][0][2] = zvals[j] * u.um\n",
    "        Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, Powers, options)\n",
    "        TrapDepths[i][j]      = TrapDepthsInKelvin[1]        \n",
    "        try:\n",
    "            if np.isnan(ExtractedTrapFrequencies[idx][0].value):\n",
    "                TrapFrequencies[i][j] = np.nan\n",
    "            else:\n",
    "                TrapFrequencies[i][j] = ExtractedTrapFrequencies[idx][0]    \n",
    "        except:\n",
    "            TrapFrequencies[i][j] = ExtractedTrapFrequencies[idx][0]                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "387ec0a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "axis = options['axis']\n",
    "if axis == 0:\n",
    "    dir = 'X'\n",
    "elif axis == 1:\n",
    "    dir = 'Y'\n",
    "else:\n",
    "    dir = 'Z'\n",
    "plt.title('Trap Frequency along ' + dir)\n",
    "\n",
    "c = ax.matshow(np.asarray(TrapFrequencies), origin='lower', cmap = 'coolwarm', aspect = 1)\n",
    "\n",
    "x_pos = np.arange(len(zvals))\n",
    "plt.xticks(x_pos, zvals)\n",
    "#plt.xticks(np.arange(min(zvals.value), max(zvals.value), 30))\n",
    "ax.xaxis.set_ticks_position('bottom')\n",
    "x_pos = np.arange(len(xvals))\n",
    "plt.yticks(x_pos, xvals)\n",
    "#plt.yticks(np.arange(min(yvals.value), max(yvals.value), 30))\n",
    "plt.xlabel('Z - shift (um)', fontsize= 12, fontweight='bold')\n",
    "plt.ylabel('X - shift (um)', fontsize= 12, fontweight='bold')\n",
    "plt.setp(ax.get_xticklabels(), rotation=45, horizontalalignment='right')\n",
    "fig.colorbar(c, ax = ax)\n",
    "plt.tight_layout()\n",
    "plt.grid(visible=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc9b993e",
   "metadata": {},
   "source": [
    "# Heating Rate in ODT Arm 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cc8af09",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_x = 27 * u.um\n",
    "w_y = 31 * u.um\n",
    "P = 2.29 * u.W\n",
    "linewidth  = 32.2 * u.MHz\n",
    "detuning   = 4.3e+14 * u.Hz \n",
    "\n",
    "HeatingRate, T_recoil, Gamma_sr, U = calculateHeatingRate(w_x, w_y, P, linewidth, detuning, wavelength = 1.064*u.um)\n",
    "\n",
    "print((U/ac.k_B).to(u.uK))\n",
    "print(Gamma_sr)\n",
    "print(HeatingRate.to(u.nK/u.s))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56bef045",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
