{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "from IPython.display import Math, display\n",
    "from matplotlib.axes import Axes\n",
    "from scipy import constants as const\n",
    "from scipy.integrate import quad\n",
    "from scipy.optimize import root_scalar\n",
    "from scipy.signal import argrelmax,argrelmin\n",
    "from tqdm import tqdm\n",
    "\n",
    "import fewfermions.analysis.units as si\n",
    "from fewfermions.simulate.traps.twod.trap import PancakeTrap\n",
    "from fewfermions.style import FIGS_PATH, setup\n",
    "\n",
    "from helpers import *\n",
    "\n",
    "colors, colors_alpha = setup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving the radial SE to see when those get populated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_power = 49 * si.uW \n",
    "\n",
    "trap: PancakeTrap = PancakeTrap(\n",
    "    power=0,  # Set pancake laser power to 0, no 2D trap\n",
    "    grad_z= 0 * si.G / si.cm,\n",
    "    grad_r=0,\n",
    "    power_tweezer= initial_power,\n",
    "    waist_tweezer=1 * si.um,\n",
    "    a=184.4*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "\n",
    "    wvl = 1064 * si.nm,\n",
    ")\n",
    "axial_width = trap.get_tweezer_rayleigh()\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": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle - 2.67062937174555 \\cdot 10^{-24} z + \\frac{5.61241215754764 \\cdot 10^{-29}}{\\pi} - \\frac{5.61241215754764 \\cdot 10^{-29} e^{\\frac{1000000000000.0 \\left(- 2 x^{2} - 2 y^{2}\\right)}{\\frac{1132096000000.0 z^{2}}{\\pi^{2}} + 1}}}{\\pi \\left(\\frac{1132096000000.0 z^{2}}{\\pi^{2}} + 1\\right)}$"
      ],
      "text/plain": [
       "-2.67062937174555e-24*z + 5.61241215754764e-29/pi - 5.61241215754764e-29*exp(1000000000000.0*(-2*x**2 - 2*y**2)/(1132096000000.0*z**2/pi**2 + 1))/(pi*(1132096000000.0*z**2/pi**2 + 1))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trap.subs(trap.get_potential())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_power = 291.5 * si.uW\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=15 * si.G / si.cm,\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",
    "axial_width = trap.get_tweezer_rayleigh()\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": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\omega_{t r}}{\\omega_{t ax}} \\approx 7.67$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "aspect_ratio = trap.get_omega_r_tweezer() / trap.get_omega_ax_tweezer()\n",
    "_aspect_ratio_latex = sp.latex(trap.omega_r_tweezer / trap.omega_ax_tweezer)\n",
    "display(Math(f\"{_aspect_ratio_latex} \\\\approx {trap.subs(aspect_ratio).evalf():.2f}\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "trap[trap.power_tweezer] = 0.52* initial_power\n",
    "# Solve the hamiltonian numerically in axial direction\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "    trap.x, (-0.5 * axial_width, 3 * axial_width), n_pot_steps, k=n_levels\n",
    ")\n",
    "\n",
    "pot_ax = trap.subs(trap.get_potential())\n",
    "pot_diff_ax = sp.diff(pot_ax, trap.x)\n",
    "pot_diff2_ax = sp.diff(pot_diff_ax, trap.x)\n",
    "pot_diff3_ax = sp.diff(pot_diff2_ax, trap.x)\n",
    "pot_diff_ax_numpy = sp.lambdify(trap.x, pot_diff_ax.subs({x: 0, y: 0}))\n",
    "pot_diff2_ax_numpy = sp.lambdify(trap.x, pot_diff2_ax.subs({x: 0, y: 0}))\n",
    "pot_diff3_ax_numpy = sp.lambdify(trap.x, pot_diff3_ax.subs({x: 0, y: 0}))\n",
    "\n",
    "barrier = root_scalar(\n",
    "    pot_diff_ax_numpy,\n",
    "    x0=1.5 * float(trap.subs(axial_width)),\n",
    "    fprime=pot_diff2_ax_numpy,\n",
    "    xtol=1e-18,\n",
    "    fprime2=pot_diff2_ax_numpy,\n",
    ").root\n",
    "\n",
    "# States that are below the potential barrier\n",
    "bound_states = energies < potential(barrier)\n",
    "\n",
    "\n",
    "# Density of states is larger on the left than on the right\n",
    "# Likely that the state in question is a true bound state\n",
    "true_bound_states = np.logical_and(\n",
    "    bound_states,\n",
    "    np.sum(states[:, coords[x] < barrier] ** 2, axis=1)\n",
    "    > np.sum(states[:, coords[x] > barrier] ** 2, axis=1),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.5217732675329976e-34\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'E / h (kHz)')"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 26468 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "width_np = float(trap.subs(axial_width))\n",
    "\n",
    "z_np = np.linspace(-0.5 * width_np, 2 * width_np, num=1000)\n",
    "\n",
    "ax: plt.Axes\n",
    "fig, ax = plt.subplots(figsize=(2.5, 2.5))\n",
    "# ax.set_title(\"Axial\")\n",
    "abs_min = np.min(potential(z_np))\n",
    "print(abs_min)\n",
    "ax.fill_between(\n",
    "    z_np / si.um,\n",
    "    potential(z_np) / const.h / si.kHz,\n",
    "    abs_min / const.h / si.kHz,\n",
    "    fc=colors_alpha[\"red\"],\n",
    "    alpha=0.5,\n",
    ")\n",
    "# ax2 = ax.twinx()\n",
    "\n",
    "for i, bound in enumerate(true_bound_states):\n",
    "    if not bound:\n",
    "        continue\n",
    "    energy = energies[i]\n",
    "    state = states[i]\n",
    "    ax.plot(\n",
    "        z_np / si.um,\n",
    "        np.where(\n",
    "            (energy > potential(z_np)) & (z_np < barrier),\n",
    "            energy / const.h / si.kHz,\n",
    "            np.nan,\n",
    "        ),\n",
    "        c=\"k\",\n",
    "        lw=0.5,\n",
    "        marker=\"None\",\n",
    "    )\n",
    "    # ax1.plot(coords[trap.z], state**2, marker=\"None\", c=\"k\")\n",
    "\n",
    "ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz, marker=\"None\")\n",
    "ax.set_title(f\"{np.sum(true_bound_states)} bound solutions, power: {trap.subs(trap.power_tweezer)/si.mW}mW\")\n",
    "ax.set_xlabel(r\"z ($\\mathrm{\\mu m}$)\")\n",
    "ax.set_ylabel(r\"E / h (kHz)\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving radial direction separatly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\omega_{t r}}{\\omega_{t ax}} \\approx 4.18$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[7.04017576e+24 1.55764305e+19 2.03219750e+14 7.77300227e+09\n",
      " 7.43763814e+05 1.75838686e+02 1.16136010e-01] s\n",
      "-9.92574367372403e-31\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27540 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from helpers import *\n",
    "\n",
    "initial_power = 49 * si.uW\n",
    "\n",
    "trap: PancakeTrap = PancakeTrap(\n",
    "    power=0,  # Set pancake laser power to 0, no 2D trap\n",
    "    grad_z= 0 * si.G / si.cm,\n",
    "    grad_r=0,\n",
    "    power_tweezer= initial_power,\n",
    "    waist_tweezer=1 * si.um,\n",
    "    a=184.4*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "\n",
    "    wvl = 1064 * si.nm,\n",
    ")\n",
    "axial_width = trap.get_tweezer_rayleigh()\n",
    "zr = float(trap.subs(trap.get_tweezer_rayleigh()))\n",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z\n",
    "\n",
    "aspect_ratio = trap.get_omega_r_tweezer() / trap.get_omega_ax_tweezer()\n",
    "_aspect_ratio_latex = sp.latex(trap.omega_r_tweezer / trap.omega_ax_tweezer)\n",
    "display(Math(f\"{_aspect_ratio_latex} \\\\approx {trap.subs(aspect_ratio).evalf():.2f}\"))\n",
    "\n",
    "\n",
    "plot_solutions(trap,100,-0.5*zr,1.9*zr,n_pot_steps=1000,display_plot=-2,state_mult=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\omega_{t r}}{\\omega_{t ax}} \\approx 7.67$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2.62536051e+67 9.76559172e+59 2.04847784e+53 1.16567274e+47\n",
      " 1.46026265e+41 3.65790780e+35 1.73924943e+30 1.52473570e+25\n",
      " 2.43003841e+20 7.02643563e+15 3.72319935e+11 3.70687856e+07\n",
      " 7.27419244e+03 3.08507686e+00 3.52499553e-03] s\n",
      "-1.235754055253792e-30\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29676 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from helpers import *\n",
    "\n",
    "initial_power = 291.5 * si.uW\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=15 * si.G / si.cm,\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",
    "axial_width = trap.get_tweezer_rayleigh()\n",
    "zr = float(trap.subs(trap.get_tweezer_rayleigh()))\n",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z\n",
    "\n",
    "aspect_ratio = trap.get_omega_r_tweezer() / trap.get_omega_ax_tweezer()\n",
    "_aspect_ratio_latex = sp.latex(trap.omega_r_tweezer / trap.omega_ax_tweezer)\n",
    "display(Math(f\"{_aspect_ratio_latex} \\\\approx {trap.subs(aspect_ratio).evalf():.2f}\"))\n",
    "\n",
    "\n",
    "plot_solutions(trap,100,-0.5*zr,1.9*zr,n_pot_steps=1000,display_plot=-2,state_mult=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.1717890943465785e-33\n",
      "26.564333266068875\n",
      "25.96489910768727\n",
      "25.382996351948655\n",
      "24.818065021051286\n",
      "24.269570728094312\n",
      "23.73700386737578\n",
      "23.21987890025091\n",
      "22.71773374263102\n",
      "22.23012926268299\n",
      "21.756648900495044\n",
      "21.2968984256903\n",
      "20.8505058545992\n",
      "20.41712155624778\n",
      "19.996418586981296\n",
      "19.588093308397557\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 18332 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "plot_solutions_ax(trap,15,-17*zr,17*zr,n_pot_steps=1000,display_plot=-2,state_mult=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.528116823253725e-34\n",
      "25.695640087867233\n",
      "23.211447531392245\n",
      "20.82326989690704\n",
      "18.533513496817722\n",
      "16.344870905888378\n",
      "14.260379578986871\n",
      "12.283498634669748\n",
      "10.418211686394796\n",
      "8.669168207811033\n",
      "7.041884072722405\n",
      "5.5430371954550886\n",
      "4.180924936038319\n",
      "2.9662175320690007\n",
      "1.9133091495166965\n",
      "1.0430571002545062\n",
      "0.3894904467789659\n",
      "0.010187209408827964\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 19396 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_solutions_rad(trap,100,-0.9*zr,0.9*zr,n_pot_steps=1000,display_plot=-2,state_mult=1e3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solve 3D SE and look at step plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\omega_{t r}}{\\omega_{t ax}} \\approx 4.18$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2.14636597e+55 3.07658108e+48 2.39971369e+42 4.91850572e+36\n",
      " 2.15913958e+31 1.85351395e+26 2.97267208e+21 8.72400278e+16\n",
      " 4.67207122e+12 4.63299621e+08 8.82595867e+04 3.48134637e+01\n",
      " 3.17758787e-02] s\n",
      "-7.363216798479622e-31\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28276 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from helpers import *\n",
    "\n",
    "initial_power = 49 * si.uW\n",
    "\n",
    "trap: PancakeTrap = PancakeTrap(\n",
    "    power=0,  # Set pancake laser power to 0, no 2D trap\n",
    "    grad_z= 0 * si.G / si.cm,\n",
    "    grad_r=0,\n",
    "    power_tweezer= initial_power,\n",
    "    waist_tweezer=1 * si.um,\n",
    "    a=184.4*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "\n",
    "    wvl = 1064 * si.nm,\n",
    ")\n",
    "axial_width = trap.get_tweezer_rayleigh()\n",
    "zr = float(trap.subs(trap.get_tweezer_rayleigh()))\n",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z\n",
    "\n",
    "aspect_ratio = trap.get_omega_r_tweezer() / trap.get_omega_ax_tweezer()\n",
    "_aspect_ratio_latex = sp.latex(trap.omega_r_tweezer / trap.omega_ax_tweezer)\n",
    "display(Math(f\"{_aspect_ratio_latex} \\\\approx {trap.subs(aspect_ratio).evalf():.2f}\"))\n",
    "\n",
    "\n",
    "plot_solutions(trap,100,-0.5*zr,1.9*zr,n_pot_steps=1000,display_plot=-2,state_mult=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_spill_steps = 100\n",
    "\n",
    "trap[trap.power_tweezer] = initial_power\n",
    "\n",
    "spill_power_factor = np.linspace(0.7, 0.52, num=n_spill_steps)\n",
    "powers = trap[trap.power_tweezer] * spill_power_factor\n",
    "t_spill = 25 * si.ms\n",
    "atom_number = np.zeros_like(powers)\n",
    "\n",
    "# Number of energy levels to compute\n",
    "# will change over time to avoid calculating too many levels\n",
    "n_levels = 30\n",
    "# Resolution of the potential when solving numerically\n",
    "n_pot_steps = 20\n",
    "\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "        (trap.x,trap.y,trap.z), ((-0.5 * axial_width,0.5 * axial_width),(-0.5 * axial_width,0.5 * axial_width),(-0.5 * axial_width, 1.8 * axial_width)), n_pot_steps, k=n_levels\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x21af2a3cc50>"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 87692 characters]",
      "text/plain": [
       "<Figure size 507.875x507.875 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(states[0,int(0.5*n_pot_steps)]**2,origin=\"lower\",extent=[np.min(coords[y]),np.max(coords[y]),np.min(coords[z]),np.max(coords[z])],cmap='viridis')\n",
    "Y ,Z = np.meshgrid(coords[y],coords[z])\n",
    "plt.contour(Y, Z, potential(coords[x][int(0.5*n_pot_steps)],Y,Z), levels=30, colors='white', linewidths=0.8)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Do 2D first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_spill_steps = 100\n",
    "\n",
    "trap[trap.power_tweezer] = initial_power\n",
    "\n",
    "spill_power_factor = np.linspace(0.7, 0.52, num=n_spill_steps)\n",
    "powers = trap[trap.power_tweezer] * spill_power_factor\n",
    "t_spill = 25 * si.ms\n",
    "atom_number = np.zeros_like(powers)\n",
    "\n",
    "# Number of energy levels to compute\n",
    "# will change over time to avoid calculating too many levels\n",
    "n_levels = 30\n",
    "# Resolution of the potential when solving numerically\n",
    "n_pot_steps = 200\n",
    "\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "        (trap.y,trap.z), ((-0.8 * axial_width,0.8 * axial_width),(-0.5 * axial_width, 1.9 * axial_width)), n_pot_steps, k=n_levels)\n",
    "\n",
    "#energies, states, potential, coords = trap.nstationary_solution(\n",
    "#        (trap.y,trap.z), ((-0.5 * axial_width,0.5 * axial_width),(-0.2 * axial_width, 1.2 * axial_width)), n_pot_steps, k=n_levels)\n",
    "\n",
    "pot_ax = trap.subs(trap.get_potential())\n",
    "pot_diff_ax = sp.diff(pot_ax, trap.z)\n",
    "pot_diff2_ax = sp.diff(pot_diff_ax, trap.z)\n",
    "pot_diff3_ax = sp.diff(pot_diff2_ax, trap.z)\n",
    "pot_diff_ax_numpy = sp.lambdify(trap.z, pot_diff_ax.subs({x: 0, y: 0}))\n",
    "pot_diff2_ax_numpy = sp.lambdify(trap.z, pot_diff2_ax.subs({x: 0, y: 0}))\n",
    "pot_diff3_ax_numpy = sp.lambdify(trap.z, pot_diff3_ax.subs({x: 0, y: 0}))\n",
    "\n",
    "barrier = root_scalar(\n",
    "    pot_diff_ax_numpy,\n",
    "    x0=1.5 * float(trap.subs(axial_width)),\n",
    "    fprime=pot_diff2_ax_numpy,\n",
    "    xtol=1e-18,\n",
    "    fprime2=pot_diff3_ax_numpy,\n",
    ").root\n",
    "\n",
    "# States that are below the potential barrier\n",
    "bound_states = energies < potential(0,barrier)\n",
    "\n",
    "\n",
    "# Density of states is larger on the left than on the right\n",
    "# Likely that the state in question is a true bound state\n",
    "true_bound_states = np.logical_and(\n",
    "    bound_states,\n",
    "    np.sum(states[:, coords[z] < barrier] ** 2, axis=(1,2))\n",
    "    > np.sum(states[:, coords[z] > barrier] ** 2, axis=(1,2)),\n",
    ")\n",
    "\n",
    "true_states = np.where(np.sum(states[:, coords[z] < barrier] ** 2, axis=(1,2))\n",
    "    > np.sum(states[:, coords[z] > barrier] ** 2, axis=(1,2)), True, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1862948 characters]",
      "text/plain": [
       "<Figure size 1800x4200 with 55 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "num_plots = np.sum(true_states)\n",
    "# Determine subplot grid layout\n",
    "cols = 4  # Number of columns\n",
    "rows = (num_plots + cols - 1) // cols  # Calculate rows dynamically\n",
    "\n",
    "# Create figure and axes\n",
    "fig, axes = plt.subplots(rows, cols, figsize=(12, 4 * rows))\n",
    "\n",
    "# Flatten the axes array for easy iteration (handles single-row/column cases)\n",
    "axes = axes.flatten()\n",
    "\n",
    "# Plot data on each subplot\n",
    "for i, ax in enumerate(axes):\n",
    "    if i < num_plots:\n",
    "        im = ax.imshow(states[true_states][i]**2,origin=\"lower\",extent=[np.min(coords[y]),np.max(coords[y]),np.min(coords[z]),np.max(coords[z])],cmap='viridis')\n",
    "        Y ,Z = np.meshgrid(coords[y],coords[z])\n",
    "        ax.contour(Y, Z, potential(Y,Z), levels=6, colors='white', linewidths=0.5)\n",
    "        fig.colorbar(im, ax=ax, orientation='vertical')\n",
    "        ax.set_title(i)\n",
    "        #if true_bound_states[i]:\n",
    "        #    ax.set_title(f\"{i}, true bound state\")\n",
    "    else:\n",
    "        # Hide unused subplots\n",
    "        ax.axis('off')\n",
    "\n",
    "# Adjust layout\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 29424 characters]",
      "text/plain": [
       "<Figure size 507.875x507.875 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(coords[z]/ si.um,potential(Y*0,Z) / const.h / si.kHz,marker=\"None\")\n",
    "plt.hlines(energies[true_bound_states] / const.h / si.kHz,np.min(coords[z])/ si.um,np.max(coords[z])/ si.um)\n",
    "#plt.vlines(barrier/si.um,-2,12)\n",
    "plt.title(f\"2D simulation: {np.sum(true_bound_states)} bound states\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[4.84434856e+47 9.50548598e+40 1.05725098e+35 3.22570673e+29\n",
      " 2.20825632e+24 3.11338011e+19 8.70485803e+14 4.78539256e+10\n",
      " 5.26285312e+06 1.21862175e+03 6.68379551e-01 1.17123191e-03] s\n",
      "-6.18167887313272e-31\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27828 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_solutions(trap,100,-0.5*zr,1.9*zr,n_pot_steps=1000,display_plot=-2,state_mult=1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\omega_{t r}}{\\omega_{t ax}} \\approx 4.18$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.21523613e+67 4.01805077e+59 7.62262425e+52 3.98990406e+46\n",
      " 4.67531538e+40 1.11396975e+35 5.12378671e+29 4.42086815e+24\n",
      " 7.06039691e+19 2.08565083e+15 1.15367133e+11 1.23014918e+07\n",
      " 2.67364424e+03 1.32463936e+00 2.00627393e-03] s\n",
      "-4.689134944192487e-31\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27992 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from helpers import *\n",
    "\n",
    "initial_power = 49 * si.uW\n",
    "\n",
    "trap: PancakeTrap = PancakeTrap(\n",
    "    power=0,  # Set pancake laser power to 0, no 2D trap\n",
    "    grad_z= -0.8 * si.G / si.cm,\n",
    "    grad_r=0,\n",
    "    power_tweezer= initial_power,\n",
    "    waist_tweezer=1 * si.um,\n",
    "    a=184.4*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "\n",
    "    wvl = 1064 * si.nm,\n",
    ")\n",
    "axial_width = trap.get_tweezer_rayleigh()\n",
    "zr = float(trap.subs(trap.get_tweezer_rayleigh()))\n",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z\n",
    "\n",
    "aspect_ratio = trap.get_omega_r_tweezer() / trap.get_omega_ax_tweezer()\n",
    "_aspect_ratio_latex = sp.latex(trap.omega_r_tweezer / trap.omega_ax_tweezer)\n",
    "display(Math(f\"{_aspect_ratio_latex} \\\\approx {trap.subs(aspect_ratio).evalf():.2f}\"))\n",
    "\n",
    "\n",
    "plot_solutions(trap,100,-0.5*zr,1.9*zr,n_pot_steps=1000,display_plot=-2,state_mult=1e3)"
   ]
  },
  {
   "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
}
