{
 "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": [
    "# Set up trap parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Erbium ###\n",
    "\n",
    "spacing_xy = 266*si.nm\n",
    "spacing_z = 532*si.nm\n",
    "\n",
    "trap: TwoSiteLattice = TwoSiteLattice(\n",
    "\n",
    "    lattice_spacing_x = spacing_xy,\n",
    "    lattice_spacing_y = spacing_xy,\n",
    "    lattice_spacing_z = spacing_z,\n",
    "    lattice_depth_x = 0,\n",
    "    lattice_depth_y = 0,\n",
    "    lattice_depth_z = 0,\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_xy = const.h**2/(8*trap.m*spacing_xy**2)\n",
    "\n",
    "trap[trap.lattice_depth_x] = 15*E_r_xy\n",
    "trap[trap.lattice_depth_y] = 15*E_r_xy\n",
    "\n",
    "#mult=4 for AR=1 / mult=0.25 for AR=2\n",
    "mult=4\n",
    "trap[trap.lattice_depth_z] = mult* 15*E_r_xy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 240.0$"
      ],
      "text/plain": [
       "240.000000000000"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "E_r_z = const.h**2/(8*trap.m*spacing_z**2)\n",
    "\n",
    "trap.subs(trap.lattice_depth_z/E_r_z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 49380 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_cutoff = -1.3*float(trap.subs(trap.lattice_spacing_x))\n",
    "right_cutoff = 1.3*float(trap.subs(trap.lattice_spacing_x))\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_x)/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": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48588 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_cutoff = -0.9*float(trap.subs(trap.lattice_spacing_y))\n",
    "right_cutoff = 0.9*float(trap.subs(trap.lattice_spacing_y))\n",
    "\n",
    "\n",
    "# Solve the hamiltonian numerically in axial direction\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "    trap.y, (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_x)/si.um}um\")\n",
    "ax.set_xlabel(r\"y ($\\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": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 56100 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_cutoff = -0.9*float(trap.subs(trap.lattice_spacing_z))\n",
    "right_cutoff = 0.9*float(trap.subs(trap.lattice_spacing_z))\n",
    "\n",
    "\n",
    "# Solve the hamiltonian numerically in axial direction\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "    trap.z, (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_x)/si.um}um\")\n",
    "ax.set_xlabel(r\"z ($\\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": "markdown",
   "metadata": {},
   "source": [
    "# Check parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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.0000\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_x))\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/2\n",
    "x_left = -a/2\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\n",
    "omega_x1, omega_x2 = trap.get_both_omega(trap.x)\n",
    "omega_y1, omega_y2 = trap.get_both_omega(trap.y)\n",
    "omega_z1, omega_z2 = trap.get_both_omega(trap.z)\n",
    "\n",
    "#recoil energy\n",
    "E_r = const.h**2/(8*trap.m*trap.lattice_spacing_x**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: {trap.subs(omega_x1)/2/np.pi:.2f} Hz\")\n",
    "print(f\"trapping frequency y: {trap.subs(omega_y1)/2/np.pi:.2f} Hz\")\n",
    "print(f\"trapping frequency z: {trap.subs(omega_z1)/2/np.pi:.2f} Hz\")\n",
    "\n",
    "print(f\"aspect ratio: {sp.sqrt(omega_x1/omega_z1):.4f}\")\n",
    "\n",
    "print(f\"tunneling distance: {tunneling_dist/si.nm:.2f} um\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Diagonalise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "files saved with ...._2025-03-27_13-50-44\n"
     ]
    }
   ],
   "source": [
    "n_grid_1D = 50\n",
    "\n",
    "n_pot_steps = [n_grid_1D,n_grid_1D,n_grid_1D]\n",
    "n_levels = 15\n",
    "\n",
    "\n",
    "left_cutoff = -1.3*float(trap.subs(trap.lattice_spacing_x))\n",
    "right_cutoff = 1.3*float(trap.subs(trap.lattice_spacing_x))\n",
    "back_cutoff = -0.9*float(trap.subs(trap.lattice_spacing_y))\n",
    "front_cutoff = 0.9*float(trap.subs(trap.lattice_spacing_y))\n",
    "bottom_cutoff = -0.9*float(trap.subs(trap.lattice_spacing_z))\n",
    "top_cutoff = 0.9*float(trap.subs(trap.lattice_spacing_z))\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": "markdown",
   "metadata": {},
   "source": [
    "# Plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 100660 characters]",
      "text/plain": [
       "<Figure size 2666.67x500 with 16 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 8\n",
    "\n",
    "states_plot = 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(n_pot_steps[0]/2), :].T,\n",
    "                   extent=[*extend[0], *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": "markdown",
   "metadata": {},
   "source": [
    "# Let's also make one for AR=2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#mult=4 for AR=1 / mult=0.25 for AR=2\n",
    "mult=0.25\n",
    "trap[trap.lattice_depth_z] = mult* 15*E_r_xy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "files saved with ...._2025-03-27_13-59-14\n"
     ]
    }
   ],
   "source": [
    "n_grid_1D = 50\n",
    "\n",
    "n_pot_steps = [n_grid_1D,n_grid_1D,n_grid_1D]\n",
    "n_levels = 15\n",
    "\n",
    "\n",
    "left_cutoff = -1.3*float(trap.subs(trap.lattice_spacing_x))\n",
    "right_cutoff = 1.3*float(trap.subs(trap.lattice_spacing_x))\n",
    "back_cutoff = -0.9*float(trap.subs(trap.lattice_spacing_y))\n",
    "front_cutoff = 0.9*float(trap.subs(trap.lattice_spacing_y))\n",
    "bottom_cutoff = -0.9*float(trap.subs(trap.lattice_spacing_z))\n",
    "top_cutoff = 0.9*float(trap.subs(trap.lattice_spacing_z))\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": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 142840 characters]",
      "text/plain": [
       "<Figure size 2666.67x500 with 16 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 8\n",
    "\n",
    "states_plot = 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(n_pot_steps[0]/2), :].T,\n",
    "                   extent=[*extend[0], *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": 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
}
