{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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",
    "from scipy import sparse\n",
    "from scipy.fft import fftn, ifftn, fftfreq\n",
    "from scipy.optimize import root_scalar\n",
    "import pickle\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 trap_numerics import functional_hamiltonian\n",
    "from boson_helpers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "trap, ham, res = import_results(9)\n",
    "\n",
    "globals().update(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "omega_r1: 7331.9Hz, omega_r2: 7331.9Hz\n",
      "omega_z1: 418.1Hz, omega_z2: 418.1Hz\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\omega_{t r}}{\\omega_{t ax}} \\approx 17.54$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\omega_{t r}}{\\omega_{t ax}} \\approx 17.54$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "omega_z1 = sp.sqrt(4*trap.power_tweezer1*trap.a/sp.pi**3/trap.m) * trap.wvl/trap.waist_tweezer1**3\n",
    "omega_z2 = sp.sqrt(4*trap.power_tweezer2*trap.a/sp.pi**3/trap.m) * trap.wvl/trap.waist_tweezer2**3\n",
    "\n",
    "omega_r1 = sp.sqrt(2*trap.power_tweezer1*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer1**2\n",
    "omega_r2 = sp.sqrt(2*trap.power_tweezer2*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer2**2\n",
    "\n",
    "print(f\"omega_r1: {float(trap.subs(omega_r1)):.1f}Hz, omega_r2: {float(trap.subs(omega_r2)):.1f}Hz\")\n",
    "print(f\"omega_z1: {float(trap.subs(omega_z1)):.1f}Hz, omega_z2: {float(trap.subs(omega_z2)):.1f}Hz\")\n",
    "\n",
    "aspect_ratio1 = omega_r1 / omega_z1\n",
    "_aspect_ratio1_latex = sp.latex(trap.omega_r_tweezer / trap.omega_ax_tweezer)\n",
    "display(Math(f\"{_aspect_ratio1_latex} \\\\approx {trap.subs(aspect_ratio1).evalf():.2f}\"))\n",
    "\n",
    "aspect_ratio2 = omega_r2 / omega_z2\n",
    "_aspect_ratio2_latex = sp.latex(trap.omega_r_tweezer / trap.omega_ax_tweezer)\n",
    "display(Math(f\"{_aspect_ratio2_latex} \\\\approx {trap.subs(aspect_ratio2).evalf():.2f}\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "GS_left, GS_right = get_localised_GS(states[0],states[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22612 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(len(y)/2), :].T,\n",
    "                   extent=[*extend[0], *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": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 33404 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "state_number = 0\n",
    "phase_mult0 = 1\n",
    "\n",
    "plt.plot(x,(phase_mult0*GS_left[:,int(len(y)/2),int(len(z)/2)]).real, label=\"real part\")\n",
    "plt.plot(x,(phase_mult0*GS_left[:,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",
    "\n",
    "state_number = 1\n",
    "phase_mult1 = 1\n",
    "\n",
    "plt.plot(x,(phase_mult1*GS_right[:,int(len(y)/2),int(len(z)/2)]).real, label=\"real part\")\n",
    "plt.plot(x,(phase_mult1*GS_right[:,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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculate J, U_s, U_dd, U:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = 0.0009612848724873431\n"
     ]
    }
   ],
   "source": [
    "print(f\"J = {get_J(GS_left, GS_right, ham, dvol)/const.h}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "U_s = 3.849482910362117\n"
     ]
    }
   ],
   "source": [
    "#scattering length\n",
    "a_s = 85* const.physical_constants[\"Bohr radius\"][0]\n",
    "#mass\n",
    "mass = float(trap.subs(trap.m))\n",
    "\n",
    "print(f\"U_s = {get_U_s(trap, GS_left, dvol)/const.h}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 24444 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#magnetic moment\n",
    "mu = float(trap.subs(trap.mu_b))\n",
    "\n",
    "#polarisation\n",
    "angles= np.deg2rad(np.linspace(0,90))\n",
    "U_dds = np.zeros_like(angles)\n",
    "polarisations = np.zeros((len(angles),3))\n",
    "polarisations[:,0] = np.sin(angles)\n",
    "polarisations[:,2] = np.cos(angles)\n",
    "\n",
    "for i, pol in enumerate(polarisations):\n",
    "    trap[trap.B_x] = pol[0]\n",
    "    trap[trap.B_y] = pol[1]\n",
    "    trap[trap.B_z] = pol[2]\n",
    "\n",
    "    U_dds[i] = get_U_dd(trap, GS_left, dvol)\n",
    "\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": "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
}
