{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Try again to calculate J and U. This time use functional Hamiltonian since that is used for diagonalisation"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "line_number = 7 # line of file in overview.txt\n",
    "\n",
    "with open(\"data/overview.txt\", \"r\") as file:\n",
    "    for current_line_number, line in enumerate(file, start=1):\n",
    "        if current_line_number == line_number:\n",
    "            timecode = line[:19]  # Get the first 19 characters\n",
    "            break  # Stop reading once we get the desired line\n",
    "\n",
    "data = np.load(f\"data/{timecode}_results.npz\",allow_pickle=True)\n",
    "for key in data.keys():\n",
    "    globals()[key] = data[key]\n",
    "\n",
    "with open(f\"data/{timecode}_trap.npz\", 'rb') as file:\n",
    "    trap = pickle.load(file)\n",
    "\n",
    "ham = sparse.load_npz(f\"data/{timecode}_hamiltonian.npz\")\n",
    "func_ham = functional_hamiltonian(dvol,pot,float(trap.subs(trap.m)))"
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "Turn state vector into real space wavefunction by adjusting normalisation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9999999999999993"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sum(np.abs(states[0]**2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "dV = np.prod(dvol)\n",
    "\n",
    "states_r = states/ np.sqrt(dV)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9999999999999993"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sum(np.abs(states_r[0]**2))*dV"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check how diagonalisation looks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 109168 characters]",
      "text/plain": [
       "<Figure size 2666.67x500 with 16 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_states = 8\n",
    "\n",
    "states_plot = states.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": "markdown",
   "metadata": {},
   "source": [
    "Find localised wave functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#rotate resulting wavefunctions into symmetric and anti-symm. states \n",
    "def f(alpha):\n",
    "\n",
    "    state_S1 = np.cos(alpha) * states[0].real + np.sin(alpha)* states[1].real\n",
    "    state_S1 /= np.sqrt(np.sum(np.abs(state_S1)**2))\n",
    "\n",
    "    state_A1 = np.sin(alpha) * states[0].real - np.cos(alpha)* states[1].real\n",
    "    state_A1 /= np.sqrt(np.sum(np.abs(state_A1)**2))\n",
    "\n",
    "    return np.sum(state_S1[int(len(x)/2):]) - np.sum(state_S1[:int(len(x)/2)])\n",
    "\n",
    "#find optimal angle\n",
    "# alpha = minimize_scalar(f,bounds=(-2*np.pi, 2*np.pi+0.1)).x\n",
    "alpha = root_scalar(f, x0=0).root\n",
    "\n",
    "state_S1 = np.cos(alpha) * states[0].real + np.sin(alpha)* states[1].real\n",
    "state_S1 /= np.sqrt(np.sum(np.abs(state_S1)**2))\n",
    "\n",
    "state_A1 = np.sin(alpha) * states[0].real - np.cos(alpha)* states[1].real\n",
    "state_A1 /= np.sqrt(np.sum(np.abs(state_A1)**2))\n",
    "\n",
    "#then transform to localised wavefunctions\n",
    "GS_LR1 = (state_S1 + state_A1)/np.sqrt(2)\n",
    "GS_LR2 = (state_S1 - state_A1)/np.sqrt(2)\n",
    "\n",
    "#asign left and right\n",
    "if abs(np.sum(GS_LR1[int(len(x)/2):])) < abs(np.sum(GS_LR1[:int(len(x)/2)])):\n",
    "    GS_left = GS_LR1\n",
    "    GS_right = GS_LR2\n",
    "else:\n",
    "    GS_left = GS_LR2\n",
    "    GS_right = GS_LR1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25960 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": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35004 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 parameters:\n",
    "\n",
    "### Start with contact interactions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9999999993379817"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#make sure we're working with state vector\n",
    "np.sum(np.abs(GS_left)**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "U_s = 3.2123021180831275 Hz\n"
     ]
    }
   ],
   "source": [
    "#scattering length\n",
    "a_s = 85* const.physical_constants[\"Bohr radius\"][0]\n",
    "\n",
    "#calculate integral over psi^4 (divide by sqrt(dV) to convert to wavefunction)\n",
    "U_s = np.sum(np.abs(GS_left/np.sqrt(dV))**4)*dV\n",
    "U_s *= 4*np.pi*const.hbar**2*a_s/float(trap.subs(trap.m))\n",
    "\n",
    "print(f\"U_s = {U_s/const.h} Hz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Then do single-particle hopping:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9999999993379817"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#make sure we're working with state vector\n",
    "np.sum(np.abs(GS_left)**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = 0.9876660072698075 Hz\n"
     ]
    }
   ],
   "source": [
    "#call Hamiltonian on one side\n",
    "H_psi = np.reshape(func_ham(np.reshape(GS_right, -1)), size)\n",
    "\n",
    "#calculate integral\n",
    "J = -(np.sum(GS_left.conj()/np.sqrt(dV)* H_psi/np.sqrt(dV))*dV)\n",
    "\n",
    "#check imaginary part of result (should vanish)\n",
    "if abs(J.real/J.imag) < 10:\n",
    "    raise Exception(f\"imaginary part of J too big: J = {J}\")\n",
    "else:\n",
    "    J = J.real\n",
    "\n",
    "print(f\"J = {J/const.h} Hz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finally dipole interactions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "angles= np.deg2rad(np.linspace(0,90))\n",
    "U_dds = np.zeros_like(angles, dtype=complex)\n",
    "\n",
    "for i, theta in enumerate(angles):\n",
    "    #compute dipolar interactions using fourier transform\n",
    "\n",
    "    #direction of polarisation (B-field) [x,y,z]\n",
    "    #theta = np.deg2rad(0)\n",
    "    pol =np.array([np.sin(theta),0,np.cos(theta)])\n",
    "\n",
    "    #density\n",
    "    rho = np.abs(GS_left/np.sqrt(dV))**2\n",
    "    # Compute Fourier transform of density\n",
    "    rho_k = fftn(rho)\n",
    "\n",
    "    #set up fourier space\n",
    "    kx = 2*np.pi*fftfreq(size[0], dvol[0])\n",
    "    ky = 2*np.pi*fftfreq(size[1], dvol[1])\n",
    "    kz = 2*np.pi*fftfreq(size[2], dvol[2])\n",
    "    KX, KY, KZ = np.meshgrid(kx, ky, kz, indexing='ij')\n",
    "    K2 = KX**2 + KY**2 + KZ**2\n",
    "    K2[K2 == 0] = np.inf  # Avoid division by zero\n",
    "\n",
    "    # Compute dipolar interaction kernel in Fourier space\n",
    "    V_k = (-1 + 3* (pol[0]*KX + pol[1]*KY + pol[2]*KZ)**2 / K2) /3\n",
    "\n",
    "    # Compute convolution in Fourier space\n",
    "    U_dd_k = V_k * rho_k\n",
    "\n",
    "    # Transform back to real space\n",
    "    U_dd_real = ifftn(U_dd_k)\n",
    "\n",
    "    # Integrate\n",
    "    U_dd = (np.sum(U_dd_real*rho) * dV) \n",
    "    #multiply by prefactor C_dd\n",
    "    U_dd *= const.mu_0*float(trap.subs(trap.mu_b))**2\n",
    "\n",
    "    #check imaginary part of result (should vanish)\n",
    "    if abs(U_dd.real/U_dd.imag) < 10:\n",
    "        raise Exception(f\"imaginary part of U_dd too big: U_dd = {U_dd}\")\n",
    "    else:\n",
    "        U_dd = U_dd.real\n",
    "\n",
    "    U_dds[i] = U_dd\n",
    "    #print(f\"U_dd = {U_dd/const.h} Hz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23612 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(np.rad2deg(angles), U_dds.real/const.h)\n",
    "#plt.plot(angles, U_dds.imag/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": [
    "---------"
   ]
  }
 ],
 "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
}
