69 lines
3.1 KiB
Python
69 lines
3.1 KiB
Python
|
import numpy as np
|
||
|
|
from scipy.sparse import spdiags
|
||
|
|
from scipy.linalg import solve_banded
|
||
|
|
from matplotlib import pyplot as plt
|
||
|
|
|
||
|
|
# Config
|
||
|
|
dx = 0.01 # space step size
|
||
|
|
h = 1.5 # space domain (max x size)
|
||
|
|
|
||
|
|
# plots a solution to file, simply for code size reduction in `solve`
|
||
|
|
def plot(x, t, C, C0, analytic, plotname):
|
||
|
|
plt.clf()
|
||
|
|
plt.plot(x, C, "o")
|
||
|
|
plt.xlim([0, h])
|
||
|
|
plt.title(plotname)
|
||
|
|
plt.xlabel("Periodic Space: x")
|
||
|
|
plt.ylabel("Wave Form: C(x, t)")
|
||
|
|
if analytic is not None:
|
||
|
|
plt.plot(x, analytic(x, t, C0))
|
||
|
|
plt.savefig(f"plots/{plotname}_{plot.i:0>5}.png")
|
||
|
|
plot.i += 1
|
||
|
|
plot.i = 0
|
||
|
|
|
||
|
|
# solver for `C_t + U C_x = 0` given initial state `C`, the current `C0` which
|
||
|
|
# incorporates `U` as well as the space and time discretization and the number
|
||
|
|
# of time steps `nt`.
|
||
|
|
def solve(C, C0, nt, plotname = None, plotskip = 1, analytic = None):
|
||
|
|
nx = C.shape[0]
|
||
|
|
if plotname is not None:
|
||
|
|
plot.i = 0 # (re)set plot counter
|
||
|
|
x = np.linspace(0, h, nx) # space grid points
|
||
|
|
fig = plt.figure(figsize = (8, 6), dpi = 100) # setup plot figure size
|
||
|
|
plot(x, 0, C, C0, analytic, plotname) # plot initial solution
|
||
|
|
# space grid point `Indices Minus 1`
|
||
|
|
im1 = np.array([nx - 1] + list(range(nx - 1)))
|
||
|
|
for t in range(nt):
|
||
|
|
# upwind update scheme
|
||
|
|
C = (1 - C0) * C + C0 * C[im1]
|
||
|
|
if plotname is not None and ((t + 1) % plotskip == 0):
|
||
|
|
plot(x, t + 1, C, C0, analytic, plotname) # plot current solution
|
||
|
|
return C
|
||
|
|
|
||
|
|
# setup initial conditions
|
||
|
|
x = np.linspace(0, h, int(h / dx))
|
||
|
|
gauss = np.exp(-10 * (4 * x - 1)**2)
|
||
|
|
square = np.array((0.1 < x) * (x < 0.3), dtype = "float")
|
||
|
|
|
||
|
|
def gauss_analytic(x, t, C0):
|
||
|
|
return np.exp(-10 * (4 * (np.mod(x - t * dx / C0, h + 1e-9)) - 1)**2)
|
||
|
|
def square_analytic(x, t, C0):
|
||
|
|
x = x - t * dx / C0
|
||
|
|
x = np.mod(x, h + 1e-9) # hack to avoid shift in the plot
|
||
|
|
return np.array((0.1 < x) * (x < 0.3), dtype = "float")
|
||
|
|
|
||
|
|
solve(gauss, 1, 200, "task02_gauss_C0_1", analytic = gauss_analytic)
|
||
|
|
solve(square, 1, 200, "task02_square_C0_1", analytic = square_analytic)
|
||
|
|
solve(gauss, 0.7, 200, "task02_gauss_C0_0.7", analytic = gauss_analytic)
|
||
|
|
solve(square, 0.7, 200, "task02_square_C0_0.7", analytic = square_analytic)
|
||
|
|
solve(gauss, 1.1, 200, "task02_gauss_C0_1.1", analytic = gauss_analytic)
|
||
|
|
solve(square, 1.1, 200, "task02_square_C0_1.1", analytic = square_analytic)
|
||
|
|
|
||
|
|
# To generate videos out of the generated plots
|
||
|
|
# $> ffmpeg -y -r 60 -i plots/task02_gauss_C0_1_%05d.png -pix_fmt yuv420p task02_gauss_C0_1.mp4
|
||
|
|
# $> ffmpeg -y -r 60 -i plots/task02_square_C0_1_%05d.png -pix_fmt yuv420p task02_square_C0_1.mp4
|
||
|
|
# $> ffmpeg -y -r 60 -i plots/task02_gauss_C0_0.7_%05d.png -pix_fmt yuv420p task02_gauss_C0_0.7.mp4
|
||
|
|
# $> ffmpeg -y -r 60 -i plots/task02_square_C0_0.7_%05d.png -pix_fmt yuv420p task02_square_C0_0.7.mp4
|
||
|
|
# $> ffmpeg -y -r 60 -i plots/task02_gauss_C0_1.1_%05d.png -pix_fmt yuv420p task02_gauss_C0_1.1.mp4
|
||
|
|
# $> ffmpeg -y -r 60 -i plots/task02_square_C0_1.1_%05d.png -pix_fmt yuv420p task02_square_C0_1.1.mp4
|