{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Troubleshooting for why \"n_levels\" changes result\n",
    "## Start with lithium1.ipynb and change all parameters one by one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [],
   "source": [
    "%config InlineBackend.figure_format = \"retina\"\n",
    "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",
    "colors, colors_alpha = setup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### original lithium1 code (adjust initial power to see 0 bound states), (n_levels=60 to check for bug)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_power = 260 * 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",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/100 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 68%|██████▊   | 68/100 [02:28<00:31,  1.00it/s]<lambdifygenerated-1230>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 4.75745677485922e-40*z/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-1231>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -1.90298270994369e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 4.75745677485922e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      "c:\\Users\\peter\\AppData\\Local\\anaconda3\\Lib\\site-packages\\scipy\\optimize\\_root_scalar.py:326: RuntimeWarning: Derivative was zero.\n",
      "  r, sol = methodc(f, x0, args=args, fprime=fprime, fprime2=fprime2, **kwargs)\n",
      " 69%|██████▉   | 69/100 [02:28<00:25,  1.21it/s]<lambdifygenerated-1234>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 4.74244902478079e-40*z/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-1235>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -1.89697960991232e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 4.74244902478079e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      " 70%|███████   | 70/100 [02:29<00:22,  1.35it/s]<lambdifygenerated-1238>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 4.72744127470237e-40*z/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-1239>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -1.89097650988095e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 4.72744127470237e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      " 71%|███████   | 71/100 [02:29<00:19,  1.52it/s]<lambdifygenerated-1242>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 4.71243352462395e-40*z/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-1243>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -1.88497340984958e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 4.71243352462395e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      "<lambdifygenerated-1243>:2: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  return -1.88497340984958e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 4.71243352462395e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      " 72%|███████▏  | 72/100 [02:30<00:16,  1.68it/s]<lambdifygenerated-1246>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 4.69742577454553e-40*z/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-1247>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -1.87897030981821e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 4.69742577454553e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      " 73%|███████▎  | 73/100 [02:30<00:15,  1.78it/s]<lambdifygenerated-1250>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 4.68241802446711e-40*z/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-1251>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -1.87296720978685e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 4.68241802446711e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      " 99%|█████████▉| 99/100 [02:43<00:00,  2.25it/s]<lambdifygenerated-1354>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 4.29221652242819e-40*z/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-1355>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -1.71688660897128e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 4.29221652242819e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      "<lambdifygenerated-1355>:2: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  return -1.71688660897128e-39*z**2/(pi**4*(z**2/pi**2 + 1.00808875956951e-11)**3) + 4.29221652242819e-40/(pi**2*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      "100%|██████████| 100/100 [02:43<00:00,  1.64s/it]\n"
     ]
    }
   ],
   "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 = 60\n",
    "# Resolution of the potential when solving numerically\n",
    "n_pot_steps = 1000\n",
    "\n",
    "for i, power in enumerate(tqdm(powers)):\n",
    "    trap[trap.power_tweezer] = power\n",
    "    # Solve the hamiltonian numerically in axial direction\n",
    "    energies, states, potential, coords = trap.nstationary_solution(\n",
    "        trap.z, (-0.5 * axial_width, 1.8 * axial_width), n_pot_steps, k=n_levels\n",
    "    )\n",
    "\n",
    "    # Determine the potential and its derivatives\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-28,\n",
    "        fprime2=pot_diff2_ax_numpy,\n",
    "    ).root\n",
    "    minimum = root_scalar(\n",
    "        pot_diff_ax_numpy,\n",
    "        x0=0,\n",
    "        fprime=pot_diff2_ax_numpy,\n",
    "        xtol=1e-28,\n",
    "        fprime2=pot_diff2_ax_numpy,\n",
    "    ).root\n",
    "    # States that are below the potential barrier\n",
    "    bound_states = energies < potential(barrier)\n",
    "\n",
    "    n_bound_states = np.sum(bound_states)\n",
    "    n_levles = n_bound_states + 3  # add 3 more levels to be safe\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)\n",
    "        > np.sum(states[:, coords[z] > barrier] ** 2, axis=1),\n",
    "    )\n",
    "\n",
    "    transmission_probability = np.full_like(energies, np.nan, dtype=float)\n",
    "    for j, energy in enumerate(energies):\n",
    "        if not true_bound_states[j]:\n",
    "            continue\n",
    "        intersect_end = root_scalar(\n",
    "            lambda x: potential(x) - energy,\n",
    "            bracket=(barrier, 3 * float(trap.subs(axial_width))),\n",
    "        ).root\n",
    "        intersect_start = root_scalar(\n",
    "            lambda x: potential(x) - energy,\n",
    "            bracket=(minimum, barrier),\n",
    "        ).root\n",
    "        barrier_interval = np.logical_and(\n",
    "            coords[z] > intersect_start, coords[z] < intersect_end\n",
    "        )\n",
    "        s = quad(\n",
    "            lambda x: np.sqrt(\n",
    "                2\n",
    "                * float(trap.subs(trap.m))\n",
    "                * np.clip(potential(x) - energy, a_min=0, a_max=None)\n",
    "            )\n",
    "            / const.hbar,\n",
    "            intersect_start,\n",
    "            intersect_end,\n",
    "        )\n",
    "        transmission_probability[j] = sp.exp(-2 * s[0])\n",
    "    tunneling_rate = (\n",
    "        transmission_probability * np.abs(energies - potential(minimum)) / const.h\n",
    "    )\n",
    "    atom_number[i] = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PolyCollection at 0x20195d18320>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 44440 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 373,
       "width": 373
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "ax: plt.Axes\n",
    "fig: plt.Figure\n",
    "fig, ax = plt.subplots(figsize=(2.5, 2.5))\n",
    "\n",
    "ax.set_xlabel(\"rel. tweezer power (a.u.)\")\n",
    "ax.set_ylabel(\"atom number\")\n",
    "ax.plot(spill_power_factor, atom_number, marker=\"None\")\n",
    "ax.fill_between(spill_power_factor, atom_number, fc=colors_alpha[\"red\"], alpha=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### changing polarisability and increase initial power (n_levels=60 to check for bug)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_power = 400 * 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=184.4*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\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",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/100 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 85%|████████▌ | 85/100 [03:33<00:13,  1.10it/s]<lambdifygenerated-3051>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 1.49145373853941e-39*z/(pi**3*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-3052>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -5.96581495415766e-39*z**2/(pi**5*(z**2/pi**2 + 1.00808875956951e-11)**3) + 1.49145373853941e-39/(pi**3*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      "c:\\Users\\peter\\AppData\\Local\\anaconda3\\Lib\\site-packages\\scipy\\optimize\\_root_scalar.py:326: RuntimeWarning: Derivative was zero.\n",
      "  r, sol = methodc(f, x0, args=args, fprime=fprime, fprime2=fprime2, **kwargs)\n",
      "<lambdifygenerated-3050>:3: RuntimeWarning: overflow encountered in scalar multiply\n",
      "  return 7.3974326386514e-29*x0 - 7.3974326386514e-29*x0/(99197614347.6731*z**2/pi**2 + 1) - 1.391101511745e-24*z\n",
      " 86%|████████▌ | 86/100 [03:34<00:11,  1.26it/s]<lambdifygenerated-3056>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -5.94592890431046e-39*z**2/(pi**5*(z**2/pi**2 + 1.00808875956951e-11)**3) + 1.48648222607762e-39/(pi**3*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      " 87%|████████▋ | 87/100 [03:34<00:09,  1.37it/s]<lambdifygenerated-3059>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 1.48151071361582e-39*z/(pi**3*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-3060>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -5.92604285446327e-39*z**2/(pi**5*(z**2/pi**2 + 1.00808875956951e-11)**3) + 1.48151071361582e-39/(pi**3*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      " 88%|████████▊ | 88/100 [03:35<00:08,  1.48it/s]<lambdifygenerated-3063>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 1.47653920115402e-39*z/(pi**3*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-3064>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -5.90615680461608e-39*z**2/(pi**5*(z**2/pi**2 + 1.00808875956951e-11)**3) + 1.47653920115402e-39/(pi**3*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      " 89%|████████▉ | 89/100 [03:35<00:07,  1.50it/s]<lambdifygenerated-3068>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return -5.88627075476889e-39*z**2/(pi**5*(z**2/pi**2 + 1.00808875956951e-11)**3) + 1.47156768869222e-39/(pi**3*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      "<lambdifygenerated-3067>:2: RuntimeWarning: overflow encountered in scalar power\n",
      "  return 1.47156768869222e-39*z/(pi**3*(z**2/pi**2 + 1.00808875956951e-11)**2) - 1.391101511745e-24\n",
      "<lambdifygenerated-3068>:2: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  return -5.88627075476889e-39*z**2/(pi**5*(z**2/pi**2 + 1.00808875956951e-11)**3) + 1.47156768869222e-39/(pi**3*(z**2/pi**2 + 1.00808875956951e-11)**2)\n",
      "100%|██████████| 100/100 [03:42<00:00,  2.23s/it]\n"
     ]
    }
   ],
   "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 = 100\n",
    "# Resolution of the potential when solving numerically\n",
    "n_pot_steps = 1000\n",
    "\n",
    "for i, power in enumerate(tqdm(powers)):\n",
    "    trap[trap.power_tweezer] = power\n",
    "    # Solve the hamiltonian numerically in axial direction\n",
    "    energies, states, potential, coords = trap.nstationary_solution(\n",
    "        trap.z, (-0.5 * axial_width, 1.8 * axial_width), n_pot_steps, k=n_levels\n",
    "    )\n",
    "\n",
    "    # Determine the potential and its derivatives\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-28,\n",
    "        fprime2=pot_diff2_ax_numpy,\n",
    "    ).root\n",
    "    minimum = root_scalar(\n",
    "        pot_diff_ax_numpy,\n",
    "        x0=0,\n",
    "        fprime=pot_diff2_ax_numpy,\n",
    "        xtol=1e-28,\n",
    "        fprime2=pot_diff2_ax_numpy,\n",
    "    ).root\n",
    "    # States that are below the potential barrier\n",
    "    bound_states = energies < potential(barrier)\n",
    "\n",
    "    n_bound_states = np.sum(bound_states)\n",
    "    n_levles = n_bound_states + 3  # add 3 more levels to be safe\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)\n",
    "        > np.sum(states[:, coords[z] > barrier] ** 2, axis=1),\n",
    "    )\n",
    "\n",
    "    transmission_probability = np.full_like(energies, np.nan, dtype=float)\n",
    "    for j, energy in enumerate(energies):\n",
    "        if not true_bound_states[j]:\n",
    "            continue\n",
    "        intersect_end = root_scalar(\n",
    "            lambda x: potential(x) - energy,\n",
    "            bracket=(barrier, 3 * float(trap.subs(axial_width))),\n",
    "        ).root\n",
    "        intersect_start = root_scalar(\n",
    "            lambda x: potential(x) - energy,\n",
    "            bracket=(minimum, barrier),\n",
    "        ).root\n",
    "        barrier_interval = np.logical_and(\n",
    "            coords[z] > intersect_start, coords[z] < intersect_end\n",
    "        )\n",
    "        s = quad(\n",
    "            lambda x: np.sqrt(\n",
    "                2\n",
    "                * float(trap.subs(trap.m))\n",
    "                * np.clip(potential(x) - energy, a_min=0, a_max=None)\n",
    "            )\n",
    "            / const.hbar,\n",
    "            intersect_start,\n",
    "            intersect_end,\n",
    "        )\n",
    "        transmission_probability[j] = sp.exp(-2 * s[0])\n",
    "    tunneling_rate = (\n",
    "        transmission_probability * np.abs(energies - potential(minimum)) / const.h\n",
    "    )\n",
    "    atom_number[i] = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PolyCollection at 0x20192a628a0>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 47348 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 373,
       "width": 375
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "ax: plt.Axes\n",
    "fig: plt.Figure\n",
    "fig, ax = plt.subplots(figsize=(2.5, 2.5))\n",
    "\n",
    "ax.set_xlabel(\"rel. tweezer power (a.u.)\")\n",
    "ax.set_ylabel(\"atom number\")\n",
    "ax.plot(spill_power_factor, atom_number, marker=\"None\")\n",
    "ax.fill_between(spill_power_factor, atom_number, fc=colors_alpha[\"red\"], alpha=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### n_levels=100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PolyCollection at 0x2019c293170>"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 46504 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 373,
       "width": 373
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "ax: plt.Axes\n",
    "fig: plt.Figure\n",
    "fig, ax = plt.subplots(figsize=(2.5, 2.5))\n",
    "\n",
    "ax.set_xlabel(\"rel. tweezer power (a.u.)\")\n",
    "ax.set_ylabel(\"atom number\")\n",
    "ax.plot(spill_power_factor, atom_number, marker=\"None\")\n",
    "ax.fill_between(spill_power_factor, atom_number, fc=colors_alpha[\"red\"], alpha=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot potential for these parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-7.030454907724839e-30\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'E / h (kHz)')"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 40772 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 373,
       "width": 373
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "trap[trap.power_tweezer] = 0.6 * initial_power\n",
    "# Solve the hamiltonian numerically in axial direction\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "    trap.z, (-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.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_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[z] < barrier] ** 2, axis=1)\n",
    "    > np.sum(states[:, coords[z] > barrier] ** 2, axis=1),\n",
    ")\n",
    "\n",
    "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_xlabel(r\"z ($\\mathrm{\\mu m}$)\")\n",
    "ax.set_ylabel(r\"E / h (kHz)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Change mass and magnetic moment, adjust initial power (n_levels=60 to check for bug)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_power = 3500 * si.uW\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",
    "    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",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  6%|▌         | 3/50 [01:01<15:57, 20.38s/it]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[3], line 75\u001b[0m\n\u001b[0;32m     68\u001b[0m     intersect_start \u001b[38;5;241m=\u001b[39m root_scalar(\n\u001b[0;32m     69\u001b[0m         \u001b[38;5;28;01mlambda\u001b[39;00m x: potential(x) \u001b[38;5;241m-\u001b[39m energy,\n\u001b[0;32m     70\u001b[0m         bracket\u001b[38;5;241m=\u001b[39m(minimum, barrier),\n\u001b[0;32m     71\u001b[0m     )\u001b[38;5;241m.\u001b[39mroot\n\u001b[0;32m     72\u001b[0m     barrier_interval \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mlogical_and(\n\u001b[0;32m     73\u001b[0m         coords[z] \u001b[38;5;241m>\u001b[39m intersect_start, coords[z] \u001b[38;5;241m<\u001b[39m intersect_end\n\u001b[0;32m     74\u001b[0m     )\n\u001b[1;32m---> 75\u001b[0m     s \u001b[38;5;241m=\u001b[39m quad(\n\u001b[0;32m     76\u001b[0m         \u001b[38;5;28;01mlambda\u001b[39;00m x: np\u001b[38;5;241m.\u001b[39msqrt(\n\u001b[0;32m     77\u001b[0m             \u001b[38;5;241m2\u001b[39m\n\u001b[0;32m     78\u001b[0m             \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mfloat\u001b[39m(trap\u001b[38;5;241m.\u001b[39msubs(trap\u001b[38;5;241m.\u001b[39mm))\n\u001b[0;32m     79\u001b[0m             \u001b[38;5;241m*\u001b[39m np\u001b[38;5;241m.\u001b[39mclip(potential(x) \u001b[38;5;241m-\u001b[39m energy, a_min\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m, a_max\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[0;32m     80\u001b[0m         )\n\u001b[0;32m     81\u001b[0m         \u001b[38;5;241m/\u001b[39m const\u001b[38;5;241m.\u001b[39mhbar,\n\u001b[0;32m     82\u001b[0m         intersect_start,\n\u001b[0;32m     83\u001b[0m         intersect_end,\n\u001b[0;32m     84\u001b[0m     )\n\u001b[0;32m     85\u001b[0m     transmission_probability[j] \u001b[38;5;241m=\u001b[39m sp\u001b[38;5;241m.\u001b[39mexp(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m s[\u001b[38;5;241m0\u001b[39m])\n\u001b[0;32m     86\u001b[0m tunneling_rate \u001b[38;5;241m=\u001b[39m (\n\u001b[0;32m     87\u001b[0m     transmission_probability \u001b[38;5;241m*\u001b[39m np\u001b[38;5;241m.\u001b[39mabs(energies \u001b[38;5;241m-\u001b[39m potential(minimum)) \u001b[38;5;241m/\u001b[39m const\u001b[38;5;241m.\u001b[39mh\n\u001b[0;32m     88\u001b[0m )\n",
      "File \u001b[1;32mc:\\Users\\peter\\AppData\\Local\\anaconda3\\Lib\\site-packages\\scipy\\integrate\\_quadpack_py.py:464\u001b[0m, in \u001b[0;36mquad\u001b[1;34m(func, a, b, args, full_output, epsabs, epsrel, limit, points, weight, wvar, wopts, maxp1, limlst, complex_func)\u001b[0m\n\u001b[0;32m    461\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m retval\n\u001b[0;32m    463\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m weight \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m--> 464\u001b[0m     retval \u001b[38;5;241m=\u001b[39m _quad(func, a, b, args, full_output, epsabs, epsrel, limit,\n\u001b[0;32m    465\u001b[0m                    points)\n\u001b[0;32m    466\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    467\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m points \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[1;32mc:\\Users\\peter\\AppData\\Local\\anaconda3\\Lib\\site-packages\\scipy\\integrate\\_quadpack_py.py:611\u001b[0m, in \u001b[0;36m_quad\u001b[1;34m(func, a, b, args, full_output, epsabs, epsrel, limit, points)\u001b[0m\n\u001b[0;32m    609\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m points \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    610\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m infbounds \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m--> 611\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m _quadpack\u001b[38;5;241m.\u001b[39m_qagse(func,a,b,args,full_output,epsabs,epsrel,limit)\n\u001b[0;32m    612\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    613\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m _quadpack\u001b[38;5;241m.\u001b[39m_qagie(func, bound, infbounds, args, full_output, \n\u001b[0;32m    614\u001b[0m                                 epsabs, epsrel, limit)\n",
      "Cell \u001b[1;32mIn[3], line 78\u001b[0m, in \u001b[0;36m<lambda>\u001b[1;34m(x)\u001b[0m\n\u001b[0;32m     68\u001b[0m     intersect_start \u001b[38;5;241m=\u001b[39m root_scalar(\n\u001b[0;32m     69\u001b[0m         \u001b[38;5;28;01mlambda\u001b[39;00m x: potential(x) \u001b[38;5;241m-\u001b[39m energy,\n\u001b[0;32m     70\u001b[0m         bracket\u001b[38;5;241m=\u001b[39m(minimum, barrier),\n\u001b[0;32m     71\u001b[0m     )\u001b[38;5;241m.\u001b[39mroot\n\u001b[0;32m     72\u001b[0m     barrier_interval \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mlogical_and(\n\u001b[0;32m     73\u001b[0m         coords[z] \u001b[38;5;241m>\u001b[39m intersect_start, coords[z] \u001b[38;5;241m<\u001b[39m intersect_end\n\u001b[0;32m     74\u001b[0m     )\n\u001b[0;32m     75\u001b[0m     s \u001b[38;5;241m=\u001b[39m quad(\n\u001b[0;32m     76\u001b[0m         \u001b[38;5;28;01mlambda\u001b[39;00m x: np\u001b[38;5;241m.\u001b[39msqrt(\n\u001b[0;32m     77\u001b[0m             \u001b[38;5;241m2\u001b[39m\n\u001b[1;32m---> 78\u001b[0m             \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mfloat\u001b[39m(trap\u001b[38;5;241m.\u001b[39msubs(trap\u001b[38;5;241m.\u001b[39mm))\n\u001b[0;32m     79\u001b[0m             \u001b[38;5;241m*\u001b[39m np\u001b[38;5;241m.\u001b[39mclip(potential(x) \u001b[38;5;241m-\u001b[39m energy, a_min\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m, a_max\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[0;32m     80\u001b[0m         )\n\u001b[0;32m     81\u001b[0m         \u001b[38;5;241m/\u001b[39m const\u001b[38;5;241m.\u001b[39mhbar,\n\u001b[0;32m     82\u001b[0m         intersect_start,\n\u001b[0;32m     83\u001b[0m         intersect_end,\n\u001b[0;32m     84\u001b[0m     )\n\u001b[0;32m     85\u001b[0m     transmission_probability[j] \u001b[38;5;241m=\u001b[39m sp\u001b[38;5;241m.\u001b[39mexp(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m s[\u001b[38;5;241m0\u001b[39m])\n\u001b[0;32m     86\u001b[0m tunneling_rate \u001b[38;5;241m=\u001b[39m (\n\u001b[0;32m     87\u001b[0m     transmission_probability \u001b[38;5;241m*\u001b[39m np\u001b[38;5;241m.\u001b[39mabs(energies \u001b[38;5;241m-\u001b[39m potential(minimum)) \u001b[38;5;241m/\u001b[39m const\u001b[38;5;241m.\u001b[39mh\n\u001b[0;32m     88\u001b[0m )\n",
      "File \u001b[1;32mc:\\users\\peter\\uni-bulk\\ferdy_chomaz\\code_jonas\\08-fewfermions-2dspilling\\fewfermions\\simulate\\traps\\twod\\trap\\__init__.py:130\u001b[0m, in \u001b[0;36mPancakeTrap.subs\u001b[1;34m(self, expr, symbolic_only)\u001b[0m\n\u001b[0;32m    127\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m expr_symbolic\n\u001b[0;32m    128\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    129\u001b[0m     \u001b[38;5;66;03m# Also substitute all non-symbolic values\u001b[39;00m\n\u001b[1;32m--> 130\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m expr_symbolic\u001b[38;5;241m.\u001b[39msubs(\n\u001b[0;32m    131\u001b[0m         [\n\u001b[0;32m    132\u001b[0m             (key, value)\n\u001b[0;32m    133\u001b[0m             \u001b[38;5;28;01mfor\u001b[39;00m key, value \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_values\u001b[38;5;241m.\u001b[39mitems()\n\u001b[0;32m    134\u001b[0m             \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(value, sp\u001b[38;5;241m.\u001b[39mExpr)\n\u001b[0;32m    135\u001b[0m         ]\n\u001b[0;32m    136\u001b[0m     )\n",
      "File \u001b[1;32mc:\\Users\\peter\\AppData\\Local\\anaconda3\\Lib\\site-packages\\sympy\\core\\basic.py:1079\u001b[0m, in \u001b[0;36mBasic.subs\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1076\u001b[0m \u001b[38;5;66;03m# skip if there is no change\u001b[39;00m\n\u001b[0;32m   1077\u001b[0m sequence \u001b[38;5;241m=\u001b[39m [(s1, s2) \u001b[38;5;28;01mfor\u001b[39;00m s1, s2 \u001b[38;5;129;01min\u001b[39;00m sequence \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m _aresame(s1, s2)]\n\u001b[1;32m-> 1079\u001b[0m simultaneous \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124msimultaneous\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[0;32m   1081\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m unordered:\n\u001b[0;32m   1082\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msorting\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _nodes, default_sort_key\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "n_spill_steps = 50\n",
    "\n",
    "trap[trap.power_tweezer] = initial_power\n",
    "\n",
    "spill_power_factor = np.linspace(0.65, 0.63, 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 = 100\n",
    "# Resolution of the potential when solving numerically\n",
    "n_pot_steps = 1000\n",
    "\n",
    "for i, power in enumerate(tqdm(powers)):\n",
    "    trap[trap.power_tweezer] = power\n",
    "    # Solve the hamiltonian numerically in axial direction\n",
    "    energies, states, potential, coords = trap.nstationary_solution(\n",
    "        trap.z, (-0.5 * axial_width, 1.8 * axial_width), n_pot_steps, k=n_levels\n",
    "    )\n",
    "\n",
    "    # Determine the potential and its derivatives\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-28,\n",
    "        fprime2=pot_diff2_ax_numpy,\n",
    "    ).root\n",
    "    minimum = root_scalar(\n",
    "        pot_diff_ax_numpy,\n",
    "        x0=0,\n",
    "        fprime=pot_diff2_ax_numpy,\n",
    "        xtol=1e-28,\n",
    "        fprime2=pot_diff2_ax_numpy,\n",
    "    ).root\n",
    "    # States that are below the potential barrier\n",
    "    bound_states = energies < potential(barrier)\n",
    "\n",
    "    n_bound_states = np.sum(bound_states)\n",
    "    n_levles = n_bound_states + 3  # add 3 more levels to be safe\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)\n",
    "        > np.sum(states[:, coords[z] > barrier] ** 2, axis=1),\n",
    "    )\n",
    "\n",
    "    transmission_probability = np.full_like(energies, np.nan, dtype=float)\n",
    "    for j, energy in enumerate(energies):\n",
    "        if not true_bound_states[j]:\n",
    "            continue\n",
    "        intersect_end = root_scalar(\n",
    "            lambda x: potential(x) - energy,\n",
    "            bracket=(barrier, 3 * float(trap.subs(axial_width))),\n",
    "        ).root\n",
    "        intersect_start = root_scalar(\n",
    "            lambda x: potential(x) - energy,\n",
    "            bracket=(minimum, barrier),\n",
    "        ).root\n",
    "        barrier_interval = np.logical_and(\n",
    "            coords[z] > intersect_start, coords[z] < intersect_end\n",
    "        )\n",
    "        s = quad(\n",
    "            lambda x: np.sqrt(\n",
    "                2\n",
    "                * float(trap.subs(trap.m))\n",
    "                * np.clip(potential(x) - energy, a_min=0, a_max=None)\n",
    "            )\n",
    "            / const.hbar,\n",
    "            intersect_start,\n",
    "            intersect_end,\n",
    "        )\n",
    "        transmission_probability[j] = sp.exp(-2 * s[0])\n",
    "    tunneling_rate = (\n",
    "        transmission_probability * np.abs(energies - potential(minimum)) / const.h\n",
    "    )\n",
    "    atom_number[i] = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PolyCollection at 0x2019acec920>"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 36596 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 373,
       "width": 373
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "ax: plt.Axes\n",
    "fig: plt.Figure\n",
    "fig, ax = plt.subplots(figsize=(2.5, 2.5))\n",
    "\n",
    "ax.set_xlabel(\"rel. tweezer power (a.u.)\")\n",
    "ax.set_ylabel(\"atom number\")\n",
    "ax.plot(spill_power_factor, atom_number, marker=\"None\")\n",
    "ax.fill_between(spill_power_factor, atom_number, fc=colors_alpha[\"red\"], alpha=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### n_levels=100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PolyCollection at 0x2019c2730b0>"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 40876 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 373,
       "width": 373
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "ax: plt.Axes\n",
    "fig: plt.Figure\n",
    "fig, ax = plt.subplots(figsize=(2.5, 2.5))\n",
    "\n",
    "ax.set_xlabel(\"rel. tweezer power (a.u.)\")\n",
    "ax.set_ylabel(\"atom number\")\n",
    "ax.plot(spill_power_factor, atom_number, marker=\"None\")\n",
    "ax.fill_between(spill_power_factor, atom_number, fc=colors_alpha[\"red\"], alpha=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot potential for these parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "-1.2215042174494931e-28\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'E / h (kHz)')"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 44244 characters]",
      "text/plain": [
       "<Figure size 375x375 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 373,
       "width": 373
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_pot_steps = 1000\n",
    "n_levels = 500\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",
    "minimum = root_scalar(\n",
    "        pot_diff_ax_numpy,\n",
    "        x0=0,\n",
    "        fprime=pot_diff2_ax_numpy,\n",
    "        xtol=1e-28,\n",
    "        fprime2=pot_diff3_ax_numpy,\n",
    "    ).root\n",
    "\n",
    "\n",
    "pot_min_ax_numpy = sp.lambdify(trap.z,pot_ax.subs({x: 0, y: 0}) - potential(minimum))\n",
    "right_cutoff = root_scalar(\n",
    "        pot_min_ax_numpy,\n",
    "        x0=barrier + np.abs(barrier - minimum),\n",
    "        fprime=pot_diff2_ax_numpy,\n",
    "        xtol=1e-28,\n",
    "        fprime2=pot_diff3_ax_numpy,\n",
    "    ).root\n",
    "right_cutoff = barrier + np.abs(barrier - minimum)\n",
    "left_cutoff = minimum - np.abs(barrier - minimum)\n",
    "\n",
    "trap[trap.power_tweezer] = 0.64 * initial_power\n",
    "# Solve the hamiltonian numerically in axial direction\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "    trap.z, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels\n",
    ")\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[z] < barrier] ** 2, axis=1)\n",
    "    > np.sum(states[:, coords[z] > barrier] ** 2, axis=1),\n",
    ")\n",
    "print(np.sum(true_bound_states))\n",
    "\n",
    "width_np = float(trap.subs(axial_width))\n",
    "\n",
    "z_np = np.linspace(left_cutoff, right_cutoff, 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",
    "    ax.plot(z_np/si.um, state**2 *600, marker=\"None\", c=\"k\")\n",
    "\n",
    "ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz, marker=\"None\")\n",
    "ax.set_xlabel(r\"z ($\\mathrm{\\mu m}$)\")\n",
    "ax.set_ylabel(r\"E / h (kHz)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal import argrelmax,argrelmin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.97696491e-06])"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z_in = np.linspace(-1.5* float(trap.subs(axial_width)),3* float(trap.subs(axial_width)),1000)\n",
    "potential_z = sp.lambdify(trap.z, pot_ax.subs({x: 0, y: 0}))(z_in)\n",
    "mi = z_in[argrelmin(test)]\n",
    "mi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x2ac1ca37680>]"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 51800 characters]",
      "text/plain": [
       "<Figure size 507.875x507.875 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 506,
       "width": 506
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(z_in,test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_power = 3500 * si.uW\n",
    "\n",
    "trap: PancakeTrap = PancakeTrap(\n",
    "    power=0,  # Set pancake laser power to 0, no 2D trap\n",
    "    grad_z= 5 * si.G / si.cm,\n",
    "    grad_r=0,\n",
    "    power_tweezer=initial_power,\n",
    "    waist_tweezer=1.838 * 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",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.LineCollection at 0x2ac1f71ecf0>"
      ]
     },
     "execution_count": 126,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 58804 characters]",
      "text/plain": [
       "<Figure size 507.875x507.875 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 506,
       "width": 506
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_pot_steps = 1000\n",
    "n_levels = 500\n",
    "\n",
    "#get potential minimum and maximum\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",
    "\n",
    "z_in = np.linspace(-1.5* float(trap.subs(axial_width)),3* float(trap.subs(axial_width)),n_pot_steps)\n",
    "potential_in = sp.lambdify(trap.z, pot_ax.subs({x: 0, y: 0}))(z_in)\n",
    "barrier = z_in[argrelmax(potential_in)][0]\n",
    "E_max = potential_in[argrelmax(potential_in)][0]\n",
    "minimum = z_in[argrelmin(potential_in)][0]\n",
    "E_min = potential_in[argrelmin(potential_in)][0]\n",
    "\n",
    "#get cutoff values for solving SE\n",
    "pot_min_ax_numpy = sp.lambdify(trap.z,pot_ax.subs({x: 0, y: 0}) - (E_min - 0.5*np.abs(E_max-E_min)) )\n",
    "right_cutoff = root_scalar(\n",
    "        pot_min_ax_numpy,\n",
    "        x0=barrier + np.abs(barrier - minimum),\n",
    "        fprime=pot_diff_ax_numpy,\n",
    "        xtol=1e-28,\n",
    "        fprime2=pot_diff2_ax_numpy,\n",
    "    ).root\n",
    "\n",
    "pot_min_ax_numpy = sp.lambdify(trap.z,pot_ax.subs({x: 0, y: 0}) - (E_min + 4*np.abs(E_max-E_min)) )\n",
    "left_cutoff = root_scalar(\n",
    "        pot_min_ax_numpy,\n",
    "        x0=minimum - np.abs(barrier - minimum),\n",
    "        fprime=pot_diff_ax_numpy,\n",
    "        xtol=1e-28,\n",
    "        fprime2=pot_diff2_ax_numpy,\n",
    "    ).root\n",
    "\n",
    "#plot results\n",
    "plt.plot(z_in,potential_in)\n",
    "plt.vlines(barrier,np.min(potential_in),np.max(potential_in))\n",
    "plt.vlines(minimum,np.min(potential_in),np.max(potential_in))\n",
    "plt.vlines(left_cutoff,np.min(potential_in),np.max(potential_in))\n",
    "plt.vlines(right_cutoff,np.min(potential_in),np.max(potential_in))\n",
    "plt.hlines(E_min,np.min(z_in),np.max(z_in))\n",
    "plt.hlines(E_max,np.min(z_in),np.max(z_in))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-8.426358471408197e-29"
      ]
     },
     "execution_count": 105,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "E_min - np.abs(E_max-E_min)"
   ]
  },
  {
   "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
}
