Coverage for pySDC/tutorial/step_9/B_paradiag_for_nonlinear_problems.py: 99%

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 B: ParaDiag for nonlinear problems 

12# 

13# ParaDiag works by diagonalizing the "top layer" of Kronecker products that make up the circularized composite 

14# collocation problem. For nonlinear problems, though, the composite collocation problem cannot be written as a 

15# matrix, so the diagonalization needs a linear operator to work with. There are two ways out: 

16# 

17# - **IMEX splitting**, where only the linear part is treated implicitly. The ParaDiag preconditioner is then made up 

18# of the linear implicit part only, which can be diagonalized just like for linear problems, and the nonlinear part 

19# comes in through the residual on the right-hand side. 

20# - **Averaging the Jacobian**. Nonlinear problems are solved with Newton, where the Jacobian is computed from the 

21# current solution and inverted in each Newton iteration. To write the ParaDiag preconditioner with Kronecker 

22# products and only diagonalize the outermost part, all steps need the same Jacobian: the average. 

23# 

24# As IMEX ParaDiag is a trivial extension of ParaDiag for linear problems, we focus on the second approach here, with 

25# the van der Pol oscillator as an example. The ParaDiag iteration then proceeds as follows: 

26# 

27# 1. compute the residual of the composite collocation problem, 

28# 2. average the solution across the steps and nodes, to compute the average Jacobian, 

29# 3. weighted FFT in time to diagonalize $E_\alpha$, 

30# 4. solve for the increment on the subproblems on the steps and nodes, inverting the averaged Jacobian from (2), 

31# 5. weighted inverse FFT in time, 

32# 6. increment the solution. 

33 

34# %% 

35import numpy as np 

36import scipy.sparse as sp 

37 

38from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit as sweeper_class 

39from pySDC.implementations.problem_classes.Van_der_Pol_implicit import vanderpol 

40 

41# setup parameters 

42L = 4 

43M = 3 

44alpha = 1e-4 

45restol = 1e-8 

46dt = 0.1 

47 

48# setup infrastructure 

49prob = vanderpol(newton_maxiter=1, mu=1e0, crash_at_maxiter=False) 

50N = prob.init[0] 

51 

52# make problem work on complex data 

53prob.init = tuple([*prob.init[:2]] + [np.dtype('complex128')]) 

54 

55# setup global solution array 

56u = np.zeros((L, M, N), dtype=complex) 

57 

58# setup collocation problem 

59sweep = sweeper_class({'num_nodes': M, 'quad_type': 'RADAU-RIGHT'}, None) 

60 

61# initial conditions 

62u[0, :, :] = prob.u_exact(t=0) 

63 

64print(f'Running ParaDiag test script for van der Pol with mu={prob.mu} and {L} time steps and {M} collocation nodes.') 

65 

66# %% [markdown] 

67# ## The matrices 

68# 

69# Those that make up the composite collocation problem, as in [Part A](A_paradiag_for_linear_problems), although we 

70# do not set up the full problem here. See [the paper](https://arxiv.org/abs/2103.12571) for their meaning. We 

71# diagonalize $Q G^{-1}$ on every step right away. 

72 

73# %% 

74I_M = sp.eye(M) 

75 

76H_M = sp.eye(M).tolil() * 0 

77H_M[:, -1] = 1 

78 

79Q = sweep.coll.Qmat[1:, 1:] 

80 

81E_alpha = sp.diags( 

82 [ 

83 -1.0, 

84 ] 

85 * (L - 1), 

86 offsets=-1, 

87).tolil() 

88E_alpha[0, -1] = -alpha 

89 

90gamma = alpha ** (-np.arange(L) / L) 

91D_alpha_diag_vals = np.fft.fft(1 / gamma * E_alpha[:, 0].toarray().flatten(), norm='backward') 

92 

93J = sp.diags(gamma) 

94J_inv = sp.diags(1 / gamma) 

95 

96G = [(D_alpha_diag_vals[l] * H_M + I_M).tocsc() for l in range(L)] # MxM 

97 

98# prepare diagonalization of QG^{-1} 

99w = [] 

100S = [] 

101S_inv = [] 

102 

103for l in range(L): 

104 # diagonalize QG^-1 matrix 

105 if M > 1: 

106 _w, _S = np.linalg.eig(Q @ sp.linalg.inv(G[l]).toarray()) 

107 else: 

108 _w, _S = np.linalg.eig(Q / (G[l].toarray())) 

109 _S_inv = np.linalg.inv(_S) 

110 w.append(_w) 

111 S.append(_S) 

112 S_inv.append(_S_inv) 

113 

114# %% [markdown] 

115# The matrix-vector product on the steps and the residual of the composite collocation problem, as in Part A: 

116 

117 

118# %% 

119def mat_vec(mat, vec): 

120 """ 

121 Matrix vector product 

122 

123 Args: 

124 mat (np.ndarray or scipy.sparse) : Matrix 

125 vec (np.ndarray) : vector 

126 

127 Returns: 

128 np.ndarray: mat @ vec 

129 """ 

130 res = np.zeros_like(vec) 

131 for l in range(vec.shape[0]): 

132 for k in range(vec.shape[0]): 

133 res[l] += mat[l, k] * vec[k] 

134 return res 

135 

136 

137def residual(_u, u0): 

138 """ 

139 Compute the residual of the composite collocation problem 

140 

141 Args: 

142 _u (np.ndarray): Current iterate 

143 u0 (np.ndarray): Initial conditions 

144 

145 Returns: 

146 np.ndarray: LMN size array with the residual 

147 """ 

148 res = _u * 0j 

149 for l in range(L): 

150 # build step local residual 

151 

152 # communicate initial conditions for each step 

153 if l == 0: 

154 res[l, ...] = u0[l, ...] 

155 else: 

156 res[l, ...] = _u[l - 1, -1, ...] 

157 

158 # evaluate and subtract integral over right hand side functions 

159 f_evals = np.array([prob.eval_f(_u[l, m], 0) for m in range(M)]) 

160 Qf = mat_vec(Q, f_evals) 

161 res[l, ...] -= _u[l] - dt * Qf 

162 

163 return res 

164 

165 

166# %% [markdown] 

167# ## The iteration 

168# 

169# The six steps from above. Each ParaDiag iteration does a single Newton iteration on every node, one call of 

170# `solve_jacobian` per node, so the number of Newton iterations per node equals the number of ParaDiag iterations. 

171 

172# %% 

173sol_paradiag = u.copy() * 0j 

174u0 = u.copy() 

175niter = 0 

176res = residual(sol_paradiag, u0) 

177while np.max(np.abs(res)) > restol: 

178 # compute all-at-once residual 

179 res = residual(sol_paradiag, u0) 

180 

181 # compute solution averaged across the L steps and M nodes. This is the difference to ParaDiag for linear problems. 

182 u_avg = prob.u_init 

183 u_avg[:] = np.mean(sol_paradiag, axis=(0, 1)) 

184 

185 # weighted FFT in time 

186 x = np.fft.fft(mat_vec(J_inv.toarray(), res), axis=0) 

187 

188 # perform local solves of "collocation problems" on the steps in parallel 

189 y = np.empty_like(x) 

190 for l in range(L): 

191 

192 # perform local solves on the collocation nodes in parallel 

193 x1 = S_inv[l] @ x[l] 

194 x2 = np.empty_like(x1) 

195 for m in range(M): 

196 x2[m, :] = prob.solve_jacobian(x1[m], w[l][m] * dt, u=u_avg, t=l * dt) 

197 z = S[l] @ x2 

198 y[l, ...] = sp.linalg.spsolve(G[l], z) 

199 

200 # inverse FFT in time and increment 

201 sol_paradiag += mat_vec(J.toarray(), np.fft.ifft(y, axis=0)) 

202 

203 res = residual(sol_paradiag, u0) 

204 niter += 1 

205 assert niter < 99, 'ParaDiag did not converge for nonlinear problem!' 

206print(f'Needed {niter} ParaDiag iterations, stopped at residual {np.max(np.abs(res)):.2e}') 

207 

208# %% [markdown] 

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

210# - Averaging the Jacobian requires communicating the average solution, which is why `average_jacobian` is off by 

211# default for linear problems. 

212# - We do a single Newton iteration per ParaDiag iteration, so the number of Newton iterations per node equals the 

213# number of ParaDiag iterations. 

214# ::: 

215# 

216# The check the tests run is the one inside the loop: ParaDiag has to converge in fewer than 99 iterations.