{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 2,
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import results from diagonalisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_line = 338\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": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line: 338\n",
      "separation: 0.7454 um\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 40316 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": 5,
   "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": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d = [745.43444684] nm\n",
      "omega_x = [50000.] Hz\n",
      "J = [10.32670994] Hz\n",
      "U_s = [1258.25386365] Hz\n",
      "E_exc = [6875.56532084] Hz\n",
      "E_barr = [40535.07317535] Hz\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 66324 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 scheme for different populations:\n",
    "\n",
    "The energy levels are shifted by the NNI ($V\\cdot n_{left}n_{right}$) and we have to account for this in ramping.\n",
    "\n",
    "The initial tilt is now shifted by the average shift of the transitions |n-1,1> -> |n,0> and |n-2,2> -> |n-1,1>, which is $V\\cdot (n-2)$\n",
    "\n",
    "## N=2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d: 745.43 nm\n",
      "omega_x: 50000.00 Hz\n",
      "J: 10.33 Hz\n",
      "limits set by: 1st excited state = 687.56 Hz, barrier = 4053.51 Hz, interactions = 4310.78 Hz\n",
      "=> max. tilt/interaction: 687.56 Hz\n",
      "J/U = 0.0300\n",
      "V: 27.94 Hz\n",
      "DeltaJ: 0.67 Hz\n",
      "ramp time: 968.363 ms\n",
      "fidelity: 0.99184\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 145800 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n = 2\n",
    "\n",
    "initial_occ=np.array([n-1,1])\n",
    "\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/n, E_max/n)       # On-site interaction strength\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",
    "Delta = ((n-2)*U[0]- (n-2)*V, n*U[0])        #tilt\n",
    "\n",
    "\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=initial_occ)\n",
    "\n",
    "d= n + 1\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 == np.sum(initial_occ)]\n",
    "labels_2particles = [basis_labels[idx] for idx in indices_2particles]\n",
    "\n",
    "#fidelity as the probability for |2,0> in the last 10% of steps\n",
    "#calculate fidelity\n",
    "fid  = np.mean(probs[int(0.9*len(tlist)):,indices_2particles[-1]])\n",
    "print(f\"fidelity: {fid:.5f}\")\n",
    "\n",
    "#alternatively with ratio Delta/U\n",
    "interactions = np.linspace(U[0],U[1],len(tlist))\n",
    "tilt = np.linspace(Delta[0],Delta[1],len(tlist))\n",
    "ratio = tilt/interactions\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(18, 5))\n",
    "\n",
    "# --- Plot 1: Probability of 2-particle basis states ---\n",
    "for idx, label in zip(indices_2particles, labels_2particles):\n",
    "    axs[0].plot(ratio[:int(1*len(tlist))], probs[:, idx][:int(1*len(tlist))], label=label)\n",
    "axs[0].set_xscale(\"log\")\n",
    "axs[0].set_xlim(0.5,2)\n",
    "axs[0].set_xlabel(r'$\\Delta /U$ []')\n",
    "axs[0].set_ylabel('Probability')\n",
    "axs[0].set_title(f'{n}Particle Basis State Probabilities')\n",
    "axs[0].legend()\n",
    "axs[0].grid(which=\"both\")\n",
    "\n",
    "# --- Plot 2: Occupation number vs time ---\n",
    "axs[1].plot(tlist, result.expect[0], label='Site 0')\n",
    "axs[1].plot(tlist, result.expect[1], label='Site 1')\n",
    "axs[1].set_xlabel('Time [s]')\n",
    "axs[1].set_ylabel('Occupation number')\n",
    "axs[1].set_title('Occupation Number Over Time')\n",
    "axs[1].legend()\n",
    "axs[1].grid(True)\n",
    "\n",
    "# --- Plot 3: Energies vs time ---\n",
    "axs[2].plot(tlist, np.linspace(Delta[0]/2/np.pi, Delta[1]/2/np.pi, N_steps), label=\"tilt\")\n",
    "\n",
    "axs[2].plot(tlist, ((1*U[0] - 0*U[0])/2/np.pi +V*(2*0 - 1*1)/2/np.pi) *np.ones(N_steps), label=\"transition |11> -> |20>\")\n",
    "\n",
    "axs[2].set_ylim(bottom=0)\n",
    "axs[2].set_xlabel('Time [s]')\n",
    "axs[2].set_ylabel('energy/h [Hz]')\n",
    "axs[2].set_title('parameter ramping')\n",
    "axs[2].legend()\n",
    "axs[2].grid(True)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## N=3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d: 745.43 nm\n",
      "omega_x: 50000.00 Hz\n",
      "J: 10.33 Hz\n",
      "limits set by: 1st excited state = 687.56 Hz, barrier = 4053.51 Hz, interactions = 4310.78 Hz\n",
      "=> max. tilt/interaction: 687.56 Hz\n",
      "J/U = 0.0451\n",
      "V: 27.94 Hz\n",
      "DeltaJ: 0.67 Hz\n",
      "ramp time: 968.363 ms\n",
      "fidelity: 0.96665\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 176584 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n = 3\n",
    "\n",
    "initial_occ=np.array([n-1,1])\n",
    "\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/n, E_max/n)       # On-site interaction strength\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",
    "Delta = ((n-2)*U[0]- (n-2)*V, n*U[0])        #tilt\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=initial_occ)\n",
    "\n",
    "d= n + 1\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 == np.sum(initial_occ)]\n",
    "labels_2particles = [basis_labels[idx] for idx in indices_2particles]\n",
    "\n",
    "#fidelity as the probability for |2,0> in the last 10% of steps\n",
    "#calculate fidelity\n",
    "fid  = np.mean(probs[int(0.9*len(tlist)):,indices_2particles[-1]])\n",
    "print(f\"fidelity: {fid:.5f}\")\n",
    "\n",
    "#alternatively with ratio Delta/U\n",
    "interactions = np.linspace(U[0],U[1],len(tlist))\n",
    "tilt = np.linspace(Delta[0],Delta[1],len(tlist))\n",
    "ratio = tilt/interactions\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(18, 5))\n",
    "\n",
    "# --- Plot 1: Probability of 2-particle basis states ---\n",
    "for idx, label in zip(indices_2particles, labels_2particles):\n",
    "    axs[0].plot(ratio[:int(1*len(tlist))], probs[:, idx][:int(1*len(tlist))], label=label)\n",
    "axs[0].set_xscale(\"log\")\n",
    "axs[0].set_xlabel(r'$\\Delta /U$ []')\n",
    "axs[0].set_ylabel('Probability')\n",
    "axs[0].set_title(f'{n}Particle Basis State Probabilities')\n",
    "axs[0].legend()\n",
    "axs[0].grid(which=\"both\")\n",
    "\n",
    "# --- Plot 2: Occupation number vs time ---\n",
    "axs[1].plot(tlist, result.expect[0], label='Site 0')\n",
    "axs[1].plot(tlist, result.expect[1], label='Site 1')\n",
    "axs[1].set_xlabel('Time [s]')\n",
    "axs[1].set_ylabel('Occupation number')\n",
    "axs[1].set_title('Occupation Number Over Time')\n",
    "axs[1].legend()\n",
    "axs[1].grid(True)\n",
    "\n",
    "# --- Plot 3: Energies vs time ---\n",
    "axs[2].plot(tlist, np.linspace(Delta[0]/2/np.pi, Delta[1]/2/np.pi, N_steps), label=\"tilt\")\n",
    "\n",
    "axs[2].plot(tlist, ((3*U[0] - 1*U[0])/2/np.pi +V*(3*0 - 2*1)/2/np.pi)*np.ones(N_steps), label=\"transition |21> -> |30>\")\n",
    "\n",
    "axs[2].set_ylim(bottom=0)\n",
    "axs[2].set_xlabel('Time [s]')\n",
    "axs[2].set_ylabel('energy/h [Hz]')\n",
    "axs[2].set_title('parameter ramping')\n",
    "axs[2].legend()\n",
    "axs[2].grid(True)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## N=4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d: 745.43 nm\n",
      "omega_x: 50000.00 Hz\n",
      "J: 10.33 Hz\n",
      "limits set by: 1st excited state = 687.56 Hz, barrier = 4053.51 Hz, interactions = 4310.78 Hz\n",
      "=> max. tilt/interaction: 687.56 Hz\n",
      "J/U = 0.0601\n",
      "V: 27.94 Hz\n",
      "DeltaJ: 0.67 Hz\n",
      "ramp time: 968.363 ms\n",
      "fidelity: 0.91251\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 205452 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n = 4\n",
    "\n",
    "initial_occ=np.array([n-1,1])\n",
    "\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/n, E_max/n)       # On-site interaction strength\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",
    "Delta = ((n-2)*U[0]- (n-2)*V, n*U[0])        #tilt\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=initial_occ)\n",
    "\n",
    "d= n + 1\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 == np.sum(initial_occ)]\n",
    "labels_2particles = [basis_labels[idx] for idx in indices_2particles]\n",
    "\n",
    "#fidelity as the probability for |2,0> in the last 10% of steps\n",
    "#calculate fidelity\n",
    "fid  = np.mean(probs[int(0.9*len(tlist)):,indices_2particles[-1]])\n",
    "print(f\"fidelity: {fid:.5f}\")\n",
    "\n",
    "#alternatively with ratio Delta/U\n",
    "interactions = np.linspace(U[0],U[1],len(tlist))\n",
    "tilt = np.linspace(Delta[0],Delta[1],len(tlist))\n",
    "ratio = tilt/interactions\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(18, 5))\n",
    "\n",
    "# --- Plot 1: Probability of 2-particle basis states ---\n",
    "for idx, label in zip(indices_2particles, labels_2particles):\n",
    "    axs[0].plot(ratio[:int(1*len(tlist))], probs[:, idx][:int(1*len(tlist))], label=label)\n",
    "axs[0].set_xscale(\"log\")\n",
    "axs[0].set_xlabel(r'$\\Delta /U$ []')\n",
    "axs[0].set_ylabel('Probability')\n",
    "axs[0].set_title(f'{n}Particle Basis State Probabilities')\n",
    "axs[0].legend()\n",
    "axs[0].grid(which=\"both\")\n",
    "\n",
    "# --- Plot 2: Occupation number vs time ---\n",
    "axs[1].plot(tlist, result.expect[0], label='Site 0')\n",
    "axs[1].plot(tlist, result.expect[1], label='Site 1')\n",
    "axs[1].set_xlabel('Time [s]')\n",
    "axs[1].set_ylabel('Occupation number')\n",
    "axs[1].set_title('Occupation Number Over Time')\n",
    "axs[1].legend()\n",
    "axs[1].grid(True)\n",
    "\n",
    "# --- Plot 3: Energies vs time ---\n",
    "axs[2].plot(tlist, np.linspace(Delta[0]/2/np.pi, Delta[1]/2/np.pi, N_steps), label=\"tilt\")\n",
    "\n",
    "axs[2].plot(tlist, ((6*U[0] - 3*U[0])/2/np.pi +V*(4*0 - 3*1)/2/np.pi)*np.ones(N_steps), label=\"transition |31> -> |40>\")\n",
    "axs[2].plot(tlist, ((3*U[0] - 2*U[0])/2/np.pi +V*(3*1 - 2*2)/2/np.pi)*np.ones(N_steps), label=\"transition |22> -> |31>\")\n",
    "\n",
    "axs[2].set_ylim(bottom=0)\n",
    "axs[2].set_xlabel('Time [s]')\n",
    "axs[2].set_ylabel('energy/h [Hz]')\n",
    "axs[2].set_title('parameter ramping')\n",
    "axs[2].legend()\n",
    "axs[2].grid(True)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## N=5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d: 745.43 nm\n",
      "omega_x: 50000.00 Hz\n",
      "J: 10.33 Hz\n",
      "limits set by: 1st excited state = 687.56 Hz, barrier = 4053.51 Hz, interactions = 4310.78 Hz\n",
      "=> max. tilt/interaction: 687.56 Hz\n",
      "J/U = 0.0751\n",
      "V: 27.94 Hz\n",
      "DeltaJ: 0.67 Hz\n",
      "ramp time: 968.363 ms\n",
      "fidelity: 0.83376\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 203020 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n = 5\n",
    "\n",
    "initial_occ=np.array([n-1,1])\n",
    "\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/n, E_max/n)       # On-site interaction strength\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",
    "Delta = ((n-2)*U[0]- (n-2)*V, n*U[0])        #tilt\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=initial_occ)\n",
    "\n",
    "d= n + 1\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 == np.sum(initial_occ)]\n",
    "labels_2particles = [basis_labels[idx] for idx in indices_2particles]\n",
    "\n",
    "#fidelity as the probability for |2,0> in the last 10% of steps\n",
    "#calculate fidelity\n",
    "fid  = np.mean(probs[int(0.9*len(tlist)):,indices_2particles[-1]])\n",
    "print(f\"fidelity: {fid:.5f}\")\n",
    "\n",
    "#alternatively with ratio Delta/U\n",
    "interactions = np.linspace(U[0],U[1],len(tlist))\n",
    "tilt = np.linspace(Delta[0],Delta[1],len(tlist))\n",
    "ratio = tilt/interactions\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(18, 5))\n",
    "\n",
    "# --- Plot 1: Probability of 2-particle basis states ---\n",
    "for idx, label in zip(indices_2particles, labels_2particles):\n",
    "    axs[0].plot(ratio[:int(1*len(tlist))], probs[:, idx][:int(1*len(tlist))], label=label)\n",
    "axs[0].set_xscale(\"log\")\n",
    "axs[0].set_xlabel(r'$\\Delta /U$ []')\n",
    "axs[0].set_ylabel('Probability')\n",
    "axs[0].set_title(f'{n}-Particle Basis State Probabilities')\n",
    "axs[0].legend()\n",
    "axs[0].grid(which=\"both\")\n",
    "\n",
    "# --- Plot 2: Occupation number vs time ---\n",
    "axs[1].plot(tlist, result.expect[0], label='Site 0')\n",
    "axs[1].plot(tlist, result.expect[1], label='Site 1')\n",
    "axs[1].set_xlabel('Time [s]')\n",
    "axs[1].set_ylabel('Occupation number')\n",
    "axs[1].set_title('Occupation Number Over Time')\n",
    "axs[1].legend()\n",
    "axs[1].grid(True)\n",
    "\n",
    "# --- Plot 3: Energies vs time ---\n",
    "axs[2].plot(tlist, np.linspace(Delta[0]/2/np.pi, Delta[1]/2/np.pi, N_steps), label=\"tilt\")\n",
    "\n",
    "axs[2].plot(tlist, ((10*U[0] - 6*U[0])/2/np.pi +V*(5*0 - 4*1)/2/np.pi)*np.ones(N_steps), label=\"transition |41> -> |50>\")\n",
    "axs[2].plot(tlist, ((6*U[0] - 4*U[0])/2/np.pi +V*(4*1 - 3*2)/2/np.pi)*np.ones(N_steps), label=\"transition |32> -> |41>\")\n",
    "\n",
    "axs[2].set_ylim(bottom=0)\n",
    "axs[2].set_xlabel('Time [s]')\n",
    "axs[2].set_ylabel('energy/h [Hz]')\n",
    "axs[2].set_title('parameter ramping')\n",
    "axs[2].legend()\n",
    "axs[2].grid(True)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## N=6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "d: 745.43 nm\n",
      "omega_x: 50000.00 Hz\n",
      "J: 10.33 Hz\n",
      "limits set by: 1st excited state = 687.56 Hz, barrier = 4053.51 Hz, interactions = 4310.78 Hz\n",
      "=> max. tilt/interaction: 687.56 Hz\n",
      "J/U = 0.0901\n",
      "V: 27.94 Hz\n",
      "DeltaJ: 0.67 Hz\n",
      "ramp time: 968.363 ms\n",
      "fidelity: 0.75204\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 214792 characters]",
      "text/plain": [
       "<Figure size 1800x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n = 6\n",
    "\n",
    "initial_occ=np.array([n-1,1])\n",
    "\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/n, E_max/n)       # On-site interaction strength\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",
    "Delta = ((n-2)*U[0]- (n-2)*V, n*U[0])        #tilt\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=initial_occ)\n",
    "\n",
    "d= n + 1\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 == np.sum(initial_occ)]\n",
    "labels_2particles = [basis_labels[idx] for idx in indices_2particles]\n",
    "\n",
    "#fidelity as the probability for |2,0> in the last 10% of steps\n",
    "#calculate fidelity\n",
    "fid  = np.mean(probs[int(0.9*len(tlist)):,indices_2particles[-1]])\n",
    "print(f\"fidelity: {fid:.5f}\")\n",
    "\n",
    "#alternatively with ratio Delta/U\n",
    "interactions = np.linspace(U[0],U[1],len(tlist))\n",
    "tilt = np.linspace(Delta[0],Delta[1],len(tlist))\n",
    "ratio = tilt/interactions\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(18, 5))\n",
    "\n",
    "# --- Plot 1: Probability of 2-particle basis states ---\n",
    "for idx, label in zip(indices_2particles, labels_2particles):\n",
    "    axs[0].plot(ratio[:int(1*len(tlist))], probs[:, idx][:int(1*len(tlist))], label=label)\n",
    "axs[0].set_xscale(\"log\")\n",
    "axs[0].set_xlabel(r'$\\Delta /U$ []')\n",
    "axs[0].set_ylabel('Probability')\n",
    "axs[0].set_title(f'{n}-Particle Basis State Probabilities')\n",
    "axs[0].legend()\n",
    "axs[0].grid(which=\"both\")\n",
    "\n",
    "# --- Plot 2: Occupation number vs time ---\n",
    "axs[1].plot(tlist, result.expect[0], label='Site 0')\n",
    "axs[1].plot(tlist, result.expect[1], label='Site 1')\n",
    "axs[1].set_xlabel('Time [s]')\n",
    "axs[1].set_ylabel('Occupation number')\n",
    "axs[1].set_title('Occupation Number Over Time')\n",
    "axs[1].legend()\n",
    "axs[1].grid(True)\n",
    "\n",
    "# --- Plot 3: Energies vs time ---\n",
    "axs[2].plot(tlist, np.linspace(Delta[0]/2/np.pi, Delta[1]/2/np.pi, N_steps), label=\"tilt\")\n",
    "\n",
    "axs[2].plot(tlist, ((15*U[0] - 10*U[0])/2/np.pi +V*(6*0 - 5*1)/2/np.pi)*np.ones(N_steps), label=\"transition |51> -> |60>\")\n",
    "axs[2].plot(tlist, ((10*U[0] - 7*U[0])/2/np.pi +V*(5*1 - 4*2)/2/np.pi)*np.ones(N_steps), label=\"transition |42> -> |51>\")\n",
    "axs[2].plot(tlist, ((7*U[0] - 6*U[0])/2/np.pi +V*(4*2 - 3*3)/2/np.pi)*np.ones(N_steps), label=\"transition |33> -> |42>\")\n",
    "\n",
    "axs[2].set_ylim(bottom=0)\n",
    "axs[2].set_xlabel('Time [s]')\n",
    "axs[2].set_ylabel('energy/h [Hz]')\n",
    "axs[2].set_title('parameter ramping')\n",
    "axs[2].legend()\n",
    "axs[2].grid(True)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  }
 ],
 "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
}
