{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "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": 3,
   "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": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "files saved under C:\\Users\\naeve\\FerDy-Repo\\merging_tweezer_code\\bosons\\data\n"
     ]
    }
   ],
   "source": [
    "n_pot_steps = [20,30,40]\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",
    "# Save the hamiltonian for external solver\n",
    "trap.nstationary_solution_export(r\"C:\\Users\\naeve\\FerDy-Repo\\merging_tweezer_code\\bosons\\data\",\n",
    "        [trap.x,trap.y,trap.z], extend, n_pot_steps, k=n_levels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now copy the hamiltonian.npz and parameters.npz files into Google colab and compute the eigenvectors. Plot results.npz here:\n",
    "### The results from cupyx don't seem to converge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\naeve\\AppData\\Local\\Temp\\ipykernel_19176\\3282607925.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"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 83916 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#state to plot\n",
    "state_number = 0\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_40_scipy.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_40.npz\")\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",
    "# 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",
    "\n",
    "ani.save(anim_path, writer=\"pillow\", fps=frames/10)  # Save as GIF\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Look at the results more in detail:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# results file containing the energies and states\n",
    "data = np.load(\"data/results_400_400_200.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_400_400_200.npz\")\n",
    "for key in data.keys():\n",
    "    globals()[key] = data[key]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# results file containing the energies and states\n",
    "data = np.load(\"data/results_200.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_200.npz\")\n",
    "for key in data.keys():\n",
    "    globals()[key] = data[key]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# results file containing the energies and states\n",
    "data = np.load(\"data/test_3D_50_50_50.npz\")\n",
    "for key in data.keys():\n",
    "    globals()[key] = data[key]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x22ad83de4e0>]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 31768 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(x,states[0,:,int(len(y)/2),int(len(z)/2)])\n",
    "plt.plot(x,1e29*pot[:,int(len(y)/2),int(len(z)/2)])\n",
    "#plt.plot(x,states[0,int(len(x)/2),:,int(len(z)/2)])\n",
    "#plt.plot(y,1e29*pot[int(len(x)/2),:,int(len(z)/2)])\n",
    "#plt.plot(z,states[0,int(len(x)/2),int(len(y)/2),:])\n",
    "#plt.plot(z,1e29*pot[int(len(x)/2),int(len(y)/2),:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now make good simulation for fixed indexing of x and y:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "files saved under C:\\Users\\naeve\\FerDy-Repo\\merging_tweezer_code\\bosons\\data\n"
     ]
    }
   ],
   "source": [
    "n_pot_steps = [50,50,50]\n",
    "n_levels = 10\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",
    "# Save the hamiltonian for external solver\n",
    "trap.nstationary_solution_export(r\"C:\\Users\\naeve\\FerDy-Repo\\merging_tweezer_code\\bosons\\data\",\n",
    "        [trap.x,trap.y,trap.z], extend, n_pot_steps, k=n_levels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# results file containing the energies and states\n",
    "data = np.load(\"data/results_200.npz\",allow_pickle=True)\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_200.npz\",allow_pickle=True)\n",
    "for key in data.keys():\n",
    "    globals()[key] = data[key]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x221809552e0>]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 52880 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "state_number = 0\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\naeve\\AppData\\Local\\Temp\\ipykernel_21952\\951173573.py:49: 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, 122976 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#state to plot\n",
    "state_number = 3\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(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",
    "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(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=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/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
}
