{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "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*1.04583875,     #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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61520 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "left_cutoff = -0.5*initial_distance-3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "right_cutoff = 0.5*initial_distance+3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "\n",
    "trap[trap.distance_tweezers] = initial_distance\n",
    "plot_solutions(trap,7,left_cutoff,right_cutoff,display_plot=-1,state_mult=40,n_pot_steps=1000,plot=True,ret_results=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of states [[2 3]] at 1.52um\n",
      "crossover of states [[0 1]] at 1.41um\n",
      "crossover of top state at 1.30um\n",
      "crossover of top state at 1.30um\n",
      "crossover of top state at 1.32um\n",
      "crossover of top state at 1.33um\n",
      "crossover of states [[1 2]] at 1.32um\n",
      "crossover of states [[3 4]] at 1.41um\n",
      "crossover of states [[1 2]] at 1.40um\n",
      "crossover of states [[1 2]] at 1.53um\n",
      "crossover of states [[3 4]] at 1.52um\n",
      "crossover of states [[1 2]] at 1.75um\n",
      "crossover of states [[3 4]] at 1.70um\n",
      "crossover of states [[2 3]] at 1.36um\n",
      "crossover of states [[2 3]] at 1.41um\n",
      "crossover of states [[2 3]] at 1.47um\n",
      "crossover of top state at 1.46um\n",
      "crossover of states [[2 3]] at 1.54um\n",
      "crossover of top state at 1.52um\n",
      "crossover of states [[2 3]] at 1.63um\n",
      "crossover of top state at 1.61um\n",
      "crossover of states [[2 3]] at 1.75um\n",
      "crossover of top state at 1.71um\n",
      "crossover of states [[2 3]] at 1.99um\n",
      "crossover of top state at 1.88um\n",
      "crossover of top state at 1.11um\n",
      "crossover of top state at 1.29um\n",
      "crossover of top state at 1.11um\n",
      "crossover of top state at 1.29um\n",
      "crossover of top state at 1.10um\n",
      "crossover of top state at 1.30um\n",
      "crossover of states [[3 4]] at 1.39um\n",
      "crossover of states [[3 4]] at 1.42um\n",
      "crossover of states [[3 4]] at 1.46um\n",
      "crossover of states [[3 4]] at 1.50um\n"
     ]
    }
   ],
   "source": [
    "power2_factors = np.linspace(1+1e-5,1.2,75)\n",
    "deltaEs = np.zeros_like(power2_factors)\n",
    "occ_numbers = np.zeros((len(power2_factors),2))\n",
    "\n",
    "for i, pow in enumerate(power2_factors):\n",
    "    trap[trap.power_tweezer2] = initial_power*pow\n",
    "\n",
    "    distances = np.linspace(initial_distance,0*si.um,500)\n",
    "    energies, states, potentials = loop_distances(trap, distances,n_levels=5,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",
    "    deltaEs[i] = deltaE_min\n",
    "    occ_numbers[i] = occ_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x202e0321100>]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 33548 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(initial_power*power2_factors/si.uW, deltaEs)\n",
    "#plt.ylim(0,1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x202e1532e70>]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 14252 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(initial_power*power2_factors/si.uW, occ_numbers[:,1],\"o\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 52304 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(distances/si.um,energies_left,color=\"blue\",label=\"left tweezer\")\n",
    "plt.plot(distances/si.um,energies_right,color=\"red\",label=\"right tweezer\")\n",
    "\n",
    "for i, ind in enumerate(index):\n",
    "    plt.axvline(distances[ind]/si.um,color=\"k\",alpha=0.7,label=\"energy crossings\")\n",
    "\n",
    "#plt.xlim(1.2,1.4)\n",
    "#plt.ylim(-1.75e-29,-1.3e-29)\n",
    "#plt.axhline(-1.55e-29)\n",
    "#plt.axhline(-1.55e-29+delta_E_min)\n",
    "\n",
    "handles, labels = plt.gca().get_legend_handles_labels()\n",
    "unique_labels = dict(zip(labels, handles))  # Remove duplicates\n",
    "plt.legend(unique_labels.values(), unique_labels.keys())\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "distances = np.linspace(initial_distance,0*si.um,500)\n",
    "energies, states, potentials = loop_distances(trap, distances,n_levels=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 60048 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot eigenenergies\n",
    "plt.axvline(float(trap.subs(trap.waist_tweezer1))/si.um, color=\"black\",label=\"tweezer waist\")\n",
    "for i in range(len(energies[0])):\n",
    "    #differentiate states that are coming from right and left tweezer\n",
    "    if np.sum(states[0,i,:int(states.shape[2]/2)]**2) < 0.5:\n",
    "        plt.plot(distances/si.um,energies[:,i],color=\"red\",label=f\"state {i}\")\n",
    "    else:\n",
    "        plt.plot(distances/si.um,energies[:,i],color=\"blue\",label=f\"state {i}\")\n",
    "\n",
    "plt.title(\"Eigenenergies of lowest states at different tweezer distances\")\n",
    "plt.xlabel(\"tweezer distance [um]\")\n",
    "plt.ylabel(\"energy [a.u.]\")\n",
    "#plt.xlim(1.8,2.3)\n",
    "#plt.ylim(-1.7e-29,-1.2e-29)\n",
    "\n",
    "plt.grid()\n",
    "#plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "crossover of top state at 1.30um\n"
     ]
    }
   ],
   "source": [
    "new_energies, new_states, new_potentials, index_top, index, swap_index = swapped_loop_distance(distances, energies, states, potentials)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.4799128624867176e-31\n",
      "[0 1]\n"
     ]
    }
   ],
   "source": [
    "#seperate energies and states corresponding to the initial tweezer they belonged to\n",
    "mask_right = np.where(np.sum(new_states[0,:,:int(new_states.shape[2]/2)]**2,axis=1) < 0.5,True, False)\n",
    "mask_left = np.logical_not(mask_right)\n",
    "\n",
    "energies_right = new_energies[:,mask_right]\n",
    "energies_left = new_energies[:,mask_left]\n",
    "\n",
    "states_right = new_states[:,mask_right,:]\n",
    "states_left = new_states[:,mask_left,:]\n",
    "\n",
    "#identify GS of shallower tweezer and calculated minimal energy difference to other states\n",
    "if mask_left[np.nanargmin(new_energies[0])]:\n",
    "    energy_GS = energies_right[:,0]\n",
    "    state_GS = states_right[:,0]\n",
    "    #smallest energy difference\n",
    "    delta_E_min = np.nanmin(np.abs(energy_GS[:,np.newaxis] - energies_left))\n",
    "    #final occupation if two GS are involved\n",
    "    occ_numbers = np.array([0,np.searchsorted(energies_left[-1],energy_GS[-1])])\n",
    "else:\n",
    "    energy_GS = energies_left[:,0]\n",
    "    state_GS = states_left[:,0]\n",
    "    #smallest energy difference\n",
    "    delta_E_min = np.nanmin(np.abs(energy_GS[:,np.newaxis] - energies_right))\n",
    "    #final occupation if two GS are involved\n",
    "    occ_numbers = np.array([0,np.searchsorted(energies_right[-1],energy_GS[-1])])\n",
    "\n",
    "print(delta_E_min)\n",
    "print(occ_numbers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48920 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(distances/si.um,energies_left,color=\"blue\",label=\"left tweezer\")\n",
    "plt.plot(distances/si.um,energies_right,color=\"red\",label=\"right tweezer\")\n",
    "\n",
    "for i, ind in enumerate(index):\n",
    "    plt.axvline(distances[ind]/si.um,color=\"k\",alpha=0.7,label=\"energy crossings\")\n",
    "\n",
    "#plt.xlim(1.2,1.4)\n",
    "#plt.ylim(-1.75e-29,-1.3e-29)\n",
    "#plt.axhline(-1.55e-29)\n",
    "#plt.axhline(-1.55e-29+delta_E_min)\n",
    "\n",
    "handles, labels = plt.gca().get_legend_handles_labels()\n",
    "unique_labels = dict(zip(labels, handles))  # Remove duplicates\n",
    "plt.legend(unique_labels.values(), unique_labels.keys())\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
