{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "from helpers_merging import *\n",
    "\n",
    "initial_power = 50* 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": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 67464 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] = 0.55*initial_distance\n",
    "trap[trap.power_tweezer2] = initial_power*2.5\n",
    "trap[trap.waist_tweezer2] = initial_waist*1.5\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": 89,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_deltaE_occ(power_tweezer1, power_tweezer2, waist_tweezer1, waist_tweezer2,\n",
    "               n_levels=10,n_pot_steps=2000, initial_distance=4*si.um):\n",
    "    \"\"\"\n",
    "    Returns the minimal energy gap and occupation numbers of the groundstate when the two tweezers merge.\n",
    "    \"\"\"\n",
    "    trap[trap.power_tweezer1] = power_tweezer1\n",
    "    trap[trap.power_tweezer2] = power_tweezer2\n",
    "    trap[trap.waist_tweezer1] = waist_tweezer1\n",
    "    trap[trap.waist_tweezer2] = waist_tweezer2\n",
    "\n",
    "    distances = np.linspace(initial_distance,0*si.um,500) #always has to have smallest distance at last index\n",
    "    energies, states, potentials = loop_distances(trap, distances,n_levels=n_levels,n_pot_steps=n_pot_steps)\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",
    "    print(index)\n",
    "\n",
    "    # Create the figure with two subplots (1 row, 2 columns)\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 4))\n",
    "\n",
    "    # Plot the full range on ax1\n",
    "    \"\"\"\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",
    "    ax1.plot(distances/si.um, energies_left[:], color=\"red\", label=\"left tweezer\")\n",
    "    ax1.plot(distances/si.um, energies_right[:], color=\"blue\", 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",
    "    \"\"\"\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",
    "    ax2.plot(distances/si.um, energies_left[:], color=\"red\", label=\"left tweezer\")\n",
    "    ax2.plot(distances/si.um, energies_right[:], color=\"blue\", 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.2, 2)  # Zooming in on the region between 1.0 and 1.5 µm\n",
    "    ax2.set_ylim(-2e-29,-0.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",
    "\n",
    "    return deltaE_min, occ_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[210 251]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 108320 characters]",
      "text/plain": [
       "<Figure size 800x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "factor_waist2 = 1.1\n",
    "\n",
    "trap[trap.power_tweezer1] = initial_power\n",
    "trap[trap.waist_tweezer1] = initial_waist\n",
    "trap[trap.waist_tweezer2] = initial_waist*factor_waist2\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",
    "#offset required for 1 level spacing of the left tweezer\n",
    "power2_offset = sp.pi*trap.waist_tweezer2**2*const.hbar*omega_r1/2/trap.a\n",
    "\n",
    "#factor to match trap depth and offset by half the level spacing\n",
    "power2_factor = factor_waist2**2 - 0.4*float(trap.subs(power2_offset))/initial_power\n",
    "\n",
    "deltaEs, occ_numbers = get_deltaE_occ(initial_power,initial_power*power2_factor,initial_waist,initial_waist*factor_waist2,n_levels=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "--------------------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/10 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[390 392]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 158552 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 1/10 [00:36<05:24, 36.02s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[400]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 158376 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 2/10 [01:12<04:50, 36.29s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[343]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 160060 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███       | 3/10 [01:47<04:10, 35.84s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[362 372 378 381]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 166868 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 4/10 [02:23<03:33, 35.62s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[389 391]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 166760 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|█████     | 5/10 [02:58<02:58, 35.66s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[401]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 165980 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 6/10 [03:34<02:23, 35.78s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 161580 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 70%|███████   | 7/10 [04:11<01:47, 35.96s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  1   7  10  12  15  18  22  27  29  31  36  40  43  46  48  52  60  62\n",
      "  65  68  72  76  79  83  87  89  91  93  96  99 104 107 111 114 119 123\n",
      " 127 130 135 138 142 145 362 372 377]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 174328 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|████████  | 8/10 [04:49<01:13, 36.71s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[385 388 391]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 170452 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|█████████ | 9/10 [05:27<00:37, 37.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[401]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 171692 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [06:06<00:00, 36.67s/it]\n"
     ]
    }
   ],
   "source": [
    "factor_waist2 = 1.06\n",
    "\n",
    "power2_factors = np.linspace(1+1e-5,1.15,10)\n",
    "deltaEs = np.zeros_like(power2_factors)\n",
    "occ_numbers = np.zeros((len(power2_factors),2))\n",
    "\n",
    "for i, pow in enumerate(tqdm(power2_factors)):\n",
    "    deltaEs[i], occ_numbers[i] = get_deltaE_occ(initial_power,initial_power*pow,initial_waist,initial_waist*factor_waist2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x1dcb64b8530>]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 29600 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(initial_power*power2_factors/si.uW, deltaEs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x1dcb4aba060>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 12120 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\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import minimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 160172 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 160188 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 159792 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 160164 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 160208 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 159692 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "factor_waist2 = 1.\n",
    "\n",
    "def objective(x):\n",
    "    \"\"\"function to minimise\"\"\"\n",
    "    deltaE_min, occ_num = get_deltaE_occ(initial_power,x[0],initial_waist,initial_waist*factor_waist2)\n",
    "\n",
    "    return -deltaE_min *1e32\n",
    "\n",
    "def constraint(x):\n",
    "    \"\"\"contraint for minimisation\"\"\"\n",
    "    deltaE_min, occ_num = get_deltaE_occ(initial_power,x[0],initial_waist,initial_waist*factor_waist2)\n",
    "\n",
    "    return occ_num[1] -1\n",
    "\n",
    "x0 = initial_power*1.01\n",
    "bou = [(initial_power, 2*initial_power)]\n",
    "con = {\"type\":\"eq\", \"fun\":constraint}\n",
    "\n",
    "result = minimize(objective, x0, bounds=bou, constraints=con)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimal solution: [0.000101]\n",
      "Objective value at optimal solution: -20.879059081175534\n",
      "Success: False\n",
      "Message: Singular matrix C in LSQ subproblem\n"
     ]
    }
   ],
   "source": [
    "print(\"Optimal solution:\", result.x)\n",
    "print(\"Objective value at optimal solution:\", result.fun)\n",
    "print(\"Success:\", result.success)\n",
    "print(\"Message:\", result.message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([101.])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result.x/si.uW"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "'float' object is not subscriptable",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[13], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m objective(\u001b[38;5;241m100.99\u001b[39m)\n",
      "Cell \u001b[1;32mIn[10], line 5\u001b[0m, in \u001b[0;36mobjective\u001b[1;34m(x)\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mobjective\u001b[39m(x):\n\u001b[0;32m      4\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"function to minimise\"\"\"\u001b[39;00m\n\u001b[1;32m----> 5\u001b[0m     deltaE_min, occ_num \u001b[38;5;241m=\u001b[39m get_deltaE_occ(initial_power,x[\u001b[38;5;241m0\u001b[39m],initial_waist,initial_waist\u001b[38;5;241m*\u001b[39mfactor_waist2)\n\u001b[0;32m      7\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;241m-\u001b[39mdeltaE_min \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m1e32\u001b[39m\n",
      "\u001b[1;31mTypeError\u001b[0m: 'float' object is not subscriptable"
     ]
    }
   ],
   "source": [
    "objective(100.99)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-5.6781927848\n"
     ]
    }
   ],
   "source": [
    "print(f\"{objective(102):.10f}\")"
   ]
  },
  {
   "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
}
