{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<center><h2>Quantum Bouncer </h2></center>\n",
    "\n",
    "\"\"\" From \"A SURVEY OF COMPUTATIONAL PHYSICS\", Python eBook Version\n",
    "   by RH Landau, MJ Paez, and CC Bordeianu\n",
    "   Copyright Princeton University Press, Princeton, 2011; Book  Copyright R Landau, \n",
    "   Oregon State Unv, MJ Paez, Univ Antioquia, C Bordeianu, Univ Bucharest, 2011.\n",
    "   Support by National Science Foundation , Oregon State Univ, Microsoft Corp\"\"\" \n",
    "   \n",
    "A neutron is  dropped in a uniform gravitational field, hits a hard floor,\n",
    "and then bounces. Quantized levels result when the time-independent Schrödinger equation is solved via path integration:\n",
    " \n",
    " $$\n",
    "-\\frac{\\hbar^2}{ 2m}\\frac{d^2\\psi(x)}{ dx^2} + mxg\\,\\psi(x) = E\\,\n",
    "\\psi(x), \\quad\n",
    "\\psi(x \\leq 0) = 0,\\quad\\mbox{(boundary\n",
    "condition)} \n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# QMCbouncer.py:       g.s. wavefunction via path integration \n",
    "\n",
    "%matplotlib notebook\n",
    "from vpython import *\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "# Parameters\n",
    "N = 100;  dt = 0.05;          g = 2.0;       h = 0.00;      maxel = 0\n",
    "path = np.zeros([101], float);  arr = path;\n",
    "prob = np.zeros([201],float) # Init                  \n",
    "trajec = canvas(width = 300, height=500,title = 'Spacetime Trajectory')\n",
    "trplot = curve(y = range(0, 100), color=color.magenta, display=trajec)\n",
    "\n",
    "def trjaxs():                               # plot axis for trajectories\n",
    "  trax=curve(pos=[vec(-97,-100,0),vec(100,-100,0)],color=color.cyan,\n",
    "             canvas=trajec)\n",
    "  curve(pos = [vec(-65, -100,0),vec(-65, 100,0)], color=color.cyan,\n",
    "        canvas=trajec)\n",
    "  label(pos = vec(-65,110,0), text = 't', box = 0,canvas=trajec)\n",
    "  label(pos = vec(-85, -110,0), text = '0', box = 0,canvas = trajec)\n",
    "  label(pos = vec(60, -110,0), text = 'x', box = 0, canvas=trajec) \n",
    "\n",
    "wvgraph = canvas(x=350, y=80, width=500, height=300, title = 'GS Prob')\n",
    "wvplot  = curve(x = range(0, 50)) # wave function plot\n",
    "wvfax   = curve(color = color.cyan)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def wvfaxs():       # plot axis for wavefunction\n",
    "   wvfax = curve(pos =[vec(-200,-155,0),vec(800,-155,0)],\n",
    "                 canvas=wvgraph,color=color.cyan)\n",
    "   curve(pos = [vec(-200,-150,0),vec(-200,400,0)],\n",
    "         canvas=wvgraph, color=color.cyan)\n",
    "   label(pos = vec(-70, 420,0),text = 'Probability', box = 0,\n",
    "        canvas=wvgraph)\n",
    "   label(pos = vec(600, -220,0),text = 'x', box = 0,canvas=wvgraph)\n",
    "   label(pos = vec(-200, -220,0),text = '0', box = 0, canvas=wvgraph)\n",
    "                       \n",
    "trjaxs();  wvfaxs()                                           # plot axes\n",
    "\n",
    "def energy (arr):                             # Function for Energy of path\n",
    "    esum = 0.                          \n",
    "    for i in range(0,N):\n",
    "      esum += 0.5*((arr[i+1]-arr[i])/dt)**2+g*(arr[i]+arr[i+1])/2                                                         \n",
    "    return esum\n",
    "\n",
    "def plotpath(path):                        # Function to plot xy trajectory\n",
    "   for j in range (0, N):                     \n",
    "       trplot.append(pos=vec(20*path[j] - 65, 2*j - 100,0))\n",
    "\n",
    "def plotwvf(prob):                         # Function to plot wave function\n",
    "    for i in range (0, 50):\n",
    "       wvplot.color = color.yellow\n",
    "       wvplot.append(pos=vec(20*i - 200,0.5*prob[i] - 150,0)) \n",
    "     \n",
    "oldE = energy(path)                                           # Initial E\n",
    "counter = 1                                       # for ea 100 iterations\n",
    "norm = 0.                                       # wavefunction is plotted\n",
    "maxx = 0.0\n",
    "\n",
    "while 1:                                                # \"Infinite\" loop\n",
    "    rate(100)\n",
    "    element = int(N*random.random() )\n",
    "    if element != 0 and element!= N:                   # Ends not allowed\n",
    "        change = ( (random.random() - 0.5)*20.)/10.  #  -1 =<random  =< 1\n",
    "        if  path[element] + change > 0.:              # No negative paths\n",
    "          path[element]   += change                  # change temporarily\n",
    "          newE = energy(path)                          # New trajectory E\n",
    "          if newE > oldE and exp( - newE + oldE) <= random.random() :   \n",
    "             path[element]   -= change                    # Link rejected\n",
    "             trplot.clear()\n",
    "             plotpath(path)   \n",
    "             trplot.visible=True\n",
    "          ele = int(path[element]*1250./100.)             # Scale changed\n",
    "          if  ele >= maxel:  maxel = ele            # Scale change 0 to N                                           \n",
    "          if  element != 0:  prob[ele]   += 1 \n",
    "          oldE = newE;                             \n",
    "    if counter%100 == 0:                    # plot wavefunction every 100\n",
    "        for  i in range(0, N):                      # max x value of path\n",
    "            if  path[i] >= maxx:  maxx = path[i]           \n",
    "        h = maxx/maxel                                       # space step\n",
    "        firstlast = h*0.5*(prob[0] +  prob[maxel])   # for trap. extremes\n",
    "        for  i in range(0, maxel + 1):  norm = norm + prob[i]      # norm\n",
    "        norm = norm*h + firstlast                             # Trap rule\n",
    "        wvplot.clear()\n",
    "        plotwvf(prob)  \n",
    "        wvplot.visible=True  # plot probability\n",
    "    counter   += 1        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}