{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import Math, display\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "from scipy import constants as const\n",
    "from scipy.optimize import minimize_scalar\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, TwoSiteLattice\n",
    "from boson_helpers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "line_number_1 = 45 #for AR=1\n",
    "state0_1 = 0\n",
    "state1_1 = 1\n",
    "\n",
    "line_number_2 = 46 #for AR=2\n",
    "state0_2 = 0\n",
    "state1_2 = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AR=1:\n",
      "J = h* 27.481287 Hz\n",
      "U_s = h* 3748.210 Hz\n",
      "\n",
      "\n",
      "AR=2:\n",
      "J = h* 27.481287 Hz\n",
      "U_s = h* 1774.199 Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 85168 characters]",
      "text/plain": [
       "<Figure size 500x900 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#import data for AR=1\n",
    "trap, func_ham, res =  import_results(line_number_1)\n",
    "GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0_1], res[\"states\"][state1_1])\n",
    "\n",
    "J_1, U_s_1, U_dds_1, angles_1, V_lrs_1, DeltaJs_1 = analyse_diagonalisation(line_number_1,\n",
    "                                                             state_nr0=state0_1,\n",
    "                                                             state_nr1=state1_1)\n",
    "print(\"AR=1:\")\n",
    "print(f\"J = h* {J_1/const.h :3f} Hz\")\n",
    "print(f\"U_s = h* {U_s_1/const.h :.3f} Hz\")\n",
    "\n",
    "#omega_left, omega_right = trap.get_both_omega(trap.x)\n",
    "#print(f\"omega_x = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "#omega_left, omega_right = trap.get_both_omega(trap.z)\n",
    "#print(f\"omega_z = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "\n",
    "print(\"\\n\")\n",
    "\n",
    "#import data for AR=2\n",
    "trap, func_ham, res =  import_results(line_number_2)\n",
    "GS_left_2, GS_right_2 = get_localised_GS(res[\"states\"][state0_2], res[\"states\"][state1_2])\n",
    "\n",
    "J_2, U_s_2, U_dds_2, angles_2, V_lrs_2, DeltaJs_2 = analyse_diagonalisation(line_number_2,\n",
    "                                                             state_nr0=state0_2,\n",
    "                                                             state_nr1=state1_2)\n",
    "print(\"AR=2:\")\n",
    "print(f\"J = h* {J_2/const.h :3f} Hz\")\n",
    "print(f\"U_s = h* {U_s_2/const.h :.3f} Hz\")\n",
    "\n",
    "#omega_left, omega_right = trap.get_both_omega(trap.x)\n",
    "#print(f\"omega_x = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "#omega_left, omega_right = trap.get_both_omega(trap.z)\n",
    "#print(f\"omega_z = 2pi* {float(trap.subs(omega_left))/2/np.pi:.2f}Hz\")\n",
    "\n",
    "\n",
    "# Create figure with subplots\n",
    "fig, axes = plt.subplots(3, 1, figsize=(5, 9), sharex=True)\n",
    "\n",
    "# Define colors\n",
    "colors = ['red', 'green']\n",
    "\n",
    "# Plot U_dd\n",
    "axes[0].plot(np.rad2deg(angles_1), U_dds_1 / const.h, color=colors[0], label=\"AR=1\")\n",
    "axes[0].plot(np.rad2deg(angles_2), U_dds_2 / const.h, color=colors[1], label=\"AR=2\")\n",
    "axes[0].axhline(0, color=\"black\", ls=\"--\")\n",
    "axes[0].set_ylabel(r\"$U_{dd} / h$ [Hz]\")\n",
    "axes[0].legend()\n",
    "\n",
    "# Plot NNI\n",
    "axes[1].plot(np.rad2deg(angles_1), V_lrs_1 / const.h, color=colors[0], label=\"AR=1\")\n",
    "axes[1].plot(np.rad2deg(angles_2), V_lrs_2 / const.h, color=colors[1], label=\"AR=2\")\n",
    "axes[1].axhline(0, color=\"black\", ls=\"--\")\n",
    "axes[1].set_ylabel(r\"$V_{lr} / h$ [Hz]\")\n",
    "axes[1].legend()\n",
    "\n",
    "# Plot DIT\n",
    "axes[2].plot(np.rad2deg(angles_1), DeltaJs_1 / const.h, color=colors[0], label=\"AR=1\")\n",
    "axes[2].plot(np.rad2deg(angles_2), DeltaJs_2 / const.h, color=colors[1], label=\"AR=2\")\n",
    "axes[2].axhline(0, color=\"black\", ls=\"--\")\n",
    "axes[2].set_xlabel(r\"$\\theta$ [deg]\")\n",
    "axes[2].set_ylabel(r\"$\\Delta J / h$ [Hz]\")\n",
    "axes[2].legend()\n",
    "\n",
    "# Customize ticks for all subplots\n",
    "for ax in axes:\n",
    "    ax.tick_params(axis='both', which='both', direction='in', length=5)  # Tick marks inside\n",
    "    ax.tick_params(axis='both', which='major', length=7)  # Major ticks longer\n",
    "    ax.tick_params(axis='both', which='minor', length=4)  # Minor ticks shorter\n",
    "    ax.minorticks_on()  # Enable minor ticks\n",
    "\n",
    "# Adjust layout and show the plot\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot wavefunctions\n",
    "## AR=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27248 characters]",
      "text/plain": [
       "<Figure size 1300x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot wavefunctions\n",
    "n_states = 2\n",
    "\n",
    "states_plot = [GS_left_1.real, GS_right_1.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=(13, 4))  # 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",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32816 characters]",
      "text/plain": [
       "<Figure size 1300x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 2\n",
    "\n",
    "states_plot = [GS_left_1.real, GS_right_1.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=(13, 4))  # 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\"][2]/2)].T,\n",
    "                   extent=[*res[\"extend\"][0], *res[\"extend\"][1]], 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(\"y\")\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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## AR=2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27556 characters]",
      "text/plain": [
       "<Figure size 1300x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot wavefunctions\n",
    "n_states = 2\n",
    "\n",
    "states_plot = [GS_left_2.real, GS_right_2.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=(13, 4))  # 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",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 29648 characters]",
      "text/plain": [
       "<Figure size 1300x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 2\n",
    "\n",
    "states_plot = [GS_left_2.real, GS_right_2.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=(13, 4))  # 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\"][2]/2)].T,\n",
    "                   extent=[*res[\"extend\"][0], *res[\"extend\"][1]], 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(\"y\")\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()"
   ]
  }
 ],
 "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
}
