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

1from pySDC.core.problem import Problem, WorkCounter 

2from pySDC.helpers.spectral_helper import SpectralHelper 

3import numpy as np 

4from pySDC.core.errors import ParameterError 

5from pySDC.helpers.fieldsIO import Rectilinear 

6 

7 

8class GenericSpectralLinear(Problem): 

9 """ 

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

11 hand side y using spectral methods. 

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

13 

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

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

16 `self.add_BC`. 

17 

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

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

20 

21 Attributes: 

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

23 work_counters (dict): Dictionary for counting work 

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

25 L (sparse matrix): Linear operator 

26 M (sparse matrix): Mass matrix 

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

28 Pl (sparse matrix): Left preconditioner 

29 Pr (sparse matrix): Right preconditioner 

30 """ 

31 

32 def setup_GPU(self): 

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 self.dtype_u = cupy_mesh 

39 

40 GPU_versions = { 

41 mesh: cupy_mesh, 

42 imex_mesh: imex_cupy_mesh, 

43 } 

44 

45 self.dtype_f = GPU_versions[self.dtype_f] 

46 

47 if self.comm is not None: 

48 from pySDC.helpers.NCCL_communicator import NCCLComm 

49 

50 if not isinstance(self.comm, NCCLComm): 

51 self.__dict__['comm'] = NCCLComm(self.comm) 

52 

53 def __init__( 

54 self, 

55 bases, 

56 components, 

57 comm=None, 

58 Dirichlet_recombination=True, 

59 left_preconditioner=True, 

60 solver_type='cached_direct', 

61 solver_args=None, 

62 preconditioner_args=None, 

63 useGPU=False, 

64 max_cached_factorizations=12, 

65 spectral_space=True, 

66 real_spectral_coefficients=False, 

67 heterogeneous=False, 

68 debug=False, 

69 ): 

70 """ 

71 Base class for problems discretized with spectral methods. 

72 

73 Args: 

74 bases (list of dictionaries): 1D Bases 

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

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

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

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

79 solver_type (str): Solver for linear systems 

80 solver_args (dict): Arguments for linear solver 

81 useGPU (bool): Run on GPU or CPU 

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

83 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 

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

85 heterogeneous (bool): If yes, perform memory intensive sparse matrix operations on CPU 

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

87 """ 

88 solver_args = {} if solver_args is None else solver_args 

89 

90 preconditioner_args = {} if preconditioner_args is None else preconditioner_args 

91 preconditioner_args['drop_tol'] = preconditioner_args.get('drop_tol', 1e-3) 

92 preconditioner_args['fill_factor'] = preconditioner_args.get('fill_factor', 100) 

93 

94 self._makeAttributeAndRegister( 

95 'max_cached_factorizations', 

96 'useGPU', 

97 'solver_type', 

98 'solver_args', 

99 'preconditioner_args', 

100 'left_preconditioner', 

101 'Dirichlet_recombination', 

102 'comm', 

103 'spectral_space', 

104 'real_spectral_coefficients', 

105 'heterogeneous', 

106 'debug', 

107 localVars=locals(), 

108 ) 

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

110 

111 if useGPU: 

112 self.setup_GPU() 

113 if self.solver_args is not None: 

114 if 'rtol' in self.solver_args.keys(): 

115 self.solver_args['tol'] = self.solver_args.pop('rtol') 

116 

117 for base in bases: 

118 self.spectral.add_axis(**base) 

119 self.spectral.add_component(components) 

120 

121 self.spectral.setup_fft(real_spectral_coefficients) 

122 

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

124 

125 self.work_counters[solver_type] = WorkCounter() 

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

127 

128 self.setup_preconditioner(Dirichlet_recombination, left_preconditioner) 

129 

130 self.cached_factorizations = {} 

131 

132 if self.heterogeneous: 

133 self.__heterogeneous_setup = False 

134 

135 self.logger.debug('Finished GenericSpectralLinear __init__') 

136 

137 def heterogeneous_setup(self): 

138 if self.heterogeneous and not self.__heterogeneous_setup: 

139 

140 CPU_only = ['BC_line_zero_matrix', 'BCs'] 

141 both = ['Pl', 'Pr', 'L', 'M'] 

142 

143 self.logger.debug(f'Starting heterogeneous setup. Moving {CPU_only} and copying {both} to CPU') 

144 

145 if self.useGPU: 

146 for key in CPU_only: 

147 self.logger.debug(f'Moving {key} to CPU') 

148 setattr(self.spectral, key, getattr(self.spectral, key).get()) 

149 

150 for key in both: 

151 self.logger.debug(f'Copying {key} to CPU') 

152 setattr(self, f'{key}_CPU', getattr(self, key).get()) 

153 else: 

154 for key in both: 

155 setattr(self, f'{key}_CPU', getattr(self, key)) 

156 

157 self.logger.debug('Done with heterogeneous setup') 

158 self.__heterogeneous_setup = True 

159 

160 def __getattr__(self, name): 

161 """ 

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

163 

164 Args: 

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

166 

167 Returns: 

168 request 

169 """ 

170 return getattr(self.spectral, name) 

171 

172 def _setup_operator(self, LHS): 

173 """ 

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

175 

176 Args: 

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

178 

179 Returns: 

180 sparse linear operator 

181 """ 

182 operator = self.spectral.get_empty_operator_matrix() 

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

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

185 return self.spectral.convert_operator_matrix_to_operator(operator) 

186 

187 def setup_L(self, LHS): 

188 """ 

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

190 

191 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: 

192 

193 ``` 

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

195 I = self.get_Id() 

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

197 self.setup_L(LHS) 

198 ``` 

199 

200 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. 

201 

202 Args: 

203 LHS (dict): Dictionary containing the equations. 

204 """ 

205 self.L = self._setup_operator(LHS) 

206 self.logger.debug('Set up L matrix') 

207 

208 def setup_M(self, LHS): 

209 ''' 

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

211 ''' 

212 diff_index = list(LHS.keys()) 

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

214 self.M = self._setup_operator(LHS) 

215 self.logger.debug('Set up M matrix') 

216 

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

218 """ 

219 Get left and right preconditioners. 

220 

221 Args: 

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

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

224 """ 

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

226 if self.useGPU: 

227 from cupy_backends.cuda.libs.cusparse import CuSparseError as MemError 

228 else: 

229 MemError = MemoryError 

230 

231 if left_preconditioner: 

232 self.logger.debug(f'Setting up left preconditioner with {N} local degrees of freedom') 

233 

234 # reverse Kronecker product 

235 if self.spectral.useGPU: 

236 import scipy.sparse as sp 

237 else: 

238 sp = self.spectral.sparse_lib 

239 

240 R = sp.lil_matrix((self.ncomponents * N,) * 2, dtype=int) 

241 

242 for j in range(self.ncomponents): 

243 for i in range(N): 

244 R[i * self.ncomponents + j, j * N + i] = 1 

245 

246 self.Pl = self.spectral.sparse_lib.csc_matrix(R, dtype=complex) 

247 

248 self.logger.debug('Finished setup of left preconditioner') 

249 else: 

250 Id = self.spectral.sparse_lib.eye(N) 

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

252 self.Pl = self._setup_operator(Pl_lhs) 

253 

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

255 self.logger.debug('Using Dirichlet recombination as right preconditioner') 

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

257 else: 

258 _Pr = self.spectral.sparse_lib.eye(N) 

259 

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

261 operator = self._setup_operator(Pr_lhs) 

262 

263 try: 

264 self.Pr = operator @ self.Pl.T 

265 except MemError: 

266 self.logger.debug('Setting up right preconditioner on CPU due to memory error') 

267 self.Pr = self.spectral.sparse_lib.csc_matrix(operator.get() @ self.Pl.T.get()) 

268 

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

270 """ 

271 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 

272 ``GenericSpectralLinear.setup_L`` and ``GenericSpectralLinear.setup_M``. 

273 

274 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. 

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

276 

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

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

279 """ 

280 

281 sp = self.spectral.sparse_lib 

282 

283 self.heterogeneous_setup() 

284 

285 if self.spectral_space: 

286 rhs_hat = rhs.copy() 

287 if u0 is not None: 

288 u0_hat = u0.copy().flatten() 

289 else: 

290 u0_hat = None 

291 else: 

292 rhs_hat = self.spectral.transform(rhs) 

293 if u0 is not None: 

294 u0_hat = self.spectral.transform(u0).flatten() 

295 else: 

296 u0_hat = None 

297 

298 # apply inverse right preconditioner to initial guess 

299 if u0_hat is not None and 'direct' not in self.solver_type: 

300 if not hasattr(self, '_Pr_inv'): 

301 self._PR_inv = self.linalg.splu(self.Pr.astype(complex)).solve 

302 u0_hat[...] = self._PR_inv(u0_hat) 

303 

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

305 rhs_hat = self.spectral.put_BCs_in_rhs_hat(rhs_hat) 

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

307 

308 if dt not in self.cached_factorizations.keys() or not self.solver_type.lower() == 'cached_direct': 

309 if self.heterogeneous: 

310 M = self.M_CPU 

311 L = self.L_CPU 

312 Pl = self.Pl_CPU 

313 Pr = self.Pr_CPU 

314 else: 

315 M = self.M 

316 L = self.L 

317 Pl = self.Pl 

318 Pr = self.Pr 

319 

320 A = M + dt * L 

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

322 

323 # if A.shape[0] < 200e20: 

324 # import matplotlib.pyplot as plt 

325 

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

327 # M = A # self.L 

328 # im = plt.spy(M) 

329 # plt.show() 

330 

331 if 'ilu' in self.solver_type.lower(): 

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

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

334 to_evict = list(self.cached_factorizations.keys())[0] 

335 self.cached_factorizations.pop(to_evict) 

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

337 iLU = self.linalg.spilu( 

338 A, **{**self.preconditioner_args, 'drop_tol': dt * self.preconditioner_args['drop_tol']} 

339 ) 

340 self.cached_factorizations[dt] = self.linalg.LinearOperator(A.shape, iLU.solve) 

341 self.logger.debug(f'Cached incomplete LU factorization for {dt=:.6f}') 

342 self.work_counters['factorizations']() 

343 M = self.cached_factorizations[dt] 

344 else: 

345 M = None 

346 info = 0 

347 

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

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

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

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

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

353 

354 if self.heterogeneous: 

355 import scipy.sparse as sp 

356 

357 cpu_decomp = sp.linalg.splu(A) 

358 

359 if self.useGPU: 

360 from cupyx.scipy.sparse.linalg import SuperLU 

361 

362 solver = SuperLU(cpu_decomp).solve 

363 else: 

364 solver = cpu_decomp.solve 

365 else: 

366 solver = self.spectral.linalg.factorized(A) 

367 

368 self.cached_factorizations[dt] = solver 

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

370 self.work_counters['factorizations']() 

371 

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

373 self.work_counters[self.solver_type]() 

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

375 

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

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

378 elif 'gmres' in self.solver_type.lower(): 

379 _sol_hat, _ = sp.linalg.gmres( 

380 A, 

381 rhs_hat, 

382 x0=u0_hat, 

383 **self.solver_args, 

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

385 callback_type='pr_norm', 

386 M=M, 

387 ) 

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

389 _sol_hat, info = sp.linalg.cg( 

390 A, rhs_hat, x0=u0_hat, **self.solver_args, callback=self.work_counters[self.solver_type] 

391 ) 

392 elif 'bicgstab' in self.solver_type.lower(): 

393 _sol_hat, info = self.linalg.bicgstab( 

394 A, 

395 rhs_hat, 

396 x0=u0_hat, 

397 **self.solver_args, 

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

399 M=M, 

400 ) 

401 else: 

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

403 

404 if info != 0: 

405 self.logger.warn(f'{self.solver_type} not converged! {info=}') 

406 

407 sol_hat = self.spectral.u_init_forward 

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

409 

410 if self.spectral_space: 

411 return sol_hat 

412 else: 

413 sol = self.spectral.u_init 

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

415 

416 if self.spectral.debug: 

417 self.spectral.check_BCs(sol) 

418 

419 return sol 

420 

421 def setUpFieldsIO(self): 

422 Rectilinear.setupMPI( 

423 comm=self.comm.commMPI if self.useGPU else self.comm, 

424 iLoc=[me.start for me in self.local_slice(False)], 

425 nLoc=[me.stop - me.start for me in self.local_slice(False)], 

426 ) 

427 

428 def getOutputFile(self, fileName): 

429 self.setUpFieldsIO() 

430 

431 coords = [me.get_1dgrid() for me in self.spectral.axes] 

432 assert np.allclose([len(me) for me in coords], self.spectral.global_shape[1:]) 

433 

434 fOut = Rectilinear(np.float64, fileName=fileName) 

435 fOut.setHeader(nVar=len(self.components), coords=coords) 

436 fOut.initialize() 

437 return fOut 

438 

439 def processSolutionForOutput(self, u): 

440 if self.spectral_space: 

441 return np.array(self.itransform(u).real) 

442 else: 

443 return np.array(u.real) 

444 

445 

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

447 """ 

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

449 

450 Args: 

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

452 """ 

453 

454 # get current level and problem description 

455 L = self.level 

456 

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

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

459 if stage in self.params.skip_residual_computation: 

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

461 return None 

462 

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

464 # assert L.status.updated 

465 

466 # compute the residual for each node 

467 

468 # build QF(u) 

469 res_norm = [] 

470 res = self.integrate() 

471 mask = L.prob.diff_mask 

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

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

474 # add tau if associated 

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

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

477 # use abs function from data type here 

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

479 

480 # find maximal residual over the nodes 

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

482 L.status.residual = max(res_norm) 

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

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

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

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

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

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

489 else: 

490 raise ParameterError( 

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

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

493 ) 

494 

495 # indicate that the residual has seen the new values 

496 L.status.updated = False 

497 

498 return None 

499 

500 

501def compute_residual_DAE_MPI(self, stage=None): 

502 """ 

503 Computation of the residual using the collocation matrix Q 

504 

505 Args: 

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

507 """ 

508 from mpi4py import MPI 

509 

510 L = self.level 

511 

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

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

514 if stage in self.params.skip_residual_computation: 

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

516 return None 

517 

518 # compute the residual for each node 

519 

520 # build QF(u) 

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

522 mask = L.prob.diff_mask 

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

524 # add tau if associated 

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

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

527 # use abs function from data type here 

528 res_norm = abs(res) 

529 

530 # find maximal residual over the nodes 

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

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

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

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

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

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

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

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

539 else: 

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

541 

542 # indicate that the residual has seen the new values 

543 L.status.updated = False 

544 

545 return None 

546 

547 

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

549 """ 

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

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

552 

553 Args: 

554 S (pySDC.Step): The current step 

555 

556 Returns: 

557 None 

558 """ 

559 u_ex = self.get_extrapolated_solution(S) 

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

561 if u_ex is not None: 

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

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

564 ) 

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

566 else: 

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