{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# With a good 3D solution, calculate the values of $J$ and $U_1$, $U_2$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "from scipy import constants as const\n",
    "from scipy import sparse\n",
    "from scipy.integrate import simpson\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "timecode = \"2025-03-19_11-05-43\"\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",
    "ham = sparse.load_npz(f\"data/{timecode}_hamiltonian.npz\")\n",
    "\n",
    "with open(f\"data/{timecode}_trap.npz\", 'rb') as file:\n",
    "    trap = pickle.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "omega_r1: 13361.1Hz, omega_r2: 13361.1Hz\n",
      "omega_z1: 1454.4Hz, omega_z2: 1454.4Hz\n"
     ]
    }
   ],
   "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\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check how it looks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 96476 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": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 18944 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(energies,\"o\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Normalise the states to give 1 over 3D integral:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "states /= np.sqrt(np.prod(dvol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0000121523186054"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#check if normalisation is correct\n",
    "\n",
    "integrand  = np.abs(states[0])**2\n",
    "\n",
    "integral_z = simpson(integrand, dx=dvol[2], axis=2)\n",
    "integral_y = simpson(integral_z, dx=dvol[1], axis=1)\n",
    "total_integral = simpson(integral_y, dx=dvol[0], axis=0)\n",
    "\n",
    "total_integral"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0000000000000013"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#just suming gives almost the same result\n",
    "np.sum(integrand)*np.prod(dvol)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Find wavefunctions that are localised in each well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 37992 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 36848 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "state_number = 0\n",
    "phase_mult0 = np.exp(-1.91*1j)\n",
    "\n",
    "plt.plot(x,(phase_mult0*states[state_number,:,int(len(y)/2),int(len(z)/2)]).real, label=\"real part\")\n",
    "plt.plot(x,(phase_mult0*states[state_number,:,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()\n",
    "\n",
    "state_number = 1\n",
    "phase_mult1 = np.exp((-0.791)*1j)\n",
    "\n",
    "plt.plot(x,(phase_mult1*states[state_number,:,int(len(y)/2),int(len(z)/2)]).real, label=\"real part\")\n",
    "plt.plot(x,(phase_mult1*states[state_number,:,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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "theta = np.deg2rad(45)\n",
    "\n",
    "GS_left = (np.cos(theta)*states[0] + np.sin(theta)*states[1])\n",
    "GS_left /= np.sqrt(np.sum(np.abs(GS_left)**2)*np.prod(dvol))\n",
    "\n",
    "GS_right = (np.sin(theta)*states[0] - np.cos(theta)*states[1])\n",
    "GS_right /= np.sqrt(np.sum(np.abs(GS_right)**2)*np.prod(dvol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "GS_left = (phase_mult0*states[0] - phase_mult1*states[1])\n",
    "GS_left /= np.sqrt(np.sum(np.abs(GS_left)**2)*np.prod(dvol))\n",
    "\n",
    "GS_right = (phase_mult0*states[0] + phase_mult1*states[1])\n",
    "GS_right /= np.sqrt(np.sum(np.abs(GS_right)**2)*np.prod(dvol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_S = states[1].real\n",
    "state_S /= np.sqrt(np.sum(np.abs(state_S)**2)*np.prod(dvol))\n",
    "\n",
    "state_A = states[0].real\n",
    "state_A /= np.sqrt(np.sum(np.abs(state_A)**2)*np.prod(dvol))\n",
    "\n",
    "\n",
    "GS_left = (state_S - state_A)\n",
    "GS_left /= np.sqrt(2)#np.sqrt(np.sum(np.abs(GS_left)**2)*np.prod(dvol))\n",
    "\n",
    "GS_right = (state_S + state_A)/np.sqrt(2)\n",
    "GS_right /= np.sqrt(2)#np.sqrt(np.sum(np.abs(GS_right)**2)*np.prod(dvol))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = np.reshape(functional_hamiltonian(dvol, pot, mass)(np.reshape(states[1], -1)),size)/energies[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = np.reshape(ham @ np.reshape(states[1], -1),size)/energies[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27500 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 29600 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(x,(states[1]-test)[:,int(len(y)/2),int(len(z)/2)].real, label=\"real part\")\n",
    "plt.plot(x,(states[1])[:,int(len(y)/2),int(len(z)/2)].real, 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()\n",
    "\n",
    "plt.plot(x,state_S[:,int(len(y)/2),int(len(z)/2)].real, label=\"real part\")\n",
    "plt.plot(x,state_S[:,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": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 24664 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 23068 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(x,GS_left[:,int(len(y)/2),int(len(z)/2)].real, label=\"real part\")\n",
    "plt.plot(x,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",
    "plt.show()\n",
    "\n",
    "plt.plot(x,GS_right[:,int(len(y)/2),int(len(z)/2)].real, label=\"real part\")\n",
    "plt.plot(x,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": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21596 characters]",
      "text/plain": [
       "<Figure size 1333.33x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#state_plot = GS_left.real\n",
    "#state_plot = GS_left.imag\n",
    "state_plot1 = np.abs(GS_left)**2\n",
    "\n",
    "state_plot2 = np.abs(GS_right)**2\n",
    "\n",
    "# Create a 1xn_states grid of subplots (1 row, n_states columns)\n",
    "fig, ax = plt.subplots(1, 2, figsize=(10/6*n_states, 5))  # Adjust the size as needed\n",
    "\n",
    "# Slice through the y-direction and rotate the x-z plane (by swapping x and z)\n",
    "im = ax[0].imshow(state_plot1[ :, int(len(y)/2), :].T,\n",
    "                extent=[*extend[0], *extend[2]], origin=\"lower\",\n",
    "                vmin=np.min(state_plot1), vmax=np.max(state_plot1))\n",
    "\n",
    "# Set labels\n",
    "ax[0].set_xlabel(\"x\")\n",
    "ax[0].set_ylabel(\"z\")\n",
    "ax[0].set_title(f\"left\")\n",
    "fig.colorbar(im, ax=ax[0])\n",
    "\n",
    "im = ax[1].imshow(state_plot2[ :, int(len(y)/2), :].T,\n",
    "                extent=[*extend[0], *extend[2]], origin=\"lower\",\n",
    "                vmin=np.min(state_plot2), vmax=np.max(state_plot2))\n",
    "\n",
    "# Set labels\n",
    "ax[1].set_xlabel(\"x\")\n",
    "ax[1].set_ylabel(\"z\")\n",
    "ax[1].set_title(f\"right\")\n",
    "fig.colorbar(im, ax=ax[1])  \n",
    "\n",
    "# Adjust layout for better spacing\n",
    "plt.tight_layout()\n",
    "\n",
    "# Show the plots\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_J(state_left, state_right, hamiltonian, dvol):\n",
    "    \"\"\"\n",
    "    Calculates the single-particle hopping amplitude between\n",
    "    the two input states, which normally should be lowest energy states of\n",
    "    the two wells.\n",
    "    \"\"\"\n",
    "    #flatten the wavefunctions for calculation\n",
    "    state_left_flat = state_left.reshape((-1,))\n",
    "    state_right_flat = state_right.reshape((-1,))\n",
    "\n",
    "    #calculate J\n",
    "    integrand  = -state_left_flat.conj() * (hamiltonian @ state_right_flat)\n",
    "\n",
    "    return np.sum(integrand)*np.prod(dvol)\n",
    "\n",
    "def get_U_s(state,mass,scatter_length, dvol):\n",
    "    \"\"\"\n",
    "    Calculates onsite interaction due to contact interaction\n",
    "    \"\"\"\n",
    "    #integrate psi**4\n",
    "    integrand = np.abs(state)**4\n",
    "    total_integral = np.sum(integrand)*np.prod(dvol)\n",
    "\n",
    "    print(total_integral)\n",
    "    #add factor and return\n",
    "    return 4*np.pi*const.hbar*scatter_length/mass * total_integral\n",
    "\n",
    "def get_U_dd(state, mu, pol, dvol):\n",
    "    \"\"\"Calculates the onsite interaction due to dipolar interactions\n",
    "        :pol (3,) array giving the direction of polarisation\n",
    "    \"\"\"\n",
    "    \n",
    "    raise Exception(\"not yet implemented!\")\n",
    "    x3D, y3D, z3D = np.meshgrid(x,y,z,indexing=\"ij\")\n",
    "    x3D_, y3D_, z3D_ = np.meshgrid(x,y,z,indexing=\"ij\")\n",
    "\n",
    "    location_vectors = np.vstack([x3D.flatten(), y3D.flatten(), z3D.flatten()]).T\n",
    "    location_vectors_ = np.vstack([x3D.flatten(), y3D.flatten(), z3D.flatten()]).T\n",
    "\n",
    "    integral = 0\n",
    "    for i, r in enumerate(location_vectors):\n",
    "        for j, r_ in enumerate(location_vectors_):\n",
    "            cos_theta = ((r[0]*r_[0] + r[1]*r_[1] + r[2]*r_[2])/(r[0]**2 + r[1]**2 + r[2]**2)/(r_[0]**2 + r_[1]**2 + r_[2]**2))\n",
    "            #integral += state[]\n",
    "\n",
    "    \"\"\"\n",
    "    def integrand(x,y,z, x_,y_,z_):\n",
    "        cos_theta = (x*x_ + y*y_ + z*z_)/((x**2 + y**2 + z**2)*(x_**2 + y_**2 + z_**2))\"\"\"\n",
    "    \n",
    "\n",
    "def get_U(state_left, state_right, hamiltonian, mass, scatter_length, mu):\n",
    "    \"\"\"Calculates the total onsite interaction\"\"\"\n",
    "    U_s = get_U_s(state_left, state_right, hamiltonian)\n",
    "    #U should be different if tweezers are not identical, use left one here\n",
    "    U_dd = get_U_dd(state_left,...)\n",
    "    return U_s + U_dd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = -0.18035688997810906Hz    imaginary part: 0.0Hz\n"
     ]
    }
   ],
   "source": [
    "J_Hz = get_J(GS_left, GS_right, ham, dvol) / const.h\n",
    "print(f\"J = {J_Hz.real}Hz    imaginary part: {J_Hz.imag}Hz\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.615142490815278e+18\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1.4049588450401188e+35"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a0 = const.physical_constants[\"Bohr radius\"][0]\n",
    "\n",
    "get_U_s(GS_left, float(trap.subs(trap.m)), 100*a0, dvol)/const.h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20720 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(x,np.abs(GS_left[:,int(len(y)/2),int(len(z)/2)])**2, label=\"square\")\n",
    "#plt.plot(x,np.abs(GS_left[:,int(len(y)/2),int(len(z)/2)])**4, label=\"power 4\")\n",
    "#plt.plot(x,pot[:,int(len(y)/2),int(len(z)/2)], label=\"potential (rescaled)\")\n",
    "plt.legend()\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
}
