# Crank-Nicolson method to solve the heat equation. from mpl_toolkits.mplot3d import axes3d import matplotlib.pyplot as pl import numpy as np import matplotlib.patches as mpatches  t0 = 0.0 t_final = 1.0 n_grid = 20 dt = t_final / n_grid dx = 1.0 / n_grid T = np.arange(t0, t_final+dt, dt) X = np.arange(0, 1+dx, dx) ax = pl.figure().add_subplot(111, projection='3d') ax.set_xlabel('t') ax.set_ylabel('x') ax.set_zlabel('U') ax.set_xlim([t0,t_final]) ax.set_ylim([0, 1]) Umax= .15 ax.set_zlim([0, Umax])  # Exact solution def exact(t, x):     return np.exp(-t)*np.sin(np.pi*x)/(np.pi*np.pi)      for t in T:     for x in X:         if t>t0:             ax.plot([t-dt, t],[x, x],[exact(t-dt, x), exact(t, x)], 'b-', linewidth=0.5)         if x>0.0:             ax.plot([t, t],[x-dx, x],[exact(t, x-dx), exact(t,x)], 'b-', linewidth=0.5)  # Crank-Nicolson method r = dx/(dt*dt)/(np.pi*np.pi) U = [[0.0 for _ in X] for _ in T] A = [[0.0 for _ in X] for _ in X] A_prime = [[0.0 for _ in X] for _ in X] b = [0.0 for _ in X] for i in range(len(X)):     if i in [0, len(X)-1]:         A[i][i] = 1.0         A_prime[i][i] = 1.0     else:         A[i][i] = 2+2*r         A[i][i-1] = -r         A[i][i+1] = -r         A_prime[i][i] = 2-2*r         A_prime[i][i-1] = r         A_prime[i][i+1] = r A = np.linalg.solve(A_prime, A) for i in range(len(T)):     if i==0:         for j in range(len(X)):             U[i][j] = exact(T[i], X[j])             b[j] = U[i][j]     else:         b = np.linalg.solve(A, b)         b[0], b[-1] = 0.0, 0.0         for j in range(len(X)):             U[i][j] = b[j]  for i in range(len(T)):     for j in range(len(X)):         if i>0:             ax.plot([T[i-1], T[i]],[X[j], X[j]],[U[i-1][j], U[i][j]], 'r-', linewidth=0.5)         if j>0:             ax.plot([T[i], T[i]],[X[j-1], X[j]],[U[i][j-1], U[i][j]], 'r-', linewidth=0.5)  bluePatch = mpatches.Patch(color='blue', label='Exact') redPatch = mpatches.Patch(color='red', label='Crank-Nicolson') pl.legend(handles = [bluePatch, redPatch], loc=2)  pl.show()