{
 "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": 3,
   "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": 4,
   "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": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 166\n",
      "separation: 0.8372 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 63340 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, 64336 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, 65036 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, 65884 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, 67496 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, 64004 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, 64204 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, 65436 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, 63868 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, 40796 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, 42632 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, 44160 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, 43452 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, 45700 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, 44532 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, 43564 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, 43136 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, 42564 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, 37520 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, 37568 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, 37340 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, 35796 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, 35164 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, 34024 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, 34100 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, 33788 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"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\scipy\\optimize\\_root_scalar.py:315: RuntimeWarning: Derivative was zero.\n",
      "  r, sol = methodc(f, x0, args=args, fprime=fprime, fprime2=None,\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35960 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, 36164 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, 36256 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, 33612 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, 32316 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, 32304 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, 32244 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, 31976 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, 35304 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, 36640 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, 33604 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, 33264 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, 32820 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": 6,
   "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": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 108592 characters]",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 109932 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": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<>:59: SyntaxWarning: invalid escape sequence '\\o'\n",
      "<>:59: SyntaxWarning: invalid escape sequence '\\o'\n",
      "C:\\Users\\naeve\\AppData\\Local\\Temp\\ipykernel_38208\\181982426.py:59: SyntaxWarning: invalid escape sequence '\\o'\n",
      "  \"\"\"\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 271236 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 3 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(3, 1, figsize=(10, 10), 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",
    "### 4 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].set_xlabel(\"tweezer spacing [nm]\")\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",
    "\"\"\"\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.savefig(\"BH_param_original.png\",dpi=300)\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
}
