{
 "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": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\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": "markdown",
   "metadata": {},
   "source": [
    "## Try how expensive 3D diagonalisation is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_pot_steps = [50,50,50]\n",
    "n_levels = 8\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",
    "\n",
    "# Solve the hamiltonian numerically\n",
    "energies, states, potential, coords = trap.nstationary_solution(\n",
    "        [trap.x,trap.y,trap.z], extend, n_pot_steps, k=n_levels)\n",
    "\n",
    "x = coords[trap.x]\n",
    "y = coords[trap.y]\n",
    "z = coords[trap.z]\n",
    "x3D,y3D,z3D = np.meshgrid(coords[trap.x],coords[trap.y],coords[trap.z],indexing=\"ij\")\n",
    "pot = potential(x3D,y3D,z3D)\n",
    "np.savez(\"data/test_3D.npz\",energies=energies, states=states, pot=pot, x=coords[trap.x],y=coords[trap.y],z=coords[trap.z])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\naeve\\AppData\\Local\\Temp\\ipykernel_18084\\1341043005.py:34: MatplotlibDeprecationWarning: The collections attribute was deprecated in Matplotlib 3.8 and will be removed in 3.10.\n",
      "  for c in contour.collections:\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 88832 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Generate spatial grid\n",
    "x = np.linspace(*extend[0], n_pot_steps[0])\n",
    "y = np.linspace(*extend[1], n_pot_steps[1])\n",
    "z = np.linspace(*extend[2], n_pot_steps[2])\n",
    "\n",
    "x3D, y3D, z3D = np.meshgrid(x, y, z,indexing=\"ij\")  # Ensure correct indexing\n",
    "\n",
    "# Compute potential (Replace with actual function)\n",
    "pot = potential(x3D, y3D, z3D)\n",
    "\n",
    "state_number = 3\n",
    "\n",
    "# Create figure and axis\n",
    "fig, ax = plt.subplots()\n",
    "im = ax.imshow(states[state_number, :, :, 0], extent=[*extend[1], *extend[0]], origin=\"lower\",\n",
    "               vmin=np.min(states[state_number]), vmax=np.max(states[state_number]))\n",
    "\n",
    "plt.xlabel(\"y\")\n",
    "plt.ylabel(\"x\")\n",
    "plt.colorbar(im)\n",
    "\n",
    "# Initialize contour as None before defining it globally\n",
    "contour = None\n",
    "\n",
    "# Animation update function\n",
    "def update(frame):\n",
    "    global contour  # Ensure we're modifying the global variable\n",
    "\n",
    "    im.set_data(states[state_number, :, :, frame])  # Update image data\n",
    "    ax.set_title(f\"z={z[frame]/si.um:.3f}um\")  # Update title\n",
    "\n",
    "    # Remove old contours if they exist\n",
    "    if contour is not None:\n",
    "        for c in contour.collections:\n",
    "            c.remove()\n",
    "\n",
    "    # Redraw contour plot\n",
    "    contour = ax.contour(pot[:, :, frame], levels=10, colors='white', linewidths=0.7, extent=[*extend[1], *extend[0]])\n",
    "\n",
    "# Create the first contour plot after defining update()\n",
    "contour = ax.contour(pot[:, :, 0], levels=10, colors='white', linewidths=0.7, extent=[*extend[1], *extend[0]])\n",
    "\n",
    "# Create animation\n",
    "frames = n_pot_steps[2]  # Number of slices\n",
    "ani = animation.FuncAnimation(fig, update, frames=frames, interval=100)\n",
    "\n",
    "ani.save(f\"animations/state{state_number}.gif\", writer=\"pillow\", fps=10)  # Save as GIF\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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, 19760 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.loglog(np.array([20**3,30**3,40**3]),np.array([1.1,13.4,101]))\n",
    "plt.xlim(5*10**3,10**5)\n",
    "plt.xlabel(\"n_pot_steps\")\n",
    "plt.ylabel(\"time\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## This is quite expensive. Let's try individual components."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make animations with this script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\naeve\\AppData\\Local\\Temp\\ipykernel_17932\\735623296.py:42: MatplotlibDeprecationWarning: The collections attribute was deprecated in Matplotlib 3.8 and will be removed in 3.10.\n",
      "  for c in contour.collections:\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 82988 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "data = np.load(\"test_3D_50_50_50.npz\")\n",
    "\n",
    "for key in data.keys():\n",
    "    globals()[key] = data[key]\n",
    "\n",
    "x3D, y3D, z3D = np.meshgrid(x, y, z)  # Ensure correct indexing\n",
    "\n",
    "extend = [(x.min(),x.max()),\n",
    "          (y.min(),y.max()),\n",
    "          (z.min(),z.max())]\n",
    "\n",
    "n_pot_steps = [len(x),len(y),len(z)]\n",
    "\n",
    "state_number = 0\n",
    "\n",
    "# Create figure and axis\n",
    "fig, ax = plt.subplots()\n",
    "im = ax.imshow(states[state_number, :, :, 0], extent=[*extend[0], *extend[1]], origin=\"lower\",\n",
    "               vmin=np.min(states[state_number]), vmax=np.max(states[state_number]))\n",
    "\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\")\n",
    "plt.colorbar(im)\n",
    "\n",
    "# Initialize contour as None before defining it globally\n",
    "contour = None\n",
    "\n",
    "# Animation update function\n",
    "def update(frame):\n",
    "    global contour  # Ensure we're modifying the global variable\n",
    "\n",
    "    im.set_data(states[state_number, :, :, frame])  # Update image data\n",
    "    ax.set_title(f\"z={z[frame]/si.um:.3f}um\")  # Update title\n",
    "\n",
    "    # Remove old contours if they exist\n",
    "    if contour is not None:\n",
    "        for c in contour.collections:\n",
    "            c.remove()\n",
    "\n",
    "    # Redraw contour plot\n",
    "    contour = ax.contour(pot[:, :, frame], levels=10, colors='white', linewidths=0.7, extent=[*extend[0], *extend[1]])\n",
    "\n",
    "# Create the first contour plot after defining update()\n",
    "contour = ax.contour(pot[:, :, 0], levels=10, colors='white', linewidths=0.7, extent=[*extend[0], *extend[1]])\n",
    "\n",
    "# Create animation\n",
    "frames = n_pot_steps[2]  # Number of slices\n",
    "ani = animation.FuncAnimation(fig, update, frames=frames, interval=100)\n",
    "\n",
    "ani.save(f\"animations/state{state_number}.gif\", writer=\"pillow\", fps=10)  # Save as GIF\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
