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 nc = self.ncomponents 

236 

237 rows = self.xp.arange(N * nc) 

238 cols = (rows % nc) * N + rows // nc 

239 

240 self.Pl = self.spectral.sparse_lib.csc_matrix( 

241 (self.xp.ones(N * nc, dtype=complex), (rows, cols)), shape=(N * nc, N * nc) 

242 ) 

243 

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

245 else: 

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

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

248 self.Pl = self._setup_operator(Pl_lhs) 

249 

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

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

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

253 else: 

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

255 

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

257 operator = self._setup_operator(Pr_lhs) 

258 

259 try: 

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

261 except MemError: 

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

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

264 

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

266 """ 

267 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 

268 ``GenericSpectralLinear.setup_L`` and ``GenericSpectralLinear.setup_M``. 

269 

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

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

272 

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

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

275 """ 

276 

277 sp = self.spectral.sparse_lib 

278 

279 self.heterogeneous_setup() 

280 

281 if self.spectral_space: 

282 rhs_hat = rhs.copy() 

283 if u0 is not None: 

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

285 else: 

286 u0_hat = None 

287 else: 

288 rhs_hat = self.spectral.transform(rhs) 

289 if u0 is not None: 

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

291 else: 

292 u0_hat = None 

293 

294 # apply inverse right preconditioner to initial guess 

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

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

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

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

299 

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

301 rhs_hat = self.spectral.put_BCs_in_rhs_hat(rhs_hat) 

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

303 

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

305 if self.heterogeneous: 

306 M = self.M_CPU 

307 L = self.L_CPU 

308 Pl = self.Pl_CPU 

309 Pr = self.Pr_CPU 

310 else: 

311 M = self.M 

312 L = self.L 

313 Pl = self.Pl 

314 Pr = self.Pr 

315 

316 A = M + dt * L 

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

318 

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

320 # import matplotlib.pyplot as plt 

321 

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

323 # M = A # self.L 

324 # im = plt.spy(M) 

325 # plt.show() 

326 

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

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

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

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

331 self.cached_factorizations.pop(to_evict) 

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

333 iLU = self.linalg.spilu( 

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

335 ) 

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

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

338 self.work_counters['factorizations']() 

339 M = self.cached_factorizations[dt] 

340 else: 

341 M = None 

342 info = 0 

343 

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

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

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

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

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

349 

350 if self.heterogeneous: 

351 import scipy.sparse as sp 

352 

353 cpu_decomp = sp.linalg.splu(A) 

354 

355 if self.useGPU: 

356 from cupyx.scipy.sparse.linalg import SuperLU 

357 

358 solver = SuperLU(cpu_decomp).solve 

359 else: 

360 solver = cpu_decomp.solve 

361 else: 

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

363 

364 self.cached_factorizations[dt] = solver 

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

366 self.work_counters['factorizations']() 

367 

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

369 self.work_counters[self.solver_type]() 

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

371 

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

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

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

375 _sol_hat, _ = sp.linalg.gmres( 

376 A, 

377 rhs_hat, 

378 x0=u0_hat, 

379 **self.solver_args, 

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

381 callback_type='pr_norm', 

382 M=M, 

383 ) 

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

385 _sol_hat, info = sp.linalg.cg( 

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

387 ) 

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

389 _sol_hat, info = self.linalg.bicgstab( 

390 A, 

391 rhs_hat, 

392 x0=u0_hat, 

393 **self.solver_args, 

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

395 M=M, 

396 ) 

397 else: 

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

399 

400 if info != 0: 

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

402 

403 sol_hat = self.spectral.u_init_forward 

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

405 

406 if self.spectral_space: 

407 return sol_hat 

408 else: 

409 sol = self.spectral.u_init 

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

411 

412 if self.spectral.debug: 

413 self.spectral.check_BCs(sol) 

414 

415 return sol 

416 

417 def setUpFieldsIO(self): 

418 Rectilinear.setupMPI( 

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

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

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

422 ) 

423 

424 def getOutputFile(self, fileName): 

425 self.setUpFieldsIO() 

426 

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

428 if self.spectral.useGPU: 

429 coords = [me.get() for me in coords] 

430 

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

432 

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

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

435 fOut.initialize() 

436 return fOut 

437 

438 def processSolutionForOutput(self, u): 

439 if self.spectral_space: 

440 u = self.itransform(u).real 

441 else: 

442 u = u.real 

443 

444 if self.spectral.useGPU: 

445 u = u.get() 

446 

447 return np.ascontiguousarray(u.view(np.ndarray)) 

448 

449 

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

451 """ 

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

453 

454 Args: 

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

456 """ 

457 

458 # get current level and problem description 

459 L = self.level 

460 

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

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

463 if stage in self.params.skip_residual_computation: 

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

465 return None 

466 

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

468 # assert L.status.updated 

469 

470 # compute the residual for each node 

471 

472 # build QF(u) 

473 res_norm = [] 

474 res = self.integrate() 

475 mask = L.prob.diff_mask 

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

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

478 # add tau if associated 

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

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

481 # use abs function from data type here 

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

483 

484 # find maximal residual over the nodes 

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

486 L.status.residual = max(res_norm) 

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

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

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

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

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

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

493 else: 

494 raise ParameterError( 

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

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

497 ) 

498 

499 # indicate that the residual has seen the new values 

500 L.status.updated = False 

501 

502 return None 

503 

504 

505def compute_residual_DAE_MPI(self, stage=None): 

506 """ 

507 Computation of the residual using the collocation matrix Q 

508 

509 Args: 

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

511 """ 

512 from mpi4py import MPI 

513 

514 L = self.level 

515 

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

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

518 if stage in self.params.skip_residual_computation: 

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

520 return None 

521 

522 # compute the residual for each node 

523 

524 # build QF(u) 

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

526 mask = L.prob.diff_mask 

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

528 # add tau if associated 

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

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

531 # use abs function from data type here 

532 res_norm = abs(res) 

533 

534 # find maximal residual over the nodes 

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

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

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

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

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

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

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

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

543 else: 

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

545 

546 # indicate that the residual has seen the new values 

547 L.status.updated = False 

548 

549 return None 

550 

551 

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

553 """ 

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

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

556 

557 Args: 

558 S (pySDC.Step): The current step 

559 

560 Returns: 

561 None 

562 """ 

563 u_ex = self.get_extrapolated_solution(S) 

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

565 if u_ex is not None: 

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

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

568 ) 

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

570 else: 

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