{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sympy as sp\n",
    "import qutip as qt\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import constants as const\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 boson_helpers import*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_hub(J,U,Delta,mu,V,DeltaJ,tau,\n",
    "              N_steps=1000, initial_occ=[1,1]):\n",
    "    \"\"\"\n",
    "    Solves and evolves (duration of tau) the Hubbard-Hamiltonian for the given\n",
    "    parameters and ramping. Input parameters as angular frequencies and U,Delta\n",
    "    as tuple of start and end value.\n",
    "    \"\"\"\n",
    "    # Maximum boson number per site (Hilbert space dimension d = Nmax+1)\n",
    "    Nmax = initial_occ[0] + initial_occ[1]\n",
    "    d = Nmax + 1\n",
    "\n",
    "    # --- Define Operators ---\n",
    "    # Single-site bosonic operators in the truncated Fock space\n",
    "    a = qt.destroy(d)\n",
    "    adag = a.dag()\n",
    "    n_op = adag * a\n",
    "\n",
    "    # --- Construct Hamiltonian ---\n",
    "    # Hopping term: -J (a† ⊗ a + a ⊗ a†)\n",
    "    H_hop = -J* (qt.tensor(adag, a) + qt.tensor(a, adag))\n",
    "\n",
    "    # Interaction term: (U/2) * [n(n-1)] on each site\n",
    "    # We subtract the identity appropriately to compute n(n-1)\n",
    "    H_int = 0.5*(qt.tensor(n_op * (n_op - qt.qeye(d)), qt.qeye(d)) +\n",
    "                    qt.tensor(qt.qeye(d), n_op * (n_op - qt.qeye(d))))\n",
    "    # Define the time-dependent function for the interaction\n",
    "    def int_func(t, args):\n",
    "        return args['U0']+ (args['U1']-args['U0'])* t/args['tau']\n",
    "    int_args = {'U0': U[0], 'U1': U[1], 'tau': tau}\n",
    "    \n",
    "    # Chemical potential term (if desired)\n",
    "    H_mu = -mu* (qt.tensor(n_op, qt.qeye(d)) + qt.tensor(qt.qeye(d), n_op))\n",
    "\n",
    "    #tilt potential\n",
    "    H_tilt = (1*qt.tensor(n_op, qt.qeye(d)) + 2*qt.tensor(qt.qeye(d), n_op))\n",
    "    # Define the time-dependent function for the tilt\n",
    "    def tilt_func(t, args):\n",
    "        return args['D0']+ (args['D1']-args['D0'])* t/args['tau']\n",
    "    tilt_args = {'D0': Delta[0], 'D1': Delta[1], 'tau': tau}\n",
    "\n",
    "    #nearest neighbour interaction\n",
    "    H_nni = V* qt.tensor(n_op, qt.qeye(d)) * qt.tensor(qt.qeye(d), n_op)\n",
    "\n",
    "    #density induced tunneling\n",
    "    H_dit = -DeltaJ* (qt.tensor(adag,a)* (qt.tensor(n_op,qt.qeye(d)) + qt.tensor(n_op,qt.qeye(d)) - 1) +\n",
    "                        qt.tensor(a,adag)* (qt.tensor(n_op,qt.qeye(d)) + qt.tensor(n_op,qt.qeye(d)) - 1))\n",
    "\n",
    "\n",
    "    # Total Hamiltonian\n",
    "    H_static = H_hop + H_mu + H_nni + H_dit\n",
    "    H = [H_static, [H_int, int_func],[H_tilt, tilt_func]]\n",
    "    \n",
    "\n",
    "\n",
    "    # --- Define Initial State ---\n",
    "    # The product Fock state |1,1> (one boson per site)\n",
    "    psi0 = qt.tensor(qt.basis(d, initial_occ[0]), qt.basis(d, initial_occ[1]))\n",
    "\n",
    "    # --- Time Evolution Setup ---\n",
    "    tlist = np.linspace(0, tau, N_steps)  # Evolve from time 0 to 5 in 101 steps\n",
    "\n",
    "    # We want to monitor the number operator expectation values at each site.\n",
    "    e_ops = [qt.tensor(n_op, qt.qeye(d)), qt.tensor(qt.qeye(d), n_op)]\n",
    "\n",
    "    # Create an Options object with store_states=True\n",
    "    #opts = qt.Options(store_states=True, nsteps=10000)\n",
    "    opts = {\"store_states\": True, \"nsteps\": 10000}\n",
    "\n",
    "    # Run the evolution and store the full state history\n",
    "    result = qt.sesolve(H, psi0, tlist, e_ops=e_ops, options=opts,\n",
    "                        args=int_args | tilt_args)\n",
    "    \n",
    "    return tlist, result"
   ]
  },
  {
   "attachments": {
    "image.png": {
     "image/png": "[encoded data, 154280 characters]"
    }
   },
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Verify simulation with results from Kaufmann paper:\n",
    "\n",
    "Neglectable interactions and J=262Hz (left), J=348Hz (right)\n",
    "\n",
    "![image.png](attachment:image.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Parameters ---\n",
    "J = 2*np.pi* 262#Hz       # Hopping amplitude\n",
    "U = (0, 0)       # On-site interaction strength\n",
    "Delta = (0, 0)        #tilt\n",
    "\n",
    "mu = 0.0      # Chemical potential (set to zero for simplicity)\n",
    "V = 0.        # nearest neighbour interaction energy (due to DDI)\n",
    "DeltaJ = 0.   # density induced tunneling\n",
    "\n",
    "tau = 2*si.ms     #time duration of simulation\n",
    "N_steps = int(1e3)      #number of timesteps\n",
    "\n",
    "tlist, result = solve_hub(J,U,Delta, mu,V,DeltaJ,tau,N_steps=N_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 79552 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "d=3\n",
    "\n",
    "# --- Extracting and Plotting State Vector Components ---\n",
    "# Create an array of states; each state is flattened into a 1D array of length d^2\n",
    "states = np.array([state.full().flatten() for state in result.states])\n",
    "# Compute probabilities: squared magnitude of each component\n",
    "probs = np.abs(states)**2\n",
    "\n",
    "# Create labels for each basis state in the two-site Fock space\n",
    "basis_labels = [f\"|{n0},{n1}>\" for n0 in range(d) for n1 in range(d)]\n",
    "# Create a list of corresponding total particle numbers for each basis state\n",
    "total_particles = [n0 + n1 for n0 in range(d) for n1 in range(d)]\n",
    "\n",
    "# Filter indices for basis states with exactly 2 particles\n",
    "indices_2particles = [idx for idx, tot in enumerate(total_particles) if tot == 2]\n",
    "labels_2particles = [basis_labels[idx] for idx in indices_2particles]\n",
    "\n",
    "# Plot the probability for each 2-particle basis state over time\n",
    "plt.figure(figsize=(10, 6))\n",
    "for idx, label in zip(indices_2particles, labels_2particles):\n",
    "    plt.plot(tlist, probs[:, idx], label=label)\n",
    "plt.xlabel('Time [s]')\n",
    "plt.ylabel('Probability')\n",
    "plt.title('Time Evolution of Basis States with 2 Particles')\n",
    "plt.legend(loc=\"upper right\", bbox_to_anchor=(1.25, 1))\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Parameters ---\n",
    "J = 2*np.pi* 348#Hz       # Hopping amplitude\n",
    "U = (0, 0)       # On-site interaction strength\n",
    "Delta = (0, 0)        #tilt\n",
    "\n",
    "mu = 0.0      # Chemical potential (set to zero for simplicity)\n",
    "V = 0.        # nearest neighbour interaction energy (due to DDI)\n",
    "DeltaJ = 0.   # density induced tunneling\n",
    "\n",
    "tau = 1.7*si.ms     #time duration of simulation\n",
    "N_steps = int(1e3)      #number of timesteps\n",
    "\n",
    "tlist, result = solve_hub(J,U,Delta, mu,V,DeltaJ,tau,N_steps=N_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 83108 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "d=3\n",
    "\n",
    "# --- Extracting and Plotting State Vector Components ---\n",
    "# Create an array of states; each state is flattened into a 1D array of length d^2\n",
    "states = np.array([state.full().flatten() for state in result.states])\n",
    "# Compute probabilities: squared magnitude of each component\n",
    "probs = np.abs(states)**2\n",
    "\n",
    "# Create labels for each basis state in the two-site Fock space\n",
    "basis_labels = [f\"|{n0},{n1}>\" for n0 in range(d) for n1 in range(d)]\n",
    "# Create a list of corresponding total particle numbers for each basis state\n",
    "total_particles = [n0 + n1 for n0 in range(d) for n1 in range(d)]\n",
    "\n",
    "# Filter indices for basis states with exactly 2 particles\n",
    "indices_2particles = [idx for idx, tot in enumerate(total_particles) if tot == 2]\n",
    "labels_2particles = [basis_labels[idx] for idx in indices_2particles]\n",
    "\n",
    "# Plot the probability for each 2-particle basis state over time\n",
    "plt.figure(figsize=(10, 6))\n",
    "for idx, label in zip(indices_2particles, labels_2particles):\n",
    "    plt.plot(tlist, probs[:, idx], label=label)\n",
    "plt.xlabel('Time [s]')\n",
    "plt.ylabel('Probability')\n",
    "plt.title('Time Evolution of Basis States with 2 Particles')\n",
    "plt.legend(loc=\"upper right\", bbox_to_anchor=(1.25, 1))\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import results from diagonalisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "#185,186,187,188 give nice examples for different spacings\n",
    "start_line = 186\n",
    "end_line = start_line\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": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 186\n",
      "separation: 0.7466 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37340 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": 19,
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d = [746.63873654] nm\n",
      "omega_x = [26000.] Hz\n",
      "J = [207.01515899] Hz\n",
      "U_s = [442.01909695] Hz\n",
      "E_exc = [3523.90755057] Hz\n",
      "E_barr = [11674.2389474] Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 65172 characters]",
      "text/plain": [
       "<Figure size 500x900 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(f\"d = {distance/si.nm} nm\")\n",
    "print(f\"omega_x = {frequencies/2/np.pi} Hz\")\n",
    "print(f\"J = {J/const.h} Hz\")\n",
    "print(f\"U_s = {U_s/const.h} Hz\")\n",
    "print(f\"E_exc = {E_exc/const.h} Hz\")\n",
    "print(f\"E_barr = {E_barr/const.h} Hz\")\n",
    "\n",
    "# Create figure with subplots\n",
    "fig, axes = plt.subplots(3, 1, figsize=(5, 9), sharex=True)\n",
    "\n",
    "for i, line in enumerate(lines):\n",
    "    # Plot U_dd\n",
    "    axes[0].plot(np.rad2deg(angles[i]), U_dds[i]/const.h,label=f\"{line}\")\n",
    "    axes[0].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[0].set_ylabel(r\"$U_{dd} / h$ [Hz]\")\n",
    "    axes[0].legend()\n",
    "\n",
    "    # Plot NNI\n",
    "    axes[1].plot(np.rad2deg(angles[i]), V_lrs[i]/const.h,label=f\"{line}\")\n",
    "    axes[1].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[1].set_ylabel(r\"$V_{lr} / h$ [Hz]\")\n",
    "    axes[1].legend()\n",
    "\n",
    "    # Plot DIT\n",
    "    axes[2].plot(np.rad2deg(angles[i]), DeltaJs[i]/const.h,label=f\"{line}\")\n",
    "    axes[2].axhline(0, color=\"black\", ls=\"--\")\n",
    "    axes[2].set_xlabel(r\"$\\theta$ [deg]\")\n",
    "    axes[2].set_ylabel(r\"$\\Delta J / h$ [Hz]\")\n",
    "    axes[2].legend()\n",
    "\n",
    "    # Customize ticks for all subplots\n",
    "    for ax in axes:\n",
    "        ax.tick_params(axis='both', which='both', direction='in', length=5)  # Tick marks inside\n",
    "        ax.tick_params(axis='both', which='major', length=7)  # Major ticks longer\n",
    "        ax.tick_params(axis='both', which='minor', length=4)  # Minor ticks shorter\n",
    "        ax.minorticks_on()  # Enable minor ticks\n",
    "\n",
    "# Adjust layout and show the plot\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ramping tilt and interactions against each other:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d: 746.64 nm\n",
      "omega_x: 26000.00 Hz\n",
      "J: 207.02 Hz\n",
      "limits set by: 1st excited state = 352.39 Hz, barrier = 1167.42 Hz, interactions = 1511.64 Hz\n",
      "=> max. tilt/interaction: 352.39 Hz\n",
      "J/U = 0.5875\n",
      "V: 34.16 Hz\n",
      "DeltaJ: 9.31 Hz\n",
      "ramp time: 48.306 ms\n"
     ]
    }
   ],
   "source": [
    "#maximum energy scale for ramp\n",
    "#assuming that staying below 25% of the 1st excited state is enough\n",
    "#and the contact scattering length can be tuned from 0 to 3.5* 85a0\n",
    "limits = np.array([E_exc[0]/10,\n",
    "                   E_barr[0]/10,\n",
    "                   U_s[0]*3.5 + np.max(U_dds)])\n",
    "E_max = np.min(limits)/const.hbar\n",
    "\n",
    "# --- Parameters ---\n",
    "J_freq = J[0]/const.hbar      # Hopping amplitude\n",
    "U = (E_max, 0)       # On-site interaction strength\n",
    "Delta = (0, E_max)        #tilt\n",
    "\n",
    "mu = 0.0      # Chemical potential (set to zero for simplicity)\n",
    "V = V_lrs.max()/const.hbar    # nearest neighbour interaction energy (due to DDI)\n",
    "DeltaJ = DeltaJs.max()/const.hbar   # density induced tunneling\n",
    "\n",
    "tau = 2*np.pi/J_freq *10    #take 10 times the time scale set by J for evolution\n",
    "N_steps = int(1e3)      #number of timesteps\n",
    "\n",
    "print(f\"d: {distance[0]/si.nm:.2f} nm\")\n",
    "print(f\"omega_x: {frequencies[0]/2/np.pi:.2f} Hz\")\n",
    "print(f\"J: {J_freq/2/np.pi:.2f} Hz\")\n",
    "print(f\"limits set by: 1st excited state = {limits[0]/const.h:.2f} Hz, barrier = {limits[1]/const.h:.2f} Hz, interactions = {limits[2]/const.h:.2f} Hz\")\n",
    "print(f\"=> max. tilt/interaction: {E_max/2/np.pi:.2f} Hz\")\n",
    "print(f\"J/U = {J_freq/U[0]:.4f}\")\n",
    "print(f\"V: {V/2/np.pi:.2f} Hz\")\n",
    "print(f\"DeltaJ: {DeltaJ/2/np.pi:.2f} Hz\")\n",
    "print(f\"ramp time: {tau/si.ms:.3f} ms\")\n",
    "\n",
    "tlist, result = solve_hub(J_freq,U,Delta, mu,V,DeltaJ,tau,N_steps=N_steps,\n",
    "                          initial_occ=[1,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fidelity: 0.42921\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 264820 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 64812 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "d=3\n",
    "\n",
    "# --- Extracting and Plotting State Vector Components ---\n",
    "# Create an array of states; each state is flattened into a 1D array of length d^2\n",
    "states = np.array([state.full().flatten() for state in result.states])\n",
    "# Compute probabilities: squared magnitude of each component\n",
    "probs = np.abs(states)**2\n",
    "#fidelity as the probability for |2,0> in the last 10% of steps\n",
    "fid = np.mean(probs[-int(N_steps/10):, 6])\n",
    "#fid = np.mean(probs[-int(N_steps/5):, 2] + probs[-int(N_steps/5):, 6])\n",
    "print(f\"fidelity: {fid:.5f}\")\n",
    "\n",
    "# Create labels for each basis state in the two-site Fock space\n",
    "basis_labels = [f\"|{n0},{n1}>\" for n0 in range(d) for n1 in range(d)]\n",
    "# Create a list of corresponding total particle numbers for each basis state\n",
    "total_particles = [n0 + n1 for n0 in range(d) for n1 in range(d)]\n",
    "\n",
    "# Filter indices for basis states with exactly 2 particles\n",
    "indices_2particles = [idx for idx, tot in enumerate(total_particles) if tot == 2]\n",
    "labels_2particles = [basis_labels[idx] for idx in indices_2particles]\n",
    "\n",
    "# Plot the probability for each 2-particle basis state over time\n",
    "plt.figure(figsize=(10, 6))\n",
    "for idx, label in zip(indices_2particles, labels_2particles):\n",
    "    plt.plot(tlist, probs[:, idx], label=label)\n",
    "plt.xlabel('Time [s]')\n",
    "plt.ylabel('Probability')\n",
    "plt.title('Time Evolution of Basis States with 2 Particles')\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# --- Plot the Results ---\n",
    "plt.figure(figsize=(8, 5))\n",
    "plt.plot(tlist, result.expect[0], label='Site 0')\n",
    "plt.plot(tlist, result.expect[1], label='Site 1')\n",
    "plt.ylim(-0.1,2.1)\n",
    "plt.xlabel('Time [s]')\n",
    "plt.ylabel('Occupation number')\n",
    "plt.title('Time Evolution of a 2-Site Bose-Hubbard Model')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check what tunneling time should be used"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57496 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tunneling_factor = np.logspace(-1,2,20)\n",
    "tunneling_time = np.zeros_like(tunneling_factor)\n",
    "fidelity = np.zeros_like(tunneling_factor)\n",
    "\n",
    "for i, fac in enumerate(tunneling_factor):\n",
    "    #maximum energy scale for ramp\n",
    "    #assuming that staying below 25% of the 1st excited state is enough\n",
    "    #and the contact scattering length can be tuned from 0 to 3.5* 85a0\n",
    "    limits = np.array([E_exc[0]/10,\n",
    "                    E_barr[0]/10,\n",
    "                    U_s[0]*3.5 + np.max(U_dds)])\n",
    "    E_max = np.min(limits)/const.hbar\n",
    "\n",
    "    # --- Parameters ---\n",
    "    J_freq = J[0]/const.hbar      # Hopping amplitude\n",
    "    U = (E_max, 0)       # On-site interaction strength\n",
    "    Delta = (0, E_max)        #tilt\n",
    "\n",
    "    mu = 0.0      # Chemical potential (set to zero for simplicity)\n",
    "    V = V_lrs.max()/const.hbar    # nearest neighbour interaction energy (due to DDI)\n",
    "    DeltaJ = DeltaJs.max()/const.hbar   # density induced tunneling\n",
    "\n",
    "    tau = 2*np.pi/J_freq *fac    #take 10 times the time scale set by J for evolution\n",
    "    N_steps = int(1e3)      #number of timesteps\n",
    "\n",
    "    # print(f\"J: {J_freq/2/np.pi:.2f} Hz\")\n",
    "    # print(f\"max. tilt: {E_max/2/np.pi:.2f} Hz, {limits/const.h}\")\n",
    "    # print(f\"V: {V/2/np.pi:.2f} Hz\")\n",
    "    # print(f\"DeltaJ: {DeltaJ/2/np.pi:.2f} Hz\")\n",
    "    # print(f\"ramp time: {tau/si.ms:.3f} ms\")\n",
    "\n",
    "    tlist, result = solve_hub(J_freq,U,Delta, mu,V,DeltaJ,tau,N_steps=N_steps)\n",
    "\n",
    "    # Create an array of states; each state is flattened into a 1D array of length d^2\n",
    "    states = np.array([state.full().flatten() for state in result.states])\n",
    "    # Compute probabilities: squared magnitude of each component\n",
    "    probs = np.abs(states)**2\n",
    "    #fidelity as the probability for |2,0> in the last 10% of steps\n",
    "    fid = np.mean(probs[-int(N_steps/10):, 6])\n",
    "    #fid = np.mean(probs[-int(N_steps/5):, 2] + probs[-int(N_steps/5):, 6])\n",
    "\n",
    "    tunneling_time[i] = tau\n",
    "    fidelity[i] = fid\n",
    "\n",
    "\n",
    "# Main plot\n",
    "fig, ax = plt.subplots(figsize=(8, 5))\n",
    "ax.plot(tunneling_factor, fidelity)\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_ylim(0,1)\n",
    "ax.set_xlabel(r\"$t/(2\\pi/J)$\")\n",
    "ax.set_ylabel(\"Fidelity\")\n",
    "ax.grid(True, which=\"both\", linestyle=\"--\", alpha=0.6)\n",
    "\n",
    "# Add top axis with tick locations from tunneling_factor and labels from tunneling_time\n",
    "secax = ax.secondary_xaxis('top')\n",
    "secax.set_xticks(tunneling_factor[::2])\n",
    "secax.set_xticklabels([f\"{t*1e3:.1f}\" for t in tunneling_time[::2]])\n",
    "secax.set_xlabel(\"Tunneling time [ms]\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How does fidelity and tunneling time depend on tweezer distance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 334\n",
    "end_line = 341\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": "code",
   "execution_count": 207,
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "[ 675.10721513 2155.68325884 3869.12998812]\n",
      "0\n",
      "0.6527479874251076\n",
      "0.0015725667494999004\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 55112 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "[ 678.44588677 2579.61125044 4026.62151098]\n",
      "1\n",
      "0.8161957078795382\n",
      "0.013025401795357784\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 53672 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "[ 681.62924036 3037.77940289 4163.20310173]\n",
      "2\n",
      "0.9112853199231185\n",
      "0.026479374771787748\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 53456 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "[ 684.65875099 3529.35014616 4282.57450356]\n",
      "3\n",
      "0.9559721248867851\n",
      "0.02712858277747444\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 53372 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "[ 687.55653208 4053.50731753 4388.77697771]\n",
      "4\n",
      "0.9796827277919958\n",
      "0.0857840295316235\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 53112 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "[ 690.34322583 4609.45542773 4484.82498964]\n",
      "5\n",
      "0.9862064304065286\n",
      "0.1947417378176338\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 52680 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "[ 693.03419726 5196.41904649 4572.79414173]\n",
      "6\n",
      "0.9933936540354984\n",
      "0.21992506315165225\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 52600 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "[ 695.64053596 5813.64214987 4654.13031727]\n",
      "7\n",
      "0.9942814878060094\n",
      "0.366761312513934\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 49520 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fidelity_d = np.zeros(n_lines)\n",
    "tunneling_time_d = np.zeros(n_lines)\n",
    "\n",
    "for j in range(n_lines):\n",
    "\n",
    "    tunneling_factor = np.logspace(-1,2,20)\n",
    "    tunneling_time = np.zeros_like(tunneling_factor)\n",
    "    fidelity = np.zeros_like(tunneling_factor)\n",
    "\n",
    "    for i, fac in enumerate(tunneling_factor):\n",
    "        #maximum energy scale for ramp\n",
    "        #assuming that staying below 25% of the 1st excited state is enough\n",
    "        #and the contact scattering length can be tuned from 0 to 3.5* 85a0\n",
    "        limits = np.array([E_exc[j]/10,\n",
    "                        E_barr[j]/10,\n",
    "                        U_s[j]*3.5 + np.max(U_dds)])\n",
    "        print(limits/const.h)\n",
    "        E_max = np.min(limits)/const.hbar\n",
    "\n",
    "        # --- Parameters ---\n",
    "        J_freq = J[j]/const.hbar      # Hopping amplitude\n",
    "        U = (E_max, 0)       # On-site interaction strength\n",
    "        Delta = (0, E_max)        #tilt\n",
    "\n",
    "        mu = 0.0      # Chemical potential (set to zero for simplicity)\n",
    "        V = V_lrs.max()/const.hbar    # nearest neighbour interaction energy (due to DDI)\n",
    "        DeltaJ = DeltaJs.max()/const.hbar   # density induced tunneling\n",
    "\n",
    "        tau = 2*np.pi/J_freq *fac    #take 10 times the time scale set by J for evolution\n",
    "        N_steps = int(1e3)      #number of timesteps\n",
    "\n",
    "        # print(f\"J: {J_freq/2/np.pi:.2f} Hz\")\n",
    "        # print(f\"max. tilt: {E_max/2/np.pi:.2f} Hz, {limits/const.h}\")\n",
    "        # print(f\"V: {V/2/np.pi:.2f} Hz\")\n",
    "        # print(f\"DeltaJ: {DeltaJ/2/np.pi:.2f} Hz\")\n",
    "        # print(f\"ramp time: {tau/si.ms:.3f} ms\")\n",
    "\n",
    "        tlist, result = solve_hub(J_freq,U,Delta, mu,V,DeltaJ,tau,N_steps=N_steps)\n",
    "\n",
    "        # Create an array of states; each state is flattened into a 1D array of length d^2\n",
    "        states = np.array([state.full().flatten() for state in result.states])\n",
    "        # Compute probabilities: squared magnitude of each component\n",
    "        probs = np.abs(states)**2\n",
    "        #fidelity as the probability for |2,0> in the last 10% of steps\n",
    "        fid = np.mean(probs[-int(N_steps/10):, 6])\n",
    "        #fid = np.mean(probs[-int(N_steps/5):, 2] + probs[-int(N_steps/5):, 6])\n",
    "\n",
    "        tunneling_time[i] = tau\n",
    "        fidelity[i] = fid\n",
    "\n",
    "\n",
    "    print(j)\n",
    "    print(fidelity[-1])\n",
    "    #maximum fidelity is at longest ramp time\n",
    "    fidelity_d[j] = fidelity[-1]\n",
    "    #interpolate results to find \n",
    "\n",
    "    mask = abs(fidelity - fidelity[-1]) < 0.01\n",
    "    tunneling_time_d[j] = np.min(tunneling_time[mask])\n",
    "    print(tunneling_time_d[j])\n",
    "\n",
    "    # Main plot\n",
    "    fig, ax = plt.subplots(figsize=(8, 5))\n",
    "    ax.plot(tunneling_factor, fidelity)\n",
    "    ax.set_xscale(\"log\")\n",
    "    ax.set_ylim(0,1)\n",
    "    ax.set_xlabel(r\"$t/(2\\pi/J)$\")\n",
    "    ax.set_ylabel(\"Fidelity\")\n",
    "    ax.grid(True, which=\"both\", linestyle=\"--\", alpha=0.6)\n",
    "\n",
    "    # Add top axis with tick locations from tunneling_factor and labels from tunneling_time\n",
    "    secax = ax.secondary_xaxis('top')\n",
    "    secax.set_xticks(tunneling_factor[::2])\n",
    "    secax.set_xticklabels([f\"{t*1e3:.1f}\" for t in tunneling_time[::2]])\n",
    "    secax.set_xlabel(\"Tunneling time [ms]\")\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 217,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 186768 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Create a figure with 4 subplots (stacked vertically)\n",
    "fig, axes = plt.subplots(2, 1, figsize=(10, 10), sharex=True)\n",
    "\n",
    "axes[0].plot(distance/si.nm, fidelity_d)\n",
    "#axes[0].set_yscale(\"log\")\n",
    "axes[0].set_ylim(0.5,1)\n",
    "axes[0].set_ylabel(\"maximum fidelity\")\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",
    "\n",
    "axes[1].plot(distance/si.nm, tunneling_time_d/si.ms)\n",
    "axes[1].set_ylabel(\"minimum ramp duration (1% of max. fidelity) [ms]\")\n",
    "axes[1].set_xlabel(\"tweezer spacing [nm]\")\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",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "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
}
