75 lines
2.2 KiB
Python
75 lines
2.2 KiB
Python
# Task 1.4
|
|
import numpy as np
|
|
import scipy, scipy.sparse, scipy.linalg
|
|
from matplotlib import pyplot as plt
|
|
|
|
# plotting helper function
|
|
def plot(C, x, t, name):
|
|
plt.clf()
|
|
plt.plot(x, C)
|
|
plt.xlim([0, x[-1]])
|
|
plt.ylim([0, 1.2])
|
|
plt.title(f"{name} - time: {t: >12.0f}")
|
|
plt.savefig(f"plots/{name}_{plot.fignr:0>5}.png")
|
|
plot.fignr += 1
|
|
|
|
# set static variable as function attribute
|
|
plot.fignr = 1
|
|
# and initialize a figure
|
|
plt.figure(figsize = (8, 6), dpi = 100)
|
|
|
|
# Config
|
|
D = 1e-6 # diffusion coefficient
|
|
h = 1 # space domain (max x size)
|
|
T = 1e5 # solution end time
|
|
nx = 50 # nr of space discretization points
|
|
nt = 2000 # nr of time points
|
|
|
|
# derived constants
|
|
dx = h / (nx - 1) # space step size
|
|
dt = T / (nt - 1) # time step size
|
|
d = dt * D / dx**2 # stability/stepsize coefficient
|
|
|
|
# setup matrices for implicit update scheme `A C^n+1 = (2 I - A) C^n = B C^n`
|
|
# The matrix `A` is represented as a `3 x (nx - 1)` matrix as expected by
|
|
# `scipy.linalg.solve_banded` representing the 3 diagonals of the tridiagonal
|
|
# matrix `A`.
|
|
A = np.zeros((3, nx - 1))
|
|
# upper minor diagonal
|
|
A[0, 2:] = -d
|
|
# main diagonal
|
|
A[1, 0] = 1
|
|
A[1, 1:] = 1 + 2 * d
|
|
# and lower minor diagonal
|
|
A[2, :-2] = -d
|
|
A[2, -2] = -2 * d
|
|
# The right hand side matrix `B = 2 I - A` which is represented as a sparse
|
|
# matrix in diag. layout which allows for fast matrix vector multiplication
|
|
B = -A
|
|
B[1, ] += 2
|
|
B = scipy.sparse.spdiags(B, (1, 0, -1), B.shape[1], B.shape[1])
|
|
|
|
# set descritized space evaluation points `x` of `C`
|
|
x = np.linspace(0, h, nx)
|
|
|
|
# Set initial solution
|
|
C = np.zeros(nx)
|
|
C[0] = 1
|
|
C[-1] = C[-3] # (0 = 0)
|
|
|
|
for n in range(nt):
|
|
# generate roughly 400 plots
|
|
if n % (nt // 400) == 0:
|
|
plot(C, x, n * dt, "task01_4")
|
|
# update solution using the implicit schema
|
|
C[:-1] = scipy.linalg.solve_banded((1, 1), A, B @ C[:-1])
|
|
# set/enforce boundary conditions
|
|
C[0] = 1 # left Dirichlet (theoretically not needed)
|
|
C[-1] = C[-3] # right Neumann condition
|
|
|
|
# plot final state
|
|
plot(C, x, T, "task01_4")
|
|
|
|
# to convert generated image sequence to video use:
|
|
# $> ffmpeg -r 60 -i plots/task01_4_%05d.png -pix_fmt yuv420p video_1_4.mp4
|