{
 "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",
    "from scipy.optimize import minimize_scalar\n",
    "from scipy import sparse\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 DoubleTweezer, TwoSiteLattice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trap parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Erbium ###\n",
    "\n",
    "initial_power = 100* si.uW\n",
    "initial_waist = 1.1*si.uW\n",
    "initial_distance = 1.266*si.um\n",
    "\n",
    "trap: DoubleTweezer = DoubleTweezer(\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_tweezer1 = initial_power,     #stationary\n",
    "    power_tweezer2 = initial_power,     #transfer tweezer\n",
    "    waist_tweezer1 = initial_waist,     #stationary\n",
    "    waist_tweezer2 = initial_waist,     #transfer tweezer\n",
    "    distance_tweezers = initial_distance,\n",
    "\n",
    "    m = 168 * const.value(\"atomic mass constant\"),\n",
    "    mu_b = 6.982806* const.value(\"Bohr magneton\"),\n",
    "    a_s = 85* const.value(\"Bohr radius\"),\n",
    "    a=180*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "    wvl = 532 * si.nm,\n",
    "\n",
    "    g = 0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate $E_r = h^2/8ma^2$ and depth of tweezers (single tweezer depth not accuarte if close together)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "E_r = const.h**2/(8*trap.m*trap.distance_tweezers**2)\n",
    "\n",
    "depth_tweezer1 = trap.a*2*trap.power_tweezer1/sp.pi/trap.waist_tweezer1**2\n",
    "depth_tweezer2 = trap.a*2*trap.power_tweezer2/sp.pi/trap.waist_tweezer2**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{2 P_{t1} a}{\\pi W_{t1}^{2}}$"
      ],
      "text/plain": [
       "2*P_t1*a/(pi*W_t1**2)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "depth_tweezer1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4196.086090663433"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "float(trap.subs(E_r))/const.h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Match depth to Lauriane's paper:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle P_{t1}\\left(V_1 = 15 E_r \\right) = 17729.10 \\mathrm{\\mu W}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "power = trap.subs(\n",
    "    sp.solve(\n",
    "        depth_tweezer1/E_r - 15,\n",
    "        trap.power_tweezer1,\n",
    "    )[0]\n",
    ").evalf()\n",
    "\n",
    "display(\n",
    "    Math(\n",
    "        f\"{sp.latex(trap.power_tweezer1)}\\\\left(V_1 = 15 E_r \\\\right)\"\n",
    "        f\" = {power/si.uW:.2f} \\\\mathrm{{\\\\mu W}}\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "15.365855450757438"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "float(trap.subs(depth_tweezer1/E_r))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate trapping frequencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "omega_z1 = sp.sqrt(4*trap.power_tweezer1*trap.a/sp.pi**3/trap.m) * trap.wvl/trap.waist_tweezer1**3\n",
    "omega_z2 = sp.sqrt(4*trap.power_tweezer2*trap.a/sp.pi**3/trap.m) * trap.wvl/trap.waist_tweezer2**3\n",
    "\n",
    "omega_r1 = sp.sqrt(2*trap.power_tweezer1*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer1**2\n",
    "omega_r2 = sp.sqrt(2*trap.power_tweezer2*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer2**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{2 \\sqrt{2} \\sqrt{P_{t1}} \\sqrt{a}}{\\sqrt{\\pi} W_{t1}^{2} \\sqrt{m}}$"
      ],
      "text/plain": [
       "2*sqrt(2)*sqrt(P_t1)*sqrt(a)/(sqrt(pi)*W_t1**2*sqrt(m))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "omega_r1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Match trapping frequency to Lauriane's paper:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle W_{t1}\\left(\\omega_x = 2\\pi \\cdot 29.0 \\right) = 40.87 \\mathrm{\\mu m}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "waist = trap.subs(\n",
    "    sp.solve(\n",
    "        omega_r1 - 2*np.pi *29.0,\n",
    "        trap.waist_tweezer1,\n",
    "    )[0]\n",
    ").evalf()\n",
    "\n",
    "display(\n",
    "    Math(\n",
    "        f\"{sp.latex(trap.waist_tweezer1)}\\\\left(\\\\omega_x = 2\\\\pi \\\\cdot 29.0 \\\\right)\"\n",
    "        f\" = {waist/si.um:.2f} \\\\mathrm{{\\\\mu m}}\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "320.25388559018154"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "float(trap.subs(omega_r1))/(2*np.pi)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analytically:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "w_0 = 134.21um\n",
      "P = 2110639.42uW\n"
     ]
    }
   ],
   "source": [
    "depth_in_Er = 15\n",
    "trap_freq = 29.0 #Hz\n",
    "\n",
    "waist_analytic = sp.sqrt(depth_in_Er /2)* const.h/trap.m/trap.distance_tweezers/(2*sp.pi*trap_freq)\n",
    "power_analytic = (2*sp.pi*trap_freq)**2 *sp.pi*trap.m*waist_analytic**4/8/trap.a \n",
    "\n",
    "print(f\"w_0 = {float(trap.subs(waist_analytic))/si.um:.2f}um\")\n",
    "print(f\"P = {float(trap.subs(power_analytic))/si.uW:.2f}uW\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Erbium ###\n",
    "\n",
    "initial_power = float(trap.subs(power_analytic))\n",
    "initial_waist = float(trap.subs(waist_analytic))\n",
    "initial_distance = 0.266*si.um\n",
    "\n",
    "trap: DoubleTweezer = DoubleTweezer(\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_tweezer1 = initial_power,     #stationary\n",
    "    power_tweezer2 = initial_power,     #transfer tweezer\n",
    "    waist_tweezer1 = initial_waist,     #stationary\n",
    "    waist_tweezer2 = initial_waist,     #transfer tweezer\n",
    "    distance_tweezers = initial_distance,\n",
    "\n",
    "    m = 168 * const.value(\"atomic mass constant\"),\n",
    "    mu_b = 6.982806* const.value(\"Bohr magneton\"),\n",
    "    a_s = 85* const.value(\"Bohr radius\"),\n",
    "    a=180*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "    wvl = 532 * si.nm,\n",
    "\n",
    "    g = 0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "15.000000000000009"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "E_r = const.h**2/(8*trap.m*trap.distance_tweezers**2)\n",
    "\n",
    "depth_tweezer1 = trap.a*2*trap.power_tweezer1/sp.pi/trap.waist_tweezer1**2\n",
    "depth_tweezer2 = trap.a*2*trap.power_tweezer2/sp.pi/trap.waist_tweezer2**2\n",
    "\n",
    "float(trap.subs(depth_tweezer1/E_r))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "29.00000000000001"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "omega_z1 = sp.sqrt(4*trap.power_tweezer1*trap.a/sp.pi**3/trap.m) * trap.wvl/trap.waist_tweezer1**3\n",
    "omega_z2 = sp.sqrt(4*trap.power_tweezer2*trap.a/sp.pi**3/trap.m) * trap.wvl/trap.waist_tweezer2**3\n",
    "\n",
    "omega_r1 = sp.sqrt(2*trap.power_tweezer1*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer1**2\n",
    "omega_r2 = sp.sqrt(2*trap.power_tweezer2*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer2**2\n",
    "\n",
    "float(trap.subs(omega_r1))/2/np.pi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## See how that looks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 45844 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_cutoff = -0.5*float(trap.subs(trap.distance_tweezers))-3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "right_cutoff = 0.5*float(trap.subs(trap.distance_tweezers))+3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "\n",
    "\n",
    "# Solve the hamiltonian numerically in axial direction\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "    trap.x, (left_cutoff, right_cutoff), 500, k=4\n",
    ")\n",
    "\n",
    "# States that are below the potential barrier\n",
    "bound_states = energies < potential(left_cutoff)\n",
    "\n",
    "\n",
    "z_np = np.linspace(left_cutoff, right_cutoff, num=500)\n",
    "ax: plt.Axes\n",
    "fig, ax = plt.subplots(figsize=(5, 5))\n",
    "ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz,color=\"cornflowerblue\" ,marker=\"None\")\n",
    "ax.set_title(f\"{np.sum(bound_states)} bound solutions, tweezer distance: {trap.subs(trap.distance_tweezers)/si.um}um\")\n",
    "ax.set_xlabel(r\"x ($\\mathrm{\\mu m}$)\")\n",
    "ax.set_ylabel(r\"E / h (kHz)\")\n",
    "abs_min = np.min(potential(z_np))\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",
    "    alpha=0.5,\n",
    "    color=\"cornflowerblue\"\n",
    ")\n",
    "\n",
    "count = 0\n",
    "for i, bound in enumerate(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)),\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 *1e1, marker=\"None\",label=f\"state {count}\")#, c=\"k\")\n",
    "    count += 1\n",
    "\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "185.24233721850484"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "float(trap.subs(E_r))/const.h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## This probably doesn't work, since the depth is not accurate when d is small\n",
    "\n",
    "# Let's properly estimate the depths and trapping frequencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Erbium ###\n",
    "\n",
    "initial_power = 5600* si.uW\n",
    "initial_waist = 22*si.uW\n",
    "initial_distance = 22.01*si.um\n",
    "\n",
    "trap: DoubleTweezer = DoubleTweezer(\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_tweezer1 = initial_power,     #stationary\n",
    "    power_tweezer2 = initial_power,     #transfer tweezer\n",
    "    waist_tweezer1 = initial_waist,     #stationary\n",
    "    waist_tweezer2 = initial_waist,     #transfer tweezer\n",
    "    distance_tweezers = initial_distance,\n",
    "\n",
    "    m = 168 * const.value(\"atomic mass constant\"),\n",
    "    mu_b = 6.982806* const.value(\"Bohr magneton\"),\n",
    "    a_s = 85* const.value(\"Bohr radius\"),\n",
    "    a=180*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "    wvl = 532 * si.nm,\n",
    "\n",
    "    g = 0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 41148 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_cutoff = -0.5*float(trap.subs(trap.distance_tweezers))-3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "right_cutoff = 0.5*float(trap.subs(trap.distance_tweezers))+3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "\n",
    "\n",
    "# Solve the hamiltonian numerically in axial direction\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "    trap.x, (left_cutoff, right_cutoff), 500, k=4\n",
    ")\n",
    "\n",
    "# States that are below the potential barrier\n",
    "bound_states = energies < potential(left_cutoff)\n",
    "\n",
    "\n",
    "z_np = np.linspace(left_cutoff, right_cutoff, num=500)\n",
    "ax: plt.Axes\n",
    "fig, ax = plt.subplots(figsize=(5, 5))\n",
    "ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz,color=\"cornflowerblue\" ,marker=\"None\")\n",
    "ax.set_title(f\"{np.sum(bound_states)} bound solutions, tweezer distance: {trap.subs(trap.distance_tweezers)/si.um}um\")\n",
    "ax.set_xlabel(r\"x ($\\mathrm{\\mu m}$)\")\n",
    "ax.set_ylabel(r\"E / h (kHz)\")\n",
    "abs_min = np.min(potential(z_np))\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",
    "    alpha=0.5,\n",
    "    color=\"cornflowerblue\"\n",
    ")\n",
    "\n",
    "count = 0\n",
    "for i, bound in enumerate(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)),\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 *1e-3, marker=\"None\",label=f\"state {count}\")#, c=\"k\")\n",
    "    count += 1\n",
    "\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'x' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[4], line 20\u001b[0m\n\u001b[0;32m     17\u001b[0m V_2 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mabs\u001b[39m(V_func(x_right) \u001b[38;5;241m-\u001b[39m V_func(\u001b[38;5;241m0\u001b[39m))\n\u001b[0;32m     19\u001b[0m \u001b[38;5;66;03m#trapping frequencies through second derivative\u001b[39;00m\n\u001b[1;32m---> 20\u001b[0m V_prime \u001b[38;5;241m=\u001b[39m sp\u001b[38;5;241m.\u001b[39mdiff(V, x)\n\u001b[0;32m     21\u001b[0m V_double_prime \u001b[38;5;241m=\u001b[39m sp\u001b[38;5;241m.\u001b[39mdiff(V_prime, x)\n\u001b[0;32m     23\u001b[0m omega_x1 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mfloat\u001b[39m(trap\u001b[38;5;241m.\u001b[39msubs(sp\u001b[38;5;241m.\u001b[39msqrt(V_double_prime\u001b[38;5;241m.\u001b[39msubs(trap\u001b[38;5;241m.\u001b[39mx,x_left)\u001b[38;5;241m.\u001b[39msubs(trap\u001b[38;5;241m.\u001b[39my,\u001b[38;5;241m0\u001b[39m)\u001b[38;5;241m.\u001b[39msubs(trap\u001b[38;5;241m.\u001b[39mz,\u001b[38;5;241m0\u001b[39m)\u001b[38;5;241m/\u001b[39mtrap\u001b[38;5;241m.\u001b[39mm)))\n",
      "\u001b[1;31mNameError\u001b[0m: name 'x' is not defined"
     ]
    }
   ],
   "source": [
    "V = trap.subs(trap.get_potential(apply_zero_offset=False))\n",
    "a = float(trap.subs(trap.distance_tweezers))\n",
    "\n",
    "#find minima of potential\n",
    "def V_func(x):\n",
    "    return float(V.subs(trap.x,x).subs(trap.y,0).subs(trap.z,0))\n",
    "x_right = minimize_scalar(V_func,bracket=[0,a/2]).x\n",
    "x_left = minimize_scalar(V_func,bracket=[-a/2,0]).x\n",
    "\n",
    "#catch case where both potentials have already merged\n",
    "tunneling_dist = abs(x_right-x_left)\n",
    "if tunneling_dist < 1e-15:\n",
    "    raise Exception(\"potential has only one minmum\")\n",
    "\n",
    "#depts of both tweezers\n",
    "V_1 = abs(V_func(x_left) - V_func(0))\n",
    "V_2 = abs(V_func(x_right) - V_func(0))\n",
    "\n",
    "#trapping frequencies through second derivative\n",
    "V_prime = sp.diff(V, x)\n",
    "V_double_prime = sp.diff(V_prime, x)\n",
    "\n",
    "omega_x1 = float(trap.subs(sp.sqrt(V_double_prime.subs(trap.x,x_left).subs(trap.y,0).subs(trap.z,0)/trap.m)))\n",
    "omega_x2 = float(trap.subs(sp.sqrt(V_double_prime.subs(trap.x,x_right).subs(trap.y,0).subs(trap.z,0)/trap.m)))\n",
    "\n",
    "#recoil energy\n",
    "E_r = const.h**2/(8*trap.m*trap.distance_tweezers**2)\n",
    "\n",
    "#print results\n",
    "print(f\"depth: {V_2/float(trap.subs(E_r)):.2f} E_r\")\n",
    "print(f\"trapping frequency: {omega_x1/2/np.pi:.2f} Hz\")\n",
    "print(f\"tunneling distance: {tunneling_dist/si.nm:.2f} um\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### It seems that the desired parameters are not realisable with the tweezer potential\n",
    "\n",
    "## Let's try to implement a lattice potential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Erbium ###\n",
    "\n",
    "initial_spacing = 0\n",
    "initial_depth = 0\n",
    "\n",
    "\n",
    "trap: TwoSiteLattice = TwoSiteLattice(\n",
    "\n",
    "    lattice_depth = initial_depth,\n",
    "    lattice_spacing = initial_spacing,\n",
    "\n",
    "\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 = 0,\n",
    "    waist_tweezer = 0,\n",
    "\n",
    "    m = 168 * const.value(\"atomic mass constant\"),\n",
    "    mu_b = 6.982806* const.value(\"Bohr magneton\"),\n",
    "    a_s = 137* const.value(\"Bohr radius\"),\n",
    "    a = 180*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "    wvl = 532 * si.nm,\n",
    "\n",
    "    g = 0,\n",
    ")\n",
    "\n",
    "E_r = const.h**2/(2*trap.m*trap.lattice_spacing**2)\n",
    "\n",
    "trap[trap.lattice_depth] = 15*E_r\n",
    "#lattice spacing is double the actual distance between sites\n",
    "trap[trap.lattice_spacing] = trap.wvl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 49280 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_cutoff = -1.3*0.5*float(trap.subs(trap.lattice_spacing))\n",
    "right_cutoff = 1.3*0.5*float(trap.subs(trap.lattice_spacing))\n",
    "\n",
    "\n",
    "# Solve the hamiltonian numerically in axial direction\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "    trap.x, (left_cutoff, right_cutoff), 500, k=4\n",
    ")\n",
    "\n",
    "# States that are below the potential barrier\n",
    "#bound_states = energies < potential(left_cutoff)\n",
    "bound_states = energies < np.inf\n",
    "\n",
    "z_np = np.linspace(left_cutoff, right_cutoff, num=500)\n",
    "ax: plt.Axes\n",
    "fig, ax = plt.subplots(figsize=(5, 5))\n",
    "ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz,color=\"cornflowerblue\" ,marker=\"None\")\n",
    "ax.set_title(f\"{np.sum(bound_states)} bound solutions, tweezer distance: {trap.subs(trap.lattice_spacing)/si.um}um\")\n",
    "ax.set_xlabel(r\"x ($\\mathrm{\\mu m}$)\")\n",
    "ax.set_ylabel(r\"E / h (kHz)\")\n",
    "abs_min = np.min(potential(z_np))\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",
    "    alpha=0.5,\n",
    "    color=\"cornflowerblue\"\n",
    ")\n",
    "\n",
    "count = 0\n",
    "for i, bound in enumerate(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)),\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 *1e2, marker=\"None\",label=f\"state {count}\")#, c=\"k\")\n",
    "    count += 1\n",
    "\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "depth: 15.00 E_r\n",
      "trapping frequency x: 32502.74 Hz\n",
      "trapping frequency y: 32502.74 Hz\n",
      "trapping frequency z: 32502.74 Hz\n",
      "aspect ratio: 1.00\n",
      "tunneling distance: 266.00 um\n"
     ]
    }
   ],
   "source": [
    "V = trap.subs(trap.get_potential(apply_zero_offset=False))\n",
    "a = float(trap.subs(trap.lattice_spacing))\n",
    "\n",
    "#find minima of potential\n",
    "def V_func(x):\n",
    "    return float(V.subs({trap.x:x, trap.y:0, trap.z:0}))\n",
    "#x_right = minimize_scalar(V_func,bracket=[1e-3*a,a/2],bounds=[0,a/2]).x\n",
    "#x_left = minimize_scalar(V_func,bracket=[-a/2,-1e-3*a],bounds=[-a/2,0]).x\n",
    "x_right = a/4\n",
    "x_left = -a/4\n",
    "\n",
    "#catch case where both potentials have already merged\n",
    "tunneling_dist = abs(x_right-x_left)\n",
    "if tunneling_dist < 1e-15:\n",
    "    raise Exception(\"potential has only one minmum\")\n",
    "\n",
    "#depts of both tweezers\n",
    "V_1 = abs(V_func(x_left) - V_func(0))\n",
    "V_2 = abs(V_func(x_right) - V_func(0)) \n",
    "\n",
    "#trapping frequencies through second derivative\n",
    "#trapping frequencies through second derivative\n",
    "V_double_prime_x = sp.diff(V, trap.x, 2)\n",
    "omega_x1 = float(trap.subs(sp.sqrt(V_double_prime_x.subs({trap.x:x_left, trap.y:0, trap.z:0})/trap.m)))\n",
    "omega_x2 = float(trap.subs(sp.sqrt(V_double_prime_x.subs({trap.x:x_right, trap.y:0, trap.z:0})/trap.m)))\n",
    "\n",
    "V_double_prime_y = sp.diff(V, trap.y, 2)\n",
    "omega_y1 = float(trap.subs(sp.sqrt(V_double_prime_y.subs({trap.x:x_left, trap.y:0, trap.z:0})/trap.m)))\n",
    "omega_y2 = float(trap.subs(sp.sqrt(V_double_prime_y.subs({trap.x:x_right, trap.y:0, trap.z:0})/trap.m)))\n",
    "\n",
    "V_double_prime_z = sp.diff(V, trap.z, 2)\n",
    "omega_z1 = float(trap.subs(sp.sqrt(V_double_prime_z.subs({trap.x:x_left, trap.y:0, trap.z:0})/trap.m)))\n",
    "omega_z2 = float(trap.subs(sp.sqrt(V_double_prime_z.subs({trap.x:x_right, trap.y:0, trap.z:0})/trap.m)))\n",
    "\n",
    "#recoil energy\n",
    "E_r = const.h**2/(2*trap.m*trap.lattice_spacing**2)\n",
    "\n",
    "#print results\n",
    "print(f\"depth: {V_2/float(trap.subs(E_r)):.2f} E_r\")\n",
    "\n",
    "print(f\"trapping frequency x: {omega_x1/2/np.pi:.2f} Hz\")\n",
    "print(f\"trapping frequency y: {omega_y1/2/np.pi:.2f} Hz\")\n",
    "print(f\"trapping frequency z: {omega_z1/2/np.pi:.2f} Hz\")\n",
    "\n",
    "print(f\"aspect ratio: {omega_x1/omega_z1:.2f}\")\n",
    "\n",
    "print(f\"tunneling distance: {tunneling_dist/si.nm:.2f} um\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now do diagonalisation for this potential:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "files saved with ...._2025-03-26_15-58-59\n"
     ]
    }
   ],
   "source": [
    "n_grid_1D = 50\n",
    "\n",
    "n_pot_steps = [n_grid_1D,n_grid_1D,n_grid_1D]\n",
    "n_levels = 8\n",
    "\n",
    "left_cutoff = -0.5*float(trap.subs(trap.lattice_spacing))*1.3\n",
    "right_cutoff = 0.5*float(trap.subs(trap.lattice_spacing))*1.3\n",
    "back_cutoff = -0.25*float(trap.subs(trap.lattice_spacing))*1.3\n",
    "front_cutoff = 0.25*float(trap.subs(trap.lattice_spacing))*1.3\n",
    "bottom_cutoff = -0.5*float(trap.subs(trap.lattice_spacing))*1.7\n",
    "top_cutoff = 0.5*float(trap.subs(trap.lattice_spacing))*1.7\n",
    "\n",
    "extend = [(left_cutoff,right_cutoff),\n",
    "          (back_cutoff,front_cutoff),\n",
    "          (bottom_cutoff,top_cutoff)]\n",
    "\n",
    "\n",
    "# Solve the hamiltonian numerically\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=True)\n",
    "\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from boson_helpers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "________________________________________________________________\n",
      "d = 266.00nm\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 36524 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = 27.481287 Hz\n",
      "U_s = 2607.617 Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24756 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i in range(25,26):\n",
    "\n",
    "    trap, func_ham, res =  import_results(i)\n",
    "\n",
    "    print(\"________________________________________________________________\")\n",
    "    print(f\"d = {float(trap.subs(trap.lattice_spacing)/2) /si.nm :.2f}nm\")\n",
    "\n",
    "    GS_left, GS_right = get_localised_GS(res[\"states\"][0], res[\"states\"][1])\n",
    "    \n",
    "    plt.plot(res[\"x\"],np.abs(GS_left[:,int(res[\"size\"][1]/2),int(res[\"size\"][2]/2)])**2, label=\"left state\")\n",
    "    plt.plot(res[\"x\"],np.abs(GS_right[:,int(res[\"size\"][1]/2),int(res[\"size\"][2]/2)])**2, label=\"right state\")\n",
    "\n",
    "    plt.xlabel(\"x\")\n",
    "    plt.ylabel(r\"$|\\Psi|^2$\")\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "    J, U_s, U_dds, angles, V_lrs = analyse_diagonalisation(i)\n",
    "\n",
    "    print(f\"J = {J/const.h :3f} Hz\")\n",
    "    print(f\"U_s = {U_s/const.h :.3f} Hz\")\n",
    "\n",
    "    plt.plot(np.rad2deg(angles), U_dds/const.h)\n",
    "\n",
    "    plt.axhline(0,color=\"black\",ls=\"--\")\n",
    "\n",
    "    plt.xlabel(\"theta [deg]\")\n",
    "    plt.ylabel(r\"$U_{dd} / h$ [Hz]\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 117244 characters]",
      "text/plain": [
       "<Figure size 2666.67x500 with 16 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 8\n",
    "\n",
    "states_plot = res[\"states\"].real\n",
    "#states_plot = np.angle(states)\n",
    "#states_plot = states.imag\n",
    "#states_plot = np.abs(states)**2\n",
    "\n",
    "# Create a 1xn_states grid of subplots (1 row, n_states columns)\n",
    "fig, axes = plt.subplots(1, n_states, figsize=(20/6*n_states, 5))  # Adjust the size as needed\n",
    "\n",
    "# Loop over the state numbers from 0 to n_states-1\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(states_plot[state_number, :, int(res[\"size\"][1]/2), :].T,\n",
    "                   extent=[*res[\"extend\"][0], *res[\"extend\"][2]], origin=\"lower\",\n",
    "                   vmin=np.min(states_plot[state_number]), vmax=np.max(states_plot[state_number]))\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()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25160 characters]",
      "text/plain": [
       "<Figure size 666.667x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 2\n",
    "\n",
    "states_plot = [GS_left.real, GS_right.real]\n",
    "#states_plot = np.angle(states)\n",
    "#states_plot = states.imag\n",
    "#states_plot = np.abs(states)**2\n",
    "\n",
    "# Create a 1xn_states grid of subplots (1 row, n_states columns)\n",
    "fig, axes = plt.subplots(1, n_states, figsize=(20/6*n_states, 5))  # Adjust the size as needed\n",
    "\n",
    "# Loop over the state numbers from 0 to n_states-1\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(states_plot[state_number][:, int(res[\"size\"][1]/2), :].T,\n",
    "                   extent=[*res[\"extend\"][0], *res[\"extend\"][2]], origin=\"lower\",\n",
    "                   vmin=np.min(states_plot[state_number]), vmax=np.max(states_plot[state_number]))\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": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38456 characters]",
      "text/plain": [
       "<Figure size 666.667x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 2\n",
    "\n",
    "states_plot = [GS_left.real, GS_right.real]\n",
    "#states_plot = np.angle(states)\n",
    "#states_plot = states.imag\n",
    "#states_plot = np.abs(states)**2\n",
    "\n",
    "# Create a 1xn_states grid of subplots (1 row, n_states columns)\n",
    "fig, axes = plt.subplots(1, n_states, figsize=(20/6*n_states, 5))  # Adjust the size as needed\n",
    "\n",
    "# Loop over the state numbers from 0 to n_states-1\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(states_plot[state_number][:, :, int(res[\"size\"][2]/2)].T,\n",
    "                   extent=[*res[\"extend\"][0], *res[\"extend\"][1]], origin=\"lower\",\n",
    "                   vmin=np.min(states_plot[state_number]), vmax=np.max(states_plot[state_number]))\n",
    "\n",
    "    # Set labels for each subplot\n",
    "    ax.set_xlabel(\"x\")\n",
    "    ax.set_ylabel(\"y\")\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": 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
}
