{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "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 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors, colors_alpha = setup()\n",
    "pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_power = 291.5 * si.uW\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",
    "axial_width = trap.get_tweezer_rayleigh()\n",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle W_{t}\\left(\\frac{\\omega_{t r}}{\\omega_{t ax}} = 7\\right) = 1.68\\mathrm{\\mu m}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "waist = trap.subs(\n",
    "    sp.solve(\n",
    "        trap.get_omega_r_tweezer() / trap.get_omega_ax_tweezer() - 7,\n",
    "        trap.waist_tweezer,\n",
    "    )[0]\n",
    ").evalf()\n",
    "\n",
    "display(\n",
    "    Math(\n",
    "        f\"{sp.latex(trap.waist_tweezer)}\\\\left({_aspect_ratio_latex} = 7\\\\right)\"\n",
    "        f\" = {waist/si.um:.2f}\\\\mathrm{{\\\\mu m}}\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:26<00:00,  3.83it/s]\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 = 30\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] = 2 * np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47772 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)\n",
    "fig.savefig(FIGS_PATH / \"twodtrap\" / \"1D Stufenplot.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "width_np = float(trap.subs(axial_width))\n",
    "\n",
    "z_np = np.linspace(-0.5 * width_np, 2 * width_np, num=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-5.719499903379515e-30\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41856 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, 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)\")\n",
    "\n",
    "fig.savefig(FIGS_PATH / \"twodtrap\" / \"1D Spilling Example.pdf\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
