{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "from IPython.display import Math, display\n",
    "from scipy import constants as const\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 boson_helpers import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get the Hubbard parameters from the results of a diagonalisation in one function call:\n",
    "\n",
    "## Show results for approaching tweezers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "________________________________________________________________\n",
      "d = 1.35um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32176 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = 0.444561 Hz\n",
      "U_s = 19.981 Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24692 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "________________________________________________________________\n",
      "d = 1.33um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 30904 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = 1.051351 Hz\n",
      "U_s = 19.591 Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24568 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "________________________________________________________________\n",
      "d = 1.31um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 31236 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = 2.448100 Hz\n",
      "U_s = 19.135 Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24712 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "________________________________________________________________\n",
      "d = 1.29um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 31008 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = 5.575860 Hz\n",
      "U_s = 18.593 Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24540 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "________________________________________________________________\n",
      "d = 1.27um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 31544 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = 12.302617 Hz\n",
      "U_s = 17.930 Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 23908 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i in range(10,15):\n",
    "\n",
    "    trap, func_ham, res =  import_results(i)\n",
    "\n",
    "    print(\"________________________________________________________________\")\n",
    "    print(f\"d = {float(trap.subs(trap.distance_tweezers)) /si.um :.2f}um\")\n",
    "\n",
    "    GS_left, GS_right = get_localised_GS(res[\"states\"][0], res[\"states\"][1])\n",
    "    \n",
    "    plt.plot(res[\"x\"],np.abs(GS_left[:,int(res[\"size\"][1]/2),int(res[\"size\"][2]/2)])**2, label=\"left state\")\n",
    "    plt.plot(res[\"x\"],np.abs(GS_right[:,int(res[\"size\"][1]/2),int(res[\"size\"][2]/2)])**2, label=\"right state\")\n",
    "\n",
    "    plt.xlabel(\"x\")\n",
    "    plt.ylabel(r\"$|\\Psi|^2$\")\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "    J, U_s, U_dds, angles = analyse_diagonalisation(i)\n",
    "\n",
    "    print(f\"J = {J/const.h :3f} Hz\")\n",
    "    print(f\"U_s = {U_s/const.h :.3f} Hz\")\n",
    "\n",
    "    plt.plot(np.rad2deg(angles), U_dds/const.h)\n",
    "\n",
    "    plt.axhline(0,color=\"black\",ls=\"--\")\n",
    "\n",
    "    plt.xlabel(\"theta [deg]\")\n",
    "    plt.ylabel(r\"$U_{dd} / h$ [Hz]\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "trap, func_ham, res =  import_results(14)\n",
    "\n",
    "GS_left, GS_right = get_localised_GS(res[\"states\"][0], res[\"states\"][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38948 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "plt.plot(res[\"x\"],(GS_left[:,int(res[\"size\"][1]/2),int(res[\"size\"][2]/2)]).real, label=\"real part\")\n",
    "plt.plot(res[\"x\"],(GS_left[:,int(res[\"size\"][1]/2),int(res[\"size\"][2]/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",
    "\n",
    "plt.plot(res[\"x\"],(GS_right[:,int(res[\"size\"][1]/2),int(res[\"size\"][2]/2)]).real, label=\"real part\")\n",
    "plt.plot(res[\"x\"],(GS_right[:,int(res[\"size\"][1]/2),int(res[\"size\"][2]/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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 36304 characters]",
      "text/plain": [
       "<Figure size 666.667x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 2\n",
    "\n",
    "states_plot = [GS_left.real, GS_right.real]\n",
    "#states_plot = np.angle(states)\n",
    "#states_plot = states.imag\n",
    "#states_plot = np.abs(states)**2\n",
    "\n",
    "# Create a 1xn_states grid of subplots (1 row, n_states columns)\n",
    "fig, axes = plt.subplots(1, n_states, figsize=(20/6*n_states, 5))  # Adjust the size as needed\n",
    "\n",
    "# Loop over the state numbers from 0 to n_states-1\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(states_plot[state_number][:, int(res[\"size\"][1]/2), :].T,\n",
    "                   extent=[*res[\"extend\"][0], *res[\"extend\"][2]], origin=\"lower\",\n",
    "                   vmin=np.min(states_plot[state_number]), vmax=np.max(states_plot[state_number]))\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": []
  }
 ],
 "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
}
