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

230 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-09-04 15:08 +0000

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 def heterogeneous_setup(self): 

136 if self.heterogeneous and not self.__heterogeneous_setup: 

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

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

139 

140 if self.useGPU: 

141 for key in CPU_only: 

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

143 

144 for key in both: 

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

146 else: 

147 for key in both: 

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

149 

150 self.__heterogeneous_setup = True 

151 

152 def __getattr__(self, name): 

153 """ 

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

155 

156 Args: 

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

158 

159 Returns: 

160 request 

161 """ 

162 return getattr(self.spectral, name) 

163 

164 def _setup_operator(self, LHS): 

165 """ 

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

167 

168 Args: 

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

170 

171 Returns: 

172 sparse linear operator 

173 """ 

174 operator = self.spectral.get_empty_operator_matrix() 

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

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

177 return self.spectral.convert_operator_matrix_to_operator(operator) 

178 

179 def setup_L(self, LHS): 

180 """ 

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

182 

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

184 

185 ``` 

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

187 I = self.get_Id() 

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

189 self.setup_L(LHS) 

190 ``` 

191 

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

193 

194 Args: 

195 LHS (dict): Dictionary containing the equations. 

196 """ 

197 self.L = self._setup_operator(LHS) 

198 

199 def setup_M(self, LHS): 

200 ''' 

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

202 ''' 

203 diff_index = list(LHS.keys()) 

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

205 self.M = self._setup_operator(LHS) 

206 

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

208 """ 

209 Get left and right preconditioners. 

210 

211 Args: 

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

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

214 """ 

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

216 

217 if left_preconditioner: 

218 # reverse Kronecker product 

219 if self.spectral.useGPU: 

220 import scipy.sparse as sp 

221 else: 

222 sp = self.spectral.sparse_lib 

223 

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

225 

226 for j in range(self.ncomponents): 

227 for i in range(N): 

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

229 

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

231 else: 

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

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

234 self.Pl = self._setup_operator(Pl_lhs) 

235 

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

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

238 else: 

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

240 

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

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

243 

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

245 """ 

246 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 

247 ``GenericSpectralLinear.setup_L`` and ``GenericSpectralLinear.setup_M``. 

248 

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

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

251 

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

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

254 """ 

255 

256 sp = self.spectral.sparse_lib 

257 

258 self.heterogeneous_setup() 

259 

260 if self.spectral_space: 

261 rhs_hat = rhs.copy() 

262 if u0 is not None: 

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

264 else: 

265 u0_hat = None 

266 else: 

267 rhs_hat = self.spectral.transform(rhs) 

268 if u0 is not None: 

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

270 else: 

271 u0_hat = None 

272 

273 # apply inverse right preconditioner to initial guess 

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

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

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

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

278 

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

280 rhs_hat = self.spectral.put_BCs_in_rhs_hat(rhs_hat) 

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

282 

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

284 if self.heterogeneous: 

285 M = self.M_CPU 

286 L = self.L_CPU 

287 Pl = self.Pl_CPU 

288 Pr = self.Pr_CPU 

289 else: 

290 M = self.M 

291 L = self.L 

292 Pl = self.Pl 

293 Pr = self.Pr 

294 

295 A = M + dt * L 

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

297 

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

299 # import matplotlib.pyplot as plt 

300 

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

302 # M = A # self.L 

303 # im = plt.spy(M) 

304 # plt.show() 

305 

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

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

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

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

310 self.cached_factorizations.pop(to_evict) 

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

312 iLU = self.linalg.spilu( 

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

314 ) 

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

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

317 self.work_counters['factorizations']() 

318 M = self.cached_factorizations[dt] 

319 else: 

320 M = None 

321 info = 0 

322 

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

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

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

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

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

328 

329 if self.heterogeneous: 

330 import scipy.sparse as sp 

331 

332 cpu_decomp = sp.linalg.splu(A) 

333 if self.useGPU: 

334 from cupyx.scipy.sparse.linalg import SuperLU 

335 

336 solver = SuperLU(cpu_decomp).solve 

337 else: 

338 solver = cpu_decomp.solve 

339 else: 

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

341 

342 self.cached_factorizations[dt] = solver 

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

344 self.work_counters['factorizations']() 

345 

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

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

348 

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

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

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

352 _sol_hat, _ = sp.linalg.gmres( 

353 A, 

354 rhs_hat, 

355 x0=u0_hat, 

356 **self.solver_args, 

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

358 callback_type='pr_norm', 

359 M=M, 

360 ) 

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

362 _sol_hat, info = sp.linalg.cg( 

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

364 ) 

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

366 _sol_hat, info = self.linalg.bicgstab( 

367 A, 

368 rhs_hat, 

369 x0=u0_hat, 

370 **self.solver_args, 

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

372 M=M, 

373 ) 

374 else: 

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

376 

377 if info != 0: 

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

379 

380 sol_hat = self.spectral.u_init_forward 

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

382 

383 if self.spectral_space: 

384 return sol_hat 

385 else: 

386 sol = self.spectral.u_init 

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

388 

389 if self.spectral.debug: 

390 self.spectral.check_BCs(sol) 

391 

392 return sol 

393 

394 def setUpFieldsIO(self): 

395 Rectilinear.setupMPI( 

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

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

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

399 ) 

400 

401 def getOutputFile(self, fileName): 

402 self.setUpFieldsIO() 

403 

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

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

406 

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

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

409 fOut.initialize() 

410 return fOut 

411 

412 def processSolutionForOutput(self, u): 

413 if self.spectral_space: 

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

415 else: 

416 return np.array(u.real) 

417 

418 

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

420 """ 

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

422 

423 Args: 

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

425 """ 

426 

427 # get current level and problem description 

428 L = self.level 

429 

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

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

432 if stage in self.params.skip_residual_computation: 

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

434 return None 

435 

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

437 # assert L.status.updated 

438 

439 # compute the residual for each node 

440 

441 # build QF(u) 

442 res_norm = [] 

443 res = self.integrate() 

444 mask = L.prob.diff_mask 

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

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

447 # add tau if associated 

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

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

450 # use abs function from data type here 

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

452 

453 # find maximal residual over the nodes 

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

455 L.status.residual = max(res_norm) 

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

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

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

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

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

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

462 else: 

463 raise ParameterError( 

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

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

466 ) 

467 

468 # indicate that the residual has seen the new values 

469 L.status.updated = False 

470 

471 return None 

472 

473 

474def compute_residual_DAE_MPI(self, stage=None): 

475 """ 

476 Computation of the residual using the collocation matrix Q 

477 

478 Args: 

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

480 """ 

481 from mpi4py import MPI 

482 

483 L = self.level 

484 

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

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

487 if stage in self.params.skip_residual_computation: 

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

489 return None 

490 

491 # compute the residual for each node 

492 

493 # build QF(u) 

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

495 mask = L.prob.diff_mask 

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

497 # add tau if associated 

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

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

500 # use abs function from data type here 

501 res_norm = abs(res) 

502 

503 # find maximal residual over the nodes 

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

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

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

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

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

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

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

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

512 else: 

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

514 

515 # indicate that the residual has seen the new values 

516 L.status.updated = False 

517 

518 return None 

519 

520 

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

522 """ 

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

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

525 

526 Args: 

527 S (pySDC.Step): The current step 

528 

529 Returns: 

530 None 

531 """ 

532 u_ex = self.get_extrapolated_solution(S) 

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

534 if u_ex is not None: 

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

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

537 ) 

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

539 else: 

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