{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lab 8: Assignment 2 Tutorial\n",
    "\n",
    "Consider the double potential well,\n",
    "\n",
    "$$V(x) =  \\alpha x^4  - 2x^2 + + \\frac{1}{\\alpha}$$\n",
    "\n",
    "where $x$ is the position of the particle, and we set $m=\\hbar=1$.\n",
    "For parts 1A, 1B, and 1C, you may set $\\alpha = 0.4$.\n",
    "For the last part, you will need to vary $\\alpha$.\n",
    "The minima of the potential are at $x_\\mathrm{min}^2 = \\pm\\frac{1}{\\alpha}$ and the barrier height is $V(0) = \\frac{1}{\\alpha}$.\n",
    "\n",
    "As discussed in lecture 9, we can approximate the potential by a harmonic potential near the minima.\n",
    "\n",
    "$$V_+(x)  = 4(x-x_\\mathrm{min})^2$$\n",
    "\n",
    "$$V_-(x)  = 4(x+x_\\mathrm{min})^2,$$\n",
    "\n",
    "which implies $\\omega = 2\\sqrt{2}$.\n",
    "The ground state wave functions for those approximate potential wells are\n",
    "\n",
    "$$\\psi_+(x) = \\left(\\frac{\\omega}{\\pi}\\right)^{1/4} \\exp\\left(- \\frac{\\omega}{2}(x-x_\\mathrm{min})^2\\right)$$\n",
    "\n",
    "$$\\psi_-(x) = \\left(\\frac{\\omega}{\\pi}\\right)^{1/4} \\exp\\left(- \\frac{\\omega}{2}(x+x_\\mathrm{min})^2\\right)$$\n",
    "\n",
    "We can approximate the ground state and first excited state as the symmetric and antisymmetric combinations of these wave functions, respectively.\n",
    "\n",
    "$$\\psi_0(x) \\approx \\psi_\\mathrm{S}(x) = \\frac{1}{\\sqrt{2}}\\left ( \\psi_+(x) + \\psi_-(x) \\right )$$\n",
    "\n",
    "$$\\psi_1(x) \\approx \\psi_\\mathrm{A}(x) = \\frac{1}{\\sqrt{2}}\\left ( \\psi_+(x) - \\psi_-(x) \\right )$$\n",
    "\n",
    "The energy gap between the ground state and first excited state is $\\Delta E = E_1 - E_0$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem 1A\n",
    "\n",
    "Demonstrate tunneling between the two wells using the Feynman path integral.\n",
    "Start with the particle in the right well\n",
    "$$\n",
    "  \\psi(x, 0) = \\psi_+(x),\n",
    "$$\n",
    "and evolve the wave function at each time step using the elementary propagator.\n",
    "\n",
    "Numerically, you may use discretization parameters similar to Assignment 1: $\\epsilon = \\Delta t = \\pi/128$, $x_0 = -4$, $x_{N_D} = +4$, and $N_D=600$.\n",
    "You will have to simulate a long enough period $T > t_\\mathrm{tunnel}$ that is longer than the tunneling time.\n",
    "As a reminder, the tunneling time is defined as the time it takes for the particle to reach the left well.\n",
    "\n",
    "Plot the mean position $\\langle x \\rangle$ as a function of time $t$.\n",
    "How can you estimate the tunneling time $t_\\mathrm{tunnel}$ from this plot?\n",
    "\n",
    "_**Hint:** You may approximate the propagator as_\n",
    "\n",
    "$$\\tilde {\\mathcal K}(x_b, \\epsilon; x_a, 0) \\sim \\exp \\left( i\\left( \\frac{1}{2}\\frac{(x_b - x_a)^2}{\\epsilon} - V\\left(\\frac{x_a+x_b}{2}\\right) \\epsilon \\right)\\right),$$\n",
    "\n",
    "_where $x_a$ and $x_b$ are the initial and final positions, respectively, and $\\epsilon$ is the time step._\n",
    "_Note, we lack the normalization factor in this approximation so you will need to normalize the wave function at each time step_\n",
    "\n",
    "$$\\psi(x, t) \\leftarrow \\frac{\\psi(x, t)}{\\sqrt{\\int_{-\\infty}^{\\infty} |\\psi(x, t)|^2 dx}}.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ALPHA = 0.4\n",
    "XMIN = np.sqrt(1 / ALPHA)\n",
    "OMEGA = np.sqrt(8 * ALPHA * XMIN ** 2)\n",
    "XSTART = XMIN\n",
    "BOXSIZE = 8\n",
    "ND = 600\n",
    "DELTAX = BOXSIZE / ND\n",
    "HBAR = 1\n",
    "T0 = 20 * np.pi\n",
    "DELTAT = np.pi / 128\n",
    "NT = int(T0 / DELTAT)\n",
    "\n",
    "t = np.linspace(0, T0, NT + 1)\n",
    "x = np.linspace(-BOXSIZE / 2, BOXSIZE / 2, ND + 1)\n",
    "\n",
    "def V(x):\n",
    "    return -2*x**2 + ALPHA*x**4 + 1/ALPHA\n",
    "\n",
    "def V_plus(x):\n",
    "    return 4*(x-1/np.sqrt(ALPHA))**2\n",
    "\n",
    "def V_minus(x):\n",
    "    return 4*(x+1/np.sqrt(ALPHA))**2\n",
    "\n",
    "def func_psi_0(x, x_start):\n",
    "    part1 = (OMEGA / np.pi) ** (1 / 4)\n",
    "    part2 = np.exp(-( OMEGA / (2 * HBAR)) * (x - x_start)**2)\n",
    "    return part1 * part2\n",
    "\n",
    "psi_plus = func_psi_0(x, XMIN) \n",
    "psi_minus = func_psi_0(x, -XMIN)\n",
    "psi_A = (psi_plus - psi_minus)/np.sqrt(2)\n",
    "psi_S = (psi_plus + psi_minus)/np.sqrt(2)\n",
    "\n",
    "plt.figure(dpi=300, figsize=(8,6))\n",
    "plt.plot(x, V(x), 'k-', label=r'$V(x)$')\n",
    "plt.plot(x[380:460], V_plus(x[380:460]), 'b--', label=r'$V_+(x)$')\n",
    "plt.plot(x[-460:-380], V_minus(x[-460:-380]), 'r-.', label=r'$V_-(x)$')\n",
    "plt.plot(x[310:530], psi_plus[310:530], 'b-', label=r'$\\psi_+(x)$')\n",
    "plt.plot(x[-530:-310], psi_minus[-530:-310], 'r-', label=r'$\\psi_-(x)$')\n",
    "plt.plot(x, psi_S, 'g-.', label=r'$\\psi_S(x)$')\n",
    "plt.plot(x, psi_A, 'g--', label=r'$\\psi_A(x)$')\n",
    "plt.ylabel(r'$V(x)$ or $\\psi(x)$')\n",
    "plt.xlabel(r'$x$')\n",
    "plt.xlim(-4, 4)\n",
    "plt.ylim(-1, 4)\n",
    "plt.legend()\n",
    "plt.savefig(\"V_all.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func_K(x_a, x_b, dt):\n",
    "    exponent = 1j * (0.5 * (x_b - x_a)**2 / dt - V((x_a + x_b) / 2) * dt)\n",
    "    return np.exp(exponent)\n",
    "\n",
    "K_dt = np.zeros((ND + 1, ND + 1), dtype=np.complex64)\n",
    "for i in range(ND + 1):\n",
    "    for j in range(ND + 1):\n",
    "        K_dt[i, j] = func_K(x[i], x[j], DELTAT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi_0 = func_psi_0(x, XSTART) \n",
    "\n",
    "print(\"Check normalization: \", np.sum(psi_0 * psi_0.conjugate()) * DELTAX)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi = [psi_0]\n",
    "\n",
    "for i in range(1, NT + 1):\n",
    "    psi_t = DELTAX * np.matmul(K_dt, psi[i-1])\n",
    "    # normalize to 1\n",
    "    psi_t /= np.sqrt(DELTAX * np.sum(psi_t * psi_t.conjugate()))\n",
    "    psi.append(psi_t)\n",
    "\n",
    "prob = []\n",
    "for i in range(NT + 1):\n",
    "    prob.append(np.real(psi[i] * psi[i].conjugate()))\n",
    "\n",
    "x_bar = np.zeros_like(t)\n",
    "for i in range(NT + 1):\n",
    "    x_bar[i] = sum(prob[i] * x * DELTAX)\n",
    "\n",
    "plt.figure(figsize=(6, 6))\n",
    "plt.plot(t, x_bar)\n",
    "plt.xlabel(r\"$t$\")\n",
    "plt.ylabel(r\"$\\langle x \\rangle$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# approximate tunneling time\n",
    "t[np.argmin(x_bar)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_interval = 190\n",
    "xmin, xmax, ymin, ymax = -BOXSIZE/2, BOXSIZE/2, -1, 4\n",
    "plt.figure(dpi=300, figsize=(8, 6))\n",
    "for i in range(0, NT // 2 + 1, plot_interval):\n",
    "    plt.plot(x, prob[i].real, label=f\"$t = {i * DELTAT:.2f}$\")\n",
    "plt.plot(x, V(x), 'k', label='$V(x)$')\n",
    "plt.xlabel(\"$x$\")\n",
    "plt.ylabel(\"$V(x)$ or $|\\psi(x)|^2$\")\n",
    "plt.xlim([xmin, xmax])\n",
    "plt.ylim([ymin, ymax])\n",
    "plt.legend(loc='upper right')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('playground')",
   "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.9.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "725838c595786ca03589bc9bd25e9f8a14247b2a3328879b9950fd22c9f29489"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}