{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use this notebook to diagonalise 3D Hamiltonian"
   ]
  },
  {
   "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 import sparse\n",
    "from scipy.optimize import minimize_scalar\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trap parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "omega_r1: 60954.0Hz, omega_r2: 60954.0Hz\n",
      "omega_z1: 10279.9Hz, omega_z2: 10279.9Hz\n"
     ]
    }
   ],
   "source": [
    "initial_power = 7000* si.uW\n",
    "initial_waist = 0.71*si.uW\n",
    "initial_distance = 0.808*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= 161 * const.value(\"atomic mass constant\"),\n",
    "    a_s= 85* const.value(\"Bohr radius\"),\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",
    "    wvl = 532 * si.nm,\n",
    "\n",
    "    g = 0,\n",
    ")\n",
    "\n",
    "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",
    "print(f\"omega_r1: {float(trap.subs(omega_r1)/2/np.pi):.1f}Hz, omega_r2: {float(trap.subs(omega_r2)/2/np.pi):.1f}Hz\")\n",
    "print(f\"omega_z1: {float(trap.subs(omega_z1)/2/np.pi):.1f}Hz, omega_z2: {float(trap.subs(omega_z2)/2/np.pi):.1f}Hz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 46988 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))-1*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))+1*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=6\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",
    "\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": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 52580 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_cutoff = -1*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "right_cutoff = 1*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.y, (left_cutoff, right_cutoff), 500, k=6\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",
    "\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\"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 *1e1, marker=\"None\",label=f\"state {count}\")#, c=\"k\")\n",
    "    count += 1\n",
    "\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\naeve\\AppData\\Local\\Temp\\ipykernel_13056\\2474033355.py:51: UserWarning: No artists with labels found to put in legend.  Note that artists whose label start with an underscore are ignored when legend() is called with no argument.\n",
      "  plt.legend()\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 42660 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_cutoff = -0.5*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "right_cutoff = 0.5*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\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=6\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",
    "\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\"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 *1e1, marker=\"None\",label=f\"state {count}\")#, c=\"k\")\n",
    "    count += 1\n",
    "\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "depth 1: 18.3716 uK\n",
      "depth 2: 18.3716 uK\n",
      "estimated tilt: 0.00 Hz\n",
      "trapping frequency x 1: 41891.68 Hz\n",
      "trapping frequency x 2: 41891.68 Hz\n",
      "trapping frequency y 1: 63870.24 Hz\n",
      "trapping frequency y 2: 63870.24 Hz\n",
      "trapping frequency z 1: 9108.96 Hz\n",
      "trapping frequency z 2: 9108.96 Hz\n",
      "tunneling distance: 604.67 um\n"
     ]
    }
   ],
   "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, trap.y:0, trap.z:0}))\n",
    "\n",
    "x_right = minimize_scalar(V_func,bracket=[0,a]).x\n",
    "x_left = minimize_scalar(V_func,bracket=[-a,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\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.distance_tweezers**2)\n",
    "\n",
    "#print results\n",
    "print(f\"depth 1: {V_1/const.k /si.uK:.4f} uK\")\n",
    "print(f\"depth 2: {V_2/const.k /si.uK:.4f} uK\")\n",
    "print(f\"estimated tilt: {abs(V_1 - V_2)/const.h:.2f} Hz\")\n",
    "\n",
    "print(f\"trapping frequency x 1: {trap.subs(omega_x1)/2/np.pi:.2f} Hz\")\n",
    "print(f\"trapping frequency x 2: {trap.subs(omega_x2)/2/np.pi:.2f} Hz\")\n",
    "print(f\"trapping frequency y 1: {trap.subs(omega_y1)/2/np.pi:.2f} Hz\")\n",
    "print(f\"trapping frequency y 2: {trap.subs(omega_y2)/2/np.pi:.2f} Hz\")\n",
    "print(f\"trapping frequency z 1: {trap.subs(omega_z1)/2/np.pi:.2f} Hz\")\n",
    "print(f\"trapping frequency z 2: {trap.subs(omega_z2)/2/np.pi:.2f} Hz\")\n",
    "\n",
    "print(f\"tunneling distance: {tunneling_dist/si.nm:.2f} um\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Diagonalisation parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "distances = [808*si.nm ]\n",
    "\n",
    "for i, dist in enumerate(distances):\n",
    "\n",
    "        trap[trap.distance_tweezers] = dist\n",
    "\n",
    "        n_grid_1D = 100  \n",
    "\n",
    "        n_pot_steps = [n_grid_1D,n_grid_1D,n_grid_1D]\n",
    "        n_levels = 6\n",
    "\n",
    "        left_cutoff = -0.5*initial_distance-1*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "        right_cutoff = 0.5*initial_distance+1*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "        back_cutoff = -1*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "        front_cutoff = 1*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "        bottom_cutoff = -0.5*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "        top_cutoff = 0.5*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "\n",
    "        extend = [(left_cutoff,right_cutoff),\n",
    "                (back_cutoff,front_cutoff),\n",
    "                (bottom_cutoff,top_cutoff)]\n",
    "\n",
    "        try:\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",
    "        except:\n",
    "                continue\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "files saved with ...._2025-03-31_16-07-07\n",
      "files saved with ...._2025-03-31_17-02-40\n",
      "files saved with ...._2025-03-31_17-56-25\n",
      "files saved with ...._2025-03-31_18-46-11\n",
      "files saved with ...._2025-03-31_19-34-16\n",
      "files saved with ...._2025-03-31_20-21-21\n",
      "files saved with ...._2025-03-31_21-07-40\n",
      "files saved with ...._2025-03-31_21-53-18\n",
      "files saved with ...._2025-03-31_22-38-09\n",
      "files saved with ...._2025-03-31_23-21-53\n"
     ]
    }
   ],
   "source": [
    "distances = np.linspace(730,850,10)*si.nm\n",
    "\n",
    "for i, dist in enumerate(distances):\n",
    "\n",
    "        trap[trap.distance_tweezers] = dist\n",
    "\n",
    "        n_grid_1D = 100  \n",
    "\n",
    "        n_pot_steps = [n_grid_1D,n_grid_1D,n_grid_1D]\n",
    "        n_levels = 6\n",
    "\n",
    "        left_cutoff = -0.5*initial_distance-1*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "        right_cutoff = 0.5*initial_distance+1*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "        back_cutoff = -1*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "        front_cutoff = 1*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "        bottom_cutoff = -0.5*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "        top_cutoff = 0.5*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "\n",
    "        extend = [(left_cutoff,right_cutoff),\n",
    "                (back_cutoff,front_cutoff),\n",
    "                (bottom_cutoff,top_cutoff)]\n",
    "\n",
    "        try:\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",
    "        except:\n",
    "                continue\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualise results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 74584 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 6\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()\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
}
