CrankNicolson.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import solve_banded

C = lambda x,t: (np.pi * np.cos(2 * np.pi * t) + 3.5) / (2 * np.pi - np.pi * np.sin(np.pi * x))
U = lambda x,t: np.cos(np.pi * x) - np.sin(2 * np.pi * t) / 2 + 2 * np.pi * x - 3.5 * t

phi = lambda x: U(x, 0)
psi_0 = lambda t: U(0, t)
psi_1 = lambda t: U(1, t)

def solve(x0, x1, t0, t1, h, tau):
  N = int((x1 - x0) / h + 0.01)
  M = int((t1 - t0) / tau + 0.01)

  h = (x1 - x0) / N
  tau = (t1 - t0) / M

  u = np.zeros((N + 1, M + 1))

  for i in range(M + 1):
    u[0, i] = psi_0(i * tau)
    u[N, i] = psi_1(i * tau)
  for i in range(N + 1):
    u[i, 0] = phi(i * h)
   
  for j in range(M):
    A = np.zeros((3, N - 1))
    
    for k in range (0, N - 1):
      a_h = tau * C(k * h, (j + 1) * tau) / (4 * h)
      
      A[0, k] = a_h
      A[1, k] = 1
      A[2, k] = -a_h

    B = np.zeros(N - 1)
    B[0] = tau * C(0, j * tau) / (4 * h) * u[0, j]
    B[N - 2] = tau * C(N * h, j * tau) / (4 * h) * u[N, j]

    for k in range(1, N):
      B[k - 1] = u[k, j] - tau * C(k * h, j * tau) * (u[k + 1, j] - u[k - 1, j]) / (4 * h)
    F = solve_banded((1, 1), A, B)
    for i in range(1, N):
      u[i, j + 1] = F[i - 1]
      
  return u

def plot(u, x0, x1, h, tau, x_plot): 
    print("x =", x_plot)

    x = np.linspace(x0, x1, u.shape[1])
    y_exact = [U(x_plot, t) for t in x]
    y_approx = [u[int(x_plot / h)][j] for j in range(x.shape[0])]

    print("x step:", h)
    print("t step:", tau)
    print("approximation error:", max(abs(np.array(y_approx) - np.array(y_exact))))
 
    plt.plot(x, y_exact, label="exact solution")
    plt.plot(x, y_approx, 'bo', label="numerical solution")
    plt.xlabel('time')
    plt.ylabel('value')
    plt.legend()
    plt.show()

h, tau = 0.1, 0.1
u = solve(0, 1, 0, 1, h, tau)
plot(u, 0, 1, h, tau, 0.5)