{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyse what happens with different powers\n",
    "\n",
    "## Now we have to deal with levels crossing over"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy.functions.special.error_functions import erf\n",
    "from fewfermions.simulate.numerics import discrete_hamiltonian\n",
    "from scipy import sparse\n",
    "\n",
    "from helpers_merging import *\n",
    "\n",
    "initial_power = 50* si.uW\n",
    "initial_waist = 1.1*si.uW\n",
    "initial_distance = 4*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*4,     #transfer tweezer\n",
    "    waist_tweezer1 = initial_waist,     #stationary\n",
    "    waist_tweezer2 = initial_waist*2,     #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",
    "left_cutoff = -0.5*initial_distance-3*float(trap.subs(trap.waist_tweezer1))\n",
    "right_cutoff = 0.5*initial_distance+3*float(trap.subs(trap.waist_tweezer2))\n",
    "\n",
    "x, y, z = trap.x, trap.y, trap.z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38816 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trap[trap.distance_tweezers] = 2.1*si.um#initial_distance\n",
    "plot_solutions(trap,7,left_cutoff,right_cutoff,display_plot=5,state_mult=40,n_pot_steps=2000,plot=True,ret_results=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 89528 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "check_crossover = False\n",
    "n_levels=6\n",
    "n_pot_steps=2000\n",
    "crossover_thresh = 0.3\n",
    "x_np = np.linspace(left_cutoff,right_cutoff,n_pot_steps)\n",
    "\n",
    "distances = np.linspace(initial_distance,0*si.um,500)\n",
    "energies = np.zeros((len(distances),n_levels))\n",
    "states = np.zeros((len(distances),n_levels,n_pot_steps))\n",
    "potential = np.zeros((len(distances),n_pot_steps))\n",
    "\n",
    "#inititalise arrays for swapping in case of crossover\n",
    "swapped_index = np.arange(n_levels)\n",
    "#diagonalise hamiltonian for all distances\n",
    "for i, dist in enumerate(distances):\n",
    "    trap[trap.distance_tweezers] = dist\n",
    "    num, ener, state, pot = plot_solutions(trap,n_levels,np.min(x_np),np.max(x_np),display_plot=0,state_mult=450,n_pot_steps=n_pot_steps,plot=False,ret_results=True)\n",
    "    \n",
    "    if check_crossover:\n",
    "        # check if there was a crossing of energies\n",
    "        if i==0:\n",
    "            prob_left = np.sum(state[:,:int(n_pot_steps/2)]**2,axis=1)\n",
    "        prob_left_new = np.sum(state[:,:int(n_pot_steps/2)]**2,axis=1)\n",
    "        #create list of indices to swap\n",
    "        if np.any(np.abs(prob_left - prob_left_new) > crossover_thresh):\n",
    "            swap_index = np.where(np.abs(prob_left - prob_left_new)[swapped_index] > crossover_thresh)[0]\n",
    "            print(f\"crossover of states {swap_index} at {dist/si.um:.2f}um\")\n",
    "            \"\"\"\n",
    "            if np.sum(np.abs(prob_left - prob_left_new) > crossover_thresh)==1:\n",
    "                #deal with crossover with state that is not calculated\n",
    "                swapped_index = np.append(swapped_index,swap_index[0])\n",
    "                energies = np.concatenate((energies,np.full((len(distances), 1), np.nan)),axis=1)\n",
    "                swapped_index = np.delete(swapped_index,swap_index)\n",
    "                #raise Exception('Crossover with state \"n_levels\". Add that level as well.')\"\"\"\n",
    "            if np.sum(np.abs(prob_left - prob_left_new) > crossover_thresh)==2:\n",
    "                swapped_index[swap_index] = swapped_index[np.roll(swap_index,1)]\n",
    "\n",
    "            print(swapped_index)\n",
    "                \n",
    "        #swap entries for energes and states\n",
    "        ener = ener[swapped_index]\n",
    "        state = state[swapped_index]\n",
    "        #update prob_left for next loop\n",
    "        prob_left = np.copy(prob_left_new)\n",
    "\n",
    "    # fill in arrays\n",
    "    energies[i,:len(ener)] = ener\n",
    "    states[i,:len(ener)] = state\n",
    "    potential[i] = pot(x_np)\n",
    "\n",
    "\n",
    "#plot eigenenergies\n",
    "plt.axvline(float(trap.subs(trap.waist_tweezer1))/si.um, color=\"black\",label=\"tweezer waist\")\n",
    "for i in range(n_levels):\n",
    "    plt.plot(distances/si.um,energies[:,i],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.1,2.7)\n",
    "#plt.ylim(-4.0e-29,-1.5e-29)\n",
    "\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 66504 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(n_levels):\n",
    "    #differentiate states that are coming from right and left tweezer\n",
    "    if np.sum(states[0,i,:int(n_pot_steps/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(0.8,2.3)\n",
    "#plt.ylim(-1.2e-28,-0.7e-28)\n",
    "\n",
    "plt.grid()\n",
    "#plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(500,)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sum(states[:,:,:int(n_pot_steps/2)]**2,axis=2)[:,i].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2000,)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_np.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal import find_peaks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  1 118 121 181 333]\n",
      "[  1 118 121 203 206 319]\n",
      "[  1 172 174 178 203 206 238 241 299]\n",
      "[  1 112 114 172 174 220 222 238 241 262 264 278]\n",
      "[  1 112 114 197 199 220 222 248 250 262 264 279 282 376]\n",
      "[  1 165 168 175 197 199 232 234 248 250 267 270 279 282 297]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49344 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(n_levels):\n",
    "    #plt.plot(energies[:,i])\n",
    "    #plt.plot(distances/si.um,energies[:,i],label=f\"state {i}\")\n",
    "    #plt.plot(np.sum(states[:,:,:int(n_pot_steps/2)]**2,axis=2)[:,i],label=f\"state {i}\")\n",
    "    #plt.plot(np.gradient(np.sum(states[:,:,:int(n_pot_steps/2)]**2,axis=2)[:,i]),label=f\"state {i}\")\n",
    "    #grad = np.abs(np.diff(np.sum(states[:,:,:int(n_pot_steps/2)]**2,axis=2)[:,i]))\n",
    "    grad = np.abs(np.gradient(np.gradient(np.gradient(energies[:,i]))))\n",
    "    plt.plot(grad)\n",
    "    index = find_peaks(grad)[0]\n",
    "    print(index)\n",
    "    for j, ind in enumerate(index):\n",
    "        plt.axvline(ind,alpha=0.8)\n",
    "    #plt.plot(distances/si.um,np.sum(np.gradient(states,axis=0),axis=2)[:,i])\n",
    "    #plt.plot(distances/si.um,np.gradient(np.gradient(np.gradient(energies[:,i]))))\n",
    "\n",
    "\n",
    "\n",
    "plt.xlabel(\"tweezer distance [um]\")\n",
    "plt.ylabel(\"probability to be on left half [a.u.]\")\n",
    "#plt.xlim(220,330)\n",
    "#plt.ylim(0,1e-3)\n",
    "\n",
    "\n",
    "plt.grid()\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[120]\n",
      "[]\n",
      "[]\n",
      "[]\n",
      "[]\n",
      "[204]\n",
      "[]\n",
      "[]\n",
      "[]\n",
      "[173 240]\n",
      "[]\n",
      "[]\n",
      "[113 221 263]\n",
      "[]\n",
      "[198 249 280]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 62648 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot eigenenergies\n",
    "for i in range(n_levels):\n",
    "    plt.plot(energies[:,i])\n",
    "    for j in range(i+1,n_levels):\n",
    "        diff = np.abs(energies[:,i] - energies[:,j])\n",
    "\n",
    "        mask = np.where(diff < 1e-1*np.max(diff),diff,1e-1*np.max(diff))\n",
    "        #plt.plot(mask)\n",
    "        index = find_peaks(-mask)[0]\n",
    "        print(index)\n",
    "        for k, ind in enumerate(index):\n",
    "            plt.axvline(ind,alpha=0.8)\n",
    "        "
   ]
  },
  {
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x1970c81f920>]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 21340 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(energies[:,-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x19708f3a000>]"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 21796 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "arr = np.abs(np.gradient(np.gradient(energies[:,-1])))\n",
    "index = find_peaks(arr,width=(1,4))[0]\n",
    "for i, ind in enumerate(index):\n",
    "    plt.axvline(ind,color=\"k\",alpha=0.5)\n",
    "plt.plot(energies[:,-1])"
   ]
  },
  {
   "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
}
