{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# I want to calculate $J,\\; U_s,\\, U_d$ for double tweezers to analyse how bosons can be merged into the GS. For this I need the wave functions of the 3D hamiltonian:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.widgets import Slider\n",
    "%matplotlib ipympl\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "from IPython.display import Math, display\n",
    "from matplotlib.axes import Axes\n",
    "from scipy import constants as const\n",
    "from scipy.integrate import quad\n",
    "from scipy.optimize import root_scalar\n",
    "from scipy.signal import argrelmax,argrelmin,find_peaks\n",
    "from tqdm import tqdm\n",
    "\n",
    "import fewfermions.analysis.units as si\n",
    "from fewfermions.simulate.traps.twod.trap import DoubleTweezer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_power = 50* si.uW\n",
    "initial_waist = 1.1*si.uW\n",
    "initial_distance = 2*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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_solutions(trap,n_levels,left_cutoff,right_cutoff,n_pot_steps=1000,display_plot=-1,state_mult=1e4,plot=True,ret_results=False):\n",
    "    \"\"\"Plot the potential and solutions for a given trap object\n",
    "        (diplay_plot=-1 for all, 0,...,n for individual ones and -2 for none)   \n",
    "    \"\"\"\n",
    "\n",
    "    # Solve the hamiltonian numerically in axial direction\n",
    "    energies, states, potential, coords = trap.nstationary_solution(\n",
    "        trap.z, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels\n",
    "    )\n",
    "\n",
    "    # States that are below the potential barrier\n",
    "    bound_states = energies < potential(left_cutoff)\n",
    "\n",
    "    if plot:\n",
    "        z_np = np.linspace(left_cutoff, right_cutoff, num=n_pot_steps)\n",
    "        ax: plt.Axes\n",
    "        fig, ax = plt.subplots(figsize=(5, 5))\n",
    "        ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz,color=\"cornflowerblue\" ,marker=\"None\")\n",
    "        ax.set_title(f\"{np.sum(bound_states)} bound solutions, tweezer distance: {trap.subs(trap.distance_tweezers)/si.um}um\")\n",
    "        ax.set_xlabel(r\"x ($\\mathrm{\\mu m}$)\")\n",
    "        ax.set_ylabel(r\"E / h (kHz)\")\n",
    "        abs_min = np.min(potential(z_np))\n",
    "        ax.fill_between(\n",
    "            z_np / si.um,\n",
    "            potential(z_np) / const.h / si.kHz,\n",
    "            abs_min / const.h / si.kHz,\n",
    "            alpha=0.5,\n",
    "            color=\"cornflowerblue\"\n",
    "        )\n",
    "\n",
    "        count = 0\n",
    "        for i, bound in enumerate(bound_states):\n",
    "            if not bound:\n",
    "                continue\n",
    "            energy = energies[i]\n",
    "            state = states[i]\n",
    "            ax.plot(\n",
    "                z_np / si.um,\n",
    "                np.where(\n",
    "                    (energy > potential(z_np)),\n",
    "                    energy / const.h / si.kHz,\n",
    "                    np.nan,\n",
    "                ),\n",
    "                c=\"k\",\n",
    "                lw=0.5,\n",
    "                marker=\"None\",\n",
    "            )\n",
    "            if display_plot == -1 or display_plot == count:\n",
    "                ax.plot(z_np/si.um, state *state_mult, marker=\"None\",label=f\"state {count}\")#, c=\"k\")\n",
    "            count += 1\n",
    "\n",
    "        plt.legend()\n",
    "\n",
    "\n",
    "    if ret_results:\n",
    "        return np.sum(bound_states), energies[bound_states], states[bound_states], potential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 59332 characters]",
      "text/plain": [
       "<Figure size 500x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#left_cutoff = -0.5*initial_distance-2*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "#right_cutoff = 0.5*initial_distance+2*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "left_cutoff = -2*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "right_cutoff = 2*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "\n",
    "trap[trap.distance_tweezers] = initial_distance\n",
    "trap[trap.power_tweezer2] = initial_power\n",
    "trap[trap.waist_tweezer2] = initial_waist\n",
    "plot_solutions(trap,4,left_cutoff,right_cutoff,display_plot=-1,state_mult=40,n_pot_steps=1000,plot=True,ret_results=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try how expensive 3D diagonalisation is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_pot_steps = [30,30,30]\n",
    "n_levels = 4\n",
    "\n",
    "left_cutoff = -0.5*initial_distance-2*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "right_cutoff = 0.5*initial_distance+2*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "back_cutoff = -2*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "front_cutoff = 2*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "bottom_cutoff = -2*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "top_cutoff = 2*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "\n",
    "extend = [(left_cutoff,right_cutoff),\n",
    "          (back_cutoff,front_cutoff),\n",
    "          (bottom_cutoff,top_cutoff)]\n",
    "\n",
    "extend = [(left_cutoff,right_cutoff),\n",
    "          (left_cutoff,right_cutoff),\n",
    "          (left_cutoff,right_cutoff)]\n",
    "\n",
    "# Solve the hamiltonian numerically in axial direction\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "        [trap.x,trap.y,trap.z], extend, n_pot_steps, k=n_levels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7f866bcc38e943919e4e2cac379fdff9",
       "version_major": 2,
       "version_minor": 0
      },
      "image/png": "[encoded data, 25632 characters]",
      "text/html": [
       "\n",
       "            <div style=\"display: inline-block;\">\n",
       "                <div class=\"jupyter-widgets widget-label\" style=\"text-align: center;\">\n",
       "                    Figure\n",
       "                </div>\n",
       "                <img src='data:image/png;base64,[encoded data, 25632 characters]' width=640.0/>\n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "Canvas(toolbar=None)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Initial slice index\n",
    "initial_slice = int(n_pot_steps[2] / 2)\n",
    "\n",
    "# Create figure and axis\n",
    "fig, ax = plt.subplots()\n",
    "plt.subplots_adjust(bottom=0.2)  # Leave space for the slider\n",
    "\n",
    "# Display initial image\n",
    "im = ax.imshow(states[0, :, :, initial_slice], extent=[*extend[0], *extend[1]], origin=\"lower\")\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\")\n",
    "plt.colorbar(im)\n",
    "\n",
    "# Create the slider axis\n",
    "ax_slider = plt.axes([0.2, 0.05, 0.65, 0.03])  # [left, bottom, width, height]\n",
    "\n",
    "# Create the slider\n",
    "slider = Slider(ax_slider, \"Slice\", 0, n_pot_steps[2] - 1, valinit=initial_slice, valfmt=\"%0.0f\")\n",
    "\n",
    "# Update function\n",
    "def update(val):\n",
    "    slice_idx = int(slider.val)  # Get current slider value\n",
    "    im.set_data(states[0, :, :, slice_idx])  # Update image\n",
    "    fig.canvas.draw_idle()  # Refresh plot\n",
    "\n",
    "# Connect slider to update function\n",
    "slider.on_changed(update)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{2 P_{t1} e^{\\frac{- 2 y^{2} - 2 \\left(\\frac{d_{t}}{2} + x\\right)^{2}}{W_{t1}^{2} \\left(1 + \\frac{\\lambda^{2} z^{2}}{\\pi^{2} W_{t1}^{4}}\\right)}}}{\\pi W_{t1}^{2} \\left(1 + \\frac{\\lambda^{2} z^{2}}{\\pi^{2} W_{t1}^{4}}\\right)} + \\frac{2 P_{t2} e^{\\frac{- 2 y^{2} - 2 \\left(- \\frac{d_{t}}{2} + x\\right)^{2}}{W_{t2}^{2} \\left(1 + \\frac{\\lambda^{2} z^{2}}{\\pi^{2} W_{t2}^{4}}\\right)}}}{\\pi W_{t2}^{2} \\left(1 + \\frac{\\lambda^{2} z^{2}}{\\pi^{2} W_{t2}^{4}}\\right)}$"
      ],
      "text/plain": [
       "2*P_t1*exp((-2*y**2 - 2*(d_t/2 + x)**2)/(W_t1**2*(1 + lambda**2*z**2/(pi**2*W_t1**4))))/(pi*W_t1**2*(1 + lambda**2*z**2/(pi**2*W_t1**4))) + 2*P_t2*exp((-2*y**2 - 2*(-d_t/2 + x)**2)/(W_t2**2*(1 + lambda**2*z**2/(pi**2*W_t2**4))))/(pi*W_t2**2*(1 + lambda**2*z**2/(pi**2*W_t2**4)))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trap.get_total_intensity(with_2dtrap=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19376 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.loglog(np.array([20,30,40]),np.array([1.1,13.4,101]))\n",
    "plt.xlim(10,100)\n",
    "plt.xlabel(\"n_pot_steps\")\n",
    "plt.ylabel(\"time\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22820 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 21276 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19988 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19824 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i in range(n_levels):\n",
    "    plt.imshow(states[i,:,:,int(n_pot_steps[2]/2)],extent=[*extend[0],*extend[1]],origin=\"lower\")\n",
    "    plt.xlabel(\"x\")\n",
    "    plt.ylabel(\"y\")\n",
    "    plt.colorbar()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 26016 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 24596 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 21936 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 22364 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i in range(n_levels):\n",
    "    plt.imshow(states[i,:,int(n_pot_steps[1]/4),:],extent=[*extend[0],*extend[2]],origin=\"lower\")\n",
    "    plt.xlabel(\"x\")\n",
    "    plt.ylabel(\"z\")\n",
    "    plt.colorbar()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23452 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 21764 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 19988 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 20032 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i in range(n_levels):\n",
    "    plt.imshow(states[i,int(n_pot_steps[0]/2),:,:],extent=[*extend[1],*extend[2]],origin=\"lower\")\n",
    "    plt.xlabel(\"y\")\n",
    "    plt.ylabel(\"z\")\n",
    "    plt.colorbar()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## This is quite expensive. Let's try individual components:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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
}
