08_01_Schroedinger.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Time dependent Schrödinger equation\n",
    "==================================\n",
    "\n",
    "We want to describe an electron wavefunction by a wavepacket\n",
    "$\\psi (x,y)$ that is a function of position $x$ and time $t$. We assume\n",
    "that the electron is initially localized around $x_0$, and model this by\n",
    "a Gaussian multiplying a plane wave:\n",
    "$$\\psi(x,t=0)=\\exp{\\left[-\\frac{1}{2}\\left(\\frac{x-x_0}{\\sigma _0}\n",
    "\\right)^2\\right ]} e^{ik_0x}\n",
    "$$\n",
    "\n",
    "This wave function does not correspond to an electron with a well\n",
    "defined momentum. However, if the width of the Gaussian $\\sigma _0$ is\n",
    "made very large, the electron gets spread over a sufficiently large\n",
    "region of space and can be considered as a plane wave with momentum\n",
    "$k_0$ with a slowly varying amplitude.\n",
    "\n",
    "The behavior of this wave packet as a function of time is described by\n",
    "the time-dependent Schröedinger equation (here in 1d):\n",
    "\n",
    "$$i\\frac{\\partial \\psi}{\\partial t}=H\\psi(x,t).$$ \n",
    "\n",
    "$H$ is the Hamiltonian operator:\n",
    "\n",
    "$$H=-\\frac{1}{2m}\\frac{\\partial ^2}{\\partial x^2}+V(x),$$ \n",
    "\n",
    "where $V(x)$ is a time independent potential. The\n",
    "Hamiltonian is chosen to be real. we have picked teh energy units such\n",
    "that $\\hbar=1$, and from now on, we will pick mass units such that\n",
    "$2m=1$ to make equations simpler.\n",
    "\n",
    "Scrhödinger’s equation is obviously a PDE, and we can use\n",
    "generalizations of the techniques learned in previous sections to solve\n",
    "it. The main observation is that this time we have to deal with complex\n",
    "numbers, and the function $\\psi (x,y)$ has real and imaginary parts:\n",
    "$$\\psi (x,t) = R(x,t)+iI(x,t).$$ However, is this section we will\n",
    "present an alternative method that makes the quantum mechanical nature\n",
    "of this problem more transparent.\n",
    "\n",
    "The time-evolution operator\n",
    "---------------------------\n",
    "\n",
    "The Scrödinger equation (\\[time\\]) can be integrated in a formal sense\n",
    "to obtain: \n",
    "$$\\psi(x,t)=U(t)\\psi(x,t=0)=e^{-iHt}\\psi(x,t).$$ \n",
    "\n",
    "From here we deduce that the wave function can be\n",
    "evolved forward in time by applying the time-evolution operator\n",
    "$U(t)=\\exp{-iHt}$: $$\\psi(t+\\Delta t)= e^{-iH\\delta t}\\psi(t).$$\n",
    "Likewise, the inverse of the time-evolution operator moves the wave\n",
    "function back in time: $$\\psi(t-\\Delta t)=e^{iH\\Delta t}\\psi(t),$$ where\n",
    "we have use the property $$U^{-1}(t)=U(-t).$$ Although it would be nice\n",
    "to have an algorithm based on the direct application of $U$, it has been\n",
    "shown that this is not stable. Hence, we apply the following relation:\n",
    "$$\\psi(t+\\Delta t)=\\psi(t-\\Delta t)+\\left[e^{-iH\\Delta t}-e^{iH\\Delta\n",
    "t}\\right]\\psi(t).$$ Now, the derivatives with recpect to $x$ can be\n",
    "approximated by \n",
    "\n",
    "$$\\begin{aligned}\n",
    "\\frac{\\partial \\psi}{\\partial t}\n",
    "&\\sim& \\frac{\\psi(x,t+\\Delta t)-\\psi(x, t)}{\\Delta t}, \\\\\n",
    "\\frac{\\partial ^2 \\psi}{\\partial x^2} &\\sim& \\frac{\\psi(x+\\Delta  %\n",
    "x,t)+\\psi(x-\\Delta x,t)-2\\psi(x,t)}{(\\Delta x)^2}.\n",
    "\\end{aligned}$$ \n",
    "\n",
    "The time evolution operator is\n",
    "approximated by: $$U(\\Delta t)=e^{-iH\\Delta t} \\sim 1+iH\\Delta t.$$\n",
    "Replacing the expression (\\[hami\\]) for $H$, we obtain:\n",
    "\n",
    "$$\\psi(x,t+\\Delta t)=\\psi(x,t)-i[(2\\alpha+\\Delta t\n",
    "V(x))\\psi(x,t)-\\alpha(\\psi(x+\\Delta x,t)+\\psi(x-\\Delta x,t))],\n",
    "$$ \n",
    "\n",
    "with $\\alpha=\\frac{\\Delta t}{(\\Delta x)^2}$. The\n",
    "probability of finding an electron at $(x,t)$ is given by\n",
    "$|\\psi(x,t)|^2$. This equations do no conserve this probability exactly,\n",
    "but the error is of the order of $(\\Delta t)^2$. The convergence can be\n",
    "determined by using smaller steps.\n",
    "\n",
    "We can write this expression explicitly for the real and imaginary parts, becoming:\n",
    "\n",
    "$$\\begin{aligned}\n",
    "\\mathrm{Im} \\psi(x, t + \\Delta t) = \\mathrm{Im} \\psi(x, t) + \\alpha \\mathrm{Re} \\psi (x + \\Delta x, t) + \\alpha \\mathrm{Re}\\psi(x − \\Delta x, t) − (2\\alpha + \\Delta t V (x)) \\mathrm{Re} \\psi(x, t) \\\\\n",
    "\\mathrm{Re} \\psi(x, t + \\Delta t) = \\mathrm{Re} \\psi(x, t) −  \\alpha \\mathrm{Im} \\psi (x + \\Delta x, t) − \\alpha \\mathrm{Im} \\psi (x − \\Delta x, t) + (2\\alpha + \\Delta tV (x)) \\mathrm{Im} \\psi(x, t)\n",
    "\\end{aligned}$$\n",
    "\n",
    "\n",
    "Notice the symmetry between these equations a: while the calculation of the imaginary part of the wave function at the later time involves a weighted average of the real part of the wave function at di erent positions from the earlier time, the calculation of the real part involves a weighted average of the imaginary part for di erent positions at the earlier time. This intermixing of the real and imaginary parts of the wave function may seem a bit strange, but remember that this situation is a direct result of our breaking up the wave function into its real and imaginary parts in the first place.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "### Exercise 8.1: Harmonic Potential \n",
    "Simulate a Gaussian wave-packet moving along the $x$ axis in a harmonic potential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Animation size has reached 20973919 bytes, exceeding the limit of 20971520.0. If you're sure you want a larger animation embedded, set the animation.embed_limit rc parameter to a larger value (in MB). This and further frames will be dropped.\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "from matplotlib import pyplot\n",
    "import math\n",
    "import matplotlib.animation as animation\n",
    "from IPython.display import HTML\n",
    "\n",
    "\n",
    "lx=20\n",
    "dx = 0.04\n",
    "nx = int(lx/dx)\n",
    "dt = dx**2/20.\n",
    "V0 = 60\n",
    "alpha = dt/dx**2\n",
    "\n",
    "fig = pyplot.figure()\n",
    "ax = pyplot.axes(xlim=(0, lx), ylim=(0, 2), xlabel='x', ylabel='|Psi|^2')\n",
    "points, = ax.plot([], [], marker='', linestyle='-', lw=3)\n",
    "\n",
    "psi0_r = np.zeros(nx+1)\n",
    "psi0_i = np.zeros(nx+1)\n",
    "\n",
    "x = np.arange(0, lx+dx, dx)\n",
    "\n",
    "#Define your potential \n",
    "V = np.zeros(nx+1)\n",
    "V = V0*(x-lx/2)**2\n",
    "\n",
    "#Initial conditions: wave packet\n",
    "sigma2 = 0.5**2\n",
    "k0 = np.pi*10.5\n",
    "x0 = lx/2\n",
    "for ix in range(0,nx):\n",
    "    psi0_r[ix] = math.exp(-0.5*((ix*dx-x0)**2)/sigma2)*math.cos(k0*ix*dx)\n",
    "    psi0_i[ix] = math.exp(-0.5*((ix*dx-x0)**2)/sigma2)*math.sin(k0*ix*dx)\n",
    "    \n",
    "def solve(i):\n",
    "    global psi0_r, psi0_i\n",
    "\n",
    "    for ix in range(1,nx-2):\n",
    "        psi0_i[ix]=psi0_i[ix]+alpha*psi0_r[ix+1]+alpha*psi0_r[ix-1]-(2.*alpha+dt*V[ix])*psi0_r[ix]\n",
    "    for ix in range(1,nx-2):\n",
    "        psi0_r[ix]=psi0_r[ix]-alpha*psi0_i[ix+1]-alpha*psi0_i[ix-1]+(2.*alpha+dt*V[ix])*psi0_i[ix]\n",
    "        \n",
    "    points.set_data(x,psi0_r**2 + psi0_i**2)\n",
    "    return (points,)\n",
    "\n",
    "\n",
    "#for i in range(2000):\n",
    "#    solve(i)\n",
    "      \n",
    "#pyplot.plot(x,psi0_r**2+psi0_i**2);\n",
    "\n",
    "\n",
    "anim = animation.FuncAnimation(fig, solve, frames = 8000, interval=50)\n",
    "\n",
    "HTML(anim.to_jshtml())            \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 8.2: Potential barrier ###\n",
    "Simulate a Gaussian wave-packet moving along the x axis passing through a potential barrier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.animation as animation\n",
    "from IPython.display import HTML\n",
    "\n",
    "\n",
    "fig = pyplot.figure()\n",
    "ax = pyplot.axes(xlim=(0, lx), ylim=(0, 2), xlabel='x', ylabel='$|\\Psi|^2$')\n",
    "points, = ax.plot([], [], marker='', linestyle='-', lw=3)\n",
    "\n",
    "x0=6\n",
    "\n",
    "for ix in range(0,nx):\n",
    "    psi0_r[ix] = math.exp(-0.5*((ix*dx-x0)**2)/sigma2)*math.cos(k0*ix*dx)\n",
    "    psi0_i[ix] = math.exp(-0.5*((ix*dx-x0)**2)/sigma2)*math.sin(k0*ix*dx)\n",
    "\n",
    "x = np.arange(0, lx+dx, dx)\n",
    "V = np.zeros(nx+1)\n",
    "\n",
    "for ix in range(nx//2-20,nx//2+20):\n",
    "    V[ix]=2000.\n",
    "                                               \n",
    "def solve(i):\n",
    "    global psi0_r, psi0_i\n",
    "\n",
    "    psi0_r[1:-1] = psi0_r[1:-1]- alpha*(psi0_i[2:]+psi0_i[:-2]-2*psi0_i[1:-1])+dt*V[1:-1]*psi0_i[1:-1] \n",
    "    psi0_i[1:-1] = psi0_i[1:-1]+ alpha*(psi0_r[2:]+psi0_r[:-2]-2*psi0_r[1:-1])-dt*V[1:-1]*psi0_r[1:-1] \n",
    "\n",
    "    points.set_data(x,psi0_r**2 + psi0_i**2)\n",
    "    return points\n",
    "\n",
    "\n",
    "#for i in range(2000):\n",
    "#    solve(i)\n",
    "      \n",
    "#pyplot.plot(x,psi0_r**2+psi0_i**2);\n",
    "\n",
    "\n",
    "anim = animation.FuncAnimation(fig, solve, frames = 2000, interval=10)\n",
    "\n",
    "HTML(anim.to_jshtml())            \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 8.2: Single-slit diffraction \n",
    "Young’s single-slit experiment consists of a wave passing though a small\n",
    "slit, which causes the emerging wavelets to intefere with eachother\n",
    "forming a diffraction pattern. In quantum mechanics, where particles are\n",
    "represented by probabilities, and probabilities by wave packets, it\n",
    "means that the same phenomenon should occur when a particle (electron,\n",
    "neutron) passes though a small slit. Consider a wave packet of initial\n",
    "width 3 incident on a slit of width 5, and plot the probability density\n",
    "$|\\psi ^2|$ as the packet crosses the slit. Generalize the\n",
    "time-evolution equation (\\[time\\_diff\\]) for 2 dimensions. Model the\n",
    "slit with a potential wall:\n",
    "$$V(x,y)=100 \\,\\,\\,\\,\\,\\ \\mathrm{for}\\,\\,x=10,|y|\\geq 2.5.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYUAAAD8CAYAAACYebj1AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAayUlEQVR4nO3df5Ac5X3n8fdnV1pJiEUsCBZJgPlZusMkxlhHbEwS2eRkUHHBSVEYKuVgG5dCEqrsqkvFOL6yXeQqhS9l5y4mMacYFeAjhFxsgiqHAcVxCv8IGKESvwxCEj+CFlkCJPTDEqxW+70/pqc9GvXM9E7PTM9Kn1eVanu6n+5+tnf0fOZ5uqdbEYGZmRnAQNkVMDOz/uFQMDOzlEPBzMxSDgUzM0s5FMzMLOVQMDOzVMtQkHSapO9L+qmkZyV9Jpl/gqQ1kjYmP0carH9dUmajpOs6/QuYmVnnqNX3FCQtABZExDpJw8ATwEeBTwA7IuIWSTcBIxHxubp1TwDWAkuASNZ9X0Ts7PhvYmZmhbXsKUTE1ohYl0zvAZ4DFgFXAncmxe6kEhT1PgKsiYgdSRCsAS7rRMXNzKzzZkylsKQzgPcCjwGjEbE1WfQzYDRjlUXAqzWvtyTzsra9AlgBMKiZ75s7I3M0ynIJYjKIWTOYnDVYdmXMABAwsP8AmpgE+XRmp+0+sP2NiDip6HZyh4KkY4FvA5+NiN2S0mUREZIK3S8jIlYCKwHmDY3GxSd/rMjmjloxMcGBmQPs+vVT2ffBSY5fuIuaP5VZaSYODrB74/GMPHKAuT95lcEZM9GAw6FTHhz7+iud2E6uUJA0k0og3B0R30lmb5O0ICK2JucdtmesOgYsrXl9KvCv7VfXmpl8+20mFp7Atg/PZ+bFu/nkf3yUs495vexqmQFwMMTjp5/J6hMv5J2Tz2Hke68w450JNGNKAxbWZS3/Gqp0CW4HnouIr9UsWg1cB9yS/Lw/Y/WHgD+ruTJpGfD5QjW2hmL8AAcWDjN+Eiw4cQ9nH/M6O8fm8OKPCvcozQr75Y9u4dy525h30l72jM5n3oITiee3oGMdCv0kz1/jg8DHgaclrU/m/QmVMPh7SdcDrwBXA0haAtwQEZ+OiB2S/hR4PFnv5ojY0dHfwJp68eGTWHfHu8quhhkLfmkXs86ZKLsa1kLLUIiIH1I5R5Tl0ozya4FP17xeBaxqt4JmZtY7PstjZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaVaPo5T0irgCmB7RJyfzLsXWJwUOR54KyIuyFj3ZWAPcBCYiIglHaq3mZl1QctQAO4AbgXuqs6IiI9VpyV9FdjVZP0PRcQb7VbQzMx6p2UoRMQjks7IWiZJwNXAhztbLTMzK0PRcwq/CmyLiI0NlgfwsKQnJK0ouC8zM+uyPMNHzVwL3NNk+SURMSbpZGCNpOcj4pGsgklorACYPThcsFpmZtaOtnsKkmYAvw3c26hMRIwlP7cD9wEXNSm7MiKWRMSSoYE57VbLzMwKKDJ89BvA8xGxJWuhpLmShqvTwDLgmQL7MzOzLmsZCpLuAf4NWCxpi6Trk0XXUDd0JGmhpAeSl6PADyU9CfwE+H8R8WDnqm5mZp2W5+qjaxvM/0TGvNeA5cn0i8B7CtbPzMx6yN9oNjOzlEPBzMxSDgUzM0s5FMzMLOVQMDOzlEPBzMxSDgUzM0s5FMzMLOVQMDOzlEPBzMxSDgUzM0s5FMzMLOVQMDOzlEPBzMxSDgUzM0s5FMzMLOVQMDOzlEPBzMxSDgUzM0u1DAVJqyRtl/RMzbwvSxqTtD75t7zBupdJ2iBpk6SbOllxMzPrvDw9hTuAyzLm/0VEXJD8e6B+oaRB4K+Ay4HzgGslnVeksmZm1l0tQyEiHgF2tLHti4BNEfFiRIwDfwdc2cZ2zMysR4qcU7hR0lPJ8NJIxvJFwKs1r7ck8zJJWiFpraS145P7C1TLzMza1W4ofAM4G7gA2Ap8tWhFImJlRCyJiCVDA3OKbs7MzNrQVihExLaIOBgRk8DfUBkqqjcGnFbz+tRknpmZ9am2QkHSgpqXvwU8k1HsceBcSWdKGgKuAVa3sz8zM+uNGa0KSLoHWArMl7QF+BKwVNIFQAAvA7+XlF0IfDMilkfEhKQbgYeAQWBVRDzbld/CzMw6omUoRMS1GbNvb1D2NWB5zesHgMMuVzUzs/7kbzSbmVnKoWBmZimHgpmZpRwKZmaWciiYmVnKoWBmZimHgpmZpRwKZmaWciiYmVnKoWBmZimHgpmZpRwKZmaWciiYmVnKoWBmZimHgpmZpRwKZmaWavmQHTOzo93+8xe1td6cZ6bfY+kdCke4he97iy2Pj5RdDTOGT3qb8WnU5LQbBFnbmE7hkOcZzauAK4DtEXF+Mu/Pgf8CjAObgU9GxFsZ674M7AEOAhMRsaRzVbd6GhhgcM8BBvfNYu/+Id6amMOCX97NlX/9ZNlVMwNg07557Ns3i5k/h4Hd+9DM/guJToRBnm32a1Dk+YvcAdwK3FUzbw3w+YiYkPQV4PPA5xqs/6GIeKNQLS0XzT2GoRdeY/TtE3lz/CRu27+UM+a/znGz3i67amYcmJzB+k3vYsaPjuGE77/KjN1vo1mzyq7WIVoFws7FQ7m2M7JhPNe++jEYWoZCRDwi6Yy6eQ/XvHwUuKqz1bJ2SEKzZjP7tT2c8n928vMXTueFhSMcnOPrCax8Ohgc99Rejnl6IwOzZqOhfA1sLzQLg7xBkLVOq3Dox2DoRN/tU8C9DZYF8LCkAP53RKxstBFJK4AVALMHhztQraOXBgeYwRDDP36FY4HKn8GsfAMDg2j2nLKrcYhGgdAoDPacPZk5f3jz4R++8oZDPykUCpK+AEwAdzcocklEjEk6GVgj6fmIeCSrYBIYKwHmDY26FeuAgT7rmpt1Q6NGPc8n8Kx1pxoGjZbXhsTOxUPTJhjaDgVJn6ByAvrSiMhsxCNiLPm5XdJ9wEVAZiiY2dGlyJU5eU4G15ap38dUegf1jf2xZ+5i70vzWu5/z9mThwVDVW1AZA0hlXnVUluhIOky4I+BX4+IfQ3KzAUGImJPMr0MuLntmprZtJfVGDdrvPOs3+5+69UHQlbP4Ngzdx3ys159WNQHQ+2+GvUcunH101TkuST1HmApMF/SFuBLVK42mkVlSAjg0Yi4QdJC4JsRsRwYBe5Lls8A/jYiHuzKb2FmfS1vQ1dWg9gqEBqFQL3acnl6E/0oz9VH12bMvr1B2deA5cn0i8B7CtXOzKzHWgXCstM3APDwvy/Ovb28vYWyewngex+Z2VGu2SWnjQKhfvpI4lAws67r9gnTnYuHDvnXrlZXGNU7EoPBoWBmPdFvX9KybP134xEzO2JVg6F+7LxRYPTDGHsjec8pTMWcZ8ZK/50dCmbWc3l7DXkbyZEN44WGjfJqJwiyTjJD/37L2cNHZtbXejnsVN+At3NZ6d6X5hW6HLXsYTb3FMys7zUadqrVqU/ew5sHDjnhvPeleS2/p9AoBNrtJZQZDA4FMzuqZQ09ZQUDHH6J6lR7BFlhUHbPoJ5DwcymjW6diM0TDDC1EKjtJfTr+YMsPqdgZtZAo+GfVuvkDYR+6yWAQ8HMppkiDWmzdYt+mq8PgyJ1KZOHj8zsqFBthKc6BFVt6LO+7dwsBKZbD6HKoWBm085UG/b6RrjR+s2+7zCVXkCjQOjnMKjy8JGZHdF63RA36yGU/W3lPBwKZjbt5G1cy7h9RqtvVu8/f1Ffh4NDwcymlak0qK2e9NYtee7W2q/B4HMKZjYtdPNRnFV5HsmZ1/DmgYbPZa6tW7+dZ3AomFkpptIgdvtTdbMwyPsozqrqF9yq26ieoG4UEFnHocywyBUKklYBVwDbI+L8ZN4JwL3AGcDLwNURsTNj3euA/5a8/O8RcWfxapvZdFXbwNdO98vtsxsFQqsH6lTvoFpdpzYc6q9cynoUZ1X19y0rGPL2FO4AbgXuqpl3E/C9iLhF0k3J68/VrpQEx5eAJUAAT0hanRUeZnbka9bAZz1jod2hn6l8Ea12vWogVBv2qTxZrf7ZzceeuavtYChTrlCIiEcknVE3+0pgaTJ9J/Cv1IUC8BFgTUTsAJC0BrgMuKet2prZtDXVT/x5yjc6mVudP5VGt1EgfHzkx7m38a2dFx8SDq2CoaraK+iHk89Frj4ajYityfTPgNGMMouAV2teb0nmHUbSCklrJa0dn9xfoFpm1m/Kaux68eCdI01HTjRHREiKgttYCawEmDc0WmhbZtZf8jwPoR3VnkB949/OsEz1rqjV5yc8/O+LWXb6Br618+JDytX2HOqXwaFPZ6u9q2qzb0S3ewuObigSCtskLYiIrZIWANszyozxiyEmgFOpDDOZ2VGoWaNX5ERzkbH52ltbZAVDVXVYKCsI4PBHdbYKhH48nwDFQmE1cB1wS/Lz/owyDwF/Jmkkeb0M+HyBfZrZNFffa2h1hU2vPz3XBgP84hzDVJ7PnLeHkKX6+/b7Jan3UPnEP1/SFipXFN0C/L2k64FXgKuTskuAGyLi0xGxQ9KfAo8nm7q5etLZzI5uU2n0uh0M9TfCq70zarvPW87bO8g6DmV+oU0R/Td8P29oNC4++WNlV8PM+lA74VDbyDZbP8+J6fovpbXSq1toPzj29SciYknR7fjeR2Y2rUy1h5F12+xG8ozz532YzsiG8Wn5TAWHgplNO0Vvj9HNBjlPGPRrIIBDwcysZ/o5DKp8QzwzO6I1uuFcllbnFOrvmpo1jNTs29R5r7gqk3sKZjbtFLllRjsnqvecPZl5G+1G86F5wJT9BbVmHApmNq0Uea5CO1ce5XmmQjUc6stOx2BwKJiZNdDOQ3bqw2G6BYPPKZjZtNGtRjSr4c4KhPoH7jT7YlvtXVHbuWtrWdxTMDOrkycQqvOm8mS2rPDpt96CQ8HMpoUiz1eYyjr1gZCn4W9UJitcdi4ealnPVuc/usnDR2bW1/I2jtWGtsgTzbICYSrqH8XZjrJ7Du4pmFnPVT8Jt/pEPNVA6IVlp29o+ZjO2jBp55LVMrmnYGY906iRL/vTcV61YVD/TOZO6Ifj4J6CmfVEPzR41pp7CmbWdd0OhE5d6ll9wE5enewl9Av3FMzsqNYsUOpPGNeGwJEYCOBQMDM7RP1N7rKCYSqB0OjZC3mfwtZrHj4ys66rfy5zq3JZev04ztphpL0vzct1eWqRS1H7RduhIGkxcG/NrLOAL0bE/6wpsxS4H3gpmfWdiLi53X2a2fSWFQ55Px3nfaRmq3UbrZ8nGGBq30WYSi+htq5lnpRvOxQiYgNwAYCkQWAMuC+j6A8i4op292NmR56iwyR5eh7N9tGo4c0KBiAzHFrJCoSpnBAvayipU8NHlwKbI+KVDm3POiAOTBAT44DKrooZEGjWbDTQuVOZRRrORsFSbbizwgEafxmt2XObp/qs5jLPLXQqFK4B7mmw7AOSngReA/4oIp7NKiRpBbACYPbgcIeqdXSKCCbfeZv9v7SIfWfNJZwJ1gcGJuDYp95k1qs70NAspP54Y+btNVQ1a/yzTIc7o9ZSRBTbgDREpcF/d0Rsq1t2HDAZEXslLQf+V0Sc22qb84ZG4+KTP1aoXkerGB/nwHGz2fGhhUx8cD+nLXrdl5hZXxg/OMirG0c57gcDHPejVxicHECDg2VXK9VqHH+qt6XIEwad7BE8OPb1JyJiSdHtdKKncDmwrj4QACJid830A5L+WtL8iHijA/u1OpP79nHgzFG2fXiEOR/Yyaf+w49ZNOutsqtlllp/ymn8w3H/ifGRsznxn7cwY+87aObMsqsFtD5P0elP/P1w+WmWToTCtTQYOpJ0CrAtIkLSRVS+F/FmB/ZpGWLiIBMnHcOB42HhvH0smvUW214YZsM/jZZdNTMuuuFlTpu9g+GRfew54RgmT5pHvLmlb0KhKu/ls+1sczooFAqS5gL/Gfi9mnk3AETEbcBVwO9LmgD2A9dE0fEqm5ItPxphw/0Lyq6GGedc/jqzzpkouxq5dSIcplMYVBUKhYj4OXBi3bzbaqZvBW4tsg8zszJNx4a9CJ+DNDOzlEPBzMxSDgUzM0s5FMzMLOVQMDOzlEPBzMxSDgUzM0s5FMzMLOVQMDOzlEPBzMxSDgUzM0s5FMzMLOVQMDOzlEPBzMxSDgUzM0s5FMzMLOVQMDOzlEPBzMxShUNB0suSnpa0XtLajOWS9JeSNkl6StKFRfdpZmbdUegZzTU+FBFvNFh2OXBu8u9XgG8kP83MrM/0YvjoSuCuqHgUOF7Sgh7s18zMpqgToRDAw5KekLQiY/ki4NWa11uSeYeQtELSWklrxyf3d6BaZmY2VZ0YProkIsYknQyskfR8RDwy1Y1ExEpgJcC8odHoQL3MzGyKCvcUImIs+bkduA+4qK7IGHBazetTk3lmZtZnCoWCpLmShqvTwDLgmbpiq4HfTa5Cej+wKyK2FtmvmZl1R9Hho1HgPknVbf1tRDwo6QaAiLgNeABYDmwC9gGfLLhPMzPrkkKhEBEvAu/JmH9bzXQAf1hkP2Zm1hv+RrOZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUcCmZmlmo7FCSdJun7kn4q6VlJn8kos1TSLknrk39fLFZdMzPrpiLPaJ4A/mtErJM0DDwhaU1E/LSu3A8i4ooC+zEzsx5pu6cQEVsjYl0yvQd4DljUqYqZmVnvdeScgqQzgPcCj2Us/oCkJyV9V9K7O7E/MzPrjiLDRwBIOhb4NvDZiNhdt3gd8K6I2CtpOfCPwLkNtrMCWAEwe3C4aLXMzKwNhXoKkmZSCYS7I+I79csjYndE7E2mHwBmSpqfta2IWBkRSyJiydDAnCLVMjOzNhW5+kjA7cBzEfG1BmVOScoh6aJkf2+2u08zM+uuIsNHHwQ+DjwtaX0y70+A0wEi4jbgKuD3JU0A+4FrIiIK7NPMzLqo7VCIiB8CalHmVuDWdvdhZma95W80m5lZyqFgZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUcCmZmlnIomJlZyqFgZmYph4KZmaUKhYKkyyRtkLRJ0k0Zy2dJujdZ/pikM4rsz8zMuqvtUJA0CPwVcDlwHnCtpPPqil0P7IyIc4C/AL7S7v7MzKz7ivQULgI2RcSLETEO/B1wZV2ZK4E7k+l/AC6VpAL7tGYGB9DbEwzuD8b3zWDH/mPY9+ZQ2bUyA2D3a7PZtW8OE/sGGXw70DvjMOgR7H6jiGhvRekq4LKI+HTy+uPAr0TEjTVlnknKbEleb07KvJGxvRXAiuTl+cAzbVWsd+YDh/0efcj17CzXs7Ncz85ZHBHDRTcyoxM16YSIWAmsBJC0NiKWlFylpqZDHcH17DTXs7Ncz86RtLYT2ynSdxsDTqt5fWoyL7OMpBnAPODNAvs0M7MuKhIKjwPnSjpT0hBwDbC6rsxq4Lpk+irgX6Ld8SozM+u6toePImJC0o3AQ8AgsCoinpV0M7A2IlYDtwPfkrQJ2EElOPJY2W69emg61BFcz05zPTvL9eycjtSx7RPNZmZ25PH1YGZmlnIomJlZqrRQmA63yJB0mqTvS/qppGclfSajzFJJuyStT/59sdf1TOrxsqSnkzocdmmaKv4yOZ5PSbqwhDourjlO6yXtlvTZujKlHE9JqyRtT75bU513gqQ1kjYmP0carHtdUmajpOuyynS5nn8u6fnk73qfpOMbrNv0PdKDen5Z0ljN33Z5g3Wbtg1druO9NfV7WdL6Buv28lhmtkNde39GRM//UTkxvRk4CxgCngTOqyvzB8BtyfQ1wL0l1HMBcGEyPQy8kFHPpcA/lXEc6+rxMjC/yfLlwHcBAe8HHiu5voPAz4B39cPxBH4NuBB4pmbe/wBuSqZvAr6Ssd4JwIvJz5FkeqTH9VwGzEimv5JVzzzvkR7U88vAH+V4XzRtG7pZx7rlXwW+2AfHMrMd6tb7s6yewrS4RUZEbI2Idcn0HuA5YFEv69BBVwJ3RcWjwPGSFpRYn0uBzRHxSol1SEXEI1SukKtV+x68E/hoxqofAdZExI6I2AmsAS7rZT0j4uGImEhePkrlO0OlanA888jTNnREszombc3VwD3d2PdUNGmHuvL+LCsUFgGv1rzewuGNbVomecPvAk7sSe0yJMNX7wUey1j8AUlPSvqupHf3tGK/EMDDkp5Q5ZYh9fIc8166hsb/4frheAKMRsTWZPpnwGhGmX47rp+i0iPM0uo90gs3JsNcqxoMd/TL8fxVYFtEbGywvJRjWdcOdeX96RPNOUg6Fvg28NmI2F23eB2VIZD3AF8H/rHX9UtcEhEXUrlr7R9K+rWS6tGSKl92/E3g/2Ys7pfjeYio9MX7+vptSV8AJoC7GxQp+z3yDeBs4AJgK5XhmX51Lc17CT0/ls3aoU6+P8sKhWlziwxJM6n8Ie6OiO/UL4+I3RGxN5l+AJgpaX6Pq0lEjCU/twP3UemG18pzzHvlcmBdRGyrX9AvxzOxrTrElvzcnlGmL46rpE8AVwC/kzQQh8nxHumqiNgWEQcjYhL4mwb7L/14Ju3NbwP3NirT62PZoB3qyvuzrFCYFrfISMYVbweei4ivNShzSvVch6SLqBzTnoaXpLmShqvTVE481t9ldjXwu6p4P7CrpuvZaw0/hfXD8axR+x68Drg/o8xDwDJJI8lwyLJkXs9Iugz4Y+A3I2JfgzJ53iNdVXcO67ca7D9P29BtvwE8H8ndnev1+lg2aYe68/7sxdnzBmfUl1M5i74Z+EIy72Yqb2yA2VSGFzYBPwHOKqGOl1Dpkj0FrE/+LQduAG5IytwIPEvlKolHgYtLqOdZyf6fTOpSPZ619RSVhyJtBp4GlpT0d59LpZGfVzOv9ONJJaS2AgeojLteT+Uc1veAjcA/AyckZZcA36xZ91PJ+3QT8MkS6rmJyrhx9T1avWpvIfBAs/dIj+v5reS99xSVBm1BfT2T14e1Db2qYzL/jur7saZsmceyUTvUlfenb3NhZmYpn2g2M7OUQ8HMzFIOBTMzSzkUzMws5VAwM7OUQ8HMzFIOBTMzS/1/4JjA7sH2NJ0AAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "from matplotlib import pyplot\n",
    "import math\n",
    "\n",
    "lx      = 20                            #Box length in x\n",
    "ly      = 20                            #Box length in y\n",
    "dx      = 0.25                          #Incremental step size in x (Increased this to decrease the time of the sim)\n",
    "dy      = dx                            #Incremental step size in y\n",
    "nx      = int(lx/dx)                    #Number of steps in x\n",
    "ny      = int(ly/dy)                    #Number of steps in y\n",
    "dt      = dx**2/20.                     #Incremental step size in time\n",
    "sigma2  = 0.5**2                        #Sigma2 Value\n",
    "k0      = np.pi*10.5                    #K0 value\n",
    "amp     = math.pow(1./2., 64)           #Amplitude (to avoid large values out of range. This was one issue)\n",
    "\n",
    "alpha   = (dt/2.)/dx**2                 #Alpha\n",
    "\n",
    "psi0_r  = np.zeros(shape=(ny+1,nx+1))   #Initialize real part of psi\n",
    "psi0_i  = np.zeros(shape=(ny+1,nx+1))   #Initialize imaginary part of psi\n",
    "V       = np.zeros(shape=(ny+1,nx+1))   #Initialize Potential\n",
    "\n",
    "#Define your potential wall\n",
    "V = np.zeros(shape=(ny+1,nx+1))\n",
    "for ix in range(nx//2-20,nx//2+20):\n",
    "    for iy in range(0,ny):\n",
    "        if(abs(iy*dy-ly/2)>2.5):\n",
    "            V[iy,ix] = 200.\n",
    "\n",
    "#Initial conditions: wave packet\n",
    "x0 = 6.\n",
    "y0 = ly/2.\n",
    "for x in range(0,nx):\n",
    "    for y in range(0,ny):\n",
    "        psi0_r[y,x] = math.exp(-0.5*((x*dx-x0)**2+(y*dy-y0)**2)/sigma2)*math.cos(k0*x*dx)\n",
    "        psi0_i[y,x] = math.exp(-0.5*((x*dx-x0)**2+(y*dy-y0)**2)/sigma2)*math.sin(k0*x*dx)\n",
    "                \n",
    "x = np.arange(0, lx+dx, dx)\n",
    "y = np.arange(0, ly+dy, dy)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "\n",
    "pyplot.contourf(X,Y,psi0_r**2+psi0_i**2)\n",
    "pyplot.contour(X,Y,V)\n",
    "\n",
    "\n",
    "#Function to solve incremental changes in psi\n",
    "def solve():\n",
    "    #Grab psi lists\n",
    "    global psi0_r, psi0_i\n",
    "\n",
    "    #Calculate Imaginary Part in all points except for last 2 (because of indice notation)\n",
    "    for x in range(1,nx-2):\n",
    "        for y in range(1,ny-2):\n",
    "            psi0_i[y,x] = psi0_i[y,x] + alpha*psi0_r[y,x+1] + alpha*psi0_r[y,x-1] - (2*alpha + dt*V[y,x])*psi0_r[y,x] + alpha*psi0_r[y+1,x] + alpha*psi0_r[y-1,x] - (2*alpha+dt*V[y,x])*psi0_r[y,x]\n",
    "\n",
    "    #Calculate Real Part in all points except for last 2 (because of indice notation)\n",
    "    for x in range(1,nx-2):\n",
    "        for y in range(1,ny-2):\n",
    "            psi0_r[y,x] = psi0_r[y,x] - alpha*psi0_i[y,x+1] - alpha*psi0_i[y,x-1] + (2*alpha + dt*V[y,x])*psi0_i[y,x] - alpha*psi0_i[y+1,x] - alpha*psi0_i[y-1,x] + (2*alpha+dt*V[y,x])*psi0_i[y,x]\n",
    "\n",
    "\n",
    "for i in range(1000):\n",
    "    solve()\n",
    "    \n",
    "pyplot.contourf(X,Y,psi0_r**2+psi0_i**2);\n",
    "#pyplot.contour(X,Y,V);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}