{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors as mcolors\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",
    "import qutip as qt\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": [
    "# Choose data to analyse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 292\n",
    "end_line = 332\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",
    "\n",
    "state_nrs[9,1] = 2\n",
    "state_nrs[18,1] = 2\n",
    "state_nrs[26,1] = 2\n",
    "\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 292\n",
      "separation: 1.5472 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 52456 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 293\n",
      "separation: 1.5597 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 53348 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 294\n",
      "separation: 1.5722 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 52420 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 295\n",
      "separation: 1.5847 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 51408 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 296\n",
      "separation: 1.5972 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 51944 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 297\n",
      "separation: 1.6097 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 51036 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 298\n",
      "separation: 1.6222 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49648 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 299\n",
      "separation: 1.6347 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49588 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 300\n",
      "separation: 1.6472 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49584 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 301\n",
      "separation: 1.4370 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 45680 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 302\n",
      "separation: 1.4495 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 47676 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 303\n",
      "separation: 1.4620 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 46652 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 304\n",
      "separation: 1.4745 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 44240 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 305\n",
      "separation: 1.4870 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 42240 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 306\n",
      "separation: 1.4995 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 40896 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 307\n",
      "separation: 1.5120 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 40112 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 308\n",
      "separation: 1.5245 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 39636 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 309\n",
      "separation: 1.5370 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 40572 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 310\n",
      "separation: 1.4316 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 44716 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 311\n",
      "separation: 1.4441 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 44144 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 312\n",
      "separation: 1.4566 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41920 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 313\n",
      "separation: 1.4691 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 39524 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 314\n",
      "separation: 1.4816 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 40064 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 315\n",
      "separation: 1.4941 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 38676 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 316\n",
      "separation: 1.5066 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37484 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 317\n",
      "separation: 1.5191 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 38144 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 318\n",
      "separation: 1.4231 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 38776 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 319\n",
      "separation: 1.4356 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 43796 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 320\n",
      "separation: 1.4481 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41380 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 321\n",
      "separation: 1.4606 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 38228 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 322\n",
      "separation: 1.4731 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 36516 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 323\n",
      "separation: 1.4856 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35548 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 324\n",
      "separation: 1.4981 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 38396 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 325\n",
      "separation: 1.5106 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 36792 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 326\n",
      "separation: 1.4304 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 48408 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 327\n",
      "separation: 1.4429 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 45128 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 328\n",
      "separation: 1.4554 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 41960 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 329\n",
      "separation: 1.4679 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35604 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 330\n",
      "separation: 1.4804 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34112 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 331\n",
      "separation: 1.4929 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 38316 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 332\n",
      "separation: 1.5054 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37028 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": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "frequencies = np.zeros_like(lines,dtype=float)\n",
    "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",
    "E_exc = np.zeros_like(lines,dtype=float)\n",
    "E_barr = np.zeros_like(lines,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",
    "    frequencies[i] = float(trap.subs(sp.sqrt(2*trap.power_tweezer1*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer1**2))\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",
    "    #estimate energies of 0th and 1st state\n",
    "    E0 = 0.5*(res[\"energies\"][0] + res[\"energies\"][1])\n",
    "    E1 = 0.5*(res[\"energies\"][2] + res[\"energies\"][3])\n",
    "    E_exc[i] = E1-E0\n",
    "\n",
    "    #estimate heigth of barrier\n",
    "    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",
    "    E_barr[i] = np.max([V_1,V_2])\n",
    "\n",
    "    # these states are basically 1 tweezer\n",
    "    if state1 !=1:\n",
    "        J[i] = np.nan\n",
    "        U_s[i] = np.nan\n",
    "        U_dds[i] = np.nan\n",
    "        V_lrs[i] = np.nan\n",
    "        DeltaJs[i] = np.nan\n",
    "        E_exc[i] = np.nan\n",
    "        E_barr[i] = np.nan\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 378924 characters]",
      "text/plain": [
       "<Figure size 1000x1500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Create a figure with 4 subplots (stacked vertically)\n",
    "fig, axes = plt.subplots(4, 1, figsize=(10, 15), sharex=True)\n",
    "\n",
    "# Get unique frequencies\n",
    "unique_freqs = np.unique(frequencies)[1:]\n",
    "\n",
    "# Define colors for consistency\n",
    "colors = [\"red\", \"blue\", \"green\", \"purple\", \"orange\"][1:]\n",
    "\n",
    "### 1️⃣ Hopping parameter J ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    axes[0].plot(distance[mask] / si.nm, J[mask] / const.h, color=color, \n",
    "                 label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "axes[0].axhspan(50, 150, alpha=0.3, color=\"gray\", label=\"desired region of 50-100Hz\")\n",
    "axes[0].set_ylim(0,500)\n",
    "axes[0].set_ylabel(r\"$J/h$ [Hz]\")\n",
    "axes[0].legend()\n",
    "# 🔹 Add top x-axis ticks\n",
    "axes[0].tick_params(axis=\"x\", labeltop=True, top=True)\n",
    "axes[0].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[0].minorticks_on()\n",
    "axes[0].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "### 2️⃣ Interaction parameter U ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    x = distance[mask] / si.nm\n",
    "    y = U_s[mask] / const.h\n",
    "    delta = U_dds[mask] / const.h\n",
    "\n",
    "    axes[1].plot(x, y, color=color, label=fr\"$U_s$ for $\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "    axes[1].fill_between(x, y + np.min(delta, axis=1), y + np.max(delta, axis=1), \n",
    "                         alpha=0.3, color=color, edgecolor=None, label=r\"$\\Delta U_{dd}$\")\n",
    "\n",
    "axes[1].set_ylabel(r\"$U/h$ [Hz]\")\n",
    "axes[1].legend()\n",
    "axes[1].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[1].minorticks_on()\n",
    "axes[1].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "### 3️⃣ First excited state energy ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    axes[2].plot(distance[mask] / si.nm, E_exc[mask] / const.h, color=color, \n",
    "                 label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "axes[2].set_ylabel(r\"$E_{exc}/h$ [Hz]\")\n",
    "axes[2].legend()\n",
    "axes[2].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[2].minorticks_on()\n",
    "axes[2].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "### 4️⃣ Barrier height ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    axes[3].plot(distance[mask] / si.nm, E_barr[mask] / const.h, color=color, \n",
    "                 label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "axes[3].set_ylabel(r\"$E_{barr}/h$ [Hz]\")\n",
    "axes[3].set_xlabel(\"tweezer spacing [nm]\")\n",
    "axes[3].set_ylim(0, 30000)\n",
    "axes[3].legend()\n",
    "axes[3].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[3].minorticks_on()\n",
    "axes[3].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "# Improve layout\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 112536 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 126080 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 95020 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 111072 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#hopping\n",
    "plt.figure(figsize=(10,5))\n",
    "for i, freq in enumerate(np.unique(frequencies)):\n",
    "    plt.plot(distance[frequencies == freq]/si.nm, J[frequencies == freq]/const.h,\n",
    "             label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "#plt.axhline(100,linestyle='dashed',color=\"black\",label=\"J=100Hz\")\n",
    "plt.axhspan(50,150,alpha=0.3, color=\"gray\",label=\"desired region of 50-100Hz\")\n",
    "\n",
    "plt.ylim(0,500)\n",
    "plt.title(\"Hubbard parameters at different trap frequencies\")\n",
    "plt.xlabel(\"tweezer spacing [nm]\")\n",
    "plt.ylabel(r\"$J/h$ [Hz]\")\n",
    "plt.legend()\n",
    "plt.grid(True, linestyle='-', linewidth=0.7)\n",
    "# Minor grid (finer)\n",
    "plt.minorticks_on()\n",
    "plt.grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=1)\n",
    "plt.show()\n",
    "\n",
    "\n",
    "#interactions\n",
    "plt.figure(figsize=(10,5))\n",
    "for i, freq in enumerate(np.unique(frequencies)):\n",
    "    mask = frequencies == freq\n",
    "    x = distance[mask] / si.nm\n",
    "    y = U_s[mask] / const.h\n",
    "    delta = U_dds[mask] / const.h\n",
    "\n",
    "    # Plot U_s curve\n",
    "    plt.plot(x, y, label=fr\"$U_s$ for $\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "    # Shade around the curve due to dipolar interactions\n",
    "    plt.fill_between(x, y + np.min(delta,axis=1), y + np.max(delta,axis=1),\n",
    "                     alpha=0.4, label=r\"$\\Delta U_{dd}$\")\n",
    "\n",
    "plt.xlabel(\"tweezer spacing [nm]\")\n",
    "plt.ylabel(r\"$U/h$ [Hz]\")\n",
    "plt.legend()\n",
    "plt.grid(True, linestyle='-', linewidth=0.7)\n",
    "# Minor grid (finer)\n",
    "plt.minorticks_on()\n",
    "plt.grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=1)\n",
    "plt.show()\n",
    "\n",
    "\n",
    "#first excited state\n",
    "plt.figure(figsize=(10,5))\n",
    "for i, freq in enumerate(np.unique(frequencies)):\n",
    "    plt.plot(distance[frequencies == freq]/si.nm, E_exc[frequencies == freq]/const.h,\n",
    "             label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "    \n",
    "\n",
    "\n",
    "plt.xlabel(\"tweezer spacing [nm]\")\n",
    "plt.ylabel(r\"$E_{exc}/h$ [Hz]\")\n",
    "plt.legend()\n",
    "plt.grid(True, linestyle='-', linewidth=0.7)\n",
    "# Minor grid (finer)\n",
    "plt.minorticks_on()\n",
    "plt.grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=1)\n",
    "plt.show()\n",
    "\n",
    "\n",
    "#height of barrier\n",
    "plt.figure(figsize=(10,5))\n",
    "for i, freq in enumerate(np.unique(frequencies)):\n",
    "    plt.plot(distance[frequencies == freq]/si.nm, E_barr[frequencies == freq]/const.h,\n",
    "             label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "plt.ylim(0,30000)\n",
    "plt.xlabel(\"tweezer spacing [nm]\")\n",
    "plt.ylabel(r\"$E_{barr}/h$ [Hz]\")\n",
    "plt.legend()\n",
    "plt.grid(True, linestyle='-', linewidth=0.7)\n",
    "# Minor grid (finer)\n",
    "plt.minorticks_on()\n",
    "plt.grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=1)\n",
    "plt.show()"
   ]
  }
 ],
 "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
}
