{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors as mcolors\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, curve_fit\n",
    "from scipy import sparse\n",
    "import qutip as qt\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": "markdown",
   "metadata": {},
   "source": [
    "# Choose data to analyse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 17\n",
    "end_line = 25\n",
    "\n",
    "n_lines = end_line - start_line + 1\n",
    "lines = np.arange(start_line,end_line+1,dtype=int)\n",
    "state_nrs = np.zeros((n_lines,2),dtype=int)\n",
    "state_nrs[:,1] = 1\n",
    "\n",
    "degeneracy = np.full(n_lines, True, dtype=bool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 17\n",
      "separation: 1.0447 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25168 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 18\n",
      "separation: 1.0709 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 26144 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 19\n",
      "separation: 1.0972 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25252 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 20\n",
      "separation: 1.1234 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25320 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 21\n",
      "separation: 1.1497 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25196 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 22\n",
      "separation: 1.1759 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 26020 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 23\n",
      "separation: 1.2022 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24812 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 24\n",
      "separation: 1.2284 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24708 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 25\n",
      "separation: 1.2547 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 24596 characters]",
      "text/plain": [
       "<Figure size 600x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "    \n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    print(f\"line: {line}\")\n",
    "    print(f\"separation: {trap.subs(trap.distance_tweezers)/si.um:.4f} um\")\n",
    "\n",
    "    #plot wavefunctions\n",
    "    n_states = 2\n",
    "\n",
    "    states_plot = [GS_left_1.real, GS_right_1.real]\n",
    "    #states_plot = [res[\"states\"][state0].real, res[\"states\"][state1].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=(6, 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": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "frequencies = np.zeros_like(lines,dtype=float)\n",
    "distance = np.zeros_like(lines,dtype=float)\n",
    "J = np.zeros_like(lines,dtype=float)\n",
    "#tilt = np.zeros_like(lines,dtype=float)\n",
    "U_s = np.zeros_like(lines,dtype=float)\n",
    "U_dds = np.zeros((len(lines),50),dtype=float)\n",
    "angles = np.zeros((len(lines),50),dtype=float)\n",
    "V_lrs = np.zeros((len(lines),50),dtype=float)\n",
    "DeltaJs = np.zeros((len(lines),50),dtype=float)\n",
    "\n",
    "E_exc = np.zeros_like(lines,dtype=float)\n",
    "E_barr = np.zeros_like(lines,dtype=float)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    state0 = state_nrs[i,0]\n",
    "    state1 = state_nrs[i,1]\n",
    "\n",
    "    trap, func_ham, res =  import_results(line)\n",
    "    GS_left_1, GS_right_1 = get_localised_GS(res[\"states\"][state0],\n",
    "                                             res[\"states\"][state1],\n",
    "                                             degenerate=degeneracy[i])\n",
    "    \n",
    "    frequencies[i] = float(trap.subs(sp.sqrt(2*trap.power_tweezer1*trap.a/sp.pi/trap.m) * 2/trap.waist_tweezer1**2))\n",
    "    distance[i] = float(trap.subs(trap.distance_tweezers))\n",
    "    J[i], U_s[i], U_dds[i], angles[i], V_lrs[i], DeltaJs[i] = analyse_diagonalisation(line,\n",
    "                                                                state_nr0=state0,\n",
    "                                                                state_nr1=state1,\n",
    "                                                                degenerate=degeneracy[i],\n",
    "                                                                n_angles=20)\n",
    "    \n",
    "    #estimate energies of 0th and 1st state\n",
    "    E0 = 0.5*(res[\"energies\"][0] + res[\"energies\"][1])\n",
    "    E1 = 0.5*(res[\"energies\"][2] + res[\"energies\"][3])\n",
    "    E_exc[i] = E1-E0\n",
    "\n",
    "    #estimate heigth of barrier\n",
    "    V = trap.subs(trap.get_potential(apply_zero_offset=False))\n",
    "    a = float(trap.subs(trap.distance_tweezers))\n",
    "\n",
    "    #find minima of potential\n",
    "    def V_func(x):\n",
    "        return float(V.subs({trap.x:x, trap.y:0, trap.z:0}))\n",
    "\n",
    "    x_right = minimize_scalar(V_func,bracket=[0,a]).x\n",
    "    x_left = minimize_scalar(V_func,bracket=[-a,0]).x\n",
    "\n",
    "    #catch case where both potentials have already merged\n",
    "    tunneling_dist = abs(x_right-x_left)\n",
    "    if tunneling_dist < 1e-15:\n",
    "        raise Exception(\"potential has only one minmum\")\n",
    "\n",
    "    #depts of both tweezers\n",
    "    V_1 = abs(V_func(x_left) - V_func(0))\n",
    "    V_2 = abs(V_func(x_right) - V_func(0)) \n",
    "\n",
    "    E_barr[i] = np.max([V_1,V_2])\n",
    "\n",
    "    # these states are basically 1 tweezer\n",
    "    if state1 !=1:\n",
    "        J[i] = np.nan\n",
    "        U_s[i] = np.nan\n",
    "        U_dds[i] = np.nan\n",
    "        V_lrs[i] = np.nan\n",
    "        DeltaJs[i] = np.nan\n",
    "        E_exc[i] = np.nan\n",
    "        E_barr[i] = np.nan\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 287296 characters]",
      "text/plain": [
       "<Figure size 1000x1500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Create a figure with 4 subplots (stacked vertically)\n",
    "fig, axes = plt.subplots(4, 1, figsize=(10, 15), sharex=True)\n",
    "\n",
    "# Get unique frequencies\n",
    "unique_freqs = np.unique(frequencies)\n",
    "\n",
    "# Define colors for consistency\n",
    "colors = [\"orange\", \"blue\", \"green\", \"purple\", \"orange\"]\n",
    "\n",
    "### 1️⃣ Hopping parameter J ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    axes[0].plot(distance[mask] / si.nm, J[mask] / const.h, color=color, \n",
    "                 label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "#axes[0].axhspan(50, 150, alpha=0.3, color=\"gray\", label=\"desired region of 50-100Hz\")\n",
    "axes[0].set_ylabel(r\"$J/h$ [Hz]\")\n",
    "axes[0].xaxis.set_label_position('top')\n",
    "axes[0].set_xlabel(\"tweezer spacing [nm]\")\n",
    "axes[0].legend()\n",
    "# 🔹 Add top x-axis ticks\n",
    "axes[0].tick_params(axis=\"x\", labeltop=True, top=True)\n",
    "axes[0].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[0].minorticks_on()\n",
    "axes[0].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "### 2️⃣ Interaction parameter U ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    x = distance[mask] / si.nm\n",
    "    y = U_s[mask] / const.h\n",
    "    delta = U_dds[mask] / const.h\n",
    "\n",
    "    axes[1].plot(x, y, color=color, label=fr\"$U_s$ for $\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "    axes[1].fill_between(x, y + np.min(delta, axis=1), y + np.max(delta, axis=1), \n",
    "                         alpha=0.3, color=color, edgecolor=None, label=r\"$\\Delta U_{dd}$\")\n",
    "\n",
    "axes[1].set_ylabel(r\"$U/h$ [Hz]\")\n",
    "axes[1].legend()\n",
    "axes[1].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[1].minorticks_on()\n",
    "axes[1].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "### 3️⃣ First excited state energy ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    axes[2].plot(distance[mask] / si.nm, E_exc[mask] / const.h, color=color, \n",
    "                 label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "axes[2].set_ylabel(r\"$E_{exc}/h$ [Hz]\")\n",
    "axes[2].legend()\n",
    "axes[2].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[2].minorticks_on()\n",
    "axes[2].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "### 4️⃣ Barrier height ###\n",
    "for i, (freq, color) in enumerate(zip(unique_freqs, colors)):\n",
    "    mask = frequencies == freq\n",
    "    axes[3].plot(distance[mask] / si.nm, E_barr[mask] / const.h, color=color, \n",
    "                 label=fr\"$\\omega_r$: {freq/2/np.pi/si.kHz:.1f} kHz\")\n",
    "\n",
    "axes[3].set_ylabel(r\"$E_{barr}/h$ [Hz]\")\n",
    "axes[3].set_xlabel(\"tweezer spacing [nm]\")\n",
    "#axes[3].set_ylim(0, 30000)\n",
    "axes[3].legend()\n",
    "axes[3].grid(True, linestyle='-', linewidth=0.7)\n",
    "axes[3].minorticks_on()\n",
    "axes[3].grid(True, which='minor', linestyle=':', linewidth=0.7, alpha=0.8)\n",
    "\n",
    "# Improve layout\n",
    "plt.tight_layout()\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
}
