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

1""" 

2This script introduces ParaDiag for nonlinear problems with the van der Pol oscillator as an example. 

3 

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

5collocation problem. 

6However, in nonlinear problems, the problem cannot be written as a matrix and therefore we cannot write the composite 

7collocation problem as a matrix. 

8There are two approaches for dealing with this. We can do IMEX splitting, where we treat only the linear part implicitly. 

9The ParaDiag preconditioner is then only made up of the linear implicit part and we can again write this as a matrix and 

10do the diagonalization just like for linear problems. The non-linear part then comes in via the residual on the right 

11hand side. 

12The second approach is to average Jacobians. The non-linear problems are solved with a Newton scheme, where the Jacobian 

13matrix is computed based on the current solution and then inverted in each Newton iteration. In order to write the 

14ParaDiag preconditioner as a matrix with Kronecker products and then only diagonalize the outermost part, we need to 

15have the same Jacobian on all steps. 

16The ParaDiag iteration then proceeds as follows: 

17 - (1) Compute residual of composite collocation problem 

18 - (2) Average the solution across the steps and nodes as preparation for computing the average Jacobian 

19 - (3) Weighted FFT in time to diagonalize E_alpha 

20 - (4) Solve for the increment by inverting the averaged Jacobian from (2) on the subproblems on the different steps 

21 and nodes. 

22 - (5) Weighted iFFT in time 

23 - (6) Increment solution 

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

25""" 

26 

27import numpy as np 

28import scipy.sparse as sp 

29import sys 

30 

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

32from pySDC.implementations.problem_classes.Van_der_Pol_implicit import vanderpol 

33 

34# setup output 

35out_file = open('data/step_9_B_out.txt', 'w') 

36 

37 

38def my_print(*args, **kwargs): 

39 for output in [sys.stdout, out_file]: 

40 print(*args, **kwargs, file=output) 

41 

42 

43# setup parameters 

44L = 4 

45M = 3 

46alpha = 1e-4 

47restol = 1e-8 

48dt = 0.1 

49 

50# setup infrastructure 

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

52N = prob.init[0] 

53 

54# make problem work on complex data 

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

56 

57# setup global solution array 

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

59 

60# setup collocation problem 

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

62 

63# initial conditions 

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

65 

66my_print( 

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

68) 

69 

70 

71""" 

72Setup matrices that make up the composite collocation problem. We do not set up the full composite collocation problem 

73here, however. See https://arxiv.org/abs/2103.12571 for the meaning of the matrices. 

74""" 

75I_M = sp.eye(M) 

76 

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

78H_M[:, -1] = 1 

79 

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

81 

82E_alpha = sp.diags( 

83 [ 

84 -1.0, 

85 ] 

86 * (L - 1), 

87 offsets=-1, 

88).tolil() 

89E_alpha[0, -1] = -alpha 

90 

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

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

93 

94J = sp.diags(gamma) 

95J_inv = sp.diags(1 / gamma) 

96 

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

98 

99# prepare diagonalization of QG^{-1} 

100w = [] 

101S = [] 

102S_inv = [] 

103 

104for l in range(L): 

105 # diagonalize QG^-1 matrix 

106 if M > 1: 

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

108 else: 

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

110 _S_inv = np.linalg.inv(_S) 

111 w.append(_w) 

112 S.append(_S) 

113 S_inv.append(_S_inv) 

114 

115""" 

116Setup functions for computing matrix-vector productions on the steps and for computing the residual of the composite 

117collocation problem 

118""" 

119 

120 

121def mat_vec(mat, vec): 

122 """ 

123 Matrix vector product 

124 

125 Args: 

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

127 vec (np.ndarray) : vector 

128 

129 Returns: 

130 np.ndarray: mat @ vec 

131 """ 

132 res = np.zeros_like(vec) 

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

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

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

136 return res 

137 

138 

139def residual(_u, u0): 

140 """ 

141 Compute the residual of the composite collocation problem 

142 

143 Args: 

144 _u (np.ndarray): Current iterate 

145 u0 (np.ndarray): Initial conditions 

146 

147 Returns: 

148 np.ndarray: LMN size array with the residual 

149 """ 

150 res = _u * 0j 

151 for l in range(L): 

152 # build step local residual 

153 

154 # communicate initial conditions for each step 

155 if l == 0: 

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

157 else: 

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

159 

160 # evaluate and subtract integral over right hand side functions 

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

162 Qf = mat_vec(Q, f_evals) 

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

164 

165 return res 

166 

167 

168# do ParaDiag 

169sol_paradiag = u.copy() * 0j 

170u0 = u.copy() 

171niter = 0 

172res = residual(sol_paradiag, u0) 

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

174 # compute all-at-once residual 

175 res = residual(sol_paradiag, u0) 

176 

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

178 u_avg = prob.u_init 

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

180 

181 # weighted FFT in time 

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

183 

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

185 y = np.empty_like(x) 

186 for l in range(L): 

187 

188 # perform local solves on the collocation nodes in parallel 

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

190 x2 = np.empty_like(x1) 

191 for m in range(M): 

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

193 z = S[l] @ x2 

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

195 

196 # inverse FFT in time and increment 

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

198 

199 res = residual(sol_paradiag, u0) 

200 niter += 1 

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

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