{
 "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, 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": "markdown",
   "metadata": {},
   "source": [
    "# Get the Hubbard parameters from the results of a diagonalisation in one function call:\n",
    "\n",
    "## Show results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 63\n",
    "end_line = 68\n",
    "\n",
    "n_lines = end_line - start_line + 1\n",
    "lines = np.arange(start_line,end_line+1,dtype=int)\n",
    "state_nrs = np.zeros((n_lines,2),dtype=int)\n",
    "state_nrs[:,1] = 1\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check at what point the tweezers separate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 63\n",
      "separation: 1.2750 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24396 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 64\n",
      "separation: 1.3000 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25468 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 65\n",
      "separation: 1.3250 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25256 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 66\n",
      "separation: 1.3500 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24984 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 67\n",
      "separation: 1.3750 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25116 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 68\n",
      "separation: 1.4000 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25032 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "    \n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    print(f\"line: {line}\")\n",
    "    print(f\"separation: {trap.subs(trap.distance_tweezers)/si.um:.4f} um\")\n",
    "\n",
    "    #plot wavefunctions\n",
    "    n_states = 2\n",
    "\n",
    "    states_plot = [GS_left_1.real, GS_right_1.real]\n",
    "    #states_plot = [res[\"states\"][state0].real, res[\"states\"][state1].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=(6, 4))  # 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",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = [10.76459135  3.9825663   1.41321814  0.48828016  0.16587758  0.05579203] Hz\n",
      "U_s = [18.16656197 18.94287365 19.55298704 20.05269908 20.47741687 20.84363256] Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 125900 characters]",
      "text/plain": [
       "<Figure size 500x900 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "distance = np.zeros_like(lines,dtype=float)\n",
    "J = np.zeros_like(lines,dtype=float)\n",
    "#tilt = np.zeros_like(lines,dtype=float)\n",
    "U_s = np.zeros_like(lines,dtype=float)\n",
    "U_dds = np.zeros((len(lines),50),dtype=float)\n",
    "angles = np.zeros((len(lines),50),dtype=float)\n",
    "V_lrs = np.zeros((len(lines),50),dtype=float)\n",
    "DeltaJs = np.zeros((len(lines),50),dtype=float)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "\n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "\n",
    "    distance[i] = float(trap.subs(trap.distance_tweezers))\n",
    "    J[i], U_s[i], U_dds[i], angles[i], V_lrs[i], DeltaJs[i] = analyse_diagonalisation(line,\n",
    "                                                                state_nr0=state0,\n",
    "                                                                state_nr1=state1,\n",
    "                                                                degenerate=degeneracy[i])\n",
    "\n",
    "    # E_left = (func_ham(GS_left_1.flatten())/GS_left_1.flatten())[0]\n",
    "    # E_right = (func_ham(GS_right_1.flatten())/GS_right_1.flatten())[0]\n",
    "    # print(E_left/const.h)\n",
    "    # print(E_right/const.h)\n",
    "    # tilt[i] = E_right - E_left\n",
    "\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.x)\n",
    "    #print(f\"omega_x = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.z)\n",
    "    #print(f\"omega_z = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "\n",
    "print(f\"J = {J/const.h} Hz\")\n",
    "print(f\"U_s = {U_s/const.h} Hz\")\n",
    "\n",
    "# Create figure with subplots\n",
    "fig, axes = plt.subplots(3, 1, figsize=(5, 9), sharex=True)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    # Plot U_dd\n",
    "    axes[0].plot(np.rad2deg(angles[i]), U_dds[i]/const.h,label=f\"{line}\")\n",
    "    axes[0].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[0].set_ylabel(r\"$U_{dd} / h$ [Hz]\")\n",
    "    axes[0].legend()\n",
    "\n",
    "    # Plot NNI\n",
    "    axes[1].plot(np.rad2deg(angles[i]), V_lrs[i]/const.h,label=f\"{line}\")\n",
    "    axes[1].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[1].set_ylabel(r\"$V_{lr} / h$ [Hz]\")\n",
    "    axes[1].legend()\n",
    "\n",
    "    # Plot DIT\n",
    "    axes[2].plot(np.rad2deg(angles[i]), DeltaJs[i]/const.h,label=f\"{line}\")\n",
    "    axes[2].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[2].set_xlabel(r\"$\\theta$ [deg]\")\n",
    "    axes[2].set_ylabel(r\"$\\Delta J / h$ [Hz]\")\n",
    "    axes[2].legend()\n",
    "\n",
    "    # Customize ticks for all subplots\n",
    "    for ax in axes:\n",
    "        ax.tick_params(axis='both', which='both', direction='in', length=5)  # Tick marks inside\n",
    "        ax.tick_params(axis='both', which='major', length=7)  # Major ticks longer\n",
    "        ax.tick_params(axis='both', which='minor', length=4)  # Minor ticks shorter\n",
    "        ax.minorticks_on()  # Enable minor ticks\n",
    "\n",
    "# Adjust layout and show the plot\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x0 = -24.85 nm\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\naeve\\AppData\\Local\\Temp\\ipykernel_25260\\3593060390.py:2: RuntimeWarning: overflow encountered in exp\n",
      "  return A*np.exp((x-d)/x0)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 20228 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def exponential(x,x0,A,d):\n",
    "    return A*np.exp((x-d)/x0)\n",
    "\n",
    "popt, pcov = curve_fit(exponential, distance, J,p0=[-25e-9, 0.1*const.h,1380e-9])\n",
    "\n",
    "print(f\"x0 = {popt[0]/si.nm:.2f} nm\")\n",
    "\n",
    "plt.plot(distance/si.nm, J/const.h,\".\")\n",
    "plt.plot(distance/si.nm, exponential(distance,*popt)/const.h)\n",
    "\n",
    "plt.ylabel(\"J [Hz]\")\n",
    "plt.xlabel(\"x [nm]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Do the same for larger trapping frequency:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 74\n",
    "end_line = 84\n",
    "\n",
    "n_lines = end_line - start_line + 1\n",
    "lines = np.arange(start_line,end_line+1,dtype=int)\n",
    "state_nrs = np.zeros((n_lines,2),dtype=int)\n",
    "state_nrs[:,1] = 1\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 74\n",
      "separation: 0.7571 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 21820 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 75\n",
      "separation: 0.7714 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 22180 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 76\n",
      "separation: 0.7857 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 21912 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 77\n",
      "separation: 0.8000 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 21880 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 78\n",
      "separation: 0.8143 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 23144 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 79\n",
      "separation: 0.8286 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29928 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 80\n",
      "separation: 0.8429 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27892 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 81\n",
      "separation: 0.8571 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 21988 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 82\n",
      "separation: 0.8714 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 23148 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 83\n",
      "separation: 0.8857 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 22712 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 84\n",
      "separation: 0.9000 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 23024 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "    \n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    print(f\"line: {line}\")\n",
    "    print(f\"separation: {trap.subs(trap.distance_tweezers)/si.um:.4f} um\")\n",
    "\n",
    "    #plot wavefunctions\n",
    "    n_states = 2\n",
    "\n",
    "    states_plot = [GS_left_1.real, GS_right_1.real]\n",
    "    #states_plot = [res[\"states\"][state0].real, res[\"states\"][state1].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=(6, 4))  # 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",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = [12.28626301  1.50269376  0.28180441  0.39235061  1.52504364 -1.60494868\n",
      " -0.47799808  0.29744892  1.7027932   2.98919735  2.78954554] Hz\n",
      "U_s = [433.34720095 475.59990671 469.97352607 499.39261094 553.13492095\n",
      " 548.50091279 510.71591276 527.06457265 592.28021292 631.4864043\n",
      " 611.36628181] Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 162296 characters]",
      "text/plain": [
       "<Figure size 500x900 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "distance = np.zeros_like(lines,dtype=float)\n",
    "J = np.zeros_like(lines,dtype=float)\n",
    "#tilt = np.zeros_like(lines,dtype=float)\n",
    "U_s = np.zeros_like(lines,dtype=float)\n",
    "U_dds = np.zeros((len(lines),50),dtype=float)\n",
    "angles = np.zeros((len(lines),50),dtype=float)\n",
    "V_lrs = np.zeros((len(lines),50),dtype=float)\n",
    "DeltaJs = np.zeros((len(lines),50),dtype=float)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "\n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    distance[i] = float(trap.subs(trap.distance_tweezers))\n",
    "    J[i], U_s[i], U_dds[i], angles[i], V_lrs[i], DeltaJs[i] = analyse_diagonalisation(line,\n",
    "                                                                state_nr0=state0,\n",
    "                                                                state_nr1=state1,\n",
    "                                                                degenerate=degeneracy[i])\n",
    "\n",
    "    # E_left = (func_ham(GS_left_1.flatten())/GS_left_1.flatten())[0]\n",
    "    # E_right = (func_ham(GS_right_1.flatten())/GS_right_1.flatten())[0]\n",
    "    # print(E_left/const.h)\n",
    "    # print(E_right/const.h)\n",
    "    # tilt[i] = E_right - E_left\n",
    "\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.x)\n",
    "    #print(f\"omega_x = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.z)\n",
    "    #print(f\"omega_z = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "\n",
    "print(f\"J = {J/const.h} Hz\")\n",
    "print(f\"U_s = {U_s/const.h} Hz\")\n",
    "\n",
    "# Create figure with subplots\n",
    "fig, axes = plt.subplots(3, 1, figsize=(5, 9), sharex=True)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    # Plot U_dd\n",
    "    axes[0].plot(np.rad2deg(angles[i]), U_dds[i]/const.h,label=f\"{line}\")\n",
    "    axes[0].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[0].set_ylabel(r\"$U_{dd} / h$ [Hz]\")\n",
    "    axes[0].legend()\n",
    "\n",
    "    # Plot NNI\n",
    "    axes[1].plot(np.rad2deg(angles[i]), V_lrs[i]/const.h,label=f\"{line}\")\n",
    "    axes[1].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[1].set_ylabel(r\"$V_{lr} / h$ [Hz]\")\n",
    "    axes[1].legend()\n",
    "\n",
    "    # Plot DIT\n",
    "    axes[2].plot(np.rad2deg(angles[i]), DeltaJs[i]/const.h,label=f\"{line}\")\n",
    "    axes[2].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[2].set_xlabel(r\"$\\theta$ [deg]\")\n",
    "    axes[2].set_ylabel(r\"$\\Delta J / h$ [Hz]\")\n",
    "    axes[2].legend()\n",
    "\n",
    "    # Customize ticks for all subplots\n",
    "    for ax in axes:\n",
    "        ax.tick_params(axis='both', which='both', direction='in', length=5)  # Tick marks inside\n",
    "        ax.tick_params(axis='both', which='major', length=7)  # Major ticks longer\n",
    "        ax.tick_params(axis='both', which='minor', length=4)  # Minor ticks shorter\n",
    "        ax.minorticks_on()  # Enable minor ticks\n",
    "\n",
    "# Adjust layout and show the plot\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x0 = -6.91 nm\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 17968 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def exponential(x,x0,A,d):\n",
    "    return A*np.exp((x-d)/x0)\n",
    "\n",
    "popt, pcov = curve_fit(exponential, distance, J,p0=[-5e-9, 0.1*const.h,770e-9])\n",
    "\n",
    "print(f\"x0 = {popt[0]/si.nm:.2f} nm\")\n",
    "\n",
    "plt.plot(distance/si.nm, J/const.h,\".\")\n",
    "plt.plot(distance/si.nm, exponential(distance,*popt)/const.h)\n",
    "\n",
    "plt.ylabel(\"J [Hz]\")\n",
    "plt.xlabel(\"x [nm]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Seems like we need to control the tweezer distance within about 10nm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 88\n",
    "end_line = 94\n",
    "\n",
    "n_lines = end_line - start_line + 1\n",
    "lines = np.arange(start_line,end_line+1,dtype=int)\n",
    "state_nrs = np.zeros((n_lines,2),dtype=int)\n",
    "state_nrs[:,1] = 1\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 88\n",
      "separation: 0.7319 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29540 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 89\n",
      "separation: 0.7428 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27340 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 90\n",
      "separation: 0.7539 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27936 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 91\n",
      "separation: 0.7652 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28104 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 92\n",
      "separation: 0.7766 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27776 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 93\n",
      "separation: 0.7882 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28124 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 94\n",
      "separation: 0.8000 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27744 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "    \n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    print(f\"line: {line}\")\n",
    "    print(f\"separation: {trap.subs(trap.distance_tweezers)/si.um:.4f} um\")\n",
    "\n",
    "    #plot wavefunctions\n",
    "    n_states = 2\n",
    "\n",
    "    states_plot = [GS_left_1.real, GS_right_1.real]\n",
    "    #states_plot = [res[\"states\"][state0].real, res[\"states\"][state1].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=(6, 4))  # 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",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = [ 2.33669250e+01  5.14466514e-01  6.52124558e-03  4.38240060e-05\n",
      " -1.47570996e-05 -2.94650800e-05 -5.31258918e-05] Hz\n",
      "U_s = [1334.81641772 1472.89943189 1567.97881363 1644.15384339 1707.63613706\n",
      " 1764.79497325 1808.30193624] Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 129732 characters]",
      "text/plain": [
       "<Figure size 500x900 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "distance = np.zeros_like(lines,dtype=float)\n",
    "J = np.zeros_like(lines,dtype=float)\n",
    "#tilt = np.zeros_like(lines,dtype=float)\n",
    "U_s = np.zeros_like(lines,dtype=float)\n",
    "U_dds = np.zeros((len(lines),50),dtype=float)\n",
    "angles = np.zeros((len(lines),50),dtype=float)\n",
    "V_lrs = np.zeros((len(lines),50),dtype=float)\n",
    "DeltaJs = np.zeros((len(lines),50),dtype=float)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "\n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    distance[i] = float(trap.subs(trap.distance_tweezers))\n",
    "    J[i], U_s[i], U_dds[i], angles[i], V_lrs[i], DeltaJs[i] = analyse_diagonalisation(line,\n",
    "                                                                state_nr0=state0,\n",
    "                                                                state_nr1=state1,\n",
    "                                                                degenerate=degeneracy[i])\n",
    "\n",
    "    # E_left = (func_ham(GS_left_1.flatten())/GS_left_1.flatten())[0]\n",
    "    # E_right = (func_ham(GS_right_1.flatten())/GS_right_1.flatten())[0]\n",
    "    # print(E_left/const.h)\n",
    "    # print(E_right/const.h)\n",
    "    # tilt[i] = E_right - E_left\n",
    "\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.x)\n",
    "    #print(f\"omega_x = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.z)\n",
    "    #print(f\"omega_z = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "\n",
    "print(f\"J = {J/const.h} Hz\")\n",
    "print(f\"U_s = {U_s/const.h} Hz\")\n",
    "\n",
    "# Create figure with subplots\n",
    "fig, axes = plt.subplots(3, 1, figsize=(5, 9), sharex=True)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    # Plot U_dd\n",
    "    axes[0].plot(np.rad2deg(angles[i]), U_dds[i]/const.h,label=f\"{line}\")\n",
    "    axes[0].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[0].set_ylabel(r\"$U_{dd} / h$ [Hz]\")\n",
    "    axes[0].legend()\n",
    "\n",
    "    # Plot NNI\n",
    "    axes[1].plot(np.rad2deg(angles[i]), V_lrs[i]/const.h,label=f\"{line}\")\n",
    "    axes[1].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[1].set_ylabel(r\"$V_{lr} / h$ [Hz]\")\n",
    "    axes[1].legend()\n",
    "\n",
    "    # Plot DIT\n",
    "    axes[2].plot(np.rad2deg(angles[i]), DeltaJs[i]/const.h,label=f\"{line}\")\n",
    "    axes[2].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[2].set_xlabel(r\"$\\theta$ [deg]\")\n",
    "    axes[2].set_ylabel(r\"$\\Delta J / h$ [Hz]\")\n",
    "    axes[2].legend()\n",
    "\n",
    "    # Customize ticks for all subplots\n",
    "    for ax in axes:\n",
    "        ax.tick_params(axis='both', which='both', direction='in', length=5)  # Tick marks inside\n",
    "        ax.tick_params(axis='both', which='major', length=7)  # Major ticks longer\n",
    "        ax.tick_params(axis='both', which='minor', length=4)  # Minor ticks shorter\n",
    "        ax.minorticks_on()  # Enable minor ticks\n",
    "\n",
    "# Adjust layout and show the plot\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x0 = -2.87 nm\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 17028 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def exponential(x,x0,A,d):\n",
    "    return A*np.exp((x-d)/x0)\n",
    "\n",
    "popt, pcov = curve_fit(exponential, distance, J,p0=[-5e-9, 0.1*const.h,750e-9])\n",
    "\n",
    "print(f\"x0 = {popt[0]/si.nm:.2f} nm\")\n",
    "\n",
    "plt.plot(distance/si.nm, J/const.h,\".\")\n",
    "plt.plot(distance/si.nm, exponential(distance,*popt)/const.h)\n",
    "\n",
    "plt.ylabel(\"J [Hz]\")\n",
    "plt.xlabel(\"x [nm]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Seems like the data is not good, since we're only using small grid.\n",
    "\n",
    "## Try to replicate the Kaufman paper with dysprosium, trap frequency 96kHz/2pi (at separation), 0.71um waist and 808nm separation:\n",
    "\n",
    "They have J~200Hz, so I would expect roughly 100Hz for Dysprosium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 96\n",
    "end_line = 96\n",
    "\n",
    "n_lines = end_line - start_line + 1\n",
    "lines = np.arange(start_line,end_line+1,dtype=int)\n",
    "state_nrs = np.zeros((n_lines,2),dtype=int)\n",
    "state_nrs[:,1] = 1\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = [0.63865845] Hz\n",
      "U_s = [269.17665202] Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 65704 characters]",
      "text/plain": [
       "<Figure size 500x900 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "distance = np.zeros_like(lines,dtype=float)\n",
    "J = np.zeros_like(lines,dtype=float)\n",
    "#tilt = np.zeros_like(lines,dtype=float)\n",
    "U_s = np.zeros_like(lines,dtype=float)\n",
    "U_dds = np.zeros((len(lines),50),dtype=float)\n",
    "angles = np.zeros((len(lines),50),dtype=float)\n",
    "V_lrs = np.zeros((len(lines),50),dtype=float)\n",
    "DeltaJs = np.zeros((len(lines),50),dtype=float)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "\n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    distance[i] = float(trap.subs(trap.distance_tweezers))\n",
    "    J[i], U_s[i], U_dds[i], angles[i], V_lrs[i], DeltaJs[i] = analyse_diagonalisation(line,\n",
    "                                                                state_nr0=state0,\n",
    "                                                                state_nr1=state1,\n",
    "                                                                degenerate=degeneracy[i])\n",
    "\n",
    "    # E_left = (func_ham(GS_left_1.flatten())/GS_left_1.flatten())[0]\n",
    "    # E_right = (func_ham(GS_right_1.flatten())/GS_right_1.flatten())[0]\n",
    "    # print(E_left/const.h)\n",
    "    # print(E_right/const.h)\n",
    "    # tilt[i] = E_right - E_left\n",
    "\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.x)\n",
    "    #print(f\"omega_x = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.z)\n",
    "    #print(f\"omega_z = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "\n",
    "print(f\"J = {J/const.h} Hz\")\n",
    "print(f\"U_s = {U_s/const.h} Hz\")\n",
    "\n",
    "# Create figure with subplots\n",
    "fig, axes = plt.subplots(3, 1, figsize=(5, 9), sharex=True)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    # Plot U_dd\n",
    "    axes[0].plot(np.rad2deg(angles[i]), U_dds[i]/const.h,label=f\"{line}\")\n",
    "    axes[0].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[0].set_ylabel(r\"$U_{dd} / h$ [Hz]\")\n",
    "    axes[0].legend()\n",
    "\n",
    "    # Plot NNI\n",
    "    axes[1].plot(np.rad2deg(angles[i]), V_lrs[i]/const.h,label=f\"{line}\")\n",
    "    axes[1].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[1].set_ylabel(r\"$V_{lr} / h$ [Hz]\")\n",
    "    axes[1].legend()\n",
    "\n",
    "    # Plot DIT\n",
    "    axes[2].plot(np.rad2deg(angles[i]), DeltaJs[i]/const.h,label=f\"{line}\")\n",
    "    axes[2].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[2].set_xlabel(r\"$\\theta$ [deg]\")\n",
    "    axes[2].set_ylabel(r\"$\\Delta J / h$ [Hz]\")\n",
    "    axes[2].legend()\n",
    "\n",
    "    # Customize ticks for all subplots\n",
    "    for ax in axes:\n",
    "        ax.tick_params(axis='both', which='both', direction='in', length=5)  # Tick marks inside\n",
    "        ax.tick_params(axis='both', which='major', length=7)  # Major ticks longer\n",
    "        ax.tick_params(axis='both', which='minor', length=4)  # Minor ticks shorter\n",
    "        ax.minorticks_on()  # Enable minor ticks\n",
    "\n",
    "# Adjust layout and show the plot\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See how sensitive this result is in terms of tweezer spacing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 98\n",
    "end_line = 106\n",
    "\n",
    "n_lines = end_line - start_line + 1\n",
    "lines = np.arange(start_line,end_line+1,dtype=int)\n",
    "state_nrs = np.zeros((n_lines,2),dtype=int)\n",
    "state_nrs[:,1] = 1\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 98\n",
      "separation: 0.7433 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35768 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 99\n",
      "separation: 0.7567 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35996 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 100\n",
      "separation: 0.7700 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35564 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 101\n",
      "separation: 0.7833 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35052 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 102\n",
      "separation: 0.7967 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 33840 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 103\n",
      "separation: 0.8100 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29352 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 104\n",
      "separation: 0.8233 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29220 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 105\n",
      "separation: 0.8367 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 29016 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 106\n",
      "separation: 0.8500 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28704 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "    \n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    print(f\"line: {line}\")\n",
    "    print(f\"separation: {trap.subs(trap.distance_tweezers)/si.um:.4f} um\")\n",
    "\n",
    "    #plot wavefunctions\n",
    "    n_states = 2\n",
    "\n",
    "    states_plot = [GS_left_1.real, GS_right_1.real]\n",
    "    #states_plot = [res[\"states\"][state0].real, res[\"states\"][state1].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=(6, 4))  # 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",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = [6.47297376e+02 2.57563001e+02 7.12839954e+01 1.54300059e+01\n",
      " 2.87856788e+00 4.86787632e-01 7.68575361e-02 1.15562296e-02\n",
      " 1.67843709e-03] Hz\n",
      "U_s = [192.250924   207.15110813 230.08270887 247.62052937 260.45180538\n",
      " 270.5648513  278.91147625 285.9887713  292.09844135] Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 145572 characters]",
      "text/plain": [
       "<Figure size 500x900 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "distance = np.zeros_like(lines,dtype=float)\n",
    "J = np.zeros_like(lines,dtype=float)\n",
    "#tilt = np.zeros_like(lines,dtype=float)\n",
    "U_s = np.zeros_like(lines,dtype=float)\n",
    "U_dds = np.zeros((len(lines),50),dtype=float)\n",
    "angles = np.zeros((len(lines),50),dtype=float)\n",
    "V_lrs = np.zeros((len(lines),50),dtype=float)\n",
    "DeltaJs = np.zeros((len(lines),50),dtype=float)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "\n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    distance[i] = float(trap.subs(trap.distance_tweezers))\n",
    "    J[i], U_s[i], U_dds[i], angles[i], V_lrs[i], DeltaJs[i] = analyse_diagonalisation(line,\n",
    "                                                                state_nr0=state0,\n",
    "                                                                state_nr1=state1,\n",
    "                                                                degenerate=degeneracy[i],\n",
    "                                                                n_angles=20)\n",
    "\n",
    "    # E_left = (func_ham(GS_left_1.flatten())/GS_left_1.flatten())[0]\n",
    "    # E_right = (func_ham(GS_right_1.flatten())/GS_right_1.flatten())[0]\n",
    "    # print(E_left/const.h)\n",
    "    # print(E_right/const.h)\n",
    "    # tilt[i] = E_right - E_left\n",
    "\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.x)\n",
    "    #print(f\"omega_x = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.z)\n",
    "    #print(f\"omega_z = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "\n",
    "print(f\"J = {J/const.h} Hz\")\n",
    "print(f\"U_s = {U_s/const.h} Hz\")\n",
    "\n",
    "# Create figure with subplots\n",
    "fig, axes = plt.subplots(3, 1, figsize=(5, 9), sharex=True)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    # Plot U_dd\n",
    "    axes[0].plot(np.rad2deg(angles[i]), U_dds[i]/const.h,label=f\"{line}\")\n",
    "    axes[0].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[0].set_ylabel(r\"$U_{dd} / h$ [Hz]\")\n",
    "    axes[0].legend()\n",
    "\n",
    "    # Plot NNI\n",
    "    axes[1].plot(np.rad2deg(angles[i]), V_lrs[i]/const.h,label=f\"{line}\")\n",
    "    axes[1].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[1].set_ylabel(r\"$V_{lr} / h$ [Hz]\")\n",
    "    axes[1].legend()\n",
    "\n",
    "    # Plot DIT\n",
    "    axes[2].plot(np.rad2deg(angles[i]), DeltaJs[i]/const.h,label=f\"{line}\")\n",
    "    axes[2].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[2].set_xlabel(r\"$\\theta$ [deg]\")\n",
    "    axes[2].set_ylabel(r\"$\\Delta J / h$ [Hz]\")\n",
    "    axes[2].legend()\n",
    "\n",
    "    # Customize ticks for all subplots\n",
    "    for ax in axes:\n",
    "        ax.tick_params(axis='both', which='both', direction='in', length=5)  # Tick marks inside\n",
    "        ax.tick_params(axis='both', which='major', length=7)  # Major ticks longer\n",
    "        ax.tick_params(axis='both', which='minor', length=4)  # Minor ticks shorter\n",
    "        ax.minorticks_on()  # Enable minor ticks\n",
    "\n",
    "# Adjust layout and show the plot\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x0 = -13.23 nm\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 22520 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def exponential(x,x0,A,d):\n",
    "    return A*np.exp((x-d)/x0)\n",
    "\n",
    "popt, pcov = curve_fit(exponential, distance, J,p0=[-5e-9, 0.1*const.h,750e-9])\n",
    "\n",
    "print(f\"x0 = {popt[0]/si.nm:.2f} nm\")\n",
    "\n",
    "plt.plot(distance/si.nm, J/const.h,\".\")\n",
    "plt.plot(distance/si.nm, exponential(distance,*popt)/const.h)\n",
    "\n",
    "plt.ylabel(\"J [Hz]\")\n",
    "plt.xlabel(\"x [nm]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now with trap frequency at separation of 60kHz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 108\n",
    "end_line = 117\n",
    "\n",
    "n_lines = end_line - start_line + 1\n",
    "lines = np.arange(start_line,end_line+1,dtype=int)\n",
    "state_nrs = np.zeros((n_lines,2),dtype=int)\n",
    "state_nrs[:,1] = 1\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 108\n",
      "separation: 0.7300 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25940 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 109\n",
      "separation: 0.7433 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 26004 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 110\n",
      "separation: 0.7567 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24908 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 111\n",
      "separation: 0.7700 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24860 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 112\n",
      "separation: 0.7833 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27960 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 113\n",
      "separation: 0.7967 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35044 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 114\n",
      "separation: 0.8100 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 31420 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 115\n",
      "separation: 0.8233 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 31064 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 116\n",
      "separation: 0.8367 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28740 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 117\n",
      "separation: 0.8500 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25540 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "    \n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    print(f\"line: {line}\")\n",
    "    print(f\"separation: {trap.subs(trap.distance_tweezers)/si.um:.4f} um\")\n",
    "\n",
    "    #plot wavefunctions\n",
    "    n_states = 2\n",
    "\n",
    "    states_plot = [GS_left_1.real, GS_right_1.real]\n",
    "    #states_plot = [res[\"states\"][state0].real, res[\"states\"][state1].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=(6, 4))  # 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",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = [3.95109882e+02 4.55461744e+00 1.58088853e-02 2.68149098e-05\n",
      " 2.84425642e-08 1.25134430e-09 3.05398467e-09 9.37039197e-09\n",
      " 8.81733011e-09 1.12933186e-08] Hz\n",
      "U_s = [1381.51049509 1682.55719597 1852.27750338 1973.65834756 1919.33117469\n",
      " 4737.22817343  792.13836296 4064.4449819  2600.64434161 2359.95270782] Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 140896 characters]",
      "text/plain": [
       "<Figure size 500x900 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "distance = np.zeros_like(lines,dtype=float)\n",
    "J = np.zeros_like(lines,dtype=float)\n",
    "#tilt = np.zeros_like(lines,dtype=float)\n",
    "U_s = np.zeros_like(lines,dtype=float)\n",
    "U_dds = np.zeros((len(lines),50),dtype=float)\n",
    "angles = np.zeros((len(lines),50),dtype=float)\n",
    "V_lrs = np.zeros((len(lines),50),dtype=float)\n",
    "DeltaJs = np.zeros((len(lines),50),dtype=float)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "\n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    distance[i] = float(trap.subs(trap.distance_tweezers))\n",
    "    J[i], U_s[i], U_dds[i], angles[i], V_lrs[i], DeltaJs[i] = analyse_diagonalisation(line,\n",
    "                                                                state_nr0=state0,\n",
    "                                                                state_nr1=state1,\n",
    "                                                                degenerate=degeneracy[i],\n",
    "                                                                n_angles=20)\n",
    "\n",
    "    # E_left = (func_ham(GS_left_1.flatten())/GS_left_1.flatten())[0]\n",
    "    # E_right = (func_ham(GS_right_1.flatten())/GS_right_1.flatten())[0]\n",
    "    # print(E_left/const.h)\n",
    "    # print(E_right/const.h)\n",
    "    # tilt[i] = E_right - E_left\n",
    "\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.x)\n",
    "    #print(f\"omega_x = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "    #omega_left, omega_right = trap.get_both_omega(trap.z)\n",
    "    #print(f\"omega_z = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "\n",
    "print(f\"J = {J/const.h} Hz\")\n",
    "print(f\"U_s = {U_s/const.h} Hz\")\n",
    "\n",
    "# Create figure with subplots\n",
    "fig, axes = plt.subplots(3, 1, figsize=(5, 9), sharex=True)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    # Plot U_dd\n",
    "    axes[0].plot(np.rad2deg(angles[i]), U_dds[i]/const.h,label=f\"{line}\")\n",
    "    axes[0].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[0].set_ylabel(r\"$U_{dd} / h$ [Hz]\")\n",
    "    axes[0].legend()\n",
    "\n",
    "    # Plot NNI\n",
    "    axes[1].plot(np.rad2deg(angles[i]), V_lrs[i]/const.h,label=f\"{line}\")\n",
    "    axes[1].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[1].set_ylabel(r\"$V_{lr} / h$ [Hz]\")\n",
    "    axes[1].legend()\n",
    "\n",
    "    # Plot DIT\n",
    "    axes[2].plot(np.rad2deg(angles[i]), DeltaJs[i]/const.h,label=f\"{line}\")\n",
    "    axes[2].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[2].set_xlabel(r\"$\\theta$ [deg]\")\n",
    "    axes[2].set_ylabel(r\"$\\Delta J / h$ [Hz]\")\n",
    "    axes[2].legend()\n",
    "\n",
    "    # Customize ticks for all subplots\n",
    "    for ax in axes:\n",
    "        ax.tick_params(axis='both', which='both', direction='in', length=5)  # Tick marks inside\n",
    "        ax.tick_params(axis='both', which='major', length=7)  # Major ticks longer\n",
    "        ax.tick_params(axis='both', which='minor', length=4)  # Minor ticks shorter\n",
    "        ax.minorticks_on()  # Enable minor ticks\n",
    "\n",
    "# Adjust layout and show the plot\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x0 = -2.99 nm\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 20344 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def exponential(x,x0,A,d):\n",
    "    return A*np.exp((x-d)/x0)\n",
    "\n",
    "popt, pcov = curve_fit(exponential, distance, J,p0=[-5e-9, 0.1*const.h,750e-9])\n",
    "\n",
    "print(f\"x0 = {popt[0]/si.nm:.2f} nm\")\n",
    "\n",
    "plt.plot(distance/si.nm, J/const.h,\".\")\n",
    "plt.plot(distance/si.nm, exponential(distance,*popt)/const.h)\n",
    "\n",
    "plt.ylabel(\"J [Hz]\")\n",
    "plt.xlabel(\"x [nm]\")\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
}
