{
 "cells": [
  {
   "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",
    "from scipy import constants as const\n",
    "\n",
    "import sys\n",
    "sys.path.append('C:/Users/naeve/FerDy-Repo/clean_diag/backend')\n",
    "\n",
    "import trap_units as si\n",
    "from 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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'data/test_3D.npz'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[3], line 25\u001b[0m\n\u001b[0;32m     23\u001b[0m x3D,y3D,z3D \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mmeshgrid(coords[trap\u001b[38;5;241m.\u001b[39mx],coords[trap\u001b[38;5;241m.\u001b[39my],coords[trap\u001b[38;5;241m.\u001b[39mz],indexing\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mij\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     24\u001b[0m pot \u001b[38;5;241m=\u001b[39m potential(x3D,y3D,z3D)\n\u001b[1;32m---> 25\u001b[0m np\u001b[38;5;241m.\u001b[39msavez(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdata/test_3D.npz\u001b[39m\u001b[38;5;124m\"\u001b[39m,energies\u001b[38;5;241m=\u001b[39menergies, states\u001b[38;5;241m=\u001b[39mstates, pot\u001b[38;5;241m=\u001b[39mpot, x\u001b[38;5;241m=\u001b[39mcoords[trap\u001b[38;5;241m.\u001b[39mx],y\u001b[38;5;241m=\u001b[39mcoords[trap\u001b[38;5;241m.\u001b[39my],z\u001b[38;5;241m=\u001b[39mcoords[trap\u001b[38;5;241m.\u001b[39mz])\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\numpy\\lib\\npyio.py:639\u001b[0m, in \u001b[0;36msavez\u001b[1;34m(file, *args, **kwds)\u001b[0m\n\u001b[0;32m    555\u001b[0m \u001b[38;5;129m@array_function_dispatch\u001b[39m(_savez_dispatcher)\n\u001b[0;32m    556\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21msavez\u001b[39m(file, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwds):\n\u001b[0;32m    557\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Save several arrays into a single file in uncompressed ``.npz`` format.\u001b[39;00m\n\u001b[0;32m    558\u001b[0m \n\u001b[0;32m    559\u001b[0m \u001b[38;5;124;03m    Provide arrays as keyword arguments to store them under the\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    637\u001b[0m \n\u001b[0;32m    638\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m--> 639\u001b[0m     _savez(file, args, kwds, \u001b[38;5;28;01mFalse\u001b[39;00m)\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\numpy\\lib\\npyio.py:736\u001b[0m, in \u001b[0;36m_savez\u001b[1;34m(file, args, kwds, compress, allow_pickle, pickle_kwargs)\u001b[0m\n\u001b[0;32m    733\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    734\u001b[0m     compression \u001b[38;5;241m=\u001b[39m zipfile\u001b[38;5;241m.\u001b[39mZIP_STORED\n\u001b[1;32m--> 736\u001b[0m zipf \u001b[38;5;241m=\u001b[39m zipfile_factory(file, mode\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mw\u001b[39m\u001b[38;5;124m\"\u001b[39m, compression\u001b[38;5;241m=\u001b[39mcompression)\n\u001b[0;32m    738\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m key, val \u001b[38;5;129;01min\u001b[39;00m namedict\u001b[38;5;241m.\u001b[39mitems():\n\u001b[0;32m    739\u001b[0m     fname \u001b[38;5;241m=\u001b[39m key \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m.npy\u001b[39m\u001b[38;5;124m'\u001b[39m\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\numpy\\lib\\npyio.py:103\u001b[0m, in \u001b[0;36mzipfile_factory\u001b[1;34m(file, *args, **kwargs)\u001b[0m\n\u001b[0;32m    101\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mzipfile\u001b[39;00m\n\u001b[0;32m    102\u001b[0m kwargs[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mallowZip64\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m--> 103\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m zipfile\u001b[38;5;241m.\u001b[39mZipFile(file, \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\\zipfile\\__init__.py:1331\u001b[0m, in \u001b[0;36mZipFile.__init__\u001b[1;34m(self, file, mode, compression, allowZip64, compresslevel, strict_timestamps, metadata_encoding)\u001b[0m\n\u001b[0;32m   1329\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m   1330\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m-> 1331\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfp \u001b[38;5;241m=\u001b[39m io\u001b[38;5;241m.\u001b[39mopen(file, filemode)\n\u001b[0;32m   1332\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mOSError\u001b[39;00m:\n\u001b[0;32m   1333\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m filemode \u001b[38;5;129;01min\u001b[39;00m modeDict:\n",
      "\u001b[1;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'data/test_3D.npz'"
     ]
    }
   ],
   "source": [
    "n_pot_steps = [30,30,30]\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\naeve\\AppData\\Local\\Temp\\ipykernel_13844\\1690377931.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, 83276 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 = 2\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\"state{state_number}.gif\", writer=\"pillow\", fps=10)  # Save as GIF\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "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": 2
}
