{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\naeve\\anaconda3\\envs\\use_tenpy\\Lib\\site-packages\\tenpy\\tools\\params.py:232: UserWarning: unused option ['Nmax'] for config BoseHubbardModel\n",
      "  warnings.warn(msg.format(keys=sorted(unused), name=self.name))\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "module 'tenpy.algorithms.tebd' has no attribute 'TEBD'",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mAttributeError\u001b[39m                            Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[2]\u001b[39m\u001b[32m, line 36\u001b[39m\n\u001b[32m     24\u001b[39m tebd_params = {\n\u001b[32m     25\u001b[39m     \u001b[33m'\u001b[39m\u001b[33mtrunc_params\u001b[39m\u001b[33m'\u001b[39m: {\n\u001b[32m     26\u001b[39m         \u001b[33m'\u001b[39m\u001b[33mchi_max\u001b[39m\u001b[33m'\u001b[39m: \u001b[32m100\u001b[39m,  \u001b[38;5;66;03m# Maximum bond dimension\u001b[39;00m\n\u001b[32m   (...)\u001b[39m\u001b[32m     32\u001b[39m     \u001b[33m'\u001b[39m\u001b[33mverbose\u001b[39m\u001b[33m'\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m,     \u001b[38;5;66;03m# Print progress\u001b[39;00m\n\u001b[32m     33\u001b[39m }\n\u001b[32m     35\u001b[39m \u001b[38;5;66;03m# Initialize the TEBD engine with the state and model\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m36\u001b[39m tebd_engine = \u001b[43mtebd\u001b[49m\u001b[43m.\u001b[49m\u001b[43mTEBD\u001b[49m(psi, model, tebd_params)\n\u001b[32m     38\u001b[39m \u001b[38;5;66;03m# Prepare lists to store observables during time evolution\u001b[39;00m\n\u001b[32m     39\u001b[39m times = []\n",
      "\u001b[31mAttributeError\u001b[39m: module 'tenpy.algorithms.tebd' has no attribute 'TEBD'"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from tenpy.models.hubbard import BoseHubbardModel\n",
    "from tenpy.networks.mps import MPS\n",
    "from tenpy.algorithms import tebd\n",
    "\n",
    "# --- Define Model Parameters ---\n",
    "# L: number of sites; Nmax: maximum number of bosons per site (should be >= max occupation)\n",
    "model_params = {\n",
    "    'L': 2,          # Two sites\n",
    "    'Nmax': 3,       # Maximum occupation per site (adjust as needed)\n",
    "    't': 1.0,        # Hopping amplitude\n",
    "    'U': 1.0,        # On-site interaction strength\n",
    "    'mu': 0.0,       # Chemical potential\n",
    "}\n",
    "model = BoseHubbardModel(model_params)\n",
    "\n",
    "# --- Prepare the Initial State ---\n",
    "# Here we use the product state [1, 1], meaning 1 boson per site.\n",
    "# The method lat.mps_sites() sets up the local Hilbert spaces.\n",
    "initial_state = [1, 1]\n",
    "psi = MPS.from_product_state(model.lat.mps_sites(), initial_state, bc=\"finite\")\n",
    "\n",
    "# --- TEBD Time Evolution Parameters ---\n",
    "tebd_params = {\n",
    "    'trunc_params': {\n",
    "        'chi_max': 100,  # Maximum bond dimension\n",
    "        'svd_min': 1e-10,  # SVD truncation threshold\n",
    "    },\n",
    "    'N_steps': 100,      # Number of time steps\n",
    "    'dt': 0.05,          # Time step size\n",
    "    'order': 2,          # Second order Trotter decomposition\n",
    "    'verbose': True,     # Print progress\n",
    "}\n",
    "\n",
    "# Initialize the TEBD engine with the state and model\n",
    "tebd_engine = tebd.TEBD(psi, model, tebd_params)\n",
    "\n",
    "# Prepare lists to store observables during time evolution\n",
    "times = []\n",
    "n0_expect = []\n",
    "n1_expect = []\n",
    "\n",
    "# --- Time Evolution Loop ---\n",
    "for step in range(tebd_params['N_steps']):\n",
    "    tebd_engine.run_one_step()  # Evolve the state by dt\n",
    "    current_time = (step + 1) * tebd_params['dt']\n",
    "    times.append(current_time)\n",
    "    n0_expect.append(psi.expectation_value(\"n\", 0))  # Number operator expectation at site 0\n",
    "    n1_expect.append(psi.expectation_value(\"n\", 1))  # Number operator expectation at site 1\n",
    "\n",
    "print(\"Time evolution complete.\")\n",
    "print(\"Times:\", times)\n",
    "print(\"Site 0 occupations:\", n0_expect)\n",
    "print(\"Site 1 occupations:\", n1_expect)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tenpy is designed for larger lattices. For my requirements QuTip is supposed to work well\n",
    "\n",
    "## Let's reproduce the results from Su's MATLAB code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 104784 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 67140 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import qutip as qt\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# --- Parameters ---\n",
    "J = 1.0       # Hopping amplitude\n",
    "U = 3.0       # On-site interaction strength\n",
    "mu = 0.0      # Chemical potential (set to zero for simplicity)\n",
    "Delta = 3.0   # energy tilt of sites\n",
    "Nmax = 2      # Maximum boson number per site (Hilbert space dimension d = Nmax+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: -t (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 = (U/2) * (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",
    "\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 = Delta * (1*qt.tensor(n_op, qt.qeye(d)) + 2*qt.tensor(qt.qeye(d), n_op))\n",
    "\n",
    "# Total Hamiltonian\n",
    "H = H_hop + H_int + H_mu + H_tilt\n",
    "\n",
    "# --- Define Initial State ---\n",
    "# The product Fock state |1,1> (one boson per site)\n",
    "psi0 = qt.tensor(qt.basis(d, 1), qt.basis(d, 1))\n",
    "\n",
    "# --- Time Evolution Setup ---\n",
    "tlist = np.linspace(0, 5, 501)  # 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)\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",
    "\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",
    "\n",
    "# Plot the probability for each basis state over time\n",
    "plt.figure(figsize=(10, 6))\n",
    "for idx, label in enumerate(basis_labels):\n",
    "    if idx==2 or idx==4 or idx==6:\n",
    "        plt.plot(tlist, probs[:, idx], label=label)\n",
    "plt.xlabel('Time')\n",
    "plt.ylabel('Probability')\n",
    "plt.title('Time Evolution of Individual Basis State Probabilities')\n",
    "plt.legend(loc=\"upper right\", bbox_to_anchor=(1.25, 1))\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.xlabel('Time')\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()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## If we add NNI and DIT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 95200 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 66636 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import qutip as qt\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# --- Parameters ---\n",
    "J = 1.0       # Hopping amplitude\n",
    "U = 10.0       # On-site interaction strength\n",
    "mu = 0.0      # Chemical potential (set to zero for simplicity)\n",
    "Delta = 10.0   # energy tilt of sites\n",
    "V = 0.        # nearest neighbour interaction energy (due to DDI)\n",
    "DeltaJ = 0.   # density induced tunneling\n",
    "\n",
    "Nmax = 2      # Maximum boson number per site (Hilbert space dimension d = Nmax+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 = (U/2) * (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",
    "\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 = Delta * (1*qt.tensor(n_op, qt.qeye(d)) + 2*qt.tensor(qt.qeye(d), n_op))\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 = H_hop + H_int + H_mu + H_tilt + H_nni + H_dit\n",
    "\n",
    "# --- Define Initial State ---\n",
    "# The product Fock state |1,1> (one boson per site)\n",
    "psi0 = qt.tensor(qt.basis(d, 1), qt.basis(d, 1))\n",
    "\n",
    "# --- Time Evolution Setup ---\n",
    "tlist = np.linspace(0, 5, 501)  # 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)\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",
    "\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')\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()\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.xlabel('Time')\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()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Make tilt time dependent:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 324,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import qutip as qt\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import constants as const\n",
    "\n",
    "#add relative path to backend\n",
    "import sys\n",
    "sys.path.append('../../clean_diag/backend')\n",
    "\n",
    "import trap_units as si"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 543,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Parameters ---\n",
    "J = 2*np.pi* 262#Hz       # Hopping amplitude\n",
    "U = 2*np.pi* 1000      # On-site interaction strength\n",
    "mu = 0.0      # Chemical potential (set to zero for simplicity)\n",
    "\n",
    "V = 0.        # nearest neighbour interaction energy (due to DDI)\n",
    "DeltaJ = 0.   # density induced tunneling\n",
    "\n",
    "# Parameters for the tilt\n",
    "Delta_start = 2*np.pi* 0   #tilt energy at start\n",
    "Delta_end = 4*U     #------\"-------- end\n",
    "\n",
    "omega_J = 2*J\n",
    "tau = 10*si.ms     #time duration of simulation\n",
    "\n",
    "N_steps = int(1e3)      #number of timesteps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 544,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1646.1945504810517\n",
      "6283.185307179586\n",
      "0.0\n",
      "25132.741228718343\n",
      "0.01\n"
     ]
    }
   ],
   "source": [
    "print(J)\n",
    "print(U)\n",
    "\n",
    "print(Delta_start)\n",
    "print(Delta_end)\n",
    "\n",
    "print(tau)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 545,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 125244 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 71892 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Nmax = 2      # Maximum boson number per site (Hilbert space dimension d = Nmax+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",
    "# Interaction term: (U/2) * [n(n-1)] on each site\n",
    "# We subtract the identity appropriately to compute n(n-1)\n",
    "H_int = (U/2) * (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",
    "# 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_op = (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['a']+ (args['b']-args['a'])* t/args['tau']\n",
    "\n",
    "\n",
    "#nearest neighbour interaction\n",
    "H_nni = V * qt.tensor(n_op, qt.qeye(d)) * qt.tensor(qt.qeye(d), n_op)\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_int + H_mu + H_nni + H_dit\n",
    "\n",
    "H = [H_static, [H_tilt_op, tilt_func]]\n",
    "tilt_args = {'a': Delta_start, 'b': Delta_end, 'tau': tau}\n",
    "\n",
    "\n",
    "# --- Define Initial State ---\n",
    "# The product Fock state |1,1> (one boson per site)\n",
    "psi0 = qt.tensor(qt.basis(d, 1), qt.basis(d, 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",
    "\n",
    "# Run the evolution and store the full state history\n",
    "result = qt.sesolve(H, psi0, tlist, e_ops=e_ops, options=opts, args=tilt_args)\n",
    "\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()\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()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fidelity should be given by probability of |2,0> (if particles are perfectly indistinguishable)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 546,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8284621706456528"
      ]
     },
     "execution_count": 546,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.mean(probs[-int(N_steps/5):,6])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "use_tenpy",
   "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
}
