{
 "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",
    "from scipy import sparse\n",
    "\n",
    "#add relative path to backend\n",
    "import sys\n",
    "sys.path.append('../../clean_diag/backend')\n",
    "\n",
    "import trap_units as si\n",
    "from twod_trap import DoubleTweezer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_power = 50* si.uW\n",
    "initial_waist = 1.1*si.uW\n",
    "initial_distance = 1.7*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": [
    "## Sparse matrix hamiltonian:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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 = -1.2*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "top_cutoff = 1.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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_number = 0\n",
    "\n",
    "\n",
    "plt.plot(x,states[state_number,:,int(len(y)/2),int(len(z)/2)])\n",
    "plt.plot(x,1e27*pot[:,int(len(y)/2),int(len(z)/2)])\n",
    "plt.show()\n",
    "\n",
    "plt.plot(y,states[state_number,int(len(x)/3),:,int(len(z)/2)])\n",
    "plt.plot(y,1e27*pot[int(len(x)/3),:,int(len(z)/2)])\n",
    "plt.show()\n",
    "\n",
    "plt.plot(z,states[state_number,int(len(x)/3),int(len(y)/2),:])\n",
    "plt.plot(z,1e27*pot[int(len(x)/3),int(len(y)/2),:])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Matrix free with LinearOperator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "files saved with ...._2025-03-18_13-22-26\n"
     ]
    }
   ],
   "source": [
    "n_pot_steps = [100,100,100]\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 = -1*np.max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
    "top_cutoff = 1*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",
    "        method=\"matrix_free\", export=True)\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10.0, 100.0)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 17808 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.loglog(np.array([30,75,100]),np.array([10/60,12,38])) #time in min\n",
    "plt.xlim(1e1,1e2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 37552 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 25204 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 31140 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "state_number = 3\n",
    "\n",
    "plt.plot(x,states[state_number,:,int(len(y)/2),int(len(z)/2)].real, label=\"real part\")\n",
    "plt.plot(x,states[state_number,:,int(len(y)/2),int(len(z)/2)].imag, label=\"imaginary part\")\n",
    "#plt.plot(x,pot[:,int(len(y)/2),int(len(z)/2)], label=\"potential (rescaled)\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "plt.plot(y,states[state_number,int(len(x)/3),:,int(len(z)/2)].real, label=\"real part\")\n",
    "plt.plot(y,states[state_number,int(len(x)/3),:,int(len(z)/2)].imag, label=\"imaginary part\")\n",
    "#plt.plot(y,1e27*pot[int(len(x)/3),:,int(len(z)/2)], label=\"potential (rescaled)\")\n",
    "plt.show()\n",
    "\n",
    "plt.plot(z,states[state_number,int(len(x)/3),int(len(y)/2),:].real, label=\"real part\")\n",
    "plt.plot(z,states[state_number,int(len(x)/3),int(len(y)/2),:].imag, label=\"imaginary part\")\n",
    "#plt.plot(z,1e27*pot[int(len(x)/3),int(len(y)/2),:], label=\"potential (rescaled)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 105872 characters]",
      "text/plain": [
       "<Figure size 2400x500 with 16 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 8\n",
    "\n",
    "# Create a 1xn_states grid of subplots (1 row, n_states columns)\n",
    "fig, axes = plt.subplots(1, n_states, figsize=(18/6*n_states, 5))  # Adjust the size as needed\n",
    "\n",
    "# Loop over the state numbers from 0 to 5\n",
    "for state_number, ax in zip(range(n_states), axes):\n",
    "    # Slice through the y-direction and rotate the x-z plane (by swapping x and z)\n",
    "    im = ax.imshow(np.abs(states[state_number, :, int(len(y)/2), :].T)**2, extent=[*extend[0], *extend[2]], origin=\"lower\",\n",
    "                   vmin=np.min(np.abs(states[state_number])**2), vmax=np.max(np.abs(states[state_number])**2))\n",
    "\n",
    "    # Set labels for each subplot\n",
    "    ax.set_xlabel(\"x\")\n",
    "    ax.set_ylabel(\"z\")\n",
    "    ax.set_title(f\"State {state_number}\")\n",
    "    fig.colorbar(im, ax=ax)  \n",
    "\n",
    "# Adjust layout for better spacing\n",
    "plt.tight_layout()\n",
    "\n",
    "# Show the plots\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#state to plot\n",
    "state_number = 1\n",
    "\n",
    "# path to store animation\n",
    "anim_path = f\"animations/state{state_number}.gif\"\n",
    "\"\"\"\n",
    "# results file containing the energies and states\n",
    "data = np.load(\"data/results_300.npz\")\n",
    "for key in data.keys():\n",
    "    globals()[key] = data[key]\n",
    "\n",
    "# parameters file for coordinates, potential, etc.\n",
    "data = np.load(\"data/parameters_300.npz\")\n",
    "for key in data.keys():\n",
    "    globals()[key] = data[key]\"\"\"\n",
    "\n",
    "x3D, y3D, z3D = np.meshgrid(x, y, z,indexing=\"ij\")  # 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",
    "# Create figure and axis\n",
    "fig, ax = plt.subplots()\n",
    "im = ax.imshow(np.abs(states[state_number, :, :, 0])**2, extent=[*extend[1], *extend[0]], origin=\"lower\",\n",
    "               vmin=np.min(np.abs(states[state_number])**2), vmax=np.max(np.abs(states[state_number])**2))\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",
    "vmin = np.min(pot)\n",
    "vmax = np.max(pot)\n",
    "\n",
    "# Animation update function\n",
    "def update(frame):\n",
    "    global contour  # Ensure we're modifying the global variable\n",
    "\n",
    "    im.set_data(np.abs(states[state_number, :, :, frame])**2) # 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=np.linspace(vmin, vmax, 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=np.linspace(vmin, vmax, 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",
    "\n",
    "ani.save(anim_path, writer=\"pillow\", fps=frames/5)  # Save as GIF\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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
}
