Presentation - 13 - Schroedinger's Equation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Schr&ouml;dinger's Equation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "source": [
    "\n",
    "### Peter Onyisi\n",
    "<img src=\"images/texas_logo.png\" width=\"400\" align=\"left\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# TACC: install if necessary\n",
    "try:\n",
    "    import numba\n",
    "except:\n",
    "    ! pip3 install --user numba numpy==1.19.1 scipy==1.5.2 tqdm==4.48.2\n",
    "    INSTALL_PATH='/home/jupyter/mydata/jupyter_packages/lib/python3.6/site-packages'\n",
    "    import sys, os\n",
    "    if os.access(INSTALL_PATH, os.R_OK) and INSTALL_PATH not in sys.path:\n",
    "        sys.path.insert(0, INSTALL_PATH)\n",
    "\n",
    "# generic\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "matplotlib.rcParams['figure.figsize'] = (8,6)\n",
    "matplotlib.rcParams['font.size'] = 16\n",
    "matplotlib.rcParams['animation.embed_limit'] = 60\n",
    "matplotlib.rcParams['xtick.top'] = True\n",
    "matplotlib.rcParams['ytick.right'] = True\n",
    "import math\n",
    "import numpy\n",
    "import scipy.integrate\n",
    "import scipy.stats\n",
    "import scipy.special\n",
    "from matplotlib import animation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Dynamics of Wave Functions\n",
    "\n",
    "So far we've discussed what wave functions might describe different particle configurations. We haven't talked about how they actually evolve with time - i.e., their *dynamics*. For classical particles, this is handled by Newton's Second Law:\n",
    "\n",
    "$$ m\\frac{d^2 x}{dt^2} = F(x) = -\\frac{dV}{dx} $$\n",
    "\n",
    "in one dimension (where $V$ is the potential energy describing $F$, assuming we're only using conservative forces); in three dimensions we would write\n",
    "\n",
    "$$ m\\frac{d^2\\vec{x}}{dt^2} = \\vec{F}(\\vec{x}) = -\\nabla V(\\vec{x}) $$\n",
    "\n",
    "The equivalent for wave functions is a *partial differential equation*, the *time-dependent Schr&ouml;dinger equation*. For one dimension it is:\n",
    "\n",
    "$$ i\\hbar \\frac{\\partial \\psi}{\\partial t} = -\\frac{\\hbar^2}{2m}\\frac{\\partial^2 \\psi}{\\partial x^2} + V(x)\\psi(x,t) $$\n",
    "\n",
    "where we have $\\hbar = h/2\\pi$ and $V(x)$ is the potential energy as a function of $x$. Here we now allow the wave function to be a function of both position $x$ and time $t$: $\\psi = \\psi(x,t)$.\n",
    "\n",
    "This is of course a potentially very complicated equation to solve (as are all partial differential equations) but let's try and understand the overall form.  Like Newton's Second Law, it tells us about how a particle will evolve given its state at a certain time.  That is - if $\\psi(x, t_0)$ is known, I can evaluate the right hand side of the equation for any value of $x$, and that tells me $\\partial \\psi/\\partial t$ for that $x$.   To make this a little more explicit:\n",
    "\n",
    "$$ \\frac{\\partial\\psi}{\\partial t} = \\frac{1}{i\\hbar}\\left[-\\frac{\\hbar^2}{2m}\\frac{\\partial^2 \\psi}{\\partial x^2} + V(x)\\psi(x,t)\\right] $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evolution of Momentum States\n",
    "\n",
    "We can solve the time-dependent Schr&ouml;dinger's equation explicitly in certain cases. Let us look at what happens to a state of definite momentum in a situation where $V(x) = 0$ everywhere, where we specify the constant $C$ at a time $t=0$:\n",
    "\n",
    "$$ \\psi(x, 0) = C \\exp(ipx/\\hbar) $$\n",
    "\n",
    "Then, we wind up with the following equation:\n",
    "\n",
    "$$\\begin{align*}\n",
    "\\frac{\\partial\\psi}{\\partial t} &= \\frac{1}{i\\hbar}\\times -\\frac{\\hbar^2}{2m}\\frac{\\partial^2}{\\partial x^2}\\psi \\\\\n",
    "&= \\frac{i\\hbar}{2m} \\frac{\\partial^2}{\\partial x^2} [C \\exp(ipx/\\hbar)] \\\\\n",
    "&= \\frac{i\\hbar}{2m} \\frac{\\partial}{\\partial x}[C (ip/\\hbar)\\exp(ipx/\\hbar)]\\\\\n",
    "&= \\frac{i\\hbar}{2m} C (ip/\\hbar)^2\\exp(ipx/\\hbar) \\\\\n",
    "&= -i \\frac{p^2}{2m\\hbar} C \\exp(ipx/\\hbar) \\\\\n",
    "&= -i \\frac{p^2}{2m\\hbar} \\psi\n",
    "\\end{align*}$$\n",
    "\n",
    "which holds for all values of $x$.  From our knowledge of derivatives, we know that the function that has the property that $\\partial f/\\partial t = K f$ is $f = f(0) \\exp(Kt)$ (where $K$ is a constant).  We also know that $p^2/2m = (mv)^2/2m = mv^2/2 $ is the kinetic energy $E$. So we conclude that\n",
    "\n",
    "$$ \\begin{align*}\n",
    "\\psi(x, t) &= \\psi(x, 0)\\exp(-i\\frac{p^2}{2m\\hbar}t)\\\\\n",
    "&= C\\exp(ipx/\\hbar)\\exp(-iEt/\\hbar)\\\\\n",
    "&= C\\exp(i(px-Et)/\\hbar)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "So the solution at all times looks very much like the solution at $t = 0$, except that the wave function is multiplied by $e^{-iEt/\\hbar}$.  You will  note this is of the form $e^{i\\phi}$, with $\\phi$ constant in $x$, so it is just a time-dependent rotation in the complex plane."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numerically Solving Partial Differential Equations\n",
    "\n",
    "In general it is difficult or impossible to find analytic solutions for Schr&ouml;dinger's equation. (You may get the opposite impression in follow-up classes, in which you will spend a lot of time finding analytic solutions! But by definition you will be looking at the cases where analytic solutions are possible.)  Thankfully the time-dependent Schr&ouml;dinger equation is well suited to having a computer approximate the solution.\n",
    "\n",
    "The general procedure is this:\n",
    "* since we don't know an analytic form for the wave function, we approximate it by its values at a set of points, assuming it doesn't change too much between the points. In other words we _discretize_ space. We aren't going to be able to see effects that happen at a shorter distance than the spacing between the points, but our computational cost will go up as we have more points, so there is a tradeoff of accuracy for time.\n",
    "* at any given point in space we can compute $V(x)\\psi(x)$ easily. We can _approximate_ $\\partial^2 \\psi/\\partial x^2$ by *finite differences*:\n",
    "$$ \\frac{\\partial^2 \\psi}{\\partial x^2}(x_n) \\approx \\frac{\\psi(x_{n+1})+\\psi(x_{n-1})-2\\psi(x_n)}{\\Delta x^2} $$\n",
    "* from this information we can determine $\\partial \\psi/\\partial t$. Conceptually we can then use this to extrapolate a small time step into the future, at all grid points:\n",
    "$$ \\psi(x_n, t+\\Delta t) \\approx \\psi(x_n, t) + \\Delta t \\frac{\\partial \\psi}{\\partial t}(x_n) $$\n",
    "\n",
    "In practice we make a couple of small modifications:\n",
    "* Because the finite difference method assumes that there are neighbor grid points on both sides, we have a problem at the edges. We \"solve\" this by having the grid wrap around, so going off the right edge takes us to the left edge.\n",
    "* There is a deep literature on how small a time step $\\Delta t$ is needed to produce numerically reliable solutions. For the method outlined above, while it works, it requires very small time steps $\\Delta t$ to produce good answers. There are more accurate ways to use $\\partial\\psi/\\partial t$ to extrapolate to the future, which rely on effectively computing approximations to $\\partial^2 \\psi/\\partial t^2$ and higher derivatives from $\\partial \\psi/\\partial t$ at different times.  Here we use one of these methods (*Runge-Kutta 5*).\n",
    "\n",
    "**For all computations here, we set *h*=1, so we don't have to deal with very small/large numbers.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RNG=(-10,10)\n",
    "PTS = 501\n",
    "grid = numpy.linspace(RNG[0],RNG[1],PTS)[:-1] # drop last point, which wraps to initial point\n",
    "deltaxsq = ((RNG[1]-RNG[0])/(PTS-1))**2\n",
    "\n",
    "# define potential energy\n",
    "V = numpy.zeros_like(grid)\n",
    "#V = numpy.single(numpy.frompyfunc(lambda x: 20*numpy.sin(x*math.pi/2)**2, 1, 1)(grid))\n",
    "#V = numpy.single(numpy.frompyfunc(lambda x: 100*x**2, 1, 1)(grid))\n",
    "\n",
    "# we are going to set h (not hbar) = 1 and m = 1\n",
    "# so that hbar = 1/(2*pi)\n",
    "\n",
    "def d2psi(psi):\n",
    "    return numpy.diff(psi,n=2,prepend=psi[-1],append=psi[0])/deltaxsq\n",
    "\n",
    "def dpsidt(t, psi):\n",
    "    global V\n",
    "    return -1.j*(2*math.pi*V*psi - d2psi(psi)/(2*2*math.pi))\n",
    "    \n",
    "P=1\n",
    "#wave packet\n",
    "psi=numpy.csingle(scipy.stats.norm(scale=1,loc=-6).pdf(grid))*numpy.exp(2j*math.pi*P*grid)\n",
    "# normalize psi\n",
    "psi /= math.sqrt(sum(psi*psi.conj()).real*(RNG[1]-RNG[0])/(PTS-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRNG=(0,4)\n",
    "TIMES=numpy.linspace(TRNG[0], TRNG[1], 100)\n",
    "rk=scipy.integrate.solve_ivp(dpsidt, TRNG, psi, t_eval=TIMES)\n",
    "print(\"Done with integration\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "liner, = plt.plot([],[],label='$\\Re[\\psi]$')\n",
    "linei, = plt.plot([],[],label='$\\Im[\\psi]$')\n",
    "linem, = plt.plot([],[],label='$|\\psi|^2$')\n",
    "plt.xlim(*RNG)\n",
    "plt.ylim(-1,1)\n",
    "plt.xlabel('$x$')\n",
    "plt.legend()\n",
    "def plotter(data):\n",
    "    global rk\n",
    "#     ax.clear()\n",
    "#     ax.axis([RNG[0], RNG[1], 0, 4])\n",
    "\n",
    "    psii = rk.y[:,data]\n",
    "#    return (plt.plot(grid, (psii*psii.conj()).real), \n",
    "#            plt.plot(grid, V), plt.plot(psi.real))\n",
    "    liner.set_data(grid, psii.real)\n",
    "    linei.set_data(grid, psii.imag)\n",
    "    linem.set_data(grid, (psii*psii.conj()).real)\n",
    "    return [liner,linei,linem]\n",
    "    #return plt.plot(grid, psi.real), plt.plot(grid, psi.imag)\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, len(TIMES))\n",
    "from IPython.display import HTML\n",
    "#ani.save('wobble.mp4', writer=animation.writers['ffmpeg'](fps=2))\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see if this makes sense. We have set the momentum $p = 1$ in these units, so we would expect the wave packet's velocity $v$ to be $v=p/m$.  If $m = 1$ then we expect $v=1$. Let's see how the peak of $|\\psi|^2$ changes with time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Where is the peak of the wavepacket vs time?\n",
    "psisq = rk.y*rk.y.conj()\n",
    "maxpos = numpy.argmax(psisq,axis=0)\n",
    "plt.plot(rk.t, numpy.take_along_axis(grid, maxpos, axis=0))\n",
    "plt.xlabel('$t$')\n",
    "plt.ylabel('$x$ of wavepacket peak')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the wavepacket does indeed move with the expected velocity.\n",
    "\n",
    "Let's do another run with higher momentum, in the opposite direction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P=-3\n",
    "#wave packet\n",
    "psi=numpy.csingle(scipy.stats.norm(scale=1,loc=6).pdf(grid))*numpy.exp(2j*math.pi*P*grid)\n",
    "# normalize psi\n",
    "psi /= math.sqrt(sum(psi*psi.conj()).real*(RNG[1]-RNG[0])/(PTS-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRNG=(0,4)\n",
    "TIMES=numpy.linspace(TRNG[0], TRNG[1], 100)\n",
    "rk=scipy.integrate.solve_ivp(dpsidt, TRNG, psi, t_eval=TIMES)\n",
    "print(\"Done with integration\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "# liner, = plt.plot([],[],label='$\\Re[\\psi]$')\n",
    "# linei, = plt.plot([],[],label='$\\Im[\\psi]$')\n",
    "linem, = plt.plot([],[],label='$|\\psi|^2$')\n",
    "plt.xlim(*RNG)\n",
    "plt.ylim(-1,1)\n",
    "plt.xlabel('$x$')\n",
    "plt.legend()\n",
    "def plotter(data):\n",
    "    global rk\n",
    "    psii = rk.y[:,data]\n",
    "#     liner.set_data(grid, psii.real)\n",
    "#     linei.set_data(grid, psii.imag)\n",
    "    linem.set_data(grid, (psii*psii.conj()).real)\n",
    "#     return [liner,linei,linem]\n",
    "    return [linem]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, len(TIMES))\n",
    "from IPython.display import HTML\n",
    "#ani.save('wobble.mp4', writer=animation.writers['ffmpeg'](fps=2))\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Where is the peak of the wavepacket vs time?\n",
    "psisq = rk.y*rk.y.conj()\n",
    "maxpos = numpy.argmax(psisq,axis=0)\n",
    "plt.plot(rk.t, numpy.take_along_axis(grid, maxpos, axis=0))\n",
    "plt.xlabel('$t$')\n",
    "plt.ylabel('$x$ of peak of wavepacket')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What happens if we set P to zero?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P=0\n",
    "#wave packet\n",
    "psi=numpy.csingle(scipy.stats.norm(scale=0.4,loc=0).pdf(grid))*numpy.exp(2j*math.pi*P*grid)\n",
    "# normalize psi\n",
    "psi /= math.sqrt(sum(psi*psi.conj()).real*(RNG[1]-RNG[0])/(PTS-1))\n",
    "TRNG=(0,10)\n",
    "TIMES=numpy.linspace(TRNG[0], TRNG[1], 100)\n",
    "rk=scipy.integrate.solve_ivp(dpsidt, TRNG, psi, t_eval=TIMES)\n",
    "print(\"Done with integration\")\n",
    "fig, ax = plt.subplots()\n",
    "# liner, = plt.plot([],[],label='$\\Re[\\psi]$')\n",
    "# linei, = plt.plot([],[],label='$\\Im[\\psi]$')\n",
    "linem, = plt.plot([],[],label='$|\\psi|^2$')\n",
    "plt.xlim(*RNG)\n",
    "plt.ylim(-0.5,1.5)\n",
    "plt.xlabel('$x$')\n",
    "plt.legend()\n",
    "def plotter(data):\n",
    "    global rk\n",
    "    psii = rk.y[:,data]\n",
    "#     liner.set_data(grid, psii.real)\n",
    "#     linei.set_data(grid, psii.imag)\n",
    "    linem.set_data(grid, (psii*psii.conj()).real)\n",
    "#     return [liner,linei,linem]\n",
    "    return [linem]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, len(TIMES))\n",
    "from IPython.display import HTML\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What is happening? The _average_ position of the wave packet stays put, but it spreads out. This is because a wave packet is _not_ a pure momentum state, so there are both negative and positive momentum parts of the wave function, and those spread out in opposite directions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get momentum distribution of wave packet, using a Fourier transform\n",
    "plt.scatter(numpy.fft.fftfreq(grid.size, d=grid[1]-grid[0]), numpy.abs(numpy.fft.fft(psi)), label='Magnitude of FFT')\n",
    "plt.xlim(-2,2)\n",
    "plt.xlabel('$p$')\n",
    "plt.ylabel('$|\\psi(p)|$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get momentum distribution of wave packet *after* evolution, using a Fourier transform\n",
    "plt.scatter(numpy.fft.fftfreq(grid.size, d=grid[1]-grid[0]), numpy.abs(numpy.fft.fft(rk.y[:,-1])), label='Magnitude of FFT')\n",
    "plt.xlim(-2,2)\n",
    "plt.xlabel('$p$')\n",
    "plt.ylabel('$|\\psi(p)|$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This means that the _magnitude_ of each of the momentum components is actually pretty much the same over time! But how can this be? Well, as always, we need to know how the _phase_ of the Fourier transform components, not just the magnitude.  _That_ changes with time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "liner, = plt.plot([],[],label='$\\Re[\\psi]$')\n",
    "linei, = plt.plot([],[],label='$\\Im[\\psi]$')\n",
    "plt.xlim(*RNG)\n",
    "plt.xlabel('$p$')\n",
    "plt.ylim(-30, 30)\n",
    "plt.xlim(-2,2)\n",
    "plt.legend()\n",
    "def plotter(data):\n",
    "    global rk\n",
    "    psii = rk.y[:,data]\n",
    "#     liner.set_data(grid, psii.real)\n",
    "#     linei.set_data(grid, psii.imag)\n",
    "    rfft=numpy.fft.fft(psii)\n",
    "    liner.set_data(numpy.fft.fftfreq(grid.size, d=grid[1]-grid[0]), \n",
    "                   rfft.real)\n",
    "    linei.set_data(numpy.fft.fftfreq(grid.size, d=grid[1]-grid[0]), \n",
    "                   rfft.imag)\n",
    "#     return [liner,linei,linem]\n",
    "    return [liner,linei]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, len(TIMES))\n",
    "from IPython.display import HTML\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is not usually the case that the magnitude of the momentum components stays constant for a wave function! (It only works here because the potential energy $V(x)$ is a constant.) But in \"free space\" - i.e. constant potential energy - this is true.\n",
    "\n",
    "## Non-trivial potentials: Simple Harmonic Oscillator\n",
    "\n",
    "Let's try something a little more interesting, perhaps. Let's consider a simple harmonic oscillator potential:\n",
    "\n",
    "$$ V(x) = \\frac{1}{2} kx^2 $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SPRING_CONSTANT=0.5\n",
    "V = 0.5*SPRING_CONSTANT*grid**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P=0\n",
    "#wave packet\n",
    "psi=numpy.csingle(scipy.stats.norm(scale=0.4,loc=1).pdf(grid))*numpy.exp(2j*math.pi*P*grid)\n",
    "# normalize psi\n",
    "psi /= math.sqrt(sum(psi*psi.conj()).real*(RNG[1]-RNG[0])/(PTS-1))\n",
    "TRNG=(0,10)\n",
    "TIMES=numpy.linspace(TRNG[0], TRNG[1], 100)\n",
    "rk=scipy.integrate.solve_ivp(dpsidt, TRNG, psi, t_eval=TIMES)\n",
    "print(\"Done with integration\")\n",
    "fig, ax = plt.subplots()\n",
    "# liner, = plt.plot([],[],label='$\\Re[\\psi]$')\n",
    "# linei, = plt.plot([],[],label='$\\Im[\\psi]$')\n",
    "linem, = plt.plot([],[],label='$|\\psi|^2$')\n",
    "plt.plot(grid, V, label='potential energy $V$')\n",
    "plt.xlim(*RNG)\n",
    "plt.ylim(-0.5,1.5)\n",
    "plt.xlabel('$x$')\n",
    "plt.legend()\n",
    "def plotter(data):\n",
    "    global rk\n",
    "    psii = rk.y[:,data]\n",
    "#     liner.set_data(grid, psii.real)\n",
    "#     linei.set_data(grid, psii.imag)\n",
    "    linem.set_data(grid, (psii*psii.conj()).real)\n",
    "#     return [liner,linei,linem]\n",
    "    return [linem]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, len(TIMES))\n",
    "from IPython.display import HTML\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Where is the peak of $|\\psi|^2$?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Where is the peak of the wavepacket vs time?\n",
    "psisq = rk.y*rk.y.conj()\n",
    "maxpos = numpy.argmax(psisq,axis=0)\n",
    "plt.plot(rk.t, numpy.take_along_axis(grid, maxpos, axis=0))\n",
    "plt.xlabel('$t$')\n",
    "plt.ylabel('$x$ position of peak of wavepacket')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the plot, the motion looks like simple harmonic motion with frequency $\\nu \\approx$ 1/9. Classically we would expect\n",
    "\n",
    "$$ \\nu = \\frac{1}{2\\pi}\\sqrt{\\frac{k}{m}} = 0.1125 \\approx 1/9$$\n",
    "\n",
    "so it seems that the _quantum_ oscillation frequency is the same as the _classical_ one."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What do the momentum magnitudes look like?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "linem, = plt.plot([],[],label='$|\\psi|$')\n",
    "plt.xlim(*RNG)\n",
    "plt.ylim(-10, 40)\n",
    "plt.xlim(-2,2)\n",
    "plt.xlabel('$p$')\n",
    "plt.legend()\n",
    "def plotter(data):\n",
    "    global rk\n",
    "    psii = rk.y[:,data]\n",
    "#     liner.set_data(grid, psii.real)\n",
    "#     linei.set_data(grid, psii.imag)\n",
    "    rfft=numpy.fft.fft(psii)\n",
    "    linem.set_data(numpy.fft.fftfreq(grid.size, d=grid[1]-grid[0]), \n",
    "                   numpy.abs(rfft))\n",
    "#     return [liner,linei,linem]\n",
    "    return [linem]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, len(TIMES))\n",
    "from IPython.display import HTML\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can imagine that we might be able to get a \"stationary state\" if we carefully chose the initial form of $\\psi$. (The technical definition of this is that $\\psi(x, t) = e^{i\\phi(t)}\\psi(x, t_0)$ for all $t$, i.e. $\\psi$ changes only by phase with time.) Playing around gives us a solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi=numpy.csingle(scipy.stats.norm(scale=0.47,loc=0).pdf(grid))*numpy.exp(2j*math.pi*P*grid)\n",
    "# normalize psi\n",
    "psi /= math.sqrt(sum(psi*psi.conj()).real*(RNG[1]-RNG[0])/(PTS-1))\n",
    "TRNG=(0,10)\n",
    "TIMES=numpy.linspace(TRNG[0], TRNG[1], 100)\n",
    "rk=scipy.integrate.solve_ivp(dpsidt, TRNG, psi, t_eval=TIMES)\n",
    "print(\"Done with integration\")\n",
    "fig, ax = plt.subplots()\n",
    "# liner, = plt.plot([],[],label='$\\Re[\\psi]$')\n",
    "# linei, = plt.plot([],[],label='$\\Im[\\psi]$')\n",
    "linem, = plt.plot([],[],label='$|\\psi|^2$')\n",
    "plt.plot(grid, V, label='potential energy $V$')\n",
    "plt.xlim(*RNG)\n",
    "plt.ylim(-0.5,1.5)\n",
    "plt.xlabel('$x$')\n",
    "plt.legend()\n",
    "def plotter(data):\n",
    "    global rk\n",
    "    psii = rk.y[:,data]\n",
    "#     liner.set_data(grid, psii.real)\n",
    "#     linei.set_data(grid, psii.imag)\n",
    "    linem.set_data(grid, (psii*psii.conj()).real)\n",
    "#     return [liner,linei,linem]\n",
    "    return [linem]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, len(TIMES))\n",
    "from IPython.display import HTML\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In fact if one solves for this (and in fact the simple harmonic oscillator is one of the systems where stationary solutions can be solved for exactly) one finds that the _scale_ argument to the normal distribution (i.e., the $\\sigma$ parameter) should be 0.4744, which is what we have found by trial and error.\n",
    "\n",
    "We won't justify the following, but it is possible to estimate the energy of a state by evaluating the term\n",
    "\n",
    "$$ \\frac{1}{\\psi}\\left[-\\frac{\\hbar^2}{2m}\\frac{\\partial^2}{\\partial x^2}\\psi + V(x)\\psi(x,t)\\right] $$\n",
    "\n",
    "If the result is _constant_ in $x$, then the state has a definite energy (and it turns out this means it's a stationary state)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Look at the energy of the simple harmonic oscillator\n",
    "# remember hbar = 1/2pi in our units\n",
    "# precision falls off where psi ~ 0 ...\n",
    "psi=numpy.csingle(scipy.stats.norm(scale=0.474425,loc=0).pdf(grid))*numpy.exp(2j*math.pi*P*grid)\n",
    "# normalize psi\n",
    "psi /= math.sqrt(sum(psi*psi.conj()).real*(RNG[1]-RNG[0])/(PTS-1))\n",
    "# fix zero values ...\n",
    "psi[psi==0]=1e-16\n",
    "eestimate = ((-d2psi(psi)/(8*math.pi**2)+V*psi)/psi).real\n",
    "plt.plot(grid, eestimate)\n",
    "plt.ylim(-0.2,0.2)\n",
    "plt.xlim(-5,5)\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel(r'$\\frac{\\hbar^2}{2m}\\frac{\\partial^2\\psi}{\\partial x^2} + V\\psi$')\n",
    "print(f'Estimated energy: {numpy.max(eestimate[(grid>-5) & (grid < 5)])}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "There are other possible stationary states, for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi=numpy.csingle(2*grid*scipy.stats.norm(scale=0.47,loc=0).pdf(grid))*numpy.exp(2j*math.pi*P*grid)\n",
    "# normalize psi\n",
    "psi /= math.sqrt(sum(psi*psi.conj()).real*(RNG[1]-RNG[0])/(PTS-1))\n",
    "TRNG=(0,10)\n",
    "TIMES=numpy.linspace(TRNG[0], TRNG[1], 100)\n",
    "rk=scipy.integrate.solve_ivp(dpsidt, TRNG, psi, t_eval=TIMES)\n",
    "print(\"Done with integration\")\n",
    "fig, ax = plt.subplots()\n",
    "# liner, = plt.plot([],[],label='$\\Re[\\psi]$')\n",
    "# linei, = plt.plot([],[],label='$\\Im[\\psi]$')\n",
    "linem, = plt.plot([],[],label='$|\\psi|^2$')\n",
    "plt.plot(grid, V, label='potential energy $V$')\n",
    "plt.xlim(*RNG)\n",
    "plt.ylim(-0.5,1.5)\n",
    "plt.xlabel('$x$')\n",
    "plt.legend()\n",
    "def plotter(data):\n",
    "    global rk\n",
    "    psii = rk.y[:,data]\n",
    "#     liner.set_data(grid, psii.real)\n",
    "#     linei.set_data(grid, psii.imag)\n",
    "    linem.set_data(grid, (psii*psii.conj()).real)\n",
    "#     return [liner,linei,linem]\n",
    "    return [linem]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, len(TIMES))\n",
    "from IPython.display import HTML\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Look at the energy of the simple harmonic oscillator\n",
    "# remember hbar = 1/2pi in our units\n",
    "# precision falls off where psi ~ 0 ...\n",
    "psi=numpy.csingle(2*grid*scipy.stats.norm(scale=0.474425,loc=0).pdf(grid))*numpy.exp(2j*math.pi*P*grid)\n",
    "# normalize psi\n",
    "psi /= math.sqrt(sum(psi*psi.conj()).real*(RNG[1]-RNG[0])/(PTS-1))\n",
    "# fix zero values ...\n",
    "psi[psi==0]=1e-16\n",
    "eestimate = ((-d2psi(psi)/(8*math.pi**2)+V*psi)/psi).real\n",
    "plt.plot(grid, eestimate)\n",
    "plt.ylim(-0.2,0.2)\n",
    "plt.xlim(-5,5)\n",
    "plt.xlabel('$x$')\n",
    "print(f'Estimated energy: {numpy.max(eestimate[(numpy.abs(psi)>1e-5) & (grid>-5) & (grid < 5)])}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple Harmonic Oscillator and Zero Point Energy\n",
    "\n",
    "Is this consistent with our simple ideas of quantization?  We expect that any sort of \"long-term\" state is going to be a stationary state, because otherwise it will evolve quickly.  So we imagine that if we are looking at something that appears stable over time, the particles inside must be in stationary states. Let's note that when we discussed the early idea that simple harmonic oscillators had \"minimum energy $h\\nu$\", this was *the minimum energy above the lowest energy state*. We can *always* redefine the energy of the lowest energy state by changing the potential energy. So what we really want to see is that the *energy difference between adjacent stationary states* is $h\\nu$. What do we see here? One state has energy $\\approx$ 0.0565, and another state we have considered has energy $\\approx$ 0.1691; the difference is 0.1126.  The frequency $\\nu$ of this oscillator is 0.1125 (from above), and we fixed $h=1$ for these computations, so indeed we have $\\Delta E = h\\nu$, to within the accuracy of our numerical computations!\n",
    "\n",
    "In fact, the full quantum solution for the simple harmonic oscillator gives the following energies for the stationary states:\n",
    "\n",
    "$$ E_n = \\left(n + \\frac{1}{2}\\right) h\\nu $$\n",
    "\n",
    "so the lowest energy state has $E_0 = h\\nu/2$ (the so-called *zero point energy*). Note that we could make this go away by shifting down $V(x)$ by $h\\nu/2$, so again the absolute scale is not really meaningful; the main point of the zero point energy is that, given the same potential, where _classical_ physics would allow states with energy $E=0$ (motionless at the bottom of the potential), _quantum_ physics doesn't agree.\n",
    "\n",
    "The precise reason that quantum mechanics doesn't allow $E=0$ is the following: we saw that a very well-localized wave function contains a giant spread of possible momenta. But non-zero momentum implies non-zero _kinetic_ energy. A state of zero momentum (and so zero kinetic energy) is a constant, hence has non-zero _potential_ energy. Since it is not possible to have a state that simultaneously has zero momentum and fixed position at zero, it is not possible to have total energy zero."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantum Tunneling\n",
    "\n",
    "In classical physics, the total energy of a particle (the sum of kinetic and potential energy) is well defined, and it is not possible for a particle to have negative kinetic energy. In quantum mechanics, because the wave function is spread out and doesn't have a single kinetic (or potential) energy, it is not as simple a story.\n",
    "\n",
    "The \"kinetic energy\" term in Schr&ouml;dinger's equation is\n",
    "\n",
    "$$ K \\sim -\\frac{1}{\\psi}\\frac{\\hbar^2}{2m}\\frac{\\partial^2\\psi}{\\partial x^2} $$\n",
    "\n",
    "but there's nothing requiring this to be positive. For _oscillating_ functions $\\psi \\sim \\exp(ipx/\\hbar)$, we have \n",
    "\n",
    "$$ K \\sim +\\frac{p^2}{2m} > 0 $$\n",
    "\n",
    "which is greater than zero. However, for _exponential_ functions $\\psi \\sim \\exp(-kx/\\hbar)$, we have instead\n",
    "\n",
    "$$ K \\sim -\\frac{k^2}{2m} < 0 $$\n",
    "\n",
    "So it is possible for a particle being described by a wave function to enter a region where, classically, its kinetic energy would be negative. In such cases the wave function must fall away exponentially the further into this region that you go.  (So the wave function does not vanish immediately when the potential energy gets too high, but rather goes smoothly to zero.)\n",
    "\n",
    "Now, if we have a tall (but narrow) potential energy barrier, this means that the wave function can _cross_ the barrier. The degree to which is passes through depends on how high and how wide the barrier is.\n",
    "\n",
    "Let's look at a simulation with a wave packet hitting a barrier. We will set $p = 3$ in our units, so the kinetic energy when $V(x) = 0$ is $p^2/2m = 3^2/2 = 4.5$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RNG=(-20,20)\n",
    "PTS = 1001\n",
    "grid = numpy.linspace(RNG[0],RNG[1],PTS)[:-1] # drop last point, which wraps to initial point\n",
    "deltaxsq = ((RNG[1]-RNG[0])/(PTS-1))**2\n",
    "P=3\n",
    "#wave packet\n",
    "psi=numpy.csingle(scipy.stats.norm(scale=0.8,loc=-7).pdf(grid))*numpy.exp(2j*math.pi*P*grid)\n",
    "# normalize psi\n",
    "psi /= math.sqrt(sum(psi*psi.conj()).real*(RNG[1]-RNG[0])/(PTS-1))\n",
    "\n",
    "V = numpy.zeros_like(grid)\n",
    "V[(grid>-0.5) & (grid<0)]=4\n",
    "\n",
    "TRNG=(0,8)\n",
    "TIMES=numpy.linspace(TRNG[0], TRNG[1], 100)\n",
    "rk=scipy.integrate.solve_ivp(dpsidt, TRNG, psi, t_eval=TIMES)\n",
    "print(\"Done with integration\")\n",
    "fig, ax = plt.subplots(figsize=(12,6))\n",
    "ax2=ax.twinx()\n",
    "# liner, = plt.plot([],[],label='$\\Re[\\psi]$')\n",
    "# linei, = plt.plot([],[],label='$\\Im[\\psi]$')\n",
    "linem, = ax.plot([],[],label='$|\\psi|^2$')\n",
    "linev,=ax2.plot(grid, V, color='r')\n",
    "plt.xlim(*RNG)\n",
    "plt.xlabel('$x$')\n",
    "ax.set_ylim(-0.,0.8)\n",
    "ax2.set_ylim(0,5)\n",
    "plt.legend([linem,linev], ['$|\\psi|^2$', '$V$'])\n",
    "def plotter(data):\n",
    "    global rk\n",
    "    psii = rk.y[:,data]\n",
    "#     liner.set_data(grid, psii.real)\n",
    "#     linei.set_data(grid, psii.imag)\n",
    "    linem.set_data(grid, (psii*psii.conj()).real)\n",
    "#     return [liner,linei,linem]\n",
    "    return [linem]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, len(TIMES))\n",
    "from IPython.display import HTML\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So it is possible for a particle to hop a barrier that classically would be impossible to cross, and the probability that it does so depends on the height and width of the barrier.  This can be used to make precisely controlled electronic devices to control the flow of electrons in a circuit, and is also required to explain many kinds of reactions (both chemical and nuclear) which need to get through an intermediate state with \"too high\" an energy.\n",
    "\n",
    "For example, the reaction $p + p \\to {}^2H + e^+ + \\nu_e$ in the Sun (or any future fusion reactor) requires the protons to cross the \"Coulomb barrier\" of the electrostatic repulsion of the two protons before they can be bound together by the strong (but very short-range) nuclear forces. This must be achieved using protons with thermal kinetic energy. Numerially the height of the barrier for proton-proton collisions is $\\approx$ 0.5 MeV, while $\\frac{3}{2}k_B T \\approx 2$ keV for $T=15\\times 10^6$ K at the center of the Sun.  Classically, then, it's almost impossible to find protons with high enough thermal energy to fuse; but quantum tunneling through the Coulomb barrier makes it far more likely that the protons will collide.\n",
    "\n",
    "[![Coulomb barrier](images/coulomb_barrier.png)](https://search.lib.utexas.edu/discovery/fulldisplay?docid=alma991057942374606011&context=L&vid=01UTAU_INST:SEARCH&lang=en&search_scope=MyInst_and_CI&adaptor=Local%20Search%20Engine&tab=Everything&query=any,contains,nuclear%20physics%20of%20stars&offset=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Particle in a Box\n",
    "One more important case is the *particle in a box*: the particle is in a flat-bottomed potential energy well with \"infinitely tall\" sides.  This corresponds to a particle trapped in a specific region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RNG=(-5,15)\n",
    "PTS = 1001\n",
    "grid = numpy.linspace(RNG[0],RNG[1],PTS)[:-1] # drop last point, which wraps to initial point\n",
    "deltaxsq = ((RNG[1]-RNG[0])/(PTS-1))**2\n",
    "\n",
    "V = numpy.single(numpy.frompyfunc(lambda x: 1000. if x < 0 or x > 10 else 0, 1, 1)(grid))\n",
    "\n",
    "psi=numpy.csingle(numpy.frompyfunc(lambda x: math.sin(x/5*8*math.pi) if 0 < x < 10 else 0, 1, 1)(grid))\n",
    "# normalize psi\n",
    "psi /= math.sqrt(sum(psi*psi.conj()).real*(RNG[1]-RNG[0])/(PTS-1))\n",
    "TRNG=(0,1)\n",
    "TIMES=numpy.linspace(TRNG[0],TRNG[1],50)\n",
    "rk=scipy.integrate.solve_ivp(dpsidt, TRNG, psi, t_eval=TIMES)\n",
    "print(\"Done with integration\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "fig, ax = plt.subplots(figsize=(12,6))\n",
    "ax2=ax.twinx()\n",
    "# liner, = plt.plot([],[],label='$\\Re[\\psi]$')\n",
    "# linei, = plt.plot([],[],label='$\\Im[\\psi]$')\n",
    "linem, = ax.plot([],[],label='$|\\psi|^2$')\n",
    "linev,=ax2.plot(grid, V, color='r')\n",
    "plt.xlim(*RNG)\n",
    "plt.xlabel('$x$')\n",
    "ax.set_ylim(-0.,0.8)\n",
    "ax2.set_ylim(0,5)\n",
    "plt.legend([linem,linev], ['$|\\psi|^2$', '$V$'])\n",
    "def plotter(data):\n",
    "    global rk\n",
    "    psii = rk.y[:,data]\n",
    "#     liner.set_data(grid, psii.real)\n",
    "#     linei.set_data(grid, psii.imag)\n",
    "    linem.set_data(grid, (psii*psii.conj()).real)\n",
    "#     return [liner,linei,linem]\n",
    "    return [linem]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, len(TIMES))\n",
    "from IPython.display import HTML\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a *stationary solution*, we want to have $\\psi(x)$ go to zero at the walls of the box (the particle should never be found outside the box (i.e. in the \"inifinte potential\" region), and the wave function is continuous). Inside the box, since the potential is constant and zero, we might expect an oscillatory solution to work (like in our first example); however $\\exp(ipx/\\hbar)$ isn't zero anywhere.  However, we see that if $\\psi_1$ and $\\psi_2$ are solutions of Schr&ouml;dinger's equation, then so is $A\\psi_1 + B\\psi_2$, because the equation is _linear_ in all the derivatives. So we can add together arbitrary solutions to get another solution. In particular, if we add together right- and left-headed waves, for example:\n",
    "\n",
    "$$\\begin{align*}\n",
    "\\psi(x) &= C \\exp(ipx/\\hbar) + C\\exp(i(-p)x/\\hbar) \\\\\n",
    "&= C[\\cos(px/\\hbar) + i\\sin(px/\\hbar) + \\cos(-px/\\hbar) + i\\sin(-px/\\hbar)] \\\\\n",
    "&= C[\\cos(px/\\hbar) + i\\sin(px/\\hbar) + \\cos(px/\\hbar) - i\\sin(px/\\hbar)] \\\\\n",
    "&= 2C\\cos(px/\\hbar)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "we can get a wave function that is zero at various locations. (Note that here, $p$ does _not_ indicate the momentum, since we've combined two wave functions of different momenta! Also, since the wave functions don't go off to infinity, the momentum is not sharply defined anyway.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(numpy.fft.fftfreq(grid.size, d=grid[1]-grid[0]), numpy.fft.fft(psi).real, label='FFT Real')\n",
    "plt.scatter(numpy.fft.fftfreq(grid.size, d=grid[1]-grid[0]), numpy.fft.fft(psi).imag, label='FFT Imag')\n",
    "plt.xlim(-2,2)\n",
    "plt.xlabel('$p$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We usually set up the problem so that the bottom of the well goes from $0$ to $L$; then what we want are the sine solutions:\n",
    "\n",
    "$$\\begin{align*}\n",
    "\\psi(x) &= \\frac{1}{i}[C \\exp(ipx/\\hbar) - C\\exp(i(-p)x/\\hbar)] \\\\\n",
    "&= \\frac{1}{i}C[\\cos(px/\\hbar) + i\\sin(px/\\hbar) - \\cos(-px/\\hbar) - i\\sin(-px/\\hbar)] \\\\\n",
    "&= \\frac{1}{i}C[\\cos(px/\\hbar) + i\\sin(px/\\hbar) - \\cos(px/\\hbar) + i\\sin(px/\\hbar)] \\\\\n",
    "&= 2C\\sin(px/\\hbar)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "$\\psi(0)=0$ naturally here, and the condition $\\psi(L) = 0$ forces\n",
    "\n",
    "$$ \\frac{pL}{\\hbar} = n\\pi $$\n",
    "$$ p = \\frac{n\\pi\\hbar}{L} $$\n",
    "\n",
    "So the solutions are\n",
    "\n",
    "$$\\psi_n(x) = \\sqrt{\\frac{2}{L}} \\sin\\left[\\frac{n\\pi}{L}x\\right] $$\n",
    "\n",
    "just as in a usual \"standing wave\" (the factor $\\sqrt{2/L}$ is introduced to make the integral of $|\\psi|^2$ equal to one.)  The total energies of the states turn out to be\n",
    "\n",
    "$$ E_n = \\frac{n^2\\pi^2\\hbar^2}{2mL^2} $$\n",
    "\n",
    "We're interested in the particle in a box as it can provide a rough model for how particles behave when they are \"confined\" in a region. So, for example, electrons in a metal can to a certain extent be modeled this way: they are free to move around the metal (it's a conductor), but it is very difficult for them to escape the metal (due to the work function).\n",
    "\n",
    "Note that we can add any number of solutions $\\psi_n$ derived above to get a new valid solution; but all valid solutions can be expressed by summing together the $\\psi_n$.  (This property means that the $\\psi_n$ form what we call a _basis_ for the solutions.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Higher dimension example\n",
    "We can also write Schr&ouml;dinger's equation in more than one dimension. The general equation is\n",
    "\n",
    "$$ i\\hbar \\frac{\\partial \\psi}{\\partial t} = -\\frac{\\hbar^2}{2m}\\nabla^2 \\psi + V(x)\\psi(x,t) $$\n",
    "\n",
    "so, for example, in 2 dimensions\n",
    "\n",
    "$$ i\\hbar \\frac{\\partial \\psi}{\\partial t} = -\\frac{\\hbar^2}{2m}\\left[\\frac{\\partial^2}{\\partial x^2} + \\frac{\\partial^2}{\\partial y^2}\\right] \\psi + V(x,y)\\psi(x,y,t) $$\n",
    "\n",
    "Let's look at a 2D example: a particle passing down a tunnel which encounters two open slits.  The walls are modeled by making a very high potential barrier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RNG_X=(-10,10)\n",
    "RNG_Y=(-5,5)\n",
    "PTS_X = 401\n",
    "PTS_Y = 61\n",
    "gridvalx = numpy.linspace(RNG_X[0],RNG_X[1],PTS_X)[:-1] # drop last point, which wraps to initial point\n",
    "gridvaly = numpy.linspace(RNG_Y[0],RNG_Y[1],PTS_Y)[:-1] # drop last point, which wraps to initial point\n",
    "gridx, gridy = numpy.meshgrid(gridvalx, gridvaly)\n",
    "\n",
    "deltaxsq = ((RNG_X[1]-RNG_X[0])/(PTS_X-1))**2\n",
    "deltaysq = ((RNG_Y[1]-RNG_Y[0])/(PTS_Y-1))**2\n",
    "\n",
    "def d2psi(psi):\n",
    "    psiloc = psi.reshape(gridx.shape)\n",
    "    derivy=numpy.diff(psiloc,n=2,prepend=psiloc[-1:,:],append=psiloc[:1,:], axis=0)/deltaysq\n",
    "    derivx=numpy.diff(psiloc,n=2,prepend=psiloc[:,-1:],append=psiloc[:,:1], axis=1)/deltaxsq\n",
    "    return derivx+derivy\n",
    "\n",
    "def dpsidt(t, psi):\n",
    "    global V\n",
    "    return (-1.j*(2*math.pi*V*psi.reshape(gridx.shape) - d2psi(psi)/(4*math.pi))).reshape(-1)\n",
    "\n",
    "# double slit with incident wave\n",
    "SLITWIDTH=0.6\n",
    "SLITCENTER=0.6\n",
    "ORIGWIDTH=4.\n",
    "V = numpy.single(numpy.frompyfunc(lambda x, y: 0 if (SLITCENTER-SLITWIDTH/2 < abs(y) < SLITCENTER+SLITWIDTH/2 and x < 0) \n",
    "                                  or x >= 0 else 20, 2, 1)(gridx, gridy))\n",
    "V[(numpy.abs(gridy)<ORIGWIDTH/2) & (gridx<-0.3)]=0\n",
    "    \n",
    "P = 2\n",
    "psi=numpy.csingle(scipy.stats.norm(scale=0.5,loc=-2).pdf(gridx))*numpy.exp(2j*math.pi*P*gridx)*numpy.cos(gridy*math.pi/2./(ORIGWIDTH/2))\n",
    "psi[(numpy.abs(gridy) > ORIGWIDTH/2)]=0\n",
    "psi[gridx>=-0.3]=0\n",
    "# normalize psi\n",
    "psi /= math.sqrt(numpy.sum(psi*psi.conj()).real*(RNG_X[1]-RNG_X[0])*(RNG_Y[1]-RNG_Y[0])/(PTS_X-1)/(PTS_Y-1))\n",
    "\n",
    "TRNG=(0,4)\n",
    "NTIMES=50\n",
    "rk=scipy.integrate.solve_ivp(dpsidt, TRNG, psi.reshape(-1), t_eval=numpy.linspace(TRNG[0],TRNG[1],NTIMES))\n",
    "psi = psi.reshape(gridx.shape)\n",
    "print(\"Done with integration\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.set_ylim(*RNG_Y)\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "ax.contourf(gridx, gridy, V,vmin=1,vmax=9, cmap=plt.cm.bone.reversed())\n",
    "psicont = ax.contourf(gridx,gridy,rk.y[:,0].real.reshape(psi.shape))\n",
    "def plotter(data):\n",
    "    global rk, psicont\n",
    "    for c in psicont.collections:\n",
    "        c.remove()\n",
    "    psii = rk.y[:,data].reshape(psi.shape)\n",
    "    psicont = ax.contour(gridx, gridy, (psii*psii.conj()).real,levels=numpy.linspace(0.,0.5,101))\n",
    "    return [psicont]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, NTIMES)\n",
    "from IPython.display import HTML\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As one might expect, we get a diffraction pattern. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# at last timeslice, look at |psi|^2 vs y at x=5\n",
    "# find x bin corresponding to x=5\n",
    "psislice=rk.y[:,-1].reshape(psi.shape)[gridx==5]\n",
    "plt.plot(gridy[gridx==5],(psislice*psislice.conj()).real)\n",
    "plt.xlabel('$y$')\n",
    "plt.ylabel('$|\\psi(5,y)|^2$')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What does passage through a single slit look like?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = numpy.single(numpy.frompyfunc(lambda x, y: 0 if (SLITCENTER-SLITWIDTH/2 < y < SLITCENTER+SLITWIDTH/2 and x < 0) \n",
    "                                  or x >= 0 else 20, 2, 1)(gridx, gridy))\n",
    "V[(numpy.abs(gridy)<ORIGWIDTH/2) & (gridx<-0.3)]=0\n",
    "    \n",
    "P = 2\n",
    "psi=numpy.csingle(scipy.stats.norm(scale=0.5,loc=-2).pdf(gridx))*numpy.exp(2j*math.pi*P*gridx)*numpy.cos(gridy*math.pi/2./(ORIGWIDTH/2))\n",
    "psi[(numpy.abs(gridy) > ORIGWIDTH/2)]=0\n",
    "psi[gridx>=-0.3]=0\n",
    "# normalize psi\n",
    "psi /= math.sqrt(numpy.sum(psi*psi.conj()).real*(RNG_X[1]-RNG_X[0])*(RNG_Y[1]-RNG_Y[0])/(PTS_X-1)/(PTS_Y-1))\n",
    "\n",
    "TRNG=(0,4)\n",
    "NTIMES=50\n",
    "rk=scipy.integrate.solve_ivp(dpsidt, TRNG, psi.reshape(-1), t_eval=numpy.linspace(TRNG[0],TRNG[1],NTIMES))\n",
    "psi = psi.reshape(gridx.shape)\n",
    "print(\"Done with integration\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.set_ylim(*RNG_Y)\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "ax.contourf(gridx, gridy, V,vmin=1,vmax=9, cmap=plt.cm.bone.reversed())\n",
    "psicont = ax.contourf(gridx,gridy,rk.y[:,0].real.reshape(psi.shape))\n",
    "def plotter(data):\n",
    "    global rk, psicont\n",
    "    for c in psicont.collections:\n",
    "        c.remove()\n",
    "    psii = rk.y[:,data].reshape(psi.shape)\n",
    "    psicont = ax.contour(gridx, gridy, (psii*psii.conj()).real,levels=numpy.linspace(0.,0.5,101))\n",
    "    return [psicont]\n",
    "\n",
    "print(\"Rendering...\")\n",
    "ani = animation.FuncAnimation(fig, plotter, NTIMES)\n",
    "from IPython.display import HTML\n",
    "plt.close()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Particle in a 3D box\n",
    "\n",
    "In three dimensions, Schr&ouml;dinger's equation is\n",
    "\n",
    "$$ i\\hbar \\frac{\\partial \\psi}{\\partial t} = -\\frac{\\hbar^2}{2m}\\left[\\frac{\\partial^2}{\\partial x^2} + \\frac{\\partial^2}{\\partial y^2} + \\frac{\\partial^2}{\\partial z^2}\\right] \\psi + V(x,y,z)\\psi(x,y,z,t) $$\n",
    "\n",
    "We can set up a cuboidal box with sides $[L_x, L_y, L_z]$ similar to the earlier example, with $V(x)=0$ inside the box and \"infinity\" outside. It turns out that in this case the differential equation is called _separable_ and we can assume the solutions are products of the solutions for each dimension, i.e.\n",
    "\n",
    "$$\\psi_{(n_x,n_y,n_z)}(x) = \\sqrt{\\frac{2}{L_x}} \\sqrt{\\frac{2}{L_y}} \\sqrt{\\frac{2}{L_z}}\\sin\\left[\\frac{n_x\\pi}{L_x}x\\right]\\sin\\left[\\frac{n_y\\pi}{L_y}y\\right]\\sin\\left[\\frac{n_z\\pi}{L_z}z\\right] $$\n",
    "\n",
    "where $n_x$, $n_y$, $n_z$ are positive integers. The energies of the states are\n",
    "\n",
    "$$ E_{(n_x,n_y,n_z)} = \\frac{\\pi^2\\hbar^2}{2m}\\left(\\frac{n_x^2}{L_x^2}+\\frac{n_y^2}{L_y^2}+\\frac{n_z^2}{L_z^2}\\right) $$\n",
    "\n",
    "If the box is a cube of side length $L$, this is just\n",
    "\n",
    "$$ E_{(n_x,n_y,n_z)} = \\frac{\\pi^2\\hbar^2}{2m}\\left(\\frac{n_x^2+n_y^2+n_z^2}{L^2}\\right) $$\n",
    "\n",
    "A thing that we will at some point want to know is how many possible basis states $\\psi_{(n_x,n_y,n_z)}$ are to be found with energy less than a certain energy $E$.  If we look at a cube of side length $L$, then there is an acceptable state at every lattice point $(n_x, n_y, n_z)$ satisfying\n",
    "\n",
    "$$ n_x^2+n_y^2+n_z^2 < \\frac{2mL^2 E}{\\pi^2\\hbar^2}$$\n",
    "\n",
    "(note that $n_x$, $n_y$, $n_z$ are positive integers).  But the number of such lattice points is approximately the same as $1/8$ times the volume of the sphere of radius $\\sqrt{n_x^2+n_y^2+n_z^2}$. So we expect that the number of states $N$ with energy $\\le E$ is\n",
    "\n",
    "$$ \\begin{align*}\n",
    "N(E) &= \\frac{1}{8} \\times \\frac{4}{3} \\pi r^3\\\\\n",
    "&=\\frac{1}{6} \\pi \\left[\\frac{2mL^2 E}{\\pi^2\\hbar^2}\\right]^{3/2}\\\\\n",
    "&=\\frac{1}{6} \\frac{(2m)^{3/2}L^3}{\\pi^2\\hbar^3} E^{3/2} = \\frac{4\\pi(2m)^{3/2}}{3}\\frac{L^3}{h^3} E^{3/2}\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Note that this result is proportional to the volume $L^3$ of the cube, and includes a $1/h^3$ factor (like we had when we talked about counting states in statistical mechanics).  We can take this particle in a box model as a model for the ideal gas (in terms of state counting) and derive the Sackur-Tetrode equation this way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}