Final: WIP: Task 4.3 is hard.

Do I need to force psi to be real?

Final/Final - Tight-binding propagation method.ipynb

@@ -104,7 +104,8 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "from matplotlib import pyplot as plt"
+    "from matplotlib import pyplot as plt\n",
+    "import scipy.linalg"
    ]
   },
   {
@@ -1251,7 +1252,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 15,
    "metadata": {
     "deletable": false,
     "nbgrader": {
@@ -1269,8 +1270,21 @@
    "outputs": [],
    "source": [
     "def getU_exact(tau, H):\n",
-    "    # YOUR CODE HERE\n",
-    "    raise NotImplementedError()"
+    "    \"\"\"\n",
+    "    Calculates the time-propagation matrix for time-step tau.\n",
+    "    \n",
+    "    Args:\n",
+    "        tau: time-step numeric to calculate the time-propagation matrix at\n",
+    "        H:   n by n array representing the hamiltonian matrix\n",
+    "\n",
+    "    Returns:\n",
+    "        The time-propagation matrix U(tau).\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    hbar = 1\n",
+    "    return scipy.linalg.expm(-1j*hbar*H*tau)\n",
+    "    \n",
+    "    # TODO: The sentence above looks faulty. U(tau) seems wrong."
    ]
   },
   {
@@ -1300,7 +1314,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 16,
    "metadata": {
     "deletable": false,
     "nbgrader": {
@@ -1318,8 +1332,37 @@
    "outputs": [],
    "source": [
     "def timePropagate(U, c0, t):\n",
-    "    # YOUR CODE HERE\n",
-    "    raise NotImplementedError()"
+    "    \n",
+    "    # The function assumes that t[j] = j*tau, so we assert that the\n",
+    "    # values are equidistant.\n",
+    "    assert np.allclose(t[1:] - t[:-1], t[1] - t[0])\n",
+    "    # And that indeed t[0] = 0*tau = 0.\n",
+    "    assert np.allclose(t[0], 0)\n",
+    "    \n",
+    "    c = np.zeros(t.shape + c0.shape, dtype=c0.dtype)\n",
+    "    c[0] = c0\n",
+    "    \n",
+    "    # TODO: Add docstring.\n",
+    "    for j in range(1, len(t)):\n",
+    "        c[j] = U@c[j - 1]\n",
+    "    \n",
+    "    return c"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tau = .1\n",
+    "t = np.arange(100)*tau\n",
+    "H = np.array([[1,2],[-2,4]])\n",
+    "U = getU_exact(tau, H)\n",
+    "\n",
+    "c0 = np.array([2,2], dtype=np.complex128)\n",
+    "\n",
+    "timePropagate(U, c0, t)"
    ]
   },
   {
@@ -1382,7 +1425,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 17,
    "metadata": {
     "deletable": false,
     "nbgrader": {
@@ -1397,10 +1440,56 @@
      "task": false
     }
    },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_2000248/327551893.py:14: ComplexWarning: Casting complex values to real discards the imaginary part\n",
+      "  c[j] = U@c[j - 1]\n"
+     ]
+    }
+   ],
+   "source": [
+    "n = 100\n",
+    "# TODO: Should one be able to pass a through everything?\n",
+    "a = 1\n",
+    "sigma = .25\n",
+    "tau = 1.5\n",
+    "t = np.arange(201)*tau\n",
+    "\n",
+    "H = TBHamiltonian(n, sigma)\n",
+    "U = getU_exact(tau, H)\n",
+    "#c0 = np.zeros(n, dtype=np.complex128)\n",
+    "c0 = np.zeros(n)\n",
+    "c0[n//2] = 1\n",
+    "\n",
+    "c = timePropagate(U, c0, t)\n",
+    "xi = atomic_positions(n, a)\n",
+    "psi = c@atomic_basis(x, xi, sigma).T\n",
+    "\n",
+    "x = np.linspace(-1, 11, 150)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
    "outputs": [],
    "source": [
-    "# YOUR CODE HERE\n",
-    "raise NotImplementedError()"
+    "# This cell can be deleted!\n",
+    "x = np.linspace(-1, 11, 150)\n",
+    "xi = atomic_positions(n, a)\n",
+    "#t = 2\n",
+    "bas = atomic_basis(x, xi, sigma)\n",
+    "psi = c@bas.T\n",
+    "#atomic_basis(x, xi, sigma).T.shape\n",
+    "#print(x.shape, xi.shape, c[t].shape, bas.shape)\n",
+    "#print(psi)\n",
+    "psi.shape\n",
+    "#psi[t][x]\n",
+    "plt.plot(x, psi[30])\n",
+    "print(psi[2].shape, x.shape)"
    ]
   },
   {
@@ -1422,10 +1511,33 @@
    },
    "outputs": [],
    "source": [
-    "# Animate here ...\n",
+    "# use matplotlib's animation package\n",
+    "import matplotlib.pylab as plt\n",
+    "import matplotlib\n",
+    "import matplotlib.animation as animation\n",
+    "# set the animation style to \"jshtml\" (for the use in Jupyter)\n",
+    "matplotlib.rcParams['animation.html'] = 'jshtml'\n",
     "\n",
-    "# YOUR CODE HERE\n",
-    "raise NotImplementedError()\n",
+    "# create a figure for the animation\n",
+    "fig = plt.figure()\n",
+    "plt.grid(True)\n",
+    "plt.xlim(-1, 11)     # fix x limits\n",
+    "plt.ylim(-1e-6, 1e-6)     # fix y limits\n",
+    "plt.xlabel('$x$')\n",
+    "plt.ylabel('$\\\\psi(t, x)$')\n",
+    "\n",
+    "# Create an empty plot object and prevent its showing (we will fill it each frame)\n",
+    "myPlot, = plt.plot([0], [0])\n",
+    "plt.close()\n",
+    "\n",
+    "# This function is called each frame to generate the animation (f is the frame number)\n",
+    "def animate(f):                   \n",
+    "    myPlot.set_data(x, psi[f])  # update plot\n",
+    "\n",
+    "# Show the animation\n",
+    "frames = np.arange(1, np.size(t))  # t is the time grid here\n",
+    "myAnimation = animation.FuncAnimation(fig, animate, frames, interval = 20)\n",
+    "myAnimation\n",
     "\n",
     "# Yann has an animation about an atomic orbital that starts\n",
     "# moving to left and right and then bounce back."