{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Since 3D simulation is numerically very expensive, let's try to solve 1D potentials and compute wave function as product:"
   ]
  },
  {
   "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 decomposing all 3D spatial directions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_pot_steps = [200,200,200]\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",
    "\n",
    "# Solve the hamiltonian numerically in all directions\n",
    "energies_x, states_x, potential_x, coords_x = trap.nstationary_solution(\n",
    "        trap.x, extend[0], n_pot_steps[0], k=n_levels)\n",
    "\n",
    "energies_y, states_y, potential_y, coords_y = trap.nstationary_solution(\n",
    "        trap.y, extend[1], n_pot_steps[1], k=n_levels)\n",
    "\n",
    "energies_z, states_z, potential_z, coords_z = trap.nstationary_solution(\n",
    "        trap.z, extend[2], n_pot_steps[2], k=n_levels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\naeve\\AppData\\Local\\Temp\\ipykernel_3752\\3500672162.py:46: MatplotlibDeprecationWarning: The collections attribute was deprecated in Matplotlib 3.8 and will be removed in 3.10.\n",
      "  for c in contour.collections:\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[4], line 59\u001b[0m\n\u001b[0;32m     56\u001b[0m frames \u001b[38;5;241m=\u001b[39m n_pot_steps[\u001b[38;5;241m2\u001b[39m]  \u001b[38;5;66;03m# Number of slices\u001b[39;00m\n\u001b[0;32m     57\u001b[0m ani \u001b[38;5;241m=\u001b[39m animation\u001b[38;5;241m.\u001b[39mFuncAnimation(fig, update, frames\u001b[38;5;241m=\u001b[39mframes, interval\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m100\u001b[39m)\n\u001b[1;32m---> 59\u001b[0m ani\u001b[38;5;241m.\u001b[39msave(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstate\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstate_number\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m_decomp_all.gif\u001b[39m\u001b[38;5;124m\"\u001b[39m, writer\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpillow\u001b[39m\u001b[38;5;124m\"\u001b[39m, fps\u001b[38;5;241m=\u001b[39mframes\u001b[38;5;241m/\u001b[39m\u001b[38;5;241m5\u001b[39m)  \u001b[38;5;66;03m# Save as GIF\u001b[39;00m\n\u001b[0;32m     61\u001b[0m plt\u001b[38;5;241m.\u001b[39mshow()\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\animation.py:1085\u001b[0m, in \u001b[0;36mAnimation.save\u001b[1;34m(self, filename, writer, fps, dpi, codec, bitrate, extra_args, metadata, extra_anim, savefig_kwargs, progress_callback)\u001b[0m\n\u001b[0;32m   1082\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m data \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\u001b[38;5;241m*\u001b[39m[a\u001b[38;5;241m.\u001b[39mnew_saved_frame_seq() \u001b[38;5;28;01mfor\u001b[39;00m a \u001b[38;5;129;01min\u001b[39;00m all_anim]):\n\u001b[0;32m   1083\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m anim, d \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(all_anim, data):\n\u001b[0;32m   1084\u001b[0m         \u001b[38;5;66;03m# TODO: See if turning off blit is really necessary\u001b[39;00m\n\u001b[1;32m-> 1085\u001b[0m         anim\u001b[38;5;241m.\u001b[39m_draw_next_frame(d, blit\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[0;32m   1086\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m progress_callback \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m   1087\u001b[0m             progress_callback(frame_number, total_frames)\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\animation.py:1121\u001b[0m, in \u001b[0;36mAnimation._draw_next_frame\u001b[1;34m(self, framedata, blit)\u001b[0m\n\u001b[0;32m   1119\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pre_draw(framedata, blit)\n\u001b[0;32m   1120\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_draw_frame(framedata)\n\u001b[1;32m-> 1121\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_post_draw(framedata, blit)\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\animation.py:1146\u001b[0m, in \u001b[0;36mAnimation._post_draw\u001b[1;34m(self, framedata, blit)\u001b[0m\n\u001b[0;32m   1144\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_blit_draw(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_drawn_artists)\n\u001b[0;32m   1145\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1146\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_fig\u001b[38;5;241m.\u001b[39mcanvas\u001b[38;5;241m.\u001b[39mdraw_idle()\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\backend_bases.py:1905\u001b[0m, in \u001b[0;36mFigureCanvasBase.draw_idle\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1903\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_is_idle_drawing:\n\u001b[0;32m   1904\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_idle_draw_cntx():\n\u001b[1;32m-> 1905\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdraw(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\backends\\backend_agg.py:387\u001b[0m, in \u001b[0;36mFigureCanvasAgg.draw\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    384\u001b[0m \u001b[38;5;66;03m# Acquire a lock on the shared font cache.\u001b[39;00m\n\u001b[0;32m    385\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtoolbar\u001b[38;5;241m.\u001b[39m_wait_cursor_for_draw_cm() \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtoolbar\n\u001b[0;32m    386\u001b[0m       \u001b[38;5;28;01melse\u001b[39;00m nullcontext()):\n\u001b[1;32m--> 387\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfigure\u001b[38;5;241m.\u001b[39mdraw(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrenderer)\n\u001b[0;32m    388\u001b[0m     \u001b[38;5;66;03m# A GUI class may be need to update a window using this draw, so\u001b[39;00m\n\u001b[0;32m    389\u001b[0m     \u001b[38;5;66;03m# don't forget to call the superclass.\u001b[39;00m\n\u001b[0;32m    390\u001b[0m     \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39mdraw()\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\artist.py:95\u001b[0m, in \u001b[0;36m_finalize_rasterization.<locals>.draw_wrapper\u001b[1;34m(artist, renderer, *args, **kwargs)\u001b[0m\n\u001b[0;32m     93\u001b[0m \u001b[38;5;129m@wraps\u001b[39m(draw)\n\u001b[0;32m     94\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdraw_wrapper\u001b[39m(artist, renderer, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m---> 95\u001b[0m     result \u001b[38;5;241m=\u001b[39m draw(artist, renderer, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m     96\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m renderer\u001b[38;5;241m.\u001b[39m_rasterizing:\n\u001b[0;32m     97\u001b[0m         renderer\u001b[38;5;241m.\u001b[39mstop_rasterizing()\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\artist.py:72\u001b[0m, in \u001b[0;36mallow_rasterization.<locals>.draw_wrapper\u001b[1;34m(artist, renderer)\u001b[0m\n\u001b[0;32m     69\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m artist\u001b[38;5;241m.\u001b[39mget_agg_filter() \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m     70\u001b[0m         renderer\u001b[38;5;241m.\u001b[39mstart_filter()\n\u001b[1;32m---> 72\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m draw(artist, renderer)\n\u001b[0;32m     73\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[0;32m     74\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m artist\u001b[38;5;241m.\u001b[39mget_agg_filter() \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\figure.py:3162\u001b[0m, in \u001b[0;36mFigure.draw\u001b[1;34m(self, renderer)\u001b[0m\n\u001b[0;32m   3159\u001b[0m             \u001b[38;5;66;03m# ValueError can occur when resizing a window.\u001b[39;00m\n\u001b[0;32m   3161\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpatch\u001b[38;5;241m.\u001b[39mdraw(renderer)\n\u001b[1;32m-> 3162\u001b[0m     mimage\u001b[38;5;241m.\u001b[39m_draw_list_compositing_images(\n\u001b[0;32m   3163\u001b[0m         renderer, \u001b[38;5;28mself\u001b[39m, artists, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msuppressComposite)\n\u001b[0;32m   3165\u001b[0m     renderer\u001b[38;5;241m.\u001b[39mclose_group(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mfigure\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m   3166\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\image.py:132\u001b[0m, in \u001b[0;36m_draw_list_compositing_images\u001b[1;34m(renderer, parent, artists, suppress_composite)\u001b[0m\n\u001b[0;32m    130\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m not_composite \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m has_images:\n\u001b[0;32m    131\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m a \u001b[38;5;129;01min\u001b[39;00m artists:\n\u001b[1;32m--> 132\u001b[0m         a\u001b[38;5;241m.\u001b[39mdraw(renderer)\n\u001b[0;32m    133\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    134\u001b[0m     \u001b[38;5;66;03m# Composite any adjacent images together\u001b[39;00m\n\u001b[0;32m    135\u001b[0m     image_group \u001b[38;5;241m=\u001b[39m []\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\artist.py:72\u001b[0m, in \u001b[0;36mallow_rasterization.<locals>.draw_wrapper\u001b[1;34m(artist, renderer)\u001b[0m\n\u001b[0;32m     69\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m artist\u001b[38;5;241m.\u001b[39mget_agg_filter() \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m     70\u001b[0m         renderer\u001b[38;5;241m.\u001b[39mstart_filter()\n\u001b[1;32m---> 72\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m draw(artist, renderer)\n\u001b[0;32m     73\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[0;32m     74\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m artist\u001b[38;5;241m.\u001b[39mget_agg_filter() \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\axes\\_base.py:3101\u001b[0m, in \u001b[0;36m_AxesBase.draw\u001b[1;34m(self, renderer)\u001b[0m\n\u001b[0;32m   3098\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m spine \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mspines\u001b[38;5;241m.\u001b[39mvalues():\n\u001b[0;32m   3099\u001b[0m         artists\u001b[38;5;241m.\u001b[39mremove(spine)\n\u001b[1;32m-> 3101\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_update_title_position(renderer)\n\u001b[0;32m   3103\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maxison:\n\u001b[0;32m   3104\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m _axis \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_axis_map\u001b[38;5;241m.\u001b[39mvalues():\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\axes\\_base.py:3045\u001b[0m, in \u001b[0;36m_AxesBase._update_title_position\u001b[1;34m(self, renderer)\u001b[0m\n\u001b[0;32m   3043\u001b[0m top \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmax\u001b[39m(top, bb\u001b[38;5;241m.\u001b[39mymax)\n\u001b[0;32m   3044\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m title\u001b[38;5;241m.\u001b[39mget_text():\n\u001b[1;32m-> 3045\u001b[0m     ax\u001b[38;5;241m.\u001b[39myaxis\u001b[38;5;241m.\u001b[39mget_tightbbox(renderer)  \u001b[38;5;66;03m# update offsetText\u001b[39;00m\n\u001b[0;32m   3046\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m ax\u001b[38;5;241m.\u001b[39myaxis\u001b[38;5;241m.\u001b[39moffsetText\u001b[38;5;241m.\u001b[39mget_text():\n\u001b[0;32m   3047\u001b[0m         bb \u001b[38;5;241m=\u001b[39m ax\u001b[38;5;241m.\u001b[39myaxis\u001b[38;5;241m.\u001b[39moffsetText\u001b[38;5;241m.\u001b[39mget_tightbbox(renderer)\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\axis.py:1372\u001b[0m, in \u001b[0;36mAxis.get_tightbbox\u001b[1;34m(self, renderer, for_layout_only)\u001b[0m\n\u001b[0;32m   1369\u001b[0m     renderer \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfigure\u001b[38;5;241m.\u001b[39m_get_renderer()\n\u001b[0;32m   1370\u001b[0m ticks_to_draw \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_update_ticks()\n\u001b[1;32m-> 1372\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_update_label_position(renderer)\n\u001b[0;32m   1374\u001b[0m \u001b[38;5;66;03m# go back to just this axis's tick labels\u001b[39;00m\n\u001b[0;32m   1375\u001b[0m tlb1, tlb2 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_ticklabel_bboxes(ticks_to_draw, renderer)\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\axis.py:2654\u001b[0m, in \u001b[0;36mYAxis._update_label_position\u001b[1;34m(self, renderer)\u001b[0m\n\u001b[0;32m   2650\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[0;32m   2652\u001b[0m \u001b[38;5;66;03m# get bounding boxes for this axis and any siblings\u001b[39;00m\n\u001b[0;32m   2653\u001b[0m \u001b[38;5;66;03m# that have been set by `fig.align_ylabels()`\u001b[39;00m\n\u001b[1;32m-> 2654\u001b[0m bboxes, bboxes2 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_tick_boxes_siblings(renderer\u001b[38;5;241m=\u001b[39mrenderer)\n\u001b[0;32m   2655\u001b[0m x, y \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlabel\u001b[38;5;241m.\u001b[39mget_position()\n\u001b[0;32m   2656\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlabel_position \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mleft\u001b[39m\u001b[38;5;124m'\u001b[39m:\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\axis.py:2206\u001b[0m, in \u001b[0;36mAxis._get_tick_boxes_siblings\u001b[1;34m(self, renderer)\u001b[0m\n\u001b[0;32m   2204\u001b[0m axis \u001b[38;5;241m=\u001b[39m ax\u001b[38;5;241m.\u001b[39m_axis_map[name]\n\u001b[0;32m   2205\u001b[0m ticks_to_draw \u001b[38;5;241m=\u001b[39m axis\u001b[38;5;241m.\u001b[39m_update_ticks()\n\u001b[1;32m-> 2206\u001b[0m tlb, tlb2 \u001b[38;5;241m=\u001b[39m axis\u001b[38;5;241m.\u001b[39m_get_ticklabel_bboxes(ticks_to_draw, renderer)\n\u001b[0;32m   2207\u001b[0m bboxes\u001b[38;5;241m.\u001b[39mextend(tlb)\n\u001b[0;32m   2208\u001b[0m bboxes2\u001b[38;5;241m.\u001b[39mextend(tlb2)\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\axis.py:1351\u001b[0m, in \u001b[0;36mAxis._get_ticklabel_bboxes\u001b[1;34m(self, ticks, renderer)\u001b[0m\n\u001b[0;32m   1349\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m renderer \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m   1350\u001b[0m     renderer \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfigure\u001b[38;5;241m.\u001b[39m_get_renderer()\n\u001b[1;32m-> 1351\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m ([tick\u001b[38;5;241m.\u001b[39mlabel1\u001b[38;5;241m.\u001b[39mget_window_extent(renderer)\n\u001b[0;32m   1352\u001b[0m          \u001b[38;5;28;01mfor\u001b[39;00m tick \u001b[38;5;129;01min\u001b[39;00m ticks \u001b[38;5;28;01mif\u001b[39;00m tick\u001b[38;5;241m.\u001b[39mlabel1\u001b[38;5;241m.\u001b[39mget_visible()],\n\u001b[0;32m   1353\u001b[0m         [tick\u001b[38;5;241m.\u001b[39mlabel2\u001b[38;5;241m.\u001b[39mget_window_extent(renderer)\n\u001b[0;32m   1354\u001b[0m          \u001b[38;5;28;01mfor\u001b[39;00m tick \u001b[38;5;129;01min\u001b[39;00m ticks \u001b[38;5;28;01mif\u001b[39;00m tick\u001b[38;5;241m.\u001b[39mlabel2\u001b[38;5;241m.\u001b[39mget_visible()])\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\text.py:961\u001b[0m, in \u001b[0;36mText.get_window_extent\u001b[1;34m(self, renderer, dpi)\u001b[0m\n\u001b[0;32m    959\u001b[0m bbox, info, descent \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_layout(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_renderer)\n\u001b[0;32m    960\u001b[0m x, y \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_unitless_position()\n\u001b[1;32m--> 961\u001b[0m x, y \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_transform()\u001b[38;5;241m.\u001b[39mtransform((x, y))\n\u001b[0;32m    962\u001b[0m bbox \u001b[38;5;241m=\u001b[39m bbox\u001b[38;5;241m.\u001b[39mtranslated(x, y)\n\u001b[0;32m    963\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m bbox\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\transforms.py:1505\u001b[0m, in \u001b[0;36mTransform.transform\u001b[1;34m(self, values)\u001b[0m\n\u001b[0;32m   1502\u001b[0m values \u001b[38;5;241m=\u001b[39m values\u001b[38;5;241m.\u001b[39mreshape((\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minput_dims))\n\u001b[0;32m   1504\u001b[0m \u001b[38;5;66;03m# Transform the values\u001b[39;00m\n\u001b[1;32m-> 1505\u001b[0m res \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransform_affine(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransform_non_affine(values))\n\u001b[0;32m   1507\u001b[0m \u001b[38;5;66;03m# Convert the result back to the shape of the input values.\u001b[39;00m\n\u001b[0;32m   1508\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m ndim \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\_api\\deprecation.py:300\u001b[0m, in \u001b[0;36mrename_parameter.<locals>.wrapper\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m    295\u001b[0m     warn_deprecated(\n\u001b[0;32m    296\u001b[0m         since, message\u001b[38;5;241m=\u001b[39m\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThe \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mold\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m parameter of \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfunc\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m() \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m    297\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhas been renamed \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mnew\u001b[38;5;132;01m!r}\u001b[39;00m\u001b[38;5;124m since Matplotlib \u001b[39m\u001b[38;5;132;01m{\u001b[39;00msince\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m; support \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m    298\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfor the old name will be dropped %(removal)s.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m    299\u001b[0m     kwargs[new] \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mpop(old)\n\u001b[1;32m--> 300\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m func(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\transforms.py:2428\u001b[0m, in \u001b[0;36mCompositeGenericTransform.transform_affine\u001b[1;34m(self, values)\u001b[0m\n\u001b[0;32m   2425\u001b[0m \u001b[38;5;129m@_api\u001b[39m\u001b[38;5;241m.\u001b[39mrename_parameter(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m3.8\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpoints\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvalues\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m   2426\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mtransform_affine\u001b[39m(\u001b[38;5;28mself\u001b[39m, values):\n\u001b[0;32m   2427\u001b[0m     \u001b[38;5;66;03m# docstring inherited\u001b[39;00m\n\u001b[1;32m-> 2428\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mget_affine()\u001b[38;5;241m.\u001b[39mtransform(values)\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\transforms.py:2455\u001b[0m, in \u001b[0;36mCompositeGenericTransform.get_affine\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m   2453\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_b\u001b[38;5;241m.\u001b[39mget_affine()\n\u001b[0;32m   2454\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 2455\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m Affine2D(np\u001b[38;5;241m.\u001b[39mdot(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_b\u001b[38;5;241m.\u001b[39mget_affine()\u001b[38;5;241m.\u001b[39mget_matrix(),\n\u001b[0;32m   2456\u001b[0m                            \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_a\u001b[38;5;241m.\u001b[39mget_affine()\u001b[38;5;241m.\u001b[39mget_matrix()))\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\matplotlib\\transforms.py:1913\u001b[0m, in \u001b[0;36mAffine2D.__init__\u001b[1;34m(self, matrix, **kwargs)\u001b[0m\n\u001b[0;32m   1910\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m matrix \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m   1911\u001b[0m     \u001b[38;5;66;03m# A bit faster than np.identity(3).\u001b[39;00m\n\u001b[0;32m   1912\u001b[0m     matrix \u001b[38;5;241m=\u001b[39m IdentityTransform\u001b[38;5;241m.\u001b[39m_mtx\n\u001b[1;32m-> 1913\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_mtx \u001b[38;5;241m=\u001b[39m matrix\u001b[38;5;241m.\u001b[39mcopy()\n\u001b[0;32m   1914\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_invalid \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 79556 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "states_x3D = np.zeros((n_levels,*n_pot_steps))\n",
    "states_y3D = np.zeros((n_levels,*n_pot_steps))\n",
    "states_z3D = np.zeros((n_levels,*n_pot_steps))\n",
    "\n",
    "for i in range(n_levels):\n",
    "    states_x3D[i], states_y3D[i], states_z3D[i] = np.meshgrid(states_x[i],states_y[i],states_z[i])\n",
    "\n",
    "states = states_x3D*states_y3D*states_z3D\n",
    "\n",
    "# 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",
    "# Compute potential\n",
    "pot_x = potential_x(x)\n",
    "pot_y = potential_y(y)\n",
    "pot_z = potential_z(z)\n",
    "\n",
    "pot_x3D, pot_y3D, pot_z3D = np.meshgrid(pot_x,pot_y,pot_z)\n",
    "pot = pot_x3D*pot_y3D*pot_z3D\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}_decomp_all.gif\", writer=\"pillow\", fps=frames/5)  # Save as GIF\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This seems to change the potential significantly and there are two minima in z direction.\n",
    "\n",
    "## Try to decompose only the y direction (WIP):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_pot_steps = [200,200,200]\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",
    "\n",
    "# Solve the hamiltonian numerically in all directions\n",
    "energies_xz, states_xz, potential_xz, coords_xz = trap.nstationary_solution(\n",
    "        [trap.x,trap.z], extend[::2], n_pot_steps[::2], k=n_levels)\n",
    "\n",
    "energies_y, states_y, potential_y, coords_y = trap.nstationary_solution(\n",
    "        trap.y, extend[1], n_pot_steps[1], k=n_levels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x21f0565daf0>"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 31548 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(states_xz[0],origin=\"lower\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x2238d9d2240>]"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 22036 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(states_y[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x2241c135430>"
      ]
     },
     "execution_count": 105,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 31548 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "states = np.expand_dims(states_xz, axis=2)\n",
    "states = np.repeat(states,n_pot_steps[2],axis=2)\n",
    "\n",
    "plt.imshow(states[0,:,100,:],origin=\"lower\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {},
   "outputs": [],
   "source": [
    "states = np.expand_dims(states_xz, axis=2)\n",
    "states = np.repeat(states,n_pot_steps[2],axis=2)\n",
    "\n",
    "for i in range(n_levels):\n",
    "    for j in range(n_pot_steps[2]):\n",
    "        states[i,:,j,:] *= states_y[i,j]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x2243954b6b0>"
      ]
     },
     "execution_count": 112,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 46340 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(states[0,100,:,:],origin=\"lower\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30304 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",
    "\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",
    "\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",
    "\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}_decomp_y.gif\", writer=\"pillow\", fps=frames/5)  # Save as GIF\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "_lambdifygenerated() missing 1 required positional argument: 'z'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[26], line 7\u001b[0m\n\u001b[0;32m      4\u001b[0m z \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mlinspace(\u001b[38;5;241m*\u001b[39mextend[\u001b[38;5;241m2\u001b[39m], n_pot_steps[\u001b[38;5;241m2\u001b[39m])\n\u001b[0;32m      6\u001b[0m \u001b[38;5;66;03m# Compute potential\u001b[39;00m\n\u001b[1;32m----> 7\u001b[0m pot_x \u001b[38;5;241m=\u001b[39m potential_x(x)\n\u001b[0;32m      8\u001b[0m pot_y \u001b[38;5;241m=\u001b[39m potential_y(y)\n\u001b[0;32m      9\u001b[0m pot_z \u001b[38;5;241m=\u001b[39m potential_z(z)\n",
      "\u001b[1;31mTypeError\u001b[0m: _lambdifygenerated() missing 1 required positional argument: 'z'"
     ]
    }
   ],
   "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",
    "# Compute potential\n",
    "pot_x = potential_x(x)\n",
    "pot_y = potential_y(y)\n",
    "pot_z = potential_z(z)\n",
    "\n",
    "pot_x3D, pot_y3D, pot_z3D = np.meshgrid(pot_x,pot_y,pot_z)\n",
    "pot = pot_x3D*pot_y3D*pot_z3D\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}_decomp_y.gif\", writer=\"pillow\", fps=frames/5)  # 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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
