{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# See how gradient spilling would look like for Li:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_power = 291.5 * si.uW\n",
    "initial_grad = 21 *si.G/si.cm\n",
    "\n",
    "#parameters for lithium setup\n",
    "wvl = 1064 * si.nm\n",
    "omega_l = 2 * np.pi * const.c / wvl\n",
    "omega_0 = 2 * np.pi * const.c / (671 * si.nm)\n",
    "gamma = 2 * np.pi * 5.8724 * const.mega\n",
    "\n",
    "trap: PancakeTrap = PancakeTrap(\n",
    "    power=0,  # Set pancake laser power to 0, no 2D trap\n",
    "    grad_z=initial_grad,\n",
    "    grad_r=0,\n",
    "    power_tweezer=initial_power,\n",
    "    waist_tweezer=1.838 * si.um,\n",
    "\n",
    "    wvl = 1064 * si.nm,\n",
    "    omega_0 = 2 * np.pi * const.c / (671 * si.nm),\n",
    "    a = (3 * sp.pi * const.c**2) / (2 * omega_0**3)* (gamma / (omega_0 - omega_l) + gamma / (omega_0 + omega_l)),\n",
    "    m = 6.0151228 * const.value(\"atomic mass constant\"),\n",
    "    mu_b = const.value(\"Bohr magneton\"),\n",
    ")\n",
    "zr = float(trap.subs(trap.get_tweezer_rayleigh()))\n",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[6.12403657e+24 1.40929010e+19 1.88588017e+14 7.28586661e+09\n",
      " 6.91489717e+05 1.58272638e+02 9.72417414e-02 2.84255990e-04] s\n",
      "-6.846373487865556e-30\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 50460 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trap[trap.power_tweezer] = 1* initial_power\n",
    "trap[trap.grad_z] = 1*initial_grad\n",
    "plot_solutions(trap,100,-1*zr,2.2*zr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/50 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<lambdifygenerated-1563>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 9.25429819739509e-40*z/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2) - 3.01926539188606e-24\n",
      "<lambdifygenerated-1564>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -3.70171927895804e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 9.25429819739509e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      "100%|██████████| 50/50 [01:42<00:00,  2.05s/it]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 33992 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trap[trap.power_tweezer] = initial_power\n",
    "trap[trap.grad_z] = initial_grad\n",
    "plot_occupation_grad(trap,100,-1.5*zr,2.2*zr,grad_fac_down=1,grad_fac_up=1.5,n_spill_steps=50,n_pot_steps=int(1e3),t_spill=25*si.ms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/50 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<lambdifygenerated-1963>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 9.25429819739509e-40*z/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2) - 3.01926539188606e-24\n",
      "<lambdifygenerated-1964>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -3.70171927895804e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 9.25429819739509e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      "100%|██████████| 50/50 [01:40<00:00,  2.01s/it]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 33984 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trap[trap.power_tweezer] = initial_power\n",
    "trap[trap.grad_z] = initial_grad\n",
    "plot_occupation_grad(trap,100,-1.5*zr,2.2*zr,grad_fac_down=1,grad_fac_up=1.5,n_spill_steps=50,n_pot_steps=int(1e3),t_spill=1*si.ms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
