{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Try to use cupy and use GPU for faster compute (It does not work on my PC, since I don't have a NVIDIA GPU):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "#import numpy as np\n",
    "import cupy 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",
    "from typing import Union\n",
    "\n",
    "import fewfermions.analysis.units as si\n",
    "from fewfermions.simulate.traps.twod.trap import DoubleTweezer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": "markdown",
   "metadata": {},
   "source": [
    "## Try with cupy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy as np\n",
    "\n",
    "n_pot_steps = [40,40,40]\n",
    "n_levels = 4\n",
    "\n",
    "left_cutoff = -0.5*initial_distance-2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "right_cutoff = 0.5*initial_distance+2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "back_cutoff = -2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "front_cutoff = 2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "bottom_cutoff = -2*max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "top_cutoff = 2*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",
    "#x3D,y3D,z3D = np.meshgrid(coords[trap.x],coords[trap.y],coords[trap.z])\n",
    "np.savez(\"data/test_3D.npz\",energies=energies, states=states, pot=potential(x3D,y3D,z3D), x=coords[trap.x],y=coords[trap.y],z=coords[trap.z])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try with regular numpy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "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": [
    "import numpy as np\n",
    "\n",
    "n_pot_steps = [40,40,40]\n",
    "n_levels = 4\n",
    "\n",
    "left_cutoff = -0.5*initial_distance-2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "right_cutoff = 0.5*initial_distance+2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "back_cutoff = -2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "front_cutoff = 2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
    "bottom_cutoff = -2*max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "top_cutoff = 2*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",
    "#x3D,y3D,z3D = np.meshgrid(coords[trap.x],coords[trap.y],coords[trap.z])\n",
    "np.savez(\"data/test_3D.npz\",energies=energies, states=states, pot=potential(x3D,y3D,z3D), x=coords[trap.x],y=coords[trap.y],z=coords[trap.z])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "ename": "CUDARuntimeError",
     "evalue": "cudaErrorInsufficientDriver: CUDA driver version is insufficient for CUDA runtime version",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mCUDARuntimeError\u001b[0m                          Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[12], line 2\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;66;03m# Generate spatial grid\u001b[39;00m\n\u001b[1;32m----> 2\u001b[0m x \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mlinspace(\u001b[38;5;241m*\u001b[39mextend[\u001b[38;5;241m0\u001b[39m], n_pot_steps[\u001b[38;5;241m0\u001b[39m])\n\u001b[0;32m      3\u001b[0m y \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mlinspace(\u001b[38;5;241m*\u001b[39mextend[\u001b[38;5;241m1\u001b[39m], n_pot_steps[\u001b[38;5;241m1\u001b[39m])\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",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\cupy\\_creation\\ranges.py:161\u001b[0m, in \u001b[0;36mlinspace\u001b[1;34m(start, stop, num, endpoint, retstep, dtype, axis)\u001b[0m\n\u001b[0;32m    159\u001b[0m scalar_stop \u001b[38;5;241m=\u001b[39m cupy\u001b[38;5;241m.\u001b[39misscalar(stop)\n\u001b[0;32m    160\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m scalar_start \u001b[38;5;129;01mand\u001b[39;00m scalar_stop:\n\u001b[1;32m--> 161\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _linspace_scalar(start, stop, num, endpoint, retstep, dtype)\n\u001b[0;32m    163\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m scalar_start:\n\u001b[0;32m    164\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28misinstance\u001b[39m(start, cupy\u001b[38;5;241m.\u001b[39mndarray) \u001b[38;5;129;01mand\u001b[39;00m start\u001b[38;5;241m.\u001b[39mdtype\u001b[38;5;241m.\u001b[39mkind \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m):\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\cupy\\_creation\\ranges.py:91\u001b[0m, in \u001b[0;36m_linspace_scalar\u001b[1;34m(start, stop, num, endpoint, retstep, dtype)\u001b[0m\n\u001b[0;32m     87\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m dtype \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m     88\u001b[0m     \u001b[38;5;66;03m# In actual implementation, only float is used\u001b[39;00m\n\u001b[0;32m     89\u001b[0m     dtype \u001b[38;5;241m=\u001b[39m dt\n\u001b[1;32m---> 91\u001b[0m ret \u001b[38;5;241m=\u001b[39m cupy\u001b[38;5;241m.\u001b[39mempty((num,), dtype\u001b[38;5;241m=\u001b[39mdt)\n\u001b[0;32m     92\u001b[0m div \u001b[38;5;241m=\u001b[39m (num \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m1\u001b[39m) \u001b[38;5;28;01mif\u001b[39;00m endpoint \u001b[38;5;28;01melse\u001b[39;00m num\n\u001b[0;32m     93\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m div \u001b[38;5;241m<\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m:\n",
      "File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\cupy\\_creation\\basic.py:32\u001b[0m, in \u001b[0;36mempty\u001b[1;34m(shape, dtype, order)\u001b[0m\n\u001b[0;32m     13\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mempty\u001b[39m(\n\u001b[0;32m     14\u001b[0m         shape: _ShapeLike,\n\u001b[0;32m     15\u001b[0m         dtype: DTypeLike \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mfloat\u001b[39m,\n\u001b[0;32m     16\u001b[0m         order: _OrderCF \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mC\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[0;32m     17\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m NDArray[Any]:\n\u001b[0;32m     18\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Returns an array without initializing the elements.\u001b[39;00m\n\u001b[0;32m     19\u001b[0m \n\u001b[0;32m     20\u001b[0m \u001b[38;5;124;03m    Args:\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     30\u001b[0m \n\u001b[0;32m     31\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m---> 32\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m cupy\u001b[38;5;241m.\u001b[39mndarray(shape, dtype, order\u001b[38;5;241m=\u001b[39morder)\n",
      "File \u001b[1;32mcupy\\\\_core\\\\core.pyx:151\u001b[0m, in \u001b[0;36mcupy._core.core.ndarray.__new__\u001b[1;34m()\u001b[0m\n",
      "File \u001b[1;32mcupy\\\\_core\\\\core.pyx:239\u001b[0m, in \u001b[0;36mcupy._core.core._ndarray_base._init\u001b[1;34m()\u001b[0m\n",
      "File \u001b[1;32mcupy\\\\cuda\\\\memory.pyx:738\u001b[0m, in \u001b[0;36mcupy.cuda.memory.alloc\u001b[1;34m()\u001b[0m\n",
      "File \u001b[1;32mcupy\\\\cuda\\\\memory.pyx:1424\u001b[0m, in \u001b[0;36mcupy.cuda.memory.MemoryPool.malloc\u001b[1;34m()\u001b[0m\n",
      "File \u001b[1;32mcupy\\\\cuda\\\\memory.pyx:1444\u001b[0m, in \u001b[0;36mcupy.cuda.memory.MemoryPool.malloc\u001b[1;34m()\u001b[0m\n",
      "File \u001b[1;32mcupy\\\\cuda\\\\device.pyx:40\u001b[0m, in \u001b[0;36mcupy.cuda.device.get_device_id\u001b[1;34m()\u001b[0m\n",
      "File \u001b[1;32mcupy_backends\\\\cuda\\\\api\\\\runtime.pyx:202\u001b[0m, in \u001b[0;36mcupy_backends.cuda.api.runtime.getDevice\u001b[1;34m()\u001b[0m\n",
      "File \u001b[1;32mcupy_backends\\\\cuda\\\\api\\\\runtime.pyx:146\u001b[0m, in \u001b[0;36mcupy_backends.cuda.api.runtime.check_status\u001b[1;34m()\u001b[0m\n",
      "\u001b[1;31mCUDARuntimeError\u001b[0m: cudaErrorInsufficientDriver: CUDA driver version is insufficient for CUDA runtime version"
     ]
    }
   ],
   "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)  # Ensure correct indexing\n",
    "\n",
    "# Compute potential (Replace with actual function)\n",
    "pot = potential(x3D, y3D, 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/test{state_number}.gif\", writer=\"pillow\", fps=frames/5)  # 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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
