{
 "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: 2169.9Hz, omega_r2: 2169.9Hz\n",
      "omega_z1: 472.4Hz, omega_z2: 472.4Hz\n"
     ]
    }
   ],
   "source": [
    "initial_power = 50* si.uW\n",
    "initial_waist = 1.1*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= 184*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
    "    wvl = 1064 * 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": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "depth 1: 0.2778 Hz\n",
      "depth 2: 0.2778 Hz\n",
      "target trap frequency x: 3000.00 HZ\n",
      "trapping frequency x 1: 1061.62 Hz\n",
      "trapping frequency x 2: 1061.62 Hz\n",
      "trapping frequency y 1: 3387.98 Hz\n",
      "trapping frequency y 2: 3387.98 Hz\n",
      "trapping frequency z 1: 546.57 Hz\n",
      "trapping frequency z 2: 546.57 Hz\n",
      "harmonic oscillator length x: 243.18 nm\n",
      "harmonic oscillator length z: 338.91 nm\n",
      "tunneling distance: 0.00 nm\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 72636 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "depth 1: 2012.3643 Hz\n",
      "depth 2: 2086.0124 Hz\n",
      "target trap frequency x: 3000.00 HZ\n",
      "trapping frequency x 1: 2022.19 Hz\n",
      "trapping frequency x 2: 2022.19 Hz\n",
      "trapping frequency y 1: 3149.18 Hz\n",
      "trapping frequency y 2: 3149.18 Hz\n",
      "trapping frequency z 1: 576.15 Hz\n",
      "trapping frequency z 2: 576.15 Hz\n",
      "harmonic oscillator length x: 176.20 nm\n",
      "harmonic oscillator length z: 330.10 nm\n",
      "tunneling distance: 916.91 nm\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 72020 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_________________________________________________________________\n",
      "depth 1: 8370.8575 Hz\n",
      "depth 2: 8284.3949 Hz\n",
      "target trap frequency x: 3000.00 HZ\n",
      "trapping frequency x 1: 2619.39 Hz\n",
      "trapping frequency x 2: 2619.39 Hz\n",
      "trapping frequency y 1: 3056.76 Hz\n",
      "trapping frequency y 2: 3056.76 Hz\n",
      "trapping frequency z 1: 619.72 Hz\n",
      "trapping frequency z 2: 619.72 Hz\n",
      "harmonic oscillator length x: 154.81 nm\n",
      "harmonic oscillator length z: 318.28 nm\n",
      "tunneling distance: 1328.43 nm\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 55916 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([3*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 -200*si.nm, dist0 +200*si.nm, 3)\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 minimum\")\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 = 25  \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=False)\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": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.04466052, 1.24466052, 1.44466052])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trap_spacings/si.uW"
   ]
  },
  {
   "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
}
