{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "from helpers_merging import *\n",
    "\n",
    "initial_power = 100* si.uW\n",
    "initial_waist = 1.1*si.uW\n",
    "initial_distance = 3*si.um\n",
    "\n",
    "trap: DoubleTweezer = DoubleTweezer(\n",
    "    power=0,  # Set pancake laser power to 0, no 2D trap\n",
    "    grad_z= 0*si.G/si.cm,\n",
    "    grad_r=0,\n",
    "    power_tweezer1 = initial_power,     #stationary\n",
    "    power_tweezer2 = initial_power,     #transfer tweezer\n",
    "    waist_tweezer1 = initial_waist,     #stationary\n",
    "    waist_tweezer2 = initial_waist*1.02,     #transfer tweezer\n",
    "    distance_tweezers = initial_distance,\n",
    "\n",
    "    a=180*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "    wvl = 532 * si.nm,\n",
    "\n",
    "    g = 0,\n",
    ")\n",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make nice plots to show how the eigenenergies behave for different tweezer powers and waists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[4 5]] at 1.48um\n",
      "crossover of states [[2 3]] at 1.39um\n",
      "crossover of states [[0 1]] at 1.30um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 132364 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 136420 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[1 2]] at 1.33um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 140328 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 139884 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#make multiple plots showing the eigenenergies for different tweezer powers\n",
    "\n",
    "power2_factors = np.linspace(1+1e-3,1.09,4)\n",
    "\n",
    "for i, pow in enumerate(power2_factors):\n",
    "    #make sure the powers and waists are set correctly\n",
    "    trap[trap.power_tweezer1] = initial_power\n",
    "    trap[trap.power_tweezer2] = initial_power*pow\n",
    "\n",
    "    trap[trap.waist_tweezer1] = initial_waist\n",
    "    trap[trap.waist_tweezer2] = initial_waist\n",
    "    \n",
    "    #calculate eigenenergies\n",
    "    distances = np.linspace(initial_distance,0*si.um,500)\n",
    "    energies, states, potentials = loop_distances(trap, distances,n_levels=6,n_pot_steps=2000)\n",
    "\n",
    "    new_energies, new_states, new_potentials, index_top, index, swap_index = swapped_loop_distance(distances, energies, states, potentials)\n",
    "    energies_left, energies_right, states_left, states_right,deltaE_min, occ_num = find_ass_tweezer(new_energies,new_states)\n",
    "\n",
    "    # Create the figure with two subplots (1 row, 2 columns)\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "    # Plot the full range on ax1\n",
    "    ax1.plot(distances/si.um, energies_left[:, 0], color=\"red\", label=\"GS left tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_right[:, 0], color=\"blue\", label=\"GS right tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_left[:, 1:], color=\"red\", linestyle=\"dotted\", label=\"left tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_right[:, 1:], color=\"blue\", linestyle=\"dotted\", label=\"right tweezer\")\n",
    "\n",
    "    for i, ind in enumerate(index):\n",
    "        ax1.axvline(distances[ind]/si.um, color=\"k\", alpha=0.7, label=\"energy crossings\")\n",
    "\n",
    "    ax1.set_xlabel(\"Tweezer distance [um]\")\n",
    "    ax1.set_ylabel(\"Eigenenergy [a.u.]\")\n",
    "    ax1.set_title(fr\"$p_l={float(trap.subs(trap.power_tweezer1))/si.uW:.1f}$ uW, $p_r={float(trap.subs(trap.power_tweezer2))/si.uW:.1f}$ uW, $w_l={float(trap.subs(trap.waist_tweezer1))/si.um:.3f}$ um, $w_r={float(trap.subs(trap.waist_tweezer2))/si.um:.3f}$ um\")\n",
    "    ax1.grid()\n",
    "\n",
    "    # Plot the zoomed-in range (1 to 1.5 µm) on ax2\n",
    "    ax2.plot(distances/si.um, energies_left[:, 0], color=\"red\", label=\"GS left tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_right[:, 0], color=\"blue\", label=\"GS right tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_left[:, 1:], color=\"red\", linestyle=\"dotted\", label=\"left tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_right[:, 1:], color=\"blue\", linestyle=\"dotted\", label=\"right tweezer\")\n",
    "\n",
    "    for i, ind in enumerate(index):\n",
    "        ax2.axvline(distances[ind]/si.um, color=\"k\", alpha=0.7, label=\"energy crossings\")\n",
    "\n",
    "    ax2.set_xlabel(\"Tweezer distance [um]\")\n",
    "    ax2.set_ylabel(\"Eigenenergy [a.u.]\")\n",
    "    ax2.set_title(\"Zoomed-in Region\")\n",
    "    ax2.set_xlim(1, 1.7)  # Zooming in on the region between 1.0 and 1.5 µm\n",
    "    ax2.set_ylim(-4e-29,-2.5e-29)\n",
    "    ax2.grid()\n",
    "\n",
    "    # Remove duplicates in the legend (labels that appear in both subplots)\n",
    "    handles, labels = ax1.get_legend_handles_labels()\n",
    "    unique_labels = dict(zip(labels, handles))  # Remove duplicates\n",
    "    ax1.legend(unique_labels.values(), unique_labels.keys())\n",
    "\n",
    "    # Show the plots\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The same can be done with the waist, but it's more sensitive here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[2 3]] at 1.46um\n",
      "crossover of states [[0 1]] at 1.35um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 134032 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 137844 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 136616 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[2 3]] at 1.83um\n",
      "crossover of states [[4 5]] at 1.74um\n",
      "crossover of top state at 1.74um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 139792 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#make multiple plots showing the eigenenergies for different tweezer waists\n",
    "\n",
    "waist2_factors = np.linspace(1+1e-4,1.07,4)\n",
    "\n",
    "for i, w0 in enumerate(waist2_factors):\n",
    "    #make sure the powers and waists are set correctly\n",
    "    trap[trap.power_tweezer1] = initial_power\n",
    "    trap[trap.power_tweezer2] = initial_power\n",
    "\n",
    "    trap[trap.waist_tweezer1] = initial_waist\n",
    "    trap[trap.waist_tweezer2] = initial_waist*w0\n",
    "    \n",
    "    #calculate eigenenergies\n",
    "    distances = np.linspace(initial_distance,0*si.um,500)\n",
    "    energies, states, potentials = loop_distances(trap, distances,n_levels=6,n_pot_steps=2000)\n",
    "\n",
    "    new_energies, new_states, new_potentials, index_top, index, swap_index = swapped_loop_distance(distances, energies, states, potentials)\n",
    "    energies_left, energies_right, states_left, states_right,deltaE_min, occ_num = find_ass_tweezer(new_energies,new_states)\n",
    "\n",
    "    # Create the figure with two subplots (1 row, 2 columns)\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "    # Plot the full range on ax1\n",
    "    ax1.plot(distances/si.um, energies_left[:, 0], color=\"red\", label=\"GS left tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_right[:, 0], color=\"blue\", label=\"GS right tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_left[:, 1:], color=\"red\", linestyle=\"dotted\", label=\"left tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_right[:, 1:], color=\"blue\", linestyle=\"dotted\", label=\"right tweezer\")\n",
    "\n",
    "    for i, ind in enumerate(index):\n",
    "        ax1.axvline(distances[ind]/si.um, color=\"k\", alpha=0.7, label=\"energy crossings\")\n",
    "\n",
    "    ax1.set_xlabel(\"Tweezer distance [um]\")\n",
    "    ax1.set_ylabel(\"Eigenenergy [a.u.]\")\n",
    "    ax1.set_title(fr\"$p_l={float(trap.subs(trap.power_tweezer1))/si.uW:.1f}$ uW, $p_r={float(trap.subs(trap.power_tweezer2))/si.uW:.1f}$ uW, $w_l={float(trap.subs(trap.waist_tweezer1))/si.um:.3f}$ um, $w_r={float(trap.subs(trap.waist_tweezer2))/si.um:.3f}$ um\")\n",
    "    ax1.grid()\n",
    "\n",
    "    # Plot the zoomed-in range (1 to 1.5 µm) on ax2\n",
    "    ax2.plot(distances/si.um, energies_left[:, 0], color=\"red\", label=\"GS left tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_right[:, 0], color=\"blue\", label=\"GS right tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_left[:, 1:], color=\"red\", linestyle=\"dotted\", label=\"left tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_right[:, 1:], color=\"blue\", linestyle=\"dotted\", label=\"right tweezer\")\n",
    "\n",
    "    for i, ind in enumerate(index):\n",
    "        ax2.axvline(distances[ind]/si.um, color=\"k\", alpha=0.7, label=\"energy crossings\")\n",
    "\n",
    "    ax2.set_xlabel(\"Tweezer distance [um]\")\n",
    "    ax2.set_ylabel(\"Eigenenergy [a.u.]\")\n",
    "    ax2.set_title(\"Zoomed-in Region\")\n",
    "    ax2.set_xlim(1, 1.7)  # Zooming in on the region between 1.0 and 1.5 µm\n",
    "    ax2.set_ylim(-3.8e-29,-2.4e-29)\n",
    "    ax2.grid()\n",
    "\n",
    "    # Remove duplicates in the legend (labels that appear in both subplots)\n",
    "    handles, labels = ax1.get_legend_handles_labels()\n",
    "    unique_labels = dict(zip(labels, handles))  # Remove duplicates\n",
    "    ax1.legend(unique_labels.values(), unique_labels.keys())\n",
    "\n",
    "    # Show the plots\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Can a difference in waist be compensated by changing the power?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[2 3]] at 1.52um\n",
      "crossover of states [[4 5]] at 1.49um\n",
      "crossover of top state at 1.30um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 137496 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[1 2]] at 1.89um\n",
      "crossover of states [[3 4]] at 1.77um\n",
      "crossover of top state at 1.69um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 140208 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 136372 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of top state at 1.41um\n",
      "crossover of states [[0 1]] at 1.38um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 140216 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#set the waists to be sligthly different and compensate by changing the power\n",
    "\n",
    "power2_factors = np.linspace(1+1e-3,1.15,4)\n",
    "\n",
    "for i, pow in enumerate(power2_factors):\n",
    "    #make sure the powers and waists are set correctly\n",
    "    trap[trap.power_tweezer1] = initial_power\n",
    "    trap[trap.power_tweezer2] = initial_power*pow\n",
    "\n",
    "    trap[trap.waist_tweezer1] = initial_waist\n",
    "    trap[trap.waist_tweezer2] = initial_waist*1.06\n",
    "    \n",
    "    #calculate eigenenergies\n",
    "    distances = np.linspace(initial_distance,0*si.um,500)\n",
    "    energies, states, potentials = loop_distances(trap, distances,n_levels=6,n_pot_steps=2000)\n",
    "\n",
    "    new_energies, new_states, new_potentials, index_top, index, swap_index = swapped_loop_distance(distances, energies, states, potentials)\n",
    "    energies_left, energies_right, states_left, states_right,deltaE_min, occ_num = find_ass_tweezer(new_energies,new_states)\n",
    "\n",
    "    # Create the figure with two subplots (1 row, 2 columns)\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "    # Plot the full range on ax1\n",
    "    ax1.plot(distances/si.um, energies_left[:, 0], color=\"red\", label=\"GS left tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_right[:, 0], color=\"blue\", label=\"GS right tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_left[:, 1:], color=\"red\", linestyle=\"dotted\", label=\"left tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_right[:, 1:], color=\"blue\", linestyle=\"dotted\", label=\"right tweezer\")\n",
    "\n",
    "    for i, ind in enumerate(index):\n",
    "        ax1.axvline(distances[ind]/si.um, color=\"k\", alpha=0.7, label=\"energy crossings\")\n",
    "\n",
    "    ax1.set_xlabel(\"Tweezer distance [um]\")\n",
    "    ax1.set_ylabel(\"Eigenenergy [a.u.]\")\n",
    "    ax1.set_title(fr\"$p_l={float(trap.subs(trap.power_tweezer1))/si.uW:.1f}$ uW, $p_r={float(trap.subs(trap.power_tweezer2))/si.uW:.1f}$ uW, $w_l={float(trap.subs(trap.waist_tweezer1))/si.um:.3f}$ um, $w_r={float(trap.subs(trap.waist_tweezer2))/si.um:.3f}$ um\")\n",
    "    ax1.grid()\n",
    "\n",
    "    # Plot the zoomed-in range (1 to 1.5 µm) on ax2\n",
    "    ax2.plot(distances/si.um, energies_left[:, 0], color=\"red\", label=\"GS left tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_right[:, 0], color=\"blue\", label=\"GS right tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_left[:, 1:], color=\"red\", linestyle=\"dotted\", label=\"left tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_right[:, 1:], color=\"blue\", linestyle=\"dotted\", label=\"right tweezer\")\n",
    "\n",
    "    for i, ind in enumerate(index):\n",
    "        ax2.axvline(distances[ind]/si.um, color=\"k\", alpha=0.7, label=\"energy crossings\")\n",
    "\n",
    "    ax2.set_xlabel(\"Tweezer distance [um]\")\n",
    "    ax2.set_ylabel(\"Eigenenergy [a.u.]\")\n",
    "    ax2.set_title(\"Zoomed-in Region\")\n",
    "    ax2.set_xlim(1, 1.7)  # Zooming in on the region between 1.0 and 1.5 µm\n",
    "    ax2.set_ylim(-4e-29,-2.5e-29)\n",
    "    ax2.grid()\n",
    "\n",
    "    # Remove duplicates in the legend (labels that appear in both subplots)\n",
    "    handles, labels = ax1.get_legend_handles_labels()\n",
    "    unique_labels = dict(zip(labels, handles))  # Remove duplicates\n",
    "    ax1.legend(unique_labels.values(), unique_labels.keys())\n",
    "\n",
    "    # Show the plots\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Seems like it works, but if the different in waists is too large:\n",
    "\n",
    "Then the energies of both tweezer fall of at different rates and they cross at all powers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[0 1]] at 1.87um\n",
      "crossover of states [[2 3]] at 1.75um\n",
      "crossover of states [[4 5]] at 1.68um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 142504 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[0 1]] at 1.44um\n",
      "crossover of top state at 1.44um\n",
      "crossover of states [[2 3]] at 1.42um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 142872 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[1 2]] at 1.62um\n",
      "crossover of states [[3 4]] at 1.58um\n",
      "crossover of top state at 1.55um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 145284 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[4 5]] at 2.39um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 142744 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#set the waists to be sligthly different and compensate by changing the power\n",
    "\n",
    "power2_factors = np.linspace(1.26,1.4,4)\n",
    "\n",
    "for i, pow in enumerate(power2_factors):\n",
    "    #make sure the powers and waists are set correctly\n",
    "    trap[trap.power_tweezer1] = initial_power\n",
    "    trap[trap.power_tweezer2] = initial_power*pow\n",
    "\n",
    "    trap[trap.waist_tweezer1] = initial_waist\n",
    "    trap[trap.waist_tweezer2] = initial_waist*1.12\n",
    "    \n",
    "    #calculate eigenenergies\n",
    "    distances = np.linspace(initial_distance,0*si.um,500)\n",
    "    energies, states, potentials = loop_distances(trap, distances,n_levels=6,n_pot_steps=2000)\n",
    "\n",
    "    new_energies, new_states, new_potentials, index_top, index, swap_index = swapped_loop_distance(distances, energies, states, potentials)\n",
    "    energies_left, energies_right, states_left, states_right,deltaE_min, occ_num = find_ass_tweezer(new_energies,new_states)\n",
    "\n",
    "    # Create the figure with two subplots (1 row, 2 columns)\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "    # Plot the full range on ax1\n",
    "    ax1.plot(distances/si.um, energies_left[:, 0], color=\"red\", label=\"GS left tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_right[:, 0], color=\"blue\", label=\"GS right tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_left[:, 1:], color=\"red\", linestyle=\"dotted\", label=\"left tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_right[:, 1:], color=\"blue\", linestyle=\"dotted\", label=\"right tweezer\")\n",
    "\n",
    "    for i, ind in enumerate(index):\n",
    "        ax1.axvline(distances[ind]/si.um, color=\"k\", alpha=0.7, label=\"energy crossings\")\n",
    "\n",
    "    ax1.set_xlabel(\"Tweezer distance [um]\")\n",
    "    ax1.set_ylabel(\"Eigenenergy [a.u.]\")\n",
    "    ax1.set_title(fr\"$p_l={float(trap.subs(trap.power_tweezer1))/si.uW:.1f}$ uW, $p_r={float(trap.subs(trap.power_tweezer2))/si.uW:.1f}$ uW, $w_l={float(trap.subs(trap.waist_tweezer1))/si.um:.3f}$ um, $w_r={float(trap.subs(trap.waist_tweezer2))/si.um:.3f}$ um\")\n",
    "    ax1.grid()\n",
    "\n",
    "    # Plot the zoomed-in range (1 to 1.5 µm) on ax2\n",
    "    ax2.plot(distances/si.um, energies_left[:, 0], color=\"red\", label=\"GS left tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_right[:, 0], color=\"blue\", label=\"GS right tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_left[:, 1:], color=\"red\", linestyle=\"dotted\", label=\"left tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_right[:, 1:], color=\"blue\", linestyle=\"dotted\", label=\"right tweezer\")\n",
    "\n",
    "    for i, ind in enumerate(index):\n",
    "        ax2.axvline(distances[ind]/si.um, color=\"k\", alpha=0.7, label=\"energy crossings\")\n",
    "\n",
    "    ax2.set_xlabel(\"Tweezer distance [um]\")\n",
    "    ax2.set_ylabel(\"Eigenenergy [a.u.]\")\n",
    "    ax2.set_title(\"Zoomed-in Region\")\n",
    "    ax2.set_xlim(1, 1.7)  # Zooming in on the region between 1.0 and 1.5 µm\n",
    "    ax2.set_ylim(-4e-29,-2.5e-29)\n",
    "    ax2.grid()\n",
    "\n",
    "    # Remove duplicates in the legend (labels that appear in both subplots)\n",
    "    handles, labels = ax1.get_legend_handles_labels()\n",
    "    unique_labels = dict(zip(labels, handles))  # Remove duplicates\n",
    "    ax1.legend(unique_labels.values(), unique_labels.keys())\n",
    "\n",
    "    # Show the plots\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
