{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import Math, display\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "from scipy import constants as const\n",
    "\n",
    "#add relative path to backend\n",
    "import sys\n",
    "sys.path.append('../../clean_diag/backend')\n",
    "\n",
    "import trap_units as si\n",
    "from twod_trap import PancakeTrap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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",
    "\n",
    "    m= 161 * const.value(\"atomic mass constant\"),\n",
    "    mu_b= 9.93 * const.value(\"Bohr magneton\" ),\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": "markdown",
   "metadata": {},
   "source": [
    "# Gradient to cancel gravity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.847 G/cm\n"
     ]
    }
   ],
   "source": [
    "grav_grad = float(trap.subs(trap.m*const.g/trap.mu_b) /si.G*si.cm)\n",
    "print(f\"{grav_grad:.3f} G/cm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Find waist for target aspect ration (1064nm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle W_{t}\\left(\\frac{\\omega_{t r}}{\\omega_{t ax}} = 9\\right) = 2.16\\mathrm{\\mu m}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "eta_target = 9\n",
    "\n",
    "waist = trap.subs(\n",
    "    sp.solve(\n",
    "        trap.get_omega_r_tweezer() / trap.get_omega_ax_tweezer() - eta_target,\n",
    "        trap.waist_tweezer,\n",
    "    )[0]\n",
    ").evalf()\n",
    "\n",
    "_aspect_ratio_latex = sp.latex(trap.omega_r_tweezer / trap.omega_ax_tweezer)\n",
    "display(\n",
    "    Math(\n",
    "        f\"{sp.latex(trap.waist_tweezer)}\\\\left({_aspect_ratio_latex} = {eta_target}\\\\right)\"\n",
    "        f\" = {waist/si.um:.2f}\\\\mathrm{{\\\\mu m}}\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# For 532nm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "trap[trap.wvl] = 532*si.nm\n",
    "trap[trap.a] = 180* (4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle W_{t}\\left(\\frac{\\omega_{t r}}{\\omega_{t ax}} = 9\\right) = 1.08\\mathrm{\\mu m}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "eta_target = 9\n",
    "\n",
    "waist = trap.subs(\n",
    "    sp.solve(\n",
    "        trap.get_omega_r_tweezer() / trap.get_omega_ax_tweezer() - eta_target,\n",
    "        trap.waist_tweezer,\n",
    "    )[0]\n",
    ").evalf()\n",
    "\n",
    "_aspect_ratio_latex = sp.latex(trap.omega_r_tweezer / trap.omega_ax_tweezer)\n",
    "display(\n",
    "    Math(\n",
    "        f\"{sp.latex(trap.waist_tweezer)}\\\\left({_aspect_ratio_latex} = {eta_target}\\\\right)\"\n",
    "        f\" = {waist/si.um:.2f}\\\\mathrm{{\\\\mu m}}\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ploting for different target aspect ratios"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33192 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wvls = np.array([532, 1064]) * si.nm\n",
    "waists = np.linspace(0.5,4) * si.um\n",
    "\n",
    "for i, wvl in enumerate(wvls):\n",
    "    eta = np.sqrt(2)*np.pi * waists/wvl\n",
    "    plt.plot(waists/si.um,eta,label=f\"{wvl/si.nm}nm\")\n",
    "\n",
    "plt.ylabel(\"eta\")\n",
    "plt.xlabel(\"waist [um]\")\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show that z-states are occupied first (532nm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 334724 characters]",
      "text/plain": [
       "<Figure size 7500x500 with 50 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "initial_power = 91 * si.uW \n",
    "\n",
    "trap: PancakeTrap = PancakeTrap(\n",
    "    power=0,  # Set pancake laser power to 0, no 2D trap\n",
    "    grad_z= -0.55 * si.G / si.cm,\n",
    "    grad_r=0,\n",
    "    power_tweezer= initial_power,\n",
    "    waist_tweezer=1.08 * si.um,\n",
    "\n",
    "    m= 161 * const.value(\"atomic mass constant\"),\n",
    "    mu_b= 9.93 * const.value(\"Bohr magneton\" ),\n",
    "    a=180*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "\n",
    "    wvl = 532 * si.nm,\n",
    ")\n",
    "axial_width = trap.get_tweezer_rayleigh()\n",
    "zr = float(trap.subs(trap.get_tweezer_rayleigh()))\n",
    "\n",
    "#### diagonalise\n",
    "\n",
    "n_grid_1D = 80\n",
    "\n",
    "n_pot_steps = [n_grid_1D,n_grid_1D,n_grid_1D]\n",
    "n_levels = 100\n",
    "\n",
    "left_cutoff = 1*np.max([float(trap.subs(trap.waist_tweezer)),float(trap.subs(trap.waist_tweezer))])\n",
    "right_cutoff = -1*np.max([float(trap.subs(trap.waist_tweezer)),float(trap.subs(trap.waist_tweezer))])\n",
    "back_cutoff = -1*np.max([float(trap.subs(trap.waist_tweezer)),float(trap.subs(trap.waist_tweezer))])\n",
    "front_cutoff = 1*np.max([float(trap.subs(trap.waist_tweezer)),float(trap.subs(trap.waist_tweezer))])\n",
    "bottom_cutoff = -0.3*np.max([float(trap.subs(trap.get_tweezer_rayleigh())),float(trap.subs(trap.get_tweezer_rayleigh()))])\n",
    "top_cutoff = 1.5*np.max([float(trap.subs(trap.get_tweezer_rayleigh())),float(trap.subs(trap.get_tweezer_rayleigh()))])\n",
    "\n",
    "extend = [(left_cutoff,right_cutoff),\n",
    "        (back_cutoff,front_cutoff),\n",
    "        (bottom_cutoff,top_cutoff)]\n",
    "\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "                        [trap.x,trap.y,trap.z], extend, n_pot_steps, k=n_levels,\n",
    "                        method=\"matrix_free\", export=False, )\n",
    "\n",
    "x = coords[trap.x]\n",
    "y = coords[trap.y]\n",
    "z = coords[trap.z]\n",
    "x3D,y3D,z3D = np.meshgrid(coords[trap.x],coords[trap.y],coords[trap.z],indexing=\"ij\")\n",
    "pot = potential(x3D,y3D,z3D)\n",
    "\n",
    "\n",
    "#### plot\n",
    "\n",
    "n_states = 25\n",
    "\n",
    "# Create a 1xn_states grid of subplots (1 row, n_states columns)\n",
    "fig, axes = plt.subplots(1, n_states, figsize=(18/6*n_states, 5))  # Adjust the size as needed\n",
    "\n",
    "# Loop over the state numbers from 0 to 5\n",
    "for state_number, ax in zip(range(n_states), axes):\n",
    "    # Slice through the y-direction and rotate the x-z plane (by swapping x and z)\n",
    "    im = ax.imshow(np.abs(states[state_number, :, int(len(y)/2), :].T)**2, extent=[*extend[0], *extend[2]], origin=\"lower\",\n",
    "                   vmin=np.min(np.abs(states[state_number])**2), vmax=np.max(np.abs(states[state_number])**2))\n",
    "\n",
    "    # Set labels for each subplot\n",
    "    ax.set_xlabel(\"x\")\n",
    "    ax.set_ylabel(\"z\")\n",
    "    ax.set_title(f\"State {state_number}\")\n",
    "    fig.colorbar(im, ax=ax)  \n",
    "\n",
    "# Adjust layout for better spacing\n",
    "plt.tight_layout()\n",
    "\n",
    "# Show the plots\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.003802379076607923\n"
     ]
    }
   ],
   "source": [
    "initial_grad = float(trap.subs(1/trap.mu_b * (3*np.sqrt(3)/4 * initial_power * trap.a/np.pi/trap.waist_tweezer**2/zr - trap.m * trap.g)))\n",
    "print(initial_grad/si.G*si.cm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 323376 characters]",
      "text/plain": [
       "<Figure size 2000x900 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#find true bound states\n",
    "\n",
    "from scipy.optimize import root_scalar\n",
    "\n",
    "axial_width = trap.get_tweezer_rayleigh()\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({trap.x: 0, trap.y: 0}))\n",
    "pot_diff2_ax_numpy = sp.lambdify(trap.z, pot_diff2_ax.subs({trap.x: 0, trap.y: 0}))\n",
    "pot_diff3_ax_numpy = sp.lambdify(trap.z, pot_diff3_ax.subs({trap.x: 0, trap.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",
    "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",
    "# States that are below the potential barrier\n",
    "bound_states = energies < 0#potential(0,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(np.abs(states[:,:,:, coords[trap.z] < barrier]) ** 2, axis=(1,2,3))\n",
    "    > np.sum(np.abs(states[:,:,:, coords[trap.z] > barrier]) ** 2, axis=(1,2,3)),\n",
    ")\n",
    "\n",
    "# true_bound_states = np.full((len(true_bound_states)), False)\n",
    "# true_bound_states[5] = True\n",
    "# true_bound_states[10] = True\n",
    "# true_bound_states[13] = True\n",
    "# true_bound_states[14] = True\n",
    "# true_bound_states[16] = True\n",
    "# true_bound_states[20] = True\n",
    "\n",
    "true_states = states[true_bound_states]\n",
    "\n",
    "\n",
    "#### plot\n",
    "\n",
    "n_states = 12\n",
    "\n",
    "# Create a 1xn_states grid of subplots (1 row, n_states columns)\n",
    "fig, axes = plt.subplots(1, n_states, figsize=(10/6*n_states, 5*1.8), sharey=True)  # Adjust the size as needed\n",
    "\n",
    "# Remove horizontal space between axes\n",
    "fig.subplots_adjust(wspace=5)\n",
    "\n",
    "#use µm\n",
    "plot_extent = [*extend[0], *extend[2]]\n",
    "plot_extent = [x/si.um for x in plot_extent]\n",
    "\n",
    "# Loop over the state numbers from 0 to 5\n",
    "for state_number, ax in zip(range(n_states), axes):\n",
    "    # Slice through the y-direction and rotate the x-z plane (by swapping x and z)\n",
    "    im = ax.imshow(np.abs(true_states[state_number, :, int(len(y)/2), :].T)**2, extent=plot_extent, origin=\"lower\",\n",
    "                   vmin=np.min(np.abs(true_states[state_number])**2), vmax=np.max(np.abs(true_states[state_number])**2))\n",
    "\n",
    "    # Set labels for each subplot\n",
    "    ax.set_xlabel(\"x [µm]\")\n",
    "    if state_number == 0:\n",
    "        ax.set_ylabel(\"z [µm]\")\n",
    "    ax.set_title(f\"State {state_number+1}\")\n",
    "    #fig.colorbar(im, ax=ax)  \n",
    "\n",
    "# Adjust layout for better spacing\n",
    "plt.tight_layout()\n",
    "\n",
    "#fig.colorbar(im, ax=axes.ravel().tolist())\n",
    "\n",
    "plt.savefig(\"z-occupation_withgrad.png\", dpi=300)\n",
    "\n",
    "# Show the plots\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x19617529160>]"
      ]
     },
     "execution_count": 168,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 24484 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(coords[trap.z],potential(0,0,coords[trap.z]))"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
