{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_hub(J,U,Delta,mu,V,DeltaJ,tau,\n",
    "              N_steps=1000):\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",
    "    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 = 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, 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",
    "    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": [
    "----------------\n",
    "\n",
    "# Ramping the Interactions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-999.9999999999999 3000.0\n",
      "1200.0 1200.0\n"
     ]
    }
   ],
   "source": [
    "# --- Parameters ---\n",
    "J = 2*np.pi* 100#Hz       # Hopping amplitude\n",
    "U = (2*np.pi* -1000, 2*np.pi* 3000)       # On-site interaction strength\n",
    "print(U[0]/2/np.pi,U[1]/2/np.pi)\n",
    "Delta_const = 2*np.pi*1200\n",
    "Delta = (Delta_const, Delta_const)        #tilt\n",
    "print(Delta[0]/2/np.pi,Delta[1]/2/np.pi)\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 = 20*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)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fidelity: 0.987\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 112228 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 72048 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",
    "fid = np.mean(probs[-int(N_steps/5):, 2] + probs[-int(N_steps/5):, 6])\n",
    "print(f\"fidelity: {fid:.3f}\")\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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ramping the tilt instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2223.7233954432318"
      ]
     },
     "execution_count": 133,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "9.93*const.value(\"Bohr magneton\")* 2*si.G/si.cm * 800*si.nm /const.h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "20836.619123327575"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1*si.uK *const.k / const.h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Parameters ---\n",
    "J = 2*np.pi* 100#Hz       # Hopping amplitude\n",
    "U_const = 2*np.pi* 1500\n",
    "U = (U_const, 0*U_const)       # On-site interaction strength\n",
    "Delta = (2*np.pi* 0, 2*np.pi* 1500)        #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 = 20*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)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fidelity: 0.974\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 84632 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 56492 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",
    "fid = np.mean(probs[-int(N_steps/5):, 6])\n",
    "#fid = np.mean(probs[-int(N_steps/5):, 2] + probs[-int(N_steps/5):, 6])\n",
    "print(f\"fidelity: {fid:.3f}\")\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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set speed instead of time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_hub_vel(J,U,Delta,mu,V,DeltaJ,tau,\n",
    "              N_steps=1000):\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 value and speed of ramping.\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 = 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['vel_U']* t\n",
    "    int_args = {'U0': U[0], 'vel_U': U[1]}\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['Delta0']+ args['vel_Delta']* t\n",
    "    tilt_args = {'Delta0': Delta[0], 'vel_Delta': Delta[1]}\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, 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",
    "    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": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-999.9999999999999 999.9999999999999\n",
      "200000.0\n",
      "0.0 0.0\n",
      "0.01\n"
     ]
    }
   ],
   "source": [
    "# --- Parameters ---\n",
    "J = 2*np.pi* 100#Hz       # Hopping amplitude\n",
    "\n",
    "U = (2*np.pi* -1000, 2*np.pi* 1000)       # On-site interaction strength\n",
    "print(U[0]/2/np.pi,U[1]/2/np.pi)\n",
    "vel_U = 2*np.pi*4000/(20*si.ms)\n",
    "print(vel_U/2/np.pi)\n",
    "\n",
    "Delta_const = 2*np.pi*0\n",
    "Delta = (Delta_const, Delta_const)        #tilt\n",
    "print(Delta[0]/2/np.pi,Delta[1]/2/np.pi)\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 = abs(U[1]-U[0])/vel_U     #time duration of simulation\n",
    "print(tau)\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)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fidelity: 0.915\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 73124 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 32300 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",
    "fid = np.mean(probs[-int(N_steps/5):, 2] + probs[-int(N_steps/5):, 6])\n",
    "print(f\"fidelity: {fid:.3f}\")\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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try to see what different speeds do:"
   ]
  },
  {
   "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
}
