Coverage for pySDC/implementations/problem_classes/generic_spectral.py: 56%

1from pySDC.core.problem import Problem, WorkCounter 

2from pySDC.helpers.spectral_helper import SpectralHelper 

3import numpy as np 

4from pySDC.core.errors import ParameterError 

5 

6 

7class GenericSpectralLinear(Problem): 

8 """ 

9 Generic class to solve problems of the form M u_t + L u = y, with mass matrix M, linear operator L and some right 

10 hand side y using spectral methods. 

11 L may contain algebraic conditions, as long as (M + dt L) is invertible. 

12 

13 Note that the `__getattr__` method is overloaded to pass requests on to the spectral helper if they are not 

14 attributes of this class itself. For instance, you can add a BC by calling `self.spectral.add_BC` or equivalently 

15 `self.add_BC`. 

16 

17 You can port problems derived from this more or less seamlessly to GPU by using the numerical libraries that are 

18 class attributes of the spectral helper. This class will automatically switch the datatype using the `setup_GPU` class method. 

19 

20 Attributes: 

21 spectral (pySDC.helpers.spectral_helper.SpectralHelper): Spectral helper 

22 work_counters (dict): Dictionary for counting work 

23 cached_factorizations (dict): Dictionary of cached matrix factorizations for solving 

24 L (sparse matrix): Linear operator 

25 M (sparse matrix): Mass matrix 

26 diff_mask (list): Mask for separating differential and algebraic terms 

27 Pl (sparse matrix): Left preconditioner 

28 Pr (sparse matrix): Right preconditioner 

29 """ 

30 

31 @classmethod 

32 def setup_GPU(cls): 

33 """switch to GPU modules""" 

34 import cupy as cp 

35 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh, imex_cupy_mesh 

36 from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh 

37 

38 cls.dtype_u = cupy_mesh 

39 

40 GPU_versions = { 

41 mesh: cupy_mesh, 

42 imex_mesh: imex_cupy_mesh, 

43 } 

44 

45 cls.dtype_f = GPU_versions[cls.dtype_f] 

46 

47 def __init__( 

48 self, 

49 bases, 

50 components, 

51 comm=None, 

52 Dirichlet_recombination=True, 

53 left_preconditioner=True, 

54 solver_type='cached_direct', 

55 solver_args=None, 

56 useGPU=False, 

57 max_cached_factorizations=12, 

58 spectral_space=True, 

59 real_spectral_coefficients=False, 

60 debug=False, 

61 ): 

62 """ 

63 Base class for problems discretized with spectral methods. 

64 

65 Args: 

66 bases (list of dictionaries): 1D Bases 

67 components (list of strings): Components of the equations 

68 comm (mpi4py.Intracomm or None): MPI communicator 

69 Dirichlet_recombination (bool): Use Dirichlet recombination in the last axis as right preconditioner 

70 left_preconditioner (bool): Reverse the Kronecker product if yes 

71 solver_type (str): Solver for linear systems 

72 solver_args (dict): Arguments for linear solver 

73 useGPU (bool): Run on GPU or CPU 

74 max_cached_factorizations (int): Number of matrix decompositions to cache before starting eviction 

75 spectral_space (bool): If yes, the solution will not be transformed back after solving and evaluating the RHS, and is expected as input in spectral space to these functions 

76 real_spectral_coefficients (bool): If yes, allow only real values in spectral space, otherwise, allow complex. 

77 debug (bool): Make additional tests at extra computational cost 

78 """ 

79 solver_args = {} if solver_args is None else solver_args 

80 self._makeAttributeAndRegister( 

81 'max_cached_factorizations', 

82 'useGPU', 

83 'solver_type', 

84 'solver_args', 

85 'left_preconditioner', 

86 'Dirichlet_recombination', 

87 'comm', 

88 'spectral_space', 

89 'real_spectral_coefficients', 

90 'debug', 

91 localVars=locals(), 

92 ) 

93 self.spectral = SpectralHelper(comm=comm, useGPU=useGPU, debug=debug) 

94 

95 if useGPU: 

96 self.setup_GPU() 

97 

98 for base in bases: 

99 self.spectral.add_axis(**base) 

100 self.spectral.add_component(components) 

101 

102 self.spectral.setup_fft(real_spectral_coefficients) 

103 

104 super().__init__(init=self.spectral.init_forward if spectral_space else self.spectral.init) 

105 

106 self.work_counters[solver_type] = WorkCounter() 

107 self.work_counters['factorizations'] = WorkCounter() 

108 

109 self.setup_preconditioner(Dirichlet_recombination, left_preconditioner) 

110 

111 self.cached_factorizations = {} 

112 

113 def __getattr__(self, name): 

114 """ 

115 Pass requests on to the helper if they are not directly attributes of this class for convenience. 

116 

117 Args: 

118 name (str): Name of the attribute you want 

119 

120 Returns: 

121 request 

122 """ 

123 return getattr(self.spectral, name) 

124 

125 def _setup_operator(self, LHS): 

126 """ 

127 Setup a sparse linear operator by adding relationships. See documentation for ``GenericSpectralLinear.setup_L`` to learn more. 

128 

129 Args: 

130 LHS (dict): Equations to be added to the operator 

131 

132 Returns: 

133 sparse linear operator 

134 """ 

135 operator = self.spectral.get_empty_operator_matrix() 

136 for line, equation in LHS.items(): 

137 self.spectral.add_equation_lhs(operator, line, equation) 

138 return self.spectral.convert_operator_matrix_to_operator(operator) 

139 

140 def setup_L(self, LHS): 

141 """ 

142 Setup the left hand side of the linear operator L and store it in ``self.L``. 

143 

144 The argument is meant to be a dictionary with the line you want to write the equation in as the key and the relationship between components as another dictionary. For instance, you can add an algebraic condition capturing a first derivative relationship between u and ux as follows: 

145 

146 ``` 

147 Dx = self.get_differentiation_matrix(axes=(0,)) 

148 I = self.get_Id() 

149 LHS = {'ux': {'u': Dx, 'ux': -I}} 

150 self.setup_L(LHS) 

151 ``` 

152 

153 If you put zero as right hand side for the solver in the line for ux, ux will contain the x-derivative of u afterwards. 

154 

155 Args: 

156 LHS (dict): Dictionary containing the equations. 

157 """ 

158 self.L = self._setup_operator(LHS) 

159 

160 def setup_M(self, LHS): 

161 ''' 

162 Setup mass matrix, see documentation of ``GenericSpectralLinear.setup_L``. 

163 ''' 

164 diff_index = list(LHS.keys()) 

165 self.diff_mask = [me in diff_index for me in self.components] 

166 self.M = self._setup_operator(LHS) 

167 

168 def setup_preconditioner(self, Dirichlet_recombination=True, left_preconditioner=True): 

169 """ 

170 Get left and right preconditioners. 

171 

172 Args: 

173 Dirichlet_recombination (bool): Basis conversion for right preconditioner. Useful for Chebychev and Ultraspherical methods. 10/10 would recommend. 

174 left_preconditioner (bool): If True, it will interleave the variables and reverse the Kronecker product 

175 """ 

176 sp = self.spectral.sparse_lib 

177 N = np.prod(self.init[0][1:]) 

178 

179 Id = sp.eye(N) 

180 Pl_lhs = {comp: {comp: Id} for comp in self.components} 

181 self.Pl = self._setup_operator(Pl_lhs) 

182 

183 if left_preconditioner: 

184 # reverse Kronecker product 

185 

186 if self.spectral.useGPU: 

187 R = self.Pl.get().tolil() * 0 

188 else: 

189 R = self.Pl.tolil() * 0 

190 

191 for j in range(self.ncomponents): 

192 for i in range(N): 

193 R[i * self.ncomponents + j, j * N + i] = 1.0 

194 

195 self.Pl = self.spectral.sparse_lib.csc_matrix(R) 

196 

197 if Dirichlet_recombination and type(self.axes[-1]).__name__ in ['ChebychevHelper, Ultraspherical']: 

198 _Pr = self.spectral.get_Dirichlet_recombination_matrix(axis=-1) 

199 else: 

200 _Pr = Id 

201 

202 Pr_lhs = {comp: {comp: _Pr} for comp in self.components} 

203 self.Pr = self._setup_operator(Pr_lhs) @ self.Pl.T 

204 

205 def solve_system(self, rhs, dt, u0=None, *args, skip_itransform=False, **kwargs): 

206 """ 

207 Do an implicit Euler step to solve M u_t + Lu = rhs, with M the mass matrix and L the linear operator as setup by 

208 ``GenericSpectralLinear.setup_L`` and ``GenericSpectralLinear.setup_M``. 

209 

210 The implicit Euler step is (M - dt L) u = M rhs. Note that M need not be invertible as long as (M + dt*L) is. 

211 This means solving with dt=0 to mimic explicit methods does not work for all problems, in particular simple DAEs. 

212 

213 Note that by putting M rhs on the right hand side, this function can only solve algebraic conditions equal to 

214 zero. If you want something else, it should be easy to overload this function. 

215 """ 

216 

217 sp = self.spectral.sparse_lib 

218 

219 if self.spectral_space: 

220 rhs_hat = rhs.copy() 

221 else: 

222 rhs_hat = self.spectral.transform(rhs) 

223 

224 rhs_hat = (self.M @ rhs_hat.flatten()).reshape(rhs_hat.shape) 

225 rhs_hat = self.spectral.put_BCs_in_rhs_hat(rhs_hat) 

226 rhs_hat = self.Pl @ rhs_hat.flatten() 

227 

228 if dt not in self.cached_factorizations.keys(): 

229 A = self.M + dt * self.L 

230 A = self.Pl @ self.spectral.put_BCs_in_matrix(A) @ self.Pr 

231 

232 # import numpy as np 

233 # if A.shape[0] < 200: 

234 # import matplotlib.pyplot as plt 

235 

236 # # M = self.spectral.put_BCs_in_matrix(self.L.copy()) 

237 # M = A # self.L 

238 # im = plt.imshow((M / abs(M)).real) 

239 # # im = plt.imshow(np.log10(abs(A.toarray())).real) 

240 # # im = plt.imshow(((A.toarray())).real) 

241 # plt.colorbar(im) 

242 # plt.show() 

243 

244 if self.solver_type.lower() == 'cached_direct': 

245 if dt not in self.cached_factorizations.keys(): 

246 if len(self.cached_factorizations) >= self.max_cached_factorizations: 

247 self.cached_factorizations.pop(list(self.cached_factorizations.keys())[0]) 

248 self.logger.debug(f'Evicted matrix factorization for {dt=:.6f} from cache') 

249 self.cached_factorizations[dt] = self.spectral.linalg.factorized(A) 

250 self.logger.debug(f'Cached matrix factorization for {dt=:.6f}') 

251 self.work_counters['factorizations']() 

252 

253 _sol_hat = self.cached_factorizations[dt](rhs_hat) 

254 self.logger.debug(f'Used cached matrix factorization for {dt=:.6f}') 

255 

256 elif self.solver_type.lower() == 'direct': 

257 _sol_hat = sp.linalg.spsolve(A, rhs_hat) 

258 elif self.solver_type.lower() == 'gmres': 

259 _sol_hat, _ = sp.linalg.gmres( 

260 A, 

261 rhs_hat, 

262 x0=u0.flatten(), 

263 **self.solver_args, 

264 callback=self.work_counters[self.solver_type], 

265 callback_type='legacy', 

266 ) 

267 elif self.solver_type.lower() == 'cg': 

268 _sol_hat, _ = sp.linalg.cg( 

269 A, rhs_hat, x0=u0.flatten(), **self.solver_args, callback=self.work_counters[self.solver_type] 

270 ) 

271 else: 

272 raise NotImplementedError(f'Solver {self.solver_type:!} not implemented in {type(self).__name__}!') 

273 

274 sol_hat = self.spectral.u_init_forward 

275 sol_hat[...] = (self.Pr @ _sol_hat).reshape(sol_hat.shape) 

276 

277 if self.spectral_space: 

278 return sol_hat 

279 else: 

280 sol = self.spectral.u_init 

281 sol[:] = self.spectral.itransform(sol_hat).real 

282 

283 if self.spectral.debug: 

284 self.spectral.check_BCs(sol) 

285 

286 return sol 

287 

288 

289def compute_residual_DAE(self, stage=''): 

290 """ 

291 Computation of the residual that does not add u_0 - u_m in algebraic equations. 

292 

293 Args: 

294 stage (str): The current stage of the step the level belongs to 

295 """ 

296 

297 # get current level and problem description 

298 L = self.level 

299 

300 # Check if we want to skip the residual computation to gain performance 

301 # Keep in mind that skipping any residual computation is likely to give incorrect outputs of the residual! 

302 if stage in self.params.skip_residual_computation: 

303 L.status.residual = 0.0 if L.status.residual is None else L.status.residual 

304 return None 

305 

306 # check if there are new values (e.g. from a sweep) 

307 # assert L.status.updated 

308 

309 # compute the residual for each node 

310 

311 # build QF(u) 

312 res_norm = [] 

313 res = self.integrate() 

314 mask = L.prob.diff_mask 

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

316 res[m][mask] += L.u[0][mask] - L.u[m + 1][mask] 

317 # add tau if associated 

318 if L.tau[m] is not None: 

319 res[m] += L.tau[m] 

320 # use abs function from data type here 

321 res_norm.append(abs(res[m])) 

322 # print(m, [abs(me) for me in res[m]], [abs(me) for me in L.u[0] - L.u[m + 1]]) 

323 

324 # find maximal residual over the nodes 

325 if L.params.residual_type == 'full_abs': 

326 L.status.residual = max(res_norm) 

327 elif L.params.residual_type == 'last_abs': 

328 L.status.residual = res_norm[-1] 

329 elif L.params.residual_type == 'full_rel': 

330 L.status.residual = max(res_norm) / abs(L.u[0]) 

331 elif L.params.residual_type == 'last_rel': 

332 L.status.residual = res_norm[-1] / abs(L.u[0]) 

333 else: 

334 raise ParameterError( 

335 f'residual_type = {L.params.residual_type} not implemented, choose ' 

336 f'full_abs, last_abs, full_rel or last_rel instead' 

337 ) 

338 

339 # indicate that the residual has seen the new values 

340 L.status.updated = False 

341 

342 return None 

343 

344 

345def compute_residual_DAE_MPI(self, stage=None): 

346 """ 

347 Computation of the residual using the collocation matrix Q 

348 

349 Args: 

350 stage (str): The current stage of the step the level belongs to 

351 """ 

352 from mpi4py import MPI 

353 

354 L = self.level 

355 

356 # Check if we want to skip the residual computation to gain performance 

357 # Keep in mind that skipping any residual computation is likely to give incorrect outputs of the residual! 

358 if stage in self.params.skip_residual_computation: 

359 L.status.residual = 0.0 if L.status.residual is None else L.status.residual 

360 return None 

361 

362 # compute the residual for each node 

363 

364 # build QF(u) 

365 res = self.integrate(last_only=L.params.residual_type[:4] == 'last') 

366 mask = L.prob.diff_mask 

367 res[mask] += L.u[0][mask] - L.u[self.rank + 1][mask] 

368 # add tau if associated 

369 if L.tau[self.rank] is not None: 

370 res += L.tau[self.rank] 

371 # use abs function from data type here 

372 res_norm = abs(res) 

373 

374 # find maximal residual over the nodes 

375 if L.params.residual_type == 'full_abs': 

376 L.status.residual = self.comm.allreduce(res_norm, op=MPI.MAX) 

377 elif L.params.residual_type == 'last_abs': 

378 L.status.residual = self.comm.bcast(res_norm, root=self.comm.size - 1) 

379 elif L.params.residual_type == 'full_rel': 

380 L.status.residual = self.comm.allreduce(res_norm / abs(L.u[0]), op=MPI.MAX) 

381 elif L.params.residual_type == 'last_rel': 

382 L.status.residual = self.comm.bcast(res_norm / abs(L.u[0]), root=self.comm.size - 1) 

383 else: 

384 raise NotImplementedError(f'residual type \"{L.params.residual_type}\" not implemented!') 

385 

386 # indicate that the residual has seen the new values 

387 L.status.updated = False 

388 

389 return None 

390 

391 

392def get_extrapolated_error_DAE(self, S, **kwargs): 

393 """ 

394 The extrapolation estimate combines values of u and f from multiple steps to extrapolate and compare to the 

395 solution obtained by the time marching scheme. This function can be used in `EstimateExtrapolationError`. 

396 

397 Args: 

398 S (pySDC.Step): The current step 

399 

400 Returns: 

401 None 

402 """ 

403 u_ex = self.get_extrapolated_solution(S) 

404 diff_mask = S.levels[0].prob.diff_mask 

405 if u_ex is not None: 

406 S.levels[0].status.error_extrapolation_estimate = ( 

407 abs((u_ex - S.levels[0].u[-1])[diff_mask]) * self.coeff.prefactor 

408 ) 

409 # print([abs(me) for me in (u_ex - S.levels[0].u[-1]) * self.coeff.prefactor]) 

410 else: 

411 S.levels[0].status.error_extrapolation_estimate = None