def start_imex(N, n, h, v0, alpha, q):
"""Start an IMEX scheme with (q-1) steps of LIRK4."""
# Construct Laplacian matrix (multiplied by Tsin2 and alpha):
L = alpha * laplacian(n, 0)
# Compute LU factorizations of LIRK4 matrices:
I = eye(n)
Tsin2 = multmat(n, lambda x: np.sin(x)**2, [-pi, pi])
Tsin2 = csc_matrix(kron(I, Tsin2))
lu = splu(Tsin2)
lua = splu(Tsin2 - 1 / 4 * h * L)
# Time-stepping loop:
U = [v0]
NU = [N(v0)]
v = v0
for i in range(1, q):
Nv = N(v)
w = Tsin2 @ v
wa = w + h * Tsin2 @ (1 / 4 * Nv)
a = lua.solve(wa)
Na = N(a)
wb = w + h * L @ (1 / 2 * a)
wb = wb + h * Tsin2 @ (-1 / 4 * Nv + Na)
b = lua.solve(wb)
Nb = N(b)
wc = w + h * L @ (17 / 50 * a - 1 / 25 * b)
wc = wc + h * Tsin2 @ (-13 / 100 * Nv + 43 / 75 * Na + 8 / 75 * Nb)
c = lua.solve(wc)
Nc = N(c)
wd = w + h * L @ (371 / 1360 * a - 137 / 2720 * b + 15 / 544 * c)
wd = wd + h * Tsin2 @ (-6 / 85 * Nv + 42 / 85 * Na + 179 / 1360 * Nb -
15 / 272 * Nc)
d = lua.solve(wd)
Nd = N(d)
we = w + h * L @ (25 / 24 * a - 49 / 48 * b + 125 / 16 * c -
85 / 12 * d)
we = we + h * Tsin2 * (79 / 24 * Na - 5 / 8 * Nb + 25 / 2 * Nc -
85 / 6 * Nd)
e = lua.solve(we)
Ne = N(e)
v = v + h * lu.solve(L @ (25 / 24 * a - 49 / 48 * b + 125 / 16 * c -
85 / 12 * d + 1 / 4 * e))
v = v + h * (25 / 24 * Na - 49 / 48 * Nb + 125 / 16 * Nc -
85 / 12 * Nd + 1 / 4 * Ne)
U = [v] + U
NU = [N(v)] + NU
return U, NU
from chebpy.trig import diffmat, multmat
# %% Solve u''(x) + sin(x)*u'(x) + 1000*cos(2x)*u(x) = f on [0,2*pi].
# Grid:
n = 10000
x = trigpts(n, [0, 2 * pi])
# Right-hand side f:
f = lambda x: 100 + 0 * x
# Assemble matrices:
start = time.time()
D1 = diffmat(n, 1, [0, 2 * pi])
D2 = diffmat(n, 2, [0, 2 * pi])
M0 = multmat(n, lambda x: 1000 * np.cos(2 * x), [0, 2 * pi])
M1 = multmat(n, lambda x: np.sin(x), [0, 2 * pi])
L = D2 + M1 @ D1 + M0
plt.figure()
plt.spy(L)
# Assemble RHS:
F = vals2coeffs(f(x))
end = time.time()
print(f'Time (setup): {end-start:.5f}s')
# Sparse solve:
start = time.time()
U = spsolve(L, F)
end = time.time()
print(f'Time (solve): {end-start:.5f}s')
# Functions:
f = lambda x: np.cos(pi * x)
h = lambda x: np.exp(np.sin(10 * pi * x))
gex = lambda x: f(x) * h(x)
# Grid:
n = 10000
x = trigpts(n)
# Compute coeffs of f:
F = vals2coeffs(f(x))
# Multiplication by h matrix in coefficient space:
start = time.time()
M = multmat(n, h)
end = time.time()
print(f'Time (setup): {end-start:.5f}s')
plt.figure()
plt.spy(M)
# Multiply:
start = time.time()
G = M @ F
end = time.time()
print(f'Time (product): {end-start:.5f}s')
# Convert to value space:
g = coeffs2vals(G)
# Error:
# Plot initial condition:
plt.figure()
plt.contourf(LAM, TT, feval(u0, LAM, TT), 40, cmap=cm.coolwarm)
plt.colorbar()
# Assemble initial condition:
U0 = vals2coeffs(feval(u0, LAM, TT))
U0 = np.reshape(U0.T, (n*n, 1))
# Assemble Laplacian:
start = time.time()
L = alpha*laplacian(n)
# LU factorization of BDF4 matrix:
I = eye(n)
Tsin2 = multmat(n, lambda x: np.sin(x)**2, [-pi, pi])
Tsin2 = csc_matrix(kron(I, Tsin2))
lu = splu(25*Tsin2 - 12*dt*L)
end = time.time()
print(f'Time (setup): {end-start:.5f}s')
# Start initial condition with LIRK4:
U, NU = start_imex(N, n, dt, U0, alpha, q)
# Time-stepping:
start = time.time()
itrmax = round(T/dt)
for itr in range(q, itrmax + 1):
# Compute new solution:
b = 48*U[0] - 36*U[1] + 16*U[2] - 3*U[3]
# Zero-mean condition:
F[n * int(n / 2) + int(n / 2)] = 0
# Assemble Laplacian (with mean condition):
start = time.time()
L = laplacian(n, 1)
end = time.time()
print(f'Time (setup): {end-start:.5f}s')
plt.figure()
plt.spy(L)
# Sparse solve:
start = time.time()
I = eye(n)
Tsin2 = multmat(n, lambda x: np.sin(x)**2, dom)
U = spsolve(L, kron(I, Tsin2) @ F)
end = time.time()
print(f'Time (solve): {end-start:.5f}s')
# Plot solution:
U = np.reshape(U, (n, n)).T
u = coeffs2vals(U)
plt.figure()
plt.contourf(LAM, TT, u, 40, cmap=cm.coolwarm)
plt.colorbar()
# Plot exact solution:
plt.figure()
plt.contourf(LAM, TT, feval(uex, LAM, TT), 40, cmap=cm.coolwarm)
plt.colorbar()