Fill in any place that says `# YOUR CODE HERE` or YOUR ANSWER HERE, as well as your name and collaborators below.
Grading for pre-lecture assignments is all or nothing. Partial credit is available for in-class assignments and checkpoints, but **only when code is commented**.


```python
NAME = ""
COLLABORATORS = ""
```

---


```python
import grading_helper as _test
```


```python
# Space for imports, utility functions, etc.

# YOUR CODE HERE
```

# Time-Dependent Schr&ouml;dinger Equation

(Adapted from textbook exercise 9.8)

Perhaps the most important partial differential equation, at least for physicists, is the Schr&ouml;dinger equation. In one dimension the Schr&ouml;dinger equation for a particle of mass $m$ with no potential energy reads

$$-\frac{\hbar^2}{2m} \frac{\partial^2\psi}{\partial x^2} = \mathrm{i}\hbar \frac{\partial\psi}{\partial t}\,.$$

In units where $\hbar/2m=1$, the equation simplifies to

$$-\frac{\partial^2\psi}{\partial x^2} = \mathrm{i}\frac{\partial\psi}{\partial t}\,.$$

For simplicity, let's put our particle in a box with impenetrable walls, so that we only have to solve the equation in a finite-sized space. The box forces the wavefunction $\psi$ to be zero at the walls, which we'll put at $x=0$ and $x=L$.

Replacing the second derivative in the Schr&ouml;dinger equation with a finite difference and applying the forward Euler method, we get the FTCS equation

$$\psi(x,t+h) = \psi(x,t) + h \frac{\mathrm{i}}{a^2} \bigl[ \psi(x+a,t) + \psi(x-a,t) - 2\psi(x,t) \bigr]\,,$$

where $a$ is the spacing of the spatial grid points and $h$ is the size of the time-step. (Be careful not to confuse the time-step $h$ with Planck's constant $\hbar$.) Performing a similar step in reverse, we get the
backward Euler method

$$\psi(x,t+h) - h \frac{\mathrm{i}}{a^2} \bigl[ \psi(x+a,t+h) + \psi(x-a,t+h) - 2\psi(x,t+h) \bigr] = \psi(x,t)\,.$$

And taking the average of these two, we get the Crank-Nicolson equation for the Schr&ouml;dinger equation:

$$\begin{align}
\psi(x,t+h) - h \frac{\mathrm{i}}{2a^2} \bigl[ &\psi(x+a,t+h) + \psi(x-a,t+h) - 2\psi(x,t+h) \bigr]\\
  &= \psi(x,t) + h \frac{\mathrm{i}}{2a^2} \bigl[ \psi(x+a,t) + \psi(x-a,t) - 2\psi(x,t) \bigr]\,.
\end{align}$$

This gives us a set of simultaneous equations, one for each grid point.

The boundary conditions on our problem tell us that $\psi=0$ at $x=0$ and $x=L$ for all $t$. In between these points we have grid points at $a$, $2a$, $3a$, and so forth. Let us arrange the values of $\psi$ at these interior points into a vector

$$ \boldsymbol{\psi}(t) = \begin{pmatrix} \psi(a,t) \\ \psi(2a,t) \\ \psi(3a,t) \\ \vdots \end{pmatrix}\,.$$

Then the Crank-Nicolson equations can be written in the form

$$\mathbf{A}\,\boldsymbol{\psi}(t+h) = \mathbf{B}\,\boldsymbol{\psi}(t)\,,$$

where the matrices $\mathbf{A}$ and $\mathbf{B}$ are both symmetric and tridiagonal:

$$\mathbf{A} = \begin{pmatrix} a_1 & a_2 \\
                         a_2 & a_1 & a_2 \\
                             & a_2 & a_1 & a_2 \\
                             &     & a_2 & a_1 \\
                             &     &     &     & \ddots
\end{pmatrix},\qquad\qquad
\mathbf{B} = \begin{pmatrix} b_1 & b_2 \\
                             b_2 & b_1 & b_2 \\
                                 & b_2 & b_1 & b_2 \\
                                 &     & b_2 & b_1 \\
                                 &     &     &     & \ddots
\end{pmatrix}\,,$$

with

$$ a_1 = 1 + \mathrm{i}\frac{h}{a^2},\qquad a_2 = - \mathrm{i}\frac{h}{2a^2},\qquad b_1 = 1 - \mathrm{i}\frac{h}{a^2},\qquad b_2 = \mathrm{i}\frac{h}{2a^2}\,.$$

The equation $\mathbf{A}\,\boldsymbol{\psi}(t+h) = \mathbf{B}\,\boldsymbol{\psi}(t)$ has precisely the form $\mathbf{A}\mathbf{x} = \mathbf{v}$ of the simultaneous equation problems we studied in Chapter 6 and can be solved using the same methods. Specifically, since the matrix $\mathbf{A}$ is tridiagonal in this case, we can use the fast tridiagonal version of Gaussian elimination.

The last piece we need are the initial conditions. Suppose that at time $t=0$ the wavefunction of the electron has the form

$$\psi(x,0) = \exp \biggl[ -\frac{(x-x_0)^2}{2\sigma^2} \biggr]e^{\mathrm{i}\kappa x}\,,$$

where

$$x_0 = \frac{L}{2},\qquad \sigma = \frac{L}{10},\qquad \kappa = \frac{50}{L}\,,$$

and $\psi=0$ on the walls at $x=0$ and $x=L$. (This expression for $\psi(x,0)$ is not normalized&mdash;there should really be an overall multiplying coefficient to make sure that the probability density for the electron integrates to unity. It's safe to drop the constant, however, because the Schr&ouml;dinger equation is linear, so the constant cancels out on both sides of the equation and plays no part in the solution. If desired, it's easy to normalize the final result.)

Write a program to perform the Crank-Nicolson method for this electron, calculating the vector $\boldsymbol{\psi}(t)$ of values of the wavefunction for each time step between $t=0$ and $t=0.1$. Use the initial wavefunction above and $N=1001$ spatial grid points, so that $a=L/1000$. Since everything scales with $L$, the easiest approach is to set $L=1$.

Your program will have to perform the following steps:

- First, given the vector $\boldsymbol{\psi}(0)$ at $t=0$, you will have to multiply by the matrix $\mathbf{B}$ to get a vector $\mathbf{v} = \mathbf{B}\,\boldsymbol{\psi}$. Because of the tridiagonal form of $\mathbf{B}$, this is fairly simple. The $i$th component of $\mathbf{v}$ is given by

 $$v_i = b_1\psi_i + b_2(\psi_{i+1}+\psi_{i-1})\,.$$

 You will also have to choose a value for the time-step $h$. A reasonable choice is $h=10^{-4}$ (giving us 1000 time steps).


- Second, you will have to solve the linear system $\mathbf{A}\,\mathbf{x}=\mathbf{v}$ for $\mathbf{x}$, which gives you the new value of $\boldsymbol{\psi}$. You could do this using a standard linear equation solver like the function `solve` in `scipy.linalg`, but since the matrix $\mathbf{A}$ is tridiagonal, a better approach would be to use `solve_banded`.


- Third, once you have the code in place to perform a single step of the calculation, extend your program to perform repeated steps and hence solve for $\psi$ at a sequence of times a separation $h$ apart. Note that the
matrix $\mathbf{A}$ is independent of time, so it doesn't change from one step to another. You can set up the matrix just once and then keep on reusing it for every step.


- Finally, make an animation of the time evolution of the probability density $|\psi(x,t)|^2$ for the particle.


```python
%%graded # 20 points

# YOUR CODE HERE
```

6338114881554b69e044801dcb1c18e87aee17ae
ipynb
Jupyter Notebook
Assignments/11.5 Checkpoint - Time Dependent Schrodinger Equation.ipynb
davidnero/phys1321-lectures
328222be3de1d68de531f0cc008a6607032e6721
BSD-4-Clause-UC
Assignments/11.5 Checkpoint - Time Dependent Schrodinger Equation.ipynb
davidnero/phys1321-lectures
328222be3de1d68de531f0cc008a6607032e6721
BSD-4-Clause-UC
Assignments/11.5 Checkpoint - Time Dependent Schrodinger Equation.ipynb
davidnero/phys1321-lectures
328222be3de1d68de531f0cc008a6607032e6721
BSD-4-Clause-UC
Qwen/Qwen-72B
1. YES
2. YES