<jupyter_start><jupyter_text>16 PDEs: Crank-Nicholson Method  See the [derivation of the Crank-Nicholson algorithm (PDF)](16_LectureNotes_CrankNicholson.pdf).<jupyter_code>import numpy as np

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D<jupyter_output><empty_output><jupyter_text>For 3D rotatable output:<jupyter_code>%matplotlib widget<jupyter_output><empty_output><jupyter_text>For HTML/nbviewer output, use inline:<jupyter_code>%matplotlib inline<jupyter_output><empty_output><jupyter_text>Crank-Nicholson implementationWe start with a straight-forward implementation that solves the implicit step with the standard linear algebra solver `numpy.linalg.solve()` for the matrix problem$$\mathsf{A}\mathbf{x} = \mathbf{b}$$where the unknowns $\mathbf{x}$ are the values of $T$ at the next time step $j+1$ _inside_ the boundaries$$\mathbf{x} = \begin{pmatrix}T_1\\ \vdots \\ T_i\\ \vdots\\T_{-2}\end{pmatrix}$$(i.e., $T_0$ and $T_{-1}$ are _not_ include).The main problem is to set up the matrices $\mathsf{M}(\eta)$ (only has to be done once) $$\mathsf{A} = \mathsf{M}(\eta) = \begin{pmatrix}\alpha &    -1 &    \\    -1 & \alpha & -1 & \\       &     -1 & \alpha & -1 \\       && \ddots & \ddots & \ddots & \\       &&                 &     -1 & \alpha\end{pmatrix}$$and the right-hand side vector $\mathbf{b}$$$\mathbf{b} = \begin{pmatrix}{\color{blue}T_{0,j+1}} + T_{0,j} + \beta T_{1,j} + T_{2,j}\\\vdots\\T_{i-1,j} + \beta T_{i,j} + T_{i+1,j}\\\vdots\\T_{-3,j} + \beta T_{-2,j} + T_{-1,j} + {\color{blue}T_{-1,j+1}}\end{pmatrix}$$In our problem, the boundaries of $\mathbf{b}$ do not change but it is easy to add time-varying boundary conditions as well by updating $\mathbf{b} = \mathbf{b}(t)$ in the loop. To construct $\mathsf{M}(\eta)$ we use `np.diagflat()`: we build matrices for the upper, central, and lower diagonal and add them:<jupyter_code>NM = 5
alpha = 99
np.diagflat(-np.ones(NM-1), -1) + alpha*np.diagflat(np.ones(NM), 0) + np.diagflat(-np.ones(NM-1), 1)<jupyter_output><empty_output><jupyter_text>The full _Crank-Nicholson_ algorithm with the standard `np.linalg.solve()` solver:<jupyter_code>import numpy as np

def CrankNicholson_T(L_rod=1, t_max=3000, Dx=0.02, Dt=2, T0=373, Tb=273,
                     step=20, verbose=True):
    Nx = int(L_rod // Dx)
    Nt = int(t_max // Dt)

    Kappa = 237 # W/(m K)
    CHeat = 900 # J/K
    rho = 2700  # kg/m^3

    eta = Kappa * Dt / (CHeat * rho * Dx**2)

    if verbose:
        print("Nx = {0}, Nt = {1}".format(Nx, Nt))
        print("eta = {0}".format(eta))

    T = np.zeros(Nx)
    T_plot = np.zeros((int(np.ceil(Nt/step)) + 1, Nx))

    # initial conditions
    T[1:-1] = T0
    # boundary conditions
    T[0] = T[-1] = Tb

    #---------------------
    # M_eta * T[1:-1, j+1] = bT
    # M_eta * xT = bT
    # Nx-2 x Nx-2 matrix: tridiagonal
    NM = Nx - 2
    alpha = 2/eta + 2
    beta = 2/eta - 2
    M_eta = np.diagflat(-np.ones(NM-1), 1) \
            + np.diagflat(alpha * np.ones(NM), 0) \
            + np.diagflat(-np.ones(NM-1), -1)
    bT = np.zeros(NM)

    t_index = 0
    T_plot[t_index, :] = T
    for jt in range(1, Nt):
        bT[:] = T[:-2] + beta*T[1:-1] + T[2:]
        # boundaries are special cases
        bT[0] += T[0]  #  + T[0,j+1]
        bT[-1] += T[-1] #  + T[-1,j+1]

        T[1:-1] = np.linalg.solve(M_eta, bT)

        if jt % step == 0 or jt == Nt-1:
            t_index += 1
            T_plot[t_index, :] = T
            if verbose:
                print("Iteration {0:5d}".format(jt), end="\r")
    else:
        if verbose:
            print("Completed {0:5d} iterations: t={1} s".format(jt, jt*Dt))
        
    parameters = (Dx, Dt, step)
    return T_plot, parameters
T_plot, (Dx, Dt, step) = CrankNicholson_T(t_max=3000, Dx=0.02, Dt=2)<jupyter_output>Nx = 49, Nt = 1500
eta = 0.4876543209876543
Completed  1499 iterations: t=2998 s<jupyter_text>Visualize<jupyter_code>def plot_T(T_plot, Dx, Dt, step):
    X, Y = np.meshgrid(range(T_plot.shape[0]), range(T_plot.shape[1]))
    Z = T_plot[X, Y]
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")
    ax.plot_wireframe(X*Dt*step, Y*Dx, Z)
    ax.set_xlabel(r"time $t$ (s)")
    ax.set_ylabel(r"position $x$ (m)")
    ax.set_zlabel(r"temperature $T$ (K)")
    fig.tight_layout()
    return ax
plot_T(T_plot, Dx, Dt, step)<jupyter_output><empty_output><jupyter_text>Note that this also works with $\eta > 1/2$: These were the settings at which leapfrog failed:<jupyter_code>T_plot, (Dx, Dt, step) = CrankNicholson_T(t_max=3000, Dx=0.01, Dt=2)
plot_T(T_plot, Dx, Dt, step)<jupyter_output>Nx = 99, Nt = 1500
eta = 1.9506172839506173
Completed  1499 iterations: t=2998 s<jupyter_text>We can also take very large time steps – depending on the problem, this can be a huge time saver, e.g., if we're only interested in long time scale behavior:<jupyter_code>T_plot, (Dx, Dt, step) = CrankNicholson_T(t_max=3000, Dx=0.02, Dt=20)
plot_T(T_plot, Dx, Dt, step)<jupyter_output>Nx = 49, Nt = 150
eta = 4.8765432098765435
Iteration    20
Iteration    40
Iteration    60
Iteration    80
Iteration   100
Iteration   120
Iteration   140
Iteration   149
Completed   149 iterations: t=2980 s<jupyter_text>Stability analysisVon-Neumann stability analysis yields:$$|\xi(k)| = \left| \frac{1 - 2\eta\sin^2\frac{k\Delta x}{2}}{1 + 2\eta\sin^2\frac{k\Delta x}{2}} \right|$$Because $\sin^2\alpha \le 1$ it follows that the amplification factors$$|\xi(k)| \le 1$$for *all* $k$. Therefore, the Crank-Nicholson scheme is *stable for all combinations of $\Delta x$  and $\Delta t$*.  Faster solutions   Using matrix inverseWe only need to calculate the matrix inverse of `M_eta` once and can then use$$\mathbf{x}_T = \mathsf{M}(\eta)^{-1} \mathbf{b}_T$$<jupyter_code>import numpy as np

def CrankNicholson_inverse_T(L_rod=1, t_max=3000, Dx=0.02, Dt=2, T0=373, Tb=273,
                     step=20, verbose=True):
    Nx = int(L_rod // Dx)
    Nt = int(t_max // Dt)

    Kappa = 237 # W/(m K)
    CHeat = 900 # J/K
    rho = 2700  # kg/m^3

    eta = Kappa * Dt / (CHeat * rho * Dx**2)
    
    if verbose:
        print("Nx = {0}, Nt = {1}".format(Nx, Nt))
        print("eta = {0}".format(eta))

    T = np.zeros(Nx)
    T_plot = np.zeros((int(np.ceil(Nt/step)) + 1, Nx))

    # initial conditions
    T[1:-1] = T0
    # boundary conditions
    T[0] = T[-1] = Tb

    #---------------------
    # M_eta * T[1:-1, j+1] = bT
    # M_eta * xT = bT
    # Nx-2 x Nx-2 matrix: tridiagonal
    NM = Nx - 2
    alpha = 2/eta + 2
    beta = 2/eta - 2
    M_eta = np.diagflat(-np.ones(NM-1), 1) \
            + np.diagflat(alpha * np.ones(NM), 0) \
            + np.diagflat(-np.ones(NM-1), -1)
    bT = np.zeros(NM)
    
    inv_M_eta = np.linalg.inv(M_eta)

    t_index = 0
    T_plot[t_index, :] = T
    for jt in range(1, Nt):
        bT[:] = T[:-2] + beta*T[1:-1] + T[2:]
        # boundaries are special cases
        bT[0] += T[0]  #  + T[0,j+1]
        bT[-1] += T[-1] #  + T[-1,j+1]

        T[1:-1] = np.dot(inv_M_eta, bT)

        if jt % step == 0 or jt == Nt-1:
            t_index += 1
            T_plot[t_index, :] = T
            if verbose:
                print("Iteration {0:5d}".format(jt), end="\r")
    else:
        if verbose:
            print("Completed {0:5d} iterations: t={1} s".format(jt, jt*Dt))
        
    parameters = (Dx, Dt, step)
    return T_plot, parameters
T_plot, (Dx, Dt, step) = CrankNicholson_inverse_T(t_max=3000, Dx=0.02, Dt=2)
plot_T(T_plot, Dx, Dt, step)<jupyter_output>Nx = 49, Nt = 1500
eta = 0.4876543209876543
Iteration    20
Iteration    40
Iteration    60
Iteration    80
Iteration   100
Iteration   120
Iteration   140
Iteration   160
Iteration   180
Iteration   200
Iteration   220
Iteration   240
Iteration   260
Iteration   280
Iteration   300
Iteration   320
Iteration   340
Iteration   360
Iteration   380
Iteration   400
Iteration   420
Iteration   440
Iteration   460
Iteration   480
Iteration   500
Iteration   520
Iteration   540
Iteration   560
Iteration   580
Iteration   600
Iteration   620
Iteration   640
Iteration   660
Iteration   680
Iteration   700
Iteration   720
Iteration   740
Iteration   760
Iteration   780
Iteration   800
Iteration   820
Iteration   840
Iteration   860
Iteration   880
Iteration   900
Iteration   920
Iteration   940
Iteration   960
Iteration   980
Iteration  1000
Iteration  1020
Iteration  1040
Iteration  1060
Iteration  1080
Iteration  1100
Iteration  1120
Iteration  1140
Iteration  1160
Iteration  1180
Iteration  1[...]<jupyter_text>Special routines for *banded* matricesTridiagonal matrices are a special (simple) case of *banded matrices*. **scipy** contains special, fast routines to solve matrix equations for banded matrices, namely [scipy.linalg.solve_banded()](http://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.solve_banded.htmlscipy.linalg.solve_banded). The only difficulty is to format the input in a form suitable for the function:<jupyter_code>import scipy.linalg

def solve_tridiagonal_banded(A, b):
    ab = extract_tridiag_ab(A)
    return scipy.linalg.solve_banded((1, 1), ab, b)

def extract_tridiag_ab(A):
    # extract diagonals and pad (as required for solve_banded())
    ud = np.insert(np.diag(A, 1), 0, 0)         # upper diagonal
    d = np.diag(A)                              # main diagonal
    ld = np.insert(np.diag(A, -1), len(d)-1, 0) # lower diagonal
    # matrix as required by solve_banded()
    ab = np.array([ud, d, ld])
    return ab<jupyter_output><empty_output><jupyter_text>Faster Crank-Nicholson:<jupyter_code>import numpy as np
import scipy.linalg

def CrankNicholson_banded_T(L_rod=1, t_max=3000, Dx=0.02, Dt=2, T0=373, Tb=273,
                     step=20, verbose=True):
    Nx = int(L_rod // Dx)
    Nt = int(t_max // Dt)

    Kappa = 237 # W/(m K)
    CHeat = 900 # J/K
    rho = 2700  # kg/m^3

    eta = Kappa * Dt / (CHeat * rho * Dx**2)
    
    if verbose:
        print("Nx = {0}, Nt = {1}".format(Nx, Nt))
        print("eta = {0}".format(eta))

    T = np.zeros(Nx)
    T_plot = np.zeros((int(np.ceil(Nt/step)) + 1, Nx))

    # initial conditions
    T[1:-1] = T0
    # boundary conditions
    T[0] = T[-1] = Tb

    #---------------------
    # M_eta * T[1:-1, j+1] = bT
    # M_eta * xT = bT
    # Nx-2 x Nx-2 matrix: tridiagonal
    NM = Nx - 2
    alpha = 2/eta + 2
    beta = 2/eta - 2
    M_eta = np.diagflat(-np.ones(NM-1), 1) \
            + np.diagflat(alpha * np.ones(NM), 0) \
            + np.diagflat(-np.ones(NM-1), -1)
    bT = np.zeros(NM)
    
    M_eta_ab = extract_tridiag_ab(M_eta)

    t_index = 0
    T_plot[t_index, :] = T
    for jt in range(1, Nt):
        bT[:] = T[:-2] + beta*T[1:-1] + T[2:]
        # boundaries are special cases
        bT[0] += T[0]  #  + T[0,j+1]
        bT[-1] += T[-1] #  + T[-1,j+1]

        T[1:-1] = scipy.linalg.solve_banded((1, 1), M_eta_ab, bT)

        if jt % step == 0 or jt == Nt-1:
            t_index += 1
            T_plot[t_index, :] = T
            if verbose:
                print("Iteration {0:5d}".format(jt), end="\r")
    else:
        if verbose:
            print("Completed {0:5d} iterations: t={1} s".format(jt, jt*Dt))
        
    parameters = (Dx, Dt, step)
    return T_plot, parameters
T_plot, (Dx, Dt, step) = CrankNicholson_banded_T(t_max=3000, Dx=0.02, Dt=2)
plot_T(T_plot, Dx, Dt, step)<jupyter_output>Nx = 49, Nt = 1500
eta = 0.4876543209876543
Iteration    20
Iteration    40
Iteration    60
Iteration    80
Iteration   100
Iteration   120
Iteration   140
Iteration   160
Iteration   180
Iteration   200
Iteration   220
Iteration   240
Iteration   260
Iteration   280
Iteration   300
Iteration   320
Iteration   340
Iteration   360
Iteration   380
Iteration   400
Iteration   420
Iteration   440
Iteration   460
Iteration   480
Iteration   500
Iteration   520
Iteration   540
Iteration   560
Iteration   580
Iteration   600
Iteration   620
Iteration   640
Iteration   660
Iteration   680
Iteration   700
Iteration   720
Iteration   740
Iteration   760
Iteration   780
Iteration   800
Iteration   820
Iteration   840
Iteration   860
Iteration   880
Iteration   900
Iteration   920
Iteration   940
Iteration   960
Iteration   980
Iteration  1000
Iteration  1020
Iteration  1040
Iteration  1060
Iteration  1080
Iteration  1100
Iteration  1120
Iteration  1140
Iteration  1160
Iteration  1180
Iteration  1[...]<jupyter_text>Benchmarking<jupyter_code>%timeit CrankNicholson_banded_T(t_max=3000, Dx=0.002, Dt=2, verbose=False)
%timeit CrankNicholson_T(t_max=3000, Dx=0.002, Dt=2, verbose=False)<jupyter_output>13.7 s ± 1.28 s per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>For the original problem, `np.linalg.solve()` is at least as fast as the banded solution, but for 10 times smaller step size (from 0.02 to 0.002) ie from 100 x 100 to 1000 x 1000 matrix, the slow-down is about<jupyter_code>13.7/0.159<jupyter_output><empty_output><jupyter_text>The inverse matrix is very competitive for this problem size but will run into limitations for larger sizes.<jupyter_code>%timeit CrankNicholson_inverse_T(t_max=3000, Dx=0.002, Dt=2, verbose=False)<jupyter_output>205 ms ± 15.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>Tridiagonal matrix algorithmThe usual way to solve the matrix problem is to use a special algorithm for tridiagonal matrices, the [Thomas algorithm](https://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm). This can be done in $\mathcal{O}(N)$ and thus is *as fast as the simple iterative scheme*! Implementation of the Thomas algorithm in Python is not difficult (see, for instance, cdhagman's answer [Stackoverflow: Optimize A*x = B solution for a tridiagonal coefficient matrix](http://stackoverflow.com/a/23133208)).<jupyter_code># left as an exercise<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13915693245302183, 'fraction_numerical': 0.06885293477472249, 'mean_word_length': 2.92818466799658, 'pattern_counts': {'":': 0, '<': 54, '<?xml version=': 0, '>': 55, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '15', 'dir': 'jupyter-structured-clean-dedup', 'id': '16920', 'license': 'CC-BY-4.0', 'n_tokens_mistral': 6014, 'n_tokens_neox': 5344, 'n_words': 1384, 'path': '16_PDEs/16_CrankNicholson.ipynb', 'repo_name': 'ASU-CompMethodsPhysics-PHY494/PHY494-resources-2020', 'size': '13783'}