Coverage for pySDC/tutorial/step_1/C_collocation_problem_setup.py: 100%

1# --- 

2# jupyter: 

3# jupytext: 

4# formats: py:percent 

5# kernelspec: 

6# display_name: Python 3 

7# name: python3 

8# --- 

9 

10# %% [markdown] 

11# # Part C: Collocation problem setup 

12# 

13# Now time enters. Written in integral form, the heat equation from [Part A](A_spatial_problem_setup) reads 

14# 

15# $$ 

16# u(t) = u_0 + \int_{t_0}^{t} A\, u(s)\, ds , 

17# $$ 

18# 

19# with $A$ the finite-difference Laplacian (times $\nu$). **Collocation** replaces the integral by a quadrature 

20# rule on $M$ nodes $t_0 + \tau_m \Delta t$ inside one time step, and demands that the equation holds exactly at 

21# these nodes. Whatever the equation, this collocation problem of one time step is what the methods in pySDC solve; 

22# the heat equation is just the example here. 

23 

24# %% 

25import matplotlib.pyplot as plt 

26import numpy as np 

27import scipy.sparse as sp 

28 

29from pySDC.core.collocation import CollBase 

30from pySDC.implementations.problem_classes.HeatEquation_ND_FD import heatNd_unforced 

31 

32prob = heatNd_unforced( 

33 nvars=1023, # number of degrees of freedom 

34 nu=0.1, # diffusion coefficient 

35 freq=4, # frequency for the test value 

36 bc='dirichlet-zero', # boundary conditions 

37) 

38 

39# %% [markdown] 

40# ## Nodes and the collocation matrix 

41# 

42# `CollBase` provides the quadrature: here three Gauss-Radau nodes, whose last node is the right end of the 

43# interval. `node_type` picks the family of nodes (Legendre here), `quad_type` which ends of the interval are nodes 

44# themselves: none (`'GAUSS'`), one (`'RADAU-LEFT'`, `'RADAU-RIGHT'`) or both (`'LOBATTO'`). 

45 

46# %% 

47# instantiate collocation class, relative to the time interval [0,1] 

48coll = CollBase(num_nodes=3, tleft=0, tright=1, node_type='LEGENDRE', quad_type='RADAU-RIGHT') 

49 

50with np.printoptions(precision=4, suppress=True): 

51 print('nodes tau_m:', coll.nodes) 

52 print('Q =') 

53 print(coll.Qmat[1:, 1:]) 

54 

55# %% [markdown] 

56# Row $m$ of the matrix $Q = (q_{mj})$ holds the quadrature weights for the integral from $0$ to $\tau_m$. Its 

57# first row and column in `coll.Qmat` belong to the initial value and are dropped. Both are relative to $[0, 1]$: 

58# multiply by $\Delta t$ for a real time step. The collocation problem for the values $u_m \approx u(t_0 + \tau_m 

59# \Delta t)$ is then 

60# 

61# $$ 

62# u_m = u_0 + \Delta t \sum_{j=1}^{M} q_{mj} A u_j 

63# \quad\Longleftrightarrow\quad 

64# (I - \Delta t\, Q \otimes A)\, \vec u = \vec u_0 . 

65# $$ 

66# 

67# ## Solving it directly 

68# 

69# The problem is linear, so we can assemble the Kronecker product and hand it to a sparse direct solver. Since the 

70# last Radau node is the end of the step, the last block of the solution is $u(\Delta t)$, which we compare with 

71# the exact solution. 

72 

73 

74# %% 

75def solve_collocation_problem(prob, coll, dt): 

76 """ 

77 Routine to build and solve the linear collocation problem 

78 

79 Args: 

80 prob: a problem instance 

81 coll: a collocation instance 

82 dt: time-step size 

83 

84 Return: 

85 the analytic error of the solved collocation problem 

86 """ 

87 

88 # shrink collocation matrix: first line and column deals with initial value, not needed here 

89 Q = coll.Qmat[1:, 1:] 

90 

91 # build system matrix M of collocation problem 

92 M = sp.eye(prob.nvars[0] * coll.num_nodes) - dt * sp.kron(Q, prob.A) 

93 

94 # get initial value at t0 = 0 

95 u0 = prob.u_exact(t=0) 

96 # fill in u0-vector as right-hand side for the collocation problem 

97 u0_coll = np.kron(np.ones(coll.num_nodes), u0) 

98 # get exact solution at Tend = dt 

99 uend = prob.u_exact(t=dt) 

100 

101 # solve collocation problem directly 

102 u_coll = sp.linalg.spsolve(M, u0_coll) 

103 

104 # compute error 

105 err = np.linalg.norm(u_coll[-prob.nvars[0] :] - uend, np.inf) 

106 

107 return err 

108 

109 

110# set time-step size (warning: the collocation matrices are relative to [0,1], see above) 

111dt = 0.1 

112err = solve_collocation_problem(prob=prob, coll=coll, dt=dt) 

113print(f'Error of the collocation problem: {err:8.6e}') 

114 

115# %% [markdown] 

116# ## Why pySDC does not do it like this 

117# 

118# The system couples all nodes: every block of $Q \otimes A$ is a scaled copy of $A$. For a small mesh the 

119# structure is easy to see: 

120 

121# %% tags=["hide-input"] 

122small = heatNd_unforced(nvars=15, nu=0.1, freq=4, bc='dirichlet-zero') 

123M_small = sp.eye(15 * coll.num_nodes) - dt * sp.kron(coll.Qmat[1:, 1:], small.A) 

124fig, ax = plt.subplots(figsize=(3.5, 3.5)) 

125ax.spy(M_small, markersize=3) 

126ax.set_title(r'$I - \Delta t\, Q \otimes A$, 15 unknowns, $M = 3$', fontsize=10) 

127fig.tight_layout() 

128 

129# %% [markdown] 

130# That is $M$ times more unknowns than a single implicit Euler step, fully coupled. Fine for a 1D heat equation, 

131# hopeless for a 3D problem with a nonlinear right-hand side. The rest of pySDC therefore never assembles this 

132# matrix: it keeps space and time separate and solves the collocation problem iteratively with SDC, starting in 

133# [Step 2](../step_2). 

134# 

135# :::{admonition} Important things to note 

136# - The collocation matrix `Q` is and will always be relative to $[0, 1]$. Use `dt` to scale it. 

137# - `u_exact` gives the solution at any time. Give your own problem classes either an exact solution or an 

138# initial condition routine. 

139# - This is where the fun with parameters starts: how many unknowns in space, how large a `dt`, how high a 

140# frequency? Try some and watch the error. 

141# ::: 

142# 

143# The check the tests run: 

144 

145# %% 

146assert err <= 4e-04, f"ERROR: did not get collocation error as expected, got {err}"