{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bf2e25af",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "omega_r1: 2126.5Hz, omega_r2: 2126.5Hz\n",
      "omega_z1: 231.5Hz, omega_z2: 231.5Hz\n",
      "eta = 9.186412090177072\n"
     ]
    }
   ],
   "source": [
    "from helpers_merging import *\n",
    "\n",
    "initial_power = 50* si.uW\n",
    "initial_waist = 1.1*si.um\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,     #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\n",
    "\n",
    "\n",
    "omega_z1 = sp.sqrt(4*trap.power_tweezer1*trap.a/sp.pi**3/trap.m) * trap.wvl/trap.waist_tweezer1**3\n",
    "omega_z2 = sp.sqrt(4*trap.power_tweezer2*trap.a/sp.pi**3/trap.m) * trap.wvl/trap.waist_tweezer2**3\n",
    "\n",
    "omega_r1 = sp.sqrt(2*trap.power_tweezer1*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer1**2\n",
    "omega_r2 = sp.sqrt(2*trap.power_tweezer2*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer2**2\n",
    "\n",
    "print(f\"omega_r1: {float(trap.subs(omega_r1)/2/np.pi):.1f}Hz, omega_r2: {float(trap.subs(omega_r2)/2/np.pi):.1f}Hz\")\n",
    "print(f\"omega_z1: {float(trap.subs(omega_z1)/2/np.pi):.1f}Hz, omega_z2: {float(trap.subs(omega_z2)/2/np.pi):.1f}Hz\")\n",
    "print(f\"eta = {float(trap.subs(omega_r1/omega_z1))}\")\n",
    "\n",
    "P_spacing = np.pi*trap.waist_tweezer1**2*const.hbar*omega_z1/2/trap.a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2ab194ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "231.48079178308387\n",
      "2126.4779442798836\n",
      "0.5214842143140083\n"
     ]
    }
   ],
   "source": [
    "print(float(trap.subs(omega_z1/2/np.pi)))\n",
    "print(float(trap.subs(omega_r1/2/np.pi)))\n",
    "print(float(trap.subs(P_spacing/si.uW))) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9d019b7e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.5338079647322577e-31"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "const.hbar * float(trap.subs(omega_z1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d4de6778",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.40901920312263e-30"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "const.hbar * float(trap.subs(omega_r1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bc808bcb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\naeve\\FerDy-Repo\\merging_tweezer_code\\fermions\\helpers_merging.py:171: ComplexWarning: Casting complex values to real discards the imaginary part\n",
      "  states[i,:len(ener)] = state\n"
     ]
    }
   ],
   "source": [
    "#make sure the powers and waists are set correctly\n",
    "trap[trap.power_tweezer1] = initial_power\n",
    "trap[trap.power_tweezer2] = initial_power+0.4*P_spacing\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.1*si.um,500)\n",
    "energies, states, potentials = loop_distances_3D(trap, distances,n_levels=12)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5168b4ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"test_dist.npy\", distances)\n",
    "np.save(\"test_energies.npy\", energies)\n",
    "np.save(\"test_states.npy\", states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6098e84e",
   "metadata": {},
   "outputs": [],
   "source": [
    "distances = np.load(\"test_dist.npy\")\n",
    "energies = np.load(\"test_energies.npy\")\n",
    "states = np.load(\"test_states.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "522efebc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x2956afa6ab0>,\n",
       " <matplotlib.lines.Line2D at 0x295750fd5e0>,\n",
       " <matplotlib.lines.Line2D at 0x2956f565ee0>,\n",
       " <matplotlib.lines.Line2D at 0x2956eeb6870>,\n",
       " <matplotlib.lines.Line2D at 0x2956eeb51f0>,\n",
       " <matplotlib.lines.Line2D at 0x2956f428ec0>,\n",
       " <matplotlib.lines.Line2D at 0x2956f12c680>,\n",
       " <matplotlib.lines.Line2D at 0x2956f12f3e0>,\n",
       " <matplotlib.lines.Line2D at 0x2956f12fc50>,\n",
       " <matplotlib.lines.Line2D at 0x2956f12d3a0>,\n",
       " <matplotlib.lines.Line2D at 0x2957528ef00>,\n",
       " <matplotlib.lines.Line2D at 0x2957528ec90>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 40752 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(distances/si.um, energies/const.h)\n",
    "#plt.xlim(0.5,3)\n",
    "#plt.ylim(-60000,-40000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9b3fd32b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 101560 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "new_energies, new_states, new_potentials, index_top, index, swap_index = swapped_loop_distance_3D(distances, energies, states, potentials)\n",
    "energies_left, energies_right, states_left, states_right,deltaE_min, occ_num = find_ass_tweezer_3D(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(0.5, 1)  # Zooming in on the region between 1.0 and 1.5 µm\n",
    "ax2.set_ylim(-5e-29,-3e-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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ae549e5",
   "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": 5
}
