{
 "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": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 119\n",
    "end_line = 139\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[0,1] = 2\n",
    "state_nrs[1,1] = 2\n",
    "state_nrs[-4,1] = 2\n",
    "\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 141\n",
    "end_line = 164\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[0,1] = 2\n",
    "state_nrs[1,1] = 2\n",
    "state_nrs[5,1] = 2\n",
    "state_nrs[10,1] = 2\n",
    "state_nrs[15,1] = 2\n",
    "\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 166\n",
    "end_line = 206\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[0,1] = 2\n",
    "state_nrs[1,1] = 2\n",
    "state_nrs[2,1] = 2\n",
    "state_nrs[3,1] = 2\n",
    "state_nrs[9,1] = 2\n",
    "state_nrs[18,1] = 2\n",
    "state_nrs[27,1] = 2\n",
    "\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 166\n",
      "separation: 0.8372 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 69460 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 167\n",
      "separation: 0.8497 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 65556 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 168\n",
      "separation: 0.8622 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 71088 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 169\n",
      "separation: 0.8747 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 70100 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 170\n",
      "separation: 0.8872 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 69404 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 171\n",
      "separation: 0.8997 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 65400 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 172\n",
      "separation: 0.9122 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 65152 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 173\n",
      "separation: 0.9247 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 65544 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 174\n",
      "separation: 0.9372 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 68364 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 175\n",
      "separation: 0.7270 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 44760 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 176\n",
      "separation: 0.7395 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 48108 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 177\n",
      "separation: 0.7520 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 45264 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 178\n",
      "separation: 0.7645 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 47460 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 179\n",
      "separation: 0.7770 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 46428 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 180\n",
      "separation: 0.7895 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 46868 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 181\n",
      "separation: 0.8020 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 45392 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 182\n",
      "separation: 0.8145 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 44000 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 183\n",
      "separation: 0.8270 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 42584 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 184\n",
      "separation: 0.7216 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 44156 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 185\n",
      "separation: 0.7341 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 39596 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 186\n",
      "separation: 0.7466 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 38604 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 187\n",
      "separation: 0.7591 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37152 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 188\n",
      "separation: 0.7716 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37116 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 189\n",
      "separation: 0.7841 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34096 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 190\n",
      "separation: 0.7966 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34108 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 191\n",
      "separation: 0.8091 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 33808 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 192\n",
      "separation: 0.7131 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35908 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 193\n",
      "separation: 0.7256 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 36136 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 194\n",
      "separation: 0.7381 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37560 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 195\n",
      "separation: 0.7506 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 36020 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 196\n",
      "separation: 0.7631 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32660 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 197\n",
      "separation: 0.7756 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 34864 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 198\n",
      "separation: 0.7881 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32248 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 199\n",
      "separation: 0.8006 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 31972 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 200\n",
      "separation: 0.7204 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 39644 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 201\n",
      "separation: 0.7329 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: 202\n",
      "separation: 0.7454 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35972 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 203\n",
      "separation: 0.7579 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 33708 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 204\n",
      "separation: 0.7704 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32972 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 205\n",
      "separation: 0.7829 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32464 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 206\n",
      "separation: 0.7954 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32132 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": 75,
   "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": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 108756 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 112468 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 89104 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 104444 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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 395748 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": 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
}
