{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# nbi:hide_in\n",
    "import nbinteract as nbi\n",
    "import numpy as np\n",
    "import scipy as sp "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# The particle in a box in 1D\n",
    "Here's a quick reminder of the particle in a box \n",
    "\n",
    "\n",
    "<div class=\"col-md-12\">\n",
    "                    </div>\n",
    "                    <div class=\"col-md-12\">    \n",
    "                    <div class=\"col-md-4\">\n",
    "                        <img src=\"tik1d.png\" alt=\"Motivation\">\n",
    "                    </div>    \n",
    "                    <div class=\"col-md-8\">\n",
    "                        <p>  Schrödinger equation (SE): $\\hat{H}\\psi = E\\psi $</p>\n",
    "                        <p class=\"\">  \n",
    "                        SE: action of energy operator $\\hat{H}$ defines the energy and wave function $\\psi$  </p>\n",
    "                        <p class=\"\"> $$ \\hat{H} = -\\frac{\\hbar^2}{2m}\\frac{d^2}{dx^2} +V  $$</p>\n",
    "                        <p class=\" margin-top-10\"> boundary conditions: $ \\psi(0) = 0 \\quad\\text{and}\\quad \\psi(L) = 0 $</p>\n",
    "                    </div>\n",
    "                    </div>                    \n",
    "                    <div class=\"col-md-6\">                   \n",
    "                        <p><b>wave function</b></p>\n",
    "                        <p> $$ \\psi_n(x) = \\sqrt{\\frac{2}{L}}\\sin{\\frac{n\\pi x}{L}}  $$ </p>\n",
    "                    </div>\n",
    "                    <div class=\"col-md-6\">                    \n",
    "                        <p><b>energies</b></p>\n",
    "                        <p> $$ E_n = \\frac{\\hbar^2\\pi^2}{2m}\\frac{n^2}{L^2}  $$ </p>\n",
    "                    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# nbi:hide_in\n",
    "def make_xgrid(start, stop, N):\n",
    "    x = np.linspace(start,stop,N)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# nbi:hide_in\n",
    "def potential(x, box_l=0.0, box_r=5.0, height=100000.0): \n",
    "    pot = np.zeros_like(x)     #isn't numpy great? This generates an array with zeros, which has the same shape like x\n",
    "    \n",
    "    for i, xvalue in enumerate(x): #we loop through all x values\n",
    "        #enumerate is really nice. \n",
    "        #It allows us to loop through an array and, for each value, gives us an index i and the actual xvalue\n",
    "        if (xvalue<box_l): #left wall\n",
    "            pot[i] = height\n",
    "        elif (xvalue> box_l and xvalue < box_r): #inside the box\n",
    "            pot[i] = 0.0\n",
    "        elif (xvalue>box_r): #right wall\n",
    "            pot[i] = height\n",
    "        else:\n",
    "\n",
    "            raise ValueError(\"COMPUTER SAYS NO. This point should never be reached.\")\n",
    "    return pot\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# nbi:hide_in\n",
    "def create_H(x, pot):\n",
    "    hbar = 1\n",
    "    m = 1\n",
    "    dx = x[1]-x[0]\n",
    "    \n",
    "    import scipy.ndimage\n",
    "    L = scipy.ndimage.laplace(np.eye(len(x)), mode='wrap')/(2.0*dx*dx) #np.eye is a diagonal unit matrix (I) with size 'x' by 'x'\n",
    "    T= -(1./2.)*((hbar**2)/m)*L\n",
    "    \n",
    "    V = np.diag(pot)   #np.diag takes a list of numbers and puts it onto the diagonal; the offdiagonals are 0\n",
    "    \n",
    "    H = T+V\n",
    "    return H"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# nbi:hide_in\n",
    "def diagonalise_H(H):\n",
    "    import scipy.linalg as la\n",
    "    E, psi = la.eigh(H)\n",
    "    return E, psi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# nbi:hide_in\n",
    "\n",
    "#### %matplotlib notebook \n",
    "%matplotlib inline \n",
    "import matplotlib.pyplot as plt\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "    \n",
    "def plot_pib1d(x, pot, E, psi, n=0):\n",
    "        \"\"\"\n",
    "        plots the energies, the wave function and the probability density with interactive handles.\n",
    "        \"\"\"\n",
    "\n",
    "        n=n-1\n",
    "        import matplotlib.gridspec as gridspec\n",
    "        plt.figure(figsize=(11, 8), dpi= 80, facecolor='w', edgecolor='k') # figsize determines the actual size of the figure\n",
    "        gs1 = gridspec.GridSpec(2, 3)\n",
    "        gs1.update(left=0.05, right=0.90, wspace=0.35,hspace=0.25)\n",
    "        ax1 = plt.subplot(gs1[: ,0])\n",
    "        ax2 = plt.subplot(gs1[0, 1:])\n",
    "        ax3 = plt.subplot(gs1[1, 1:])\n",
    "\n",
    "\n",
    "        ax1.plot(x, pot, color='gray')\n",
    "        ax1.plot(np.ones(len(E)), E, lw=0.0, marker='_',ms=4000, color='blue')\n",
    "        ax1.plot(1, E[n], lw=0.0, marker='_',ms=4000, color='red')\n",
    "        ax1.set_ylim(-1,21)\n",
    "        ax2.plot(x, psi[:,n])  \n",
    "        ax3.plot(x, psi[:,n]*psi[:,n])\n",
    "\n",
    "        #figure out boundaries\n",
    "        start=0; stop=0\n",
    "        for i, p in enumerate(pot):\n",
    "            if (p>-0.001 and p<0.001):\n",
    "                start=x[i-1]\n",
    "                break\n",
    "        for i, p in enumerate(pot):\n",
    "            if (p>-0.001 and p<0.001):\n",
    "                stop=x[i+1]\n",
    "\n",
    "        ax2.axvline(start, color='gray')\n",
    "        ax2.axvline(stop, color='gray')  \n",
    "        ax3.axvline(start, color='gray')\n",
    "        ax3.axvline(stop, color='gray')\n",
    "        ax3.axhline(0.0,color='gray')\n",
    "\n",
    "        #Labeling of x and y axes\n",
    "        ax1.set_xlabel('x',fontsize=14)\n",
    "        ax1.set_ylabel('Energy [eV]',fontsize=14)\n",
    "        ax2.set_ylabel(r'wave function $\\psi(x)$',fontsize=14)\n",
    "        ax3.set_ylabel(r'density $|\\psi|^2$(x)',fontsize=14)\n",
    "        ax2.set_xlabel(r'x',fontsize=14)\n",
    "        ax3.set_xlabel(r'x',fontsize=14)\n",
    "\n",
    "        #Show the final result\n",
    "        plt.show()\n",
    "        return plot_pib1d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# nbi:hide_in\n",
    "def pib1d(N=100, box_l=0.0, box_r=5.0, height=100000.0, n=0):\n",
    "    \"\"\"\n",
    "    calculates and visualises the energies and wave functions of the 1d particle in a box\n",
    "    \n",
    "    This program takes the number of grid points and the start and end point of the box as arguments and \n",
    "    returns a visualisation.\n",
    "    \"\"\"\n",
    "    \n",
    "    #make an x axis\n",
    "    x = make_xgrid(box_l-1.0,box_r+1.0,N)\n",
    "    #create a box potential on that axis\n",
    "    pot = potential(x,box_l, box_r, height)\n",
    "    #calculate Hamiltonian\n",
    "    H = create_H(x,pot)\n",
    "    #diagonalise Hamiltonian, solve for E and psi\n",
    "    E, psi = diagonalise_H(H)\n",
    "    plot_pib1d(x, pot, E, psi, n)\n",
    "    \n",
    "    #If the function doesn't have a return value and simply executes a number of commands, \n",
    "    #we simply return the logical value true when everything ran smoothly to the end\n",
    "    return pib1d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b7345182ac7947a28b2cb206613ac5de",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=5.0, description='box_r', max=10.0, min=1.0), FloatSlider(value=10000.…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.pib1d(N=100, box_l=0.0, box_r=5.0, height=100000.0, n=0)>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# nbi:hide_in\n",
    "#plot results\n",
    "N = 100\n",
    "box_l = 0.0\n",
    "box_r = 5.0\n",
    "height = 100000.0\n",
    "#pib1d(N,box_l,box_r, height)\n",
    "interact(pib1d, N=fixed(N), box_l=fixed(box_l), box_r=(1.0,10.0), height=(0.,10000,10), n=range(1,21))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## The particle in a slanted box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# nbi:hide_in\n",
    "\n",
    "#potential function\n",
    "def tilted_potential(x, box_l=0.0, box_r=5.0, height=100000.0, grad=1.0): \n",
    "    pot = np.zeros_like(x)     #isn't numpy great? This generates an array with zeros, which has the same shape like x\n",
    "    \n",
    "    for i, xvalue in enumerate(x): #we loop through all x values\n",
    "        #enumerate is really nice. \n",
    "        #It allows us to loop through an array and, for each value, gives us an index i and the actual xvalue\n",
    "        if (xvalue<box_l): #left wall\n",
    "            pot[i] = height\n",
    "        elif (xvalue> box_l and xvalue < box_r): #inside the box\n",
    "            pot[i] = xvalue*grad\n",
    "        elif (xvalue>box_r): #right wall\n",
    "            pot[i] = height\n",
    "        else:\n",
    "            raise ValueError(\"COMPUTER SAYS NO. This point should never be reached.\")\n",
    "    return pot\n",
    "\n",
    "\n",
    "#new main program where we replaced the potential\n",
    "def tilted_pib(N,box_l=0.0,box_r=5.0, height=10000.0,n=1, grad=1.0):\n",
    "        \n",
    "    x = make_xgrid(box_l-1.0,box_r+1.0,N)\n",
    "    ######ONLY THE FOLLOWING LINES ARE DIFFERENT\n",
    "    x0 = (box_l+box_r)/2  #define the centre of the harmonic oscillator in the middle of the box\n",
    "    pot = tilted_potential(x, box_l, box_r, height, grad)\n",
    "    ######END OF DIFFERENCE##########\n",
    "    H = create_H(x,pot)\n",
    "    E, psi = diagonalise_H(H)\n",
    "    plot_pib1d(x, pot, E, psi, n)\n",
    "    \n",
    "    return tilted_pib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38b90a66eebe4366b53a0d237eac6737",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(Dropdown(description='n', options=(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.tilted_pib(N, box_l=0.0, box_r=5.0, height=10000.0, n=1, grad=1.0)>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# nbi:hide_in\n",
    "#plot results\n",
    "N = 100\n",
    "box_l = 0.0\n",
    "box_r = 5.0\n",
    "height = 100000.0\n",
    "#pib1d(N,box_l,box_r, height)\n",
    "interact(tilted_pib, N=fixed(N), box_l=fixed(box_l), box_r=fixed(box_r), height=fixed(height), n=range(1,20), grad=(0.,10.))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## The Harmonic oscillator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# nbi:hide_in\n",
    "\n",
    "#potential function\n",
    "def harmonic_potential(x, x0=2.5, k=10.):\n",
    "    \"\"\"\n",
    "    Function calculates harmonic potential on x-axis centred around x0 with spring constant k\n",
    "    \"\"\"\n",
    "    return 0.5*k*(x-x0)*(x-x0)\n",
    "\n",
    "\n",
    "#new main program where we replaced the potential\n",
    "def HO1D(N,box_l,box_r, k=10.0, n=0):\n",
    "        \n",
    "    x = make_xgrid(box_l-1.0,box_r+1.0,N)\n",
    "    ######ONLY THE FOLLOWING LINES ARE DIFFERENT\n",
    "    x0 = (box_l+box_r)/2  #define the centre of the harmonic oscillator in the middle of the box\n",
    "    pot = harmonic_potential(x, x0, k)\n",
    "    ######END OF DIFFERENCE##########\n",
    "    H = create_H(x,pot)\n",
    "    E, psi = diagonalise_H(H)\n",
    "    plot_pib1d(x, pot, E, psi, n)\n",
    "    \n",
    "    return HO1D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bdc6616f09ac4377ae3deaaea35e96fe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=10.0, description='k', min=1.0), Dropdown(description='n', options=(1,…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.HO1D(N, box_l, box_r, k=10.0, n=0)>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# nbi:hide_in\n",
    "#Number of grid points\n",
    "N = 100\n",
    "box_l = 0.0\n",
    "box_r = 5.0\n",
    "height = 10.0\n",
    "interact(HO1D, N=fixed(N),box_l=fixed(box_l),box_r = fixed(box_r),k=(1.,100.), n=range(1,20))"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}