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()