{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1fbc20b4",
   "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,
   "id": "f411f1d9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 400780 characters]",
      "text/plain": [
       "<Figure size 4500x500 with 30 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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= -2.89 * 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 = 100\n",
    "\n",
    "n_pot_steps = [n_grid_1D,n_grid_1D,n_grid_1D]\n",
    "n_levels = 15\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.5*np.max([float(trap.subs(trap.get_tweezer_rayleigh())),float(trap.subs(trap.get_tweezer_rayleigh()))])\n",
    "top_cutoff = 0.5*np.max([float(trap.subs(trap.get_tweezer_rayleigh())),float(trap.subs(trap.get_tweezer_rayleigh()))])\n",
    "# left_cutoff = 0.3*np.max([float(trap.subs(trap.waist_tweezer)),float(trap.subs(trap.waist_tweezer))])\n",
    "# right_cutoff = -0.3*np.max([float(trap.subs(trap.waist_tweezer)),float(trap.subs(trap.waist_tweezer))])\n",
    "# back_cutoff = -0.3*np.max([float(trap.subs(trap.waist_tweezer)),float(trap.subs(trap.waist_tweezer))])\n",
    "# front_cutoff = 0.3*np.max([float(trap.subs(trap.waist_tweezer)),float(trap.subs(trap.waist_tweezer))])\n",
    "# bottom_cutoff = -0.2*np.max([float(trap.subs(trap.get_tweezer_rayleigh())),float(trap.subs(trap.get_tweezer_rayleigh()))])\n",
    "# top_cutoff = 0.2*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 = 15\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": 36,
   "id": "953ea738",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 320628 characters]",
      "text/plain": [
       "<Figure size 3300x500 with 22 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#### plot\n",
    "\n",
    "n_states = 11\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": 4,
   "id": "b121d7ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "244.368 Hz\n",
      "1.08288036422384 µK\n"
     ]
    }
   ],
   "source": [
    "print(f\"{float(trap.subs(trap.get_omega_ax_tweezer()))/2/np.pi:.3f} Hz\")\n",
    "print(f\"{trap.subs(trap.a *2*trap.power_tweezer/(np.pi*trap.waist_tweezer**2)/ const.k) /si.uK} µK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2c381835",
   "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, curve_fit\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\n",
    "from boson_helpers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "29ad7f21",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 84788 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ind_1 = 0\n",
    "\n",
    "V_lrs_max = 0\n",
    "\n",
    "for j in range(1,10):    \n",
    "    ind_2 = j\n",
    "\n",
    "    #compute volume element\n",
    "    dvol = np.array([])\n",
    "    for coord in coords.values():\n",
    "        dvol = np.append(dvol,np.abs(np.mean(np.diff(coord))))\n",
    "\n",
    "    #compute interaction energies\n",
    "    U_s = get_U_s(trap, states[ind_1], dvol)\n",
    "\n",
    "    angles= np.deg2rad(np.linspace(0,90))\n",
    "    U_dds = np.zeros_like(angles)\n",
    "    V_lrs = np.zeros_like(angles)\n",
    "    polarisations = np.zeros((len(angles),3))\n",
    "    polarisations[:,0] = np.sin(angles)\n",
    "    polarisations[:,2] = np.cos(angles)\n",
    "\n",
    "    for i, pol in enumerate(polarisations):\n",
    "        trap[trap.B_x] = pol[0]\n",
    "        trap[trap.B_y] = pol[1]\n",
    "        trap[trap.B_z] = pol[2]\n",
    "\n",
    "        U_dds[i] = get_U_dd(trap, states[ind_1], dvol)\n",
    "        V_lrs[i] = get_NNI(trap, states[ind_1], states[ind_2], dvol)\n",
    "\n",
    "    V_lrs_max += V_lrs[0]\n",
    "    plt.plot(np.rad2deg(angles), V_lrs/const.h, label=f\"j = {j}\")\n",
    "plt.axhline(0, linestyle=\"dashed\", color=\"grey\")\n",
    "\n",
    "plt.xlim(0,90)\n",
    "\n",
    "plt.xlabel(\"angle between dipoles and z-axis [°]\")\n",
    "plt.ylabel(r\"$V_{0j}/h$ [Hz]\")\n",
    "\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "\n",
    "plt.savefig(\"NNI.png\", dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "a2926b78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-20300.06710322, -20069.67924244, -19841.26638092, -19614.82002624,\n",
       "       -19390.33165075, -19167.79269105, -18947.19454576, -18728.52857117,\n",
       "       -18511.78616692, -18296.95801305, -18212.12095438, -18084.03959135,\n",
       "       -18212.12095404, -17993.98793879, -17993.9879383 ])"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "energies/const.h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "5a2e3bde",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 2.30387861e+02,  2.28412862e+02,  2.26446355e+02,  2.24488375e+02,\n",
       "        2.22538960e+02,  2.20598145e+02,  2.18665975e+02,  2.16742404e+02,\n",
       "        2.14828154e+02,  8.48370587e+01,  1.28081363e+02, -1.28081363e+02,\n",
       "        2.18133015e+02,  4.89512402e-07])"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.diff(energies)/const.h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e51617f7",
   "metadata": {},
   "source": [
    "The maximum ratio of spacing and interaction between 0th and 1st excited state is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "d47f267d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "230.3878607786473"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.diff(energies)[0]/const.h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "a378e0f4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-55.4863825353665"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "V_lrs_max/const.h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "9e256f3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.24083900231478284"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "V_lrs_max/np.diff(energies)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b281209",
   "metadata": {},
   "source": [
    "The dipolar interactions can be turned off entirely by choosing the following angle between magnetic field and z-axis:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "fdc4c0a7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "56.93877551020409"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.rad2deg(angles[np.argmin(np.abs(np.abs(V_lrs)))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "913212bb",
   "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": 5
}
