import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
%matplotlib notebook
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
#---------------------
# set up M_eta
raise NotImplementedError
t_index = 0
T_plot[t_index, :] = T
for jt in range(1, Nt):
# solve M_eta * T(j+1) = bT
raise NotImplementedError
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)
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)
Try different $\Delta x$ and $\Delta t$.
We only need to calculate the matrix inverse of M_eta
once and can then use
T_plot, (Dx, Dt, step) = CrankNicholson_inverse_T(t_max=3000, Dx=0.02, Dt=2)
plot_T(T_plot, Dx, Dt, step)
The usual way to solve the matrix problem is to use a special algorithm for tridiagonal matrices, the Thomas 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).
Tridiagonal 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(). The only difficulty is to format the input in a form suitable for the function:
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
Faster Crank-Nicholson using banded matrices:
T_plot, (Dx, Dt, step) = CrankNicholson_banded_T(t_max=3000, Dx=0.02, Dt=2)
plot_T(T_plot, Dx, Dt, step)
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 32/0.25 ~ 120.
%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)
%timeit CrankNicholson_inverse_T(t_max=3000, Dx=0.002, Dt=2, verbose=False)