{
 "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 import sparse\n",
    "from scipy.optimize import minimize_scalar\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trap parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "omega_r1: 60954.0Hz, omega_r2: 60954.0Hz\n",
      "omega_z1: 10279.9Hz, omega_z2: 10279.9Hz\n"
     ]
    }
   ],
   "source": [
    "initial_power = 7000* si.uW\n",
    "initial_waist = 0.71*si.uW\n",
    "initial_distance = 0.808*si.um\n",
    "\n",
    "trap: DoubleTweezer = DoubleTweezer(\n",
    "    power=0,  # Set pancake laser power to 0, no 2D trap\n",
    "    grad_z= 0*si.G/si.cm,\n",
    "    grad_r=0,\n",
    "    power_tweezer1 = initial_power,     #stationary\n",
    "    power_tweezer2 = initial_power,     #transfer tweezer\n",
    "    waist_tweezer1 = initial_waist,     #stationary\n",
    "    waist_tweezer2 = initial_waist,     #transfer tweezer\n",
    "    distance_tweezers = initial_distance,\n",
    "\n",
    "    m= 161 * const.value(\"atomic mass constant\"),\n",
    "    a_s= 85* const.value(\"Bohr radius\"),\n",
    "    mu_b= 9.93 * const.value(\"Bohr magneton\" ),\n",
    "    a= 180*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "    wvl = 532 * si.nm,\n",
    "\n",
    "    g = 0,\n",
    ")\n",
    "\n",
    "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)/2/np.pi):.1f}Hz, omega_r2: {float(trap.subs(omega_r2)/2/np.pi):.1f}Hz\")\n",
    "print(f\"omega_z1: {float(trap.subs(omega_z1)/2/np.pi):.1f}Hz, omega_z2: {float(trap.subs(omega_z2)/2/np.pi):.1f}Hz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "depth 1: 21556.8326 Hz\n",
      "depth 2: 21556.8326 Hz\n",
      "target trap frequency x: 50000.00 HZ\n",
      "trapping frequency x 1: 19829.31 Hz\n",
      "trapping frequency x 2: 19829.31 Hz\n",
      "trapping frequency y 1: 54173.78 Hz\n",
      "trapping frequency y 2: 54173.78 Hz\n",
      "trapping frequency z 1: 6879.63 Hz\n",
      "trapping frequency z 2: 6879.63 Hz\n",
      "harmonic oscillator length x: 56.27 nm\n",
      "harmonic oscillator length z: 95.53 nm\n",
      "tunneling distance: 322.96 nm\n",
      "files saved with ...._2025-04-08_11-19-26\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 146132 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "depth 1: 25796.1125 Hz\n",
      "depth 2: 25796.1125 Hz\n",
      "target trap frequency x: 50000.00 HZ\n",
      "trapping frequency x 1: 20683.59 Hz\n",
      "trapping frequency x 2: 20683.59 Hz\n",
      "trapping frequency y 1: 54094.80 Hz\n",
      "trapping frequency y 2: 54094.80 Hz\n",
      "trapping frequency z 1: 6906.51 Hz\n",
      "trapping frequency z 2: 6906.51 Hz\n",
      "harmonic oscillator length x: 55.09 nm\n",
      "harmonic oscillator length z: 95.34 nm\n",
      "tunneling distance: 337.85 nm\n",
      "files saved with ...._2025-04-08_11-46-56\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 146004 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "depth 1: 30377.7940 Hz\n",
      "depth 2: 30377.7940 Hz\n",
      "target trap frequency x: 50000.00 HZ\n",
      "trapping frequency x 1: 21489.04 Hz\n",
      "trapping frequency x 2: 21489.04 Hz\n",
      "trapping frequency y 1: 54017.30 Hz\n",
      "trapping frequency y 2: 54017.30 Hz\n",
      "trapping frequency z 1: 6932.81 Hz\n",
      "trapping frequency z 2: 6932.81 Hz\n",
      "harmonic oscillator length x: 54.05 nm\n",
      "harmonic oscillator length z: 95.16 nm\n",
      "tunneling distance: 352.03 nm\n",
      "files saved with ...._2025-04-08_12-13-41\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 139576 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "depth 1: 35293.5015 Hz\n",
      "depth 2: 35293.5015 Hz\n",
      "target trap frequency x: 50000.00 HZ\n",
      "trapping frequency x 1: 22249.92 Hz\n",
      "trapping frequency x 2: 22249.92 Hz\n",
      "trapping frequency y 1: 53941.27 Hz\n",
      "trapping frequency y 2: 53941.27 Hz\n",
      "trapping frequency z 1: 6958.47 Hz\n",
      "trapping frequency z 2: 6958.47 Hz\n",
      "harmonic oscillator length x: 53.12 nm\n",
      "harmonic oscillator length z: 94.98 nm\n",
      "tunneling distance: 365.54 nm\n",
      "files saved with ...._2025-04-08_12-37-42\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 125640 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "depth 1: 40535.0732 Hz\n",
      "depth 2: 40535.0732 Hz\n",
      "target trap frequency x: 50000.00 HZ\n",
      "trapping frequency x 1: 22973.02 Hz\n",
      "trapping frequency x 2: 22973.02 Hz\n",
      "trapping frequency y 1: 53866.66 Hz\n",
      "trapping frequency y 2: 53866.66 Hz\n",
      "trapping frequency z 1: 6983.63 Hz\n",
      "trapping frequency z 2: 6983.63 Hz\n",
      "harmonic oscillator length x: 52.28 nm\n",
      "harmonic oscillator length z: 94.81 nm\n",
      "tunneling distance: 378.50 nm\n",
      "files saved with ...._2025-04-08_13-01-43\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 118428 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "depth 1: 46094.5543 Hz\n",
      "depth 2: 46094.5543 Hz\n",
      "target trap frequency x: 50000.00 HZ\n",
      "trapping frequency x 1: 23661.80 Hz\n",
      "trapping frequency x 2: 23661.80 Hz\n",
      "trapping frequency y 1: 53793.46 Hz\n",
      "trapping frequency y 2: 53793.46 Hz\n",
      "trapping frequency z 1: 7008.27 Hz\n",
      "trapping frequency z 2: 7008.27 Hz\n",
      "harmonic oscillator length x: 51.51 nm\n",
      "harmonic oscillator length z: 94.65 nm\n",
      "tunneling distance: 390.95 nm\n",
      "files saved with ...._2025-04-08_13-25-28\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 109908 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "depth 1: 51964.1905 Hz\n",
      "depth 2: 51964.1905 Hz\n",
      "target trap frequency x: 50000.00 HZ\n",
      "trapping frequency x 1: 24318.22 Hz\n",
      "trapping frequency x 2: 24318.22 Hz\n",
      "trapping frequency y 1: 53721.64 Hz\n",
      "trapping frequency y 2: 53721.64 Hz\n",
      "trapping frequency z 1: 7032.34 Hz\n",
      "trapping frequency z 2: 7032.34 Hz\n",
      "harmonic oscillator length x: 50.81 nm\n",
      "harmonic oscillator length z: 94.48 nm\n",
      "tunneling distance: 402.92 nm\n",
      "files saved with ...._2025-04-08_13-49-09\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 104920 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "depth 1: 58136.4215 Hz\n",
      "depth 2: 58136.4215 Hz\n",
      "target trap frequency x: 50000.00 HZ\n",
      "trapping frequency x 1: 24945.69 Hz\n",
      "trapping frequency x 2: 24945.69 Hz\n",
      "trapping frequency y 1: 53651.16 Hz\n",
      "trapping frequency y 2: 53651.16 Hz\n",
      "trapping frequency z 1: 7055.91 Hz\n",
      "trapping frequency z 2: 7055.91 Hz\n",
      "harmonic oscillator length x: 50.17 nm\n",
      "harmonic oscillator length z: 94.33 nm\n",
      "tunneling distance: 414.48 nm\n",
      "files saved with ...._2025-04-08_14-13-18\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 99744 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "depth 1: 64603.8751 Hz\n",
      "depth 2: 64603.8751 Hz\n",
      "target trap frequency x: 50000.00 HZ\n",
      "trapping frequency x 1: 25546.76 Hz\n",
      "trapping frequency x 2: 25546.76 Hz\n",
      "trapping frequency y 1: 53582.00 Hz\n",
      "trapping frequency y 2: 53582.00 Hz\n",
      "trapping frequency z 1: 7078.98 Hz\n",
      "trapping frequency z 2: 7078.98 Hz\n",
      "harmonic oscillator length x: 49.57 nm\n",
      "harmonic oscillator length z: 94.17 nm\n",
      "tunneling distance: 425.64 nm\n",
      "files saved with ...._2025-04-08_14-38-30\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 97412 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "#define trap frequencies to look at\n",
    "trap_freqs = np.array([50*si.kHz])*2*np.pi\n",
    "#find good separation for all frequencies\n",
    "\n",
    "#calculate corresponding power\n",
    "trap_powers = np.zeros_like(trap_freqs)\n",
    "for i, freq in enumerate(trap_freqs):\n",
    "    trap_powers[i] = float(trap.subs(trap.m*trap.waist_tweezer1**4*sp.pi*freq**2/8/trap.a))\n",
    "    harm_osc_len_est = float(trap.subs(sp.sqrt(const.hbar/trap.m/freq)))\n",
    "\n",
    "    trap[trap.power_tweezer1] = trap_powers[i]\n",
    "    trap[trap.power_tweezer2] = trap_powers[i]\n",
    "    \n",
    "    dist0 = float(trap.subs(trap.waist_tweezer1)) + harm_osc_len_est\n",
    "    trap_spacings = np.linspace(dist0 -10*si.nm, dist0 +10*si.nm, 9)\n",
    "    #loop over distances\n",
    "    for j, dist in enumerate(trap_spacings):\n",
    "\n",
    "        trap[trap.distance_tweezers] = dist\n",
    "\n",
    "        \n",
    "        ##################\n",
    "        #print trap parameters\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",
    "            continue\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",
    "        #trapping frequencies\n",
    "        omega_x1, omega_x2 = trap.get_both_omega(trap.x)\n",
    "        omega_y1, omega_y2 = trap.get_both_omega(trap.y)\n",
    "        omega_z1, omega_z2 = trap.get_both_omega(trap.z)\n",
    "\n",
    "        harm_osc_len_x = float(trap.subs(sp.sqrt(const.hbar/trap.m/omega_x1)))\n",
    "        harm_osc_len_z = float(trap.subs(sp.sqrt(const.hbar/trap.m/omega_z1)))\n",
    "\n",
    "        #print results\n",
    "        print(f\"depth 1: {V_1/const.h:.4f} Hz\")\n",
    "        print(f\"depth 2: {V_2/const.h:.4f} Hz\")\n",
    "        #print(f\"estimated tilt: {abs(V_1 - V_2)/const.h:.2f} Hz\")\n",
    "\n",
    "        print(f\"target trap frequency x: {freq/2/np.pi:.2f} HZ\")\n",
    "        print(f\"trapping frequency x 1: {trap.subs(omega_x1)/2/np.pi:.2f} Hz\")\n",
    "        print(f\"trapping frequency x 2: {trap.subs(omega_x2)/2/np.pi:.2f} Hz\")\n",
    "        print(f\"trapping frequency y 1: {trap.subs(omega_y1)/2/np.pi:.2f} Hz\")\n",
    "        print(f\"trapping frequency y 2: {trap.subs(omega_y2)/2/np.pi:.2f} Hz\")\n",
    "        print(f\"trapping frequency z 1: {trap.subs(omega_z1)/2/np.pi:.2f} Hz\")\n",
    "        print(f\"trapping frequency z 2: {trap.subs(omega_z2)/2/np.pi:.2f} Hz\")\n",
    "\n",
    "        print(f\"harmonic oscillator length x: {harm_osc_len_x/si.nm:.2f} nm\")\n",
    "        print(f\"harmonic oscillator length z: {harm_osc_len_z/si.nm:.2f} nm\")\n",
    "        print(f\"tunneling distance: {tunneling_dist/si.nm:.2f} nm\")\n",
    "\n",
    "        #############\n",
    "        #test 3D diag\n",
    "\n",
    "        n_grid_1D = 75  \n",
    "\n",
    "        n_pot_steps = [n_grid_1D,n_grid_1D,n_grid_1D]\n",
    "        n_levels = 6\n",
    "\n",
    "        left_cutoff = -0.5*tunneling_dist-5*harm_osc_len_x\n",
    "        right_cutoff = 0.5*tunneling_dist+5*harm_osc_len_x\n",
    "        back_cutoff = -5*harm_osc_len_x\n",
    "        front_cutoff = 5*harm_osc_len_x\n",
    "        bottom_cutoff = -5*harm_osc_len_z\n",
    "        top_cutoff = 5*harm_osc_len_z\n",
    "\n",
    "        extend = [(left_cutoff,right_cutoff),\n",
    "                (back_cutoff,front_cutoff),\n",
    "                (bottom_cutoff,top_cutoff)]\n",
    "        \n",
    "        # Solve the hamiltonian numerically\n",
    "        energies, states, potential, coords = trap.nstationary_solution(\n",
    "                [trap.x,trap.y,trap.z], extend, n_pot_steps, k=n_levels,\n",
    "                method=\"matrix_free\", export=True)\n",
    "\n",
    "        x = coords[trap.x]\n",
    "        y = coords[trap.y]\n",
    "        z = coords[trap.z]\n",
    "        x3D,y3D,z3D = np.meshgrid(coords[trap.x],coords[trap.y],coords[trap.z],indexing=\"ij\")\n",
    "        pot = potential(x3D,y3D,z3D)\n",
    "\n",
    "        n_states = 6\n",
    "\n",
    "        # Create a 1xn_states grid of subplots (1 row, n_states columns)\n",
    "        fig, axes = plt.subplots(1, n_states, figsize=(18/6*n_states, 5))  # Adjust the size as needed\n",
    "\n",
    "        # Loop over the state numbers from 0 to 5\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(np.abs(states[state_number, :, int(len(y)/2), :].T)**2, extent=[*extend[0], *extend[2]], origin=\"lower\",\n",
    "                        vmin=np.min(np.abs(states[state_number])**2), vmax=np.max(np.abs(states[state_number])**2))\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",
    "\n",
    "        print(\"_________________________________________________________________\")"
   ]
  },
  {
   "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
}
