Coverage for pySDC/implementations/sweeper_classes/ParaDiagSweepers.py: 89%

1""" 

2These sweepers are made for use with ParaDiag. They can be used to some degree with SDC as well, but unless you know what you are doing, you probably want another sweeper. 

3""" 

4 

5from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit 

6from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order 

7import numpy as np 

8import scipy.sparse as sp 

9 

10 

11class QDiagonalization(generic_implicit): 

12 """ 

13 Sweeper solving the collocation problem directly via diagonalization of Q. Mainly made for ParaDiag. 

14 Can be reconfigured for use with SDC. 

15 

16 Note that the initial conditions for the collocation problem are generally stored in node zero in pySDC. However, 

17 this sweeper is intended for ParaDiag, where a node-local residual is needed as a right hand side for this sweeper 

18 rather than a step local one. Therefore, this sweeper has an option `ignore_ic`. If true, the value in node zero 

19 will only be used in computing the step-local residual, but not in the solves. If false, the values on the nodes 

20 will be ignored in the solves and the node-zero value will be used as initial conditions. When using this as a time- 

21 parallel algorithm outside ParaDiag, you should set this parameter to false, which is not the default! 

22 

23 Similarly, in ParaDiag, the solution is in Fourier space right after the solve. It therefore makes little sense to 

24 evaluate the right hand side directly after. By default, this is not done! Set `update_f_evals=True` in the 

25 parameters if you want to use this sweeper in SDC. With `ignore_ic`, the solve gives only an increment, which is 

26 added to the solution later, so there is nothing to evaluate the right hand side at and `update_f_evals` is ignored. 

27 """ 

28 

29 def __init__(self, params, level): 

30 """ 

31 Initialization routine for the custom sweeper 

32 

33 Args: 

34 params: parameters for the sweeper 

35 level (pySDC.Level.level): the level that uses this sweeper 

36 """ 

37 if 'G_inv' not in params.keys(): 

38 params['G_inv'] = np.eye(params['num_nodes']) 

39 params['update_f_evals'] = params.get('update_f_evals', False) 

40 params['ignore_ic'] = params.get('ignore_ic', True) 

41 

42 super().__init__(params, level) 

43 

44 self.set_G_inv(self.params.G_inv) 

45 

46 def set_G_inv(self, G_inv): 

47 """ 

48 In ParaDiag, QG^{-1} is diagonalized. This function stores the G_inv matrix and computes and stores the diagonalization. 

49 """ 

50 self.params.G_inv = G_inv 

51 self.w, self.S, self.S_inv = self.computeDiagonalization(A=self.coll.Qmat[1:, 1:] @ self.params.G_inv) 

52 

53 @staticmethod 

54 def computeDiagonalization(A): 

55 """ 

56 Compute diagonalization of dense matrix A = S diag(w) S^-1 

57 

58 Args: 

59 A (numpy.ndarray): dense matrix to diagonalize 

60 

61 Returns: 

62 numpy.array: Diagonal entries of the diagonalized matrix w 

63 numpy.ndarray: Matrix of eigenvectors S 

64 numpy.ndarray: Inverse of S 

65 """ 

66 w, S = np.linalg.eig(A) 

67 S_inv = np.linalg.inv(S) 

68 assert np.allclose(S @ np.diag(w) @ S_inv, A) 

69 return w, S, S_inv 

70 

71 def mat_vec(self, mat, vec): 

72 """ 

73 Compute matrix-vector multiplication. Vector can be list. 

74 

75 Args: 

76 mat: Matrix 

77 vec: Vector 

78 

79 Returns: 

80 list: mat @ vec 

81 """ 

82 assert mat.shape[1] == len(vec) 

83 result = [] 

84 for m in range(mat.shape[0]): 

85 result.append(self.level.prob.u_init) 

86 for j in range(mat.shape[1]): 

87 result[-1] += mat[m, j] * vec[j] 

88 return result 

89 

90 def update_nodes(self): 

91 """ 

92 Update the u- and f-values at the collocation nodes -> corresponds to a single sweep over all nodes 

93 

94 Returns: 

95 None 

96 """ 

97 

98 L = self.level 

99 P = L.prob 

100 M = self.coll.num_nodes 

101 

102 if L.tau[0] is not None: 

103 raise NotImplementedError('This sweeper does not work with multi-level SDC') 

104 

105 # perform local solves on the collocation nodes, can be parallelized! 

106 if self.params.ignore_ic: 

107 x1 = self.mat_vec(self.S_inv, [self.level.residual[m] for m in range(M)]) 

108 else: 

109 x1 = self.mat_vec(self.S_inv, [self.level.u[0] for _ in range(M)]) 

110 

111 # get averaged state over all nodes for constructing the Jacobian 

112 u_avg = P.u_init 

113 if not any(me is None for me in L.u_avg): 

114 for m in range(M): 

115 u_avg += L.u_avg[m] / M 

116 

117 x2 = [] 

118 for m in range(M): 

119 x2.append(P.solve_jacobian(x1[m], self.w[m] * L.dt, u=u_avg, t=L.time + L.dt * self.coll.nodes[m])) 

120 z = self.mat_vec(self.S, x2) 

121 y = self.mat_vec(self.params.G_inv, z) 

122 

123 # update solution and evaluate right hand side 

124 for m in range(M): 

125 if self.params.ignore_ic: 

126 L.increment[m] = y[m] 

127 else: 

128 L.u[m + 1] = y[m] 

129 if self.params.update_f_evals: 

130 L.f[m + 1] = P.eval_f(L.u[m + 1], L.time + L.dt * self.coll.nodes[m]) 

131 

132 L.status.updated = True 

133 return None 

134 

135 def eval_f_at_all_nodes(self): 

136 """ 

137 Evaluate the right hand side at all collocation nodes, since the solve does not do this by default. 

138 

139 Returns: 

140 None 

141 """ 

142 L = self.level 

143 P = self.level.prob 

144 for m in range(self.coll.num_nodes): 

145 L.f[m + 1] = P.eval_f(L.u[m + 1], L.time + L.dt * self.coll.nodes[m]) 

146 

147 def get_residual(self): 

148 """ 

149 This function computes and returns the "spatially extended" residual, not the norm of the residual! 

150 

151 Returns: 

152 pySDC.datatype: Spatially extended residual 

153 

154 """ 

155 self.eval_f_at_all_nodes() 

156 

157 # start with integral dt*Q*F 

158 residual = self.integrate() 

159 

160 # subtract u and add u0 to arrive at r = dt*Q*F - u + u0 

161 for m in range(self.coll.num_nodes): 

162 residual[m] -= self.level.u[m + 1] 

163 residual[m] += self.level.u[0] 

164 

165 return residual 

166 

167 def compute_residual(self, *args, **kwargs): 

168 """ 

169 Evaluate the right hand side at all nodes and then compute the residual as in the base class. 

170 

171 Args: 

172 *args: passed on to the base class method, typically the stage (str) 

173 **kwargs: passed on to the base class method 

174 

175 Returns: 

176 None 

177 """ 

178 self.eval_f_at_all_nodes() 

179 return super().compute_residual(*args, **kwargs) 

180 

181 

182class QDiagonalizationIMEX(QDiagonalization): 

183 """ 

184 Variant of QDiagonalization for ParaDiag on problems whose right hand side has implicit and explicit parts. 

185 Note that it will not work with SDC. 

186 """ 

187 

188 integrate = imex_1st_order.integrate