Coverage for pySDC/implementations/problem_classes/generic_spectral.py: 84%
254 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +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
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.
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`.
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` method.
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 """
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
38 self.dtype_u = cupy_mesh
40 GPU_versions = {
41 mesh: cupy_mesh,
42 imex_mesh: imex_cupy_mesh,
43 }
45 self.dtype_f = GPU_versions[self.dtype_f]
47 if self.comm is not None:
48 from pySDC.helpers.NCCL_communicator import NCCLComm
50 if not isinstance(self.comm, NCCLComm):
51 self.__dict__['comm'] = NCCLComm(self.comm)
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.
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 preconditioner_args (dict): Arguments for the incomplete LU preconditioner of the `ilu` solver types,
82 passed to `spilu`. `drop_tol` (default 1e-3) is scaled by the step size, `fill_factor` defaults to 100
83 useGPU (bool): Run on GPU or CPU
84 max_cached_factorizations (int): Number of matrix decompositions to cache before starting eviction
85 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
86 real_spectral_coefficients (bool): If yes, allow only real values in spectral space, otherwise, allow complex.
87 heterogeneous (bool): If yes, perform memory intensive sparse matrix operations on CPU
88 debug (bool): Make additional tests at extra computational cost
89 """
90 solver_args = {} if solver_args is None else solver_args
92 preconditioner_args = {} if preconditioner_args is None else preconditioner_args
93 preconditioner_args['drop_tol'] = preconditioner_args.get('drop_tol', 1e-3)
94 preconditioner_args['fill_factor'] = preconditioner_args.get('fill_factor', 100)
96 self._makeAttributeAndRegister(
97 'max_cached_factorizations',
98 'useGPU',
99 'solver_type',
100 'solver_args',
101 'preconditioner_args',
102 'left_preconditioner',
103 'Dirichlet_recombination',
104 'comm',
105 'spectral_space',
106 'real_spectral_coefficients',
107 'heterogeneous',
108 'debug',
109 localVars=locals(),
110 )
111 self.spectral = SpectralHelper(comm=comm, useGPU=useGPU, debug=debug)
113 if useGPU:
114 self.setup_GPU()
115 if self.solver_args is not None:
116 if 'rtol' in self.solver_args.keys():
117 self.solver_args['tol'] = self.solver_args.pop('rtol')
119 for base in bases:
120 self.spectral.add_axis(**base)
121 self.spectral.add_component(components)
123 self.spectral.setup_fft(real_spectral_coefficients)
125 super().__init__(init=self.spectral.init_forward if spectral_space else self.spectral.init)
127 self.work_counters[solver_type] = WorkCounter()
128 self.work_counters['factorizations'] = WorkCounter()
130 self.setup_preconditioner(Dirichlet_recombination, left_preconditioner)
132 self.cached_factorizations = {}
133 self._restricted_BCs_hat = {}
135 if self.heterogeneous:
136 self.__heterogeneous_setup = False
138 self.logger.debug('Finished GenericSpectralLinear __init__')
140 def heterogeneous_setup(self):
141 """
142 Prepare the operators for heterogeneous runs, once and only if `heterogeneous` is set. On GPU, the BC matrices
143 `BC_line_zero_matrix` and `BCs` of the spectral helper are moved to the CPU, and CPU copies of `Pl`, `Pr`, `L`
144 and `M` are stored with the suffix `_CPU`. On CPU, the `_CPU` attributes refer to the operators themselves.
146 Returns:
147 None
148 """
149 if self.heterogeneous and not self.__heterogeneous_setup:
151 CPU_only = ['BC_line_zero_matrix', 'BCs']
152 both = ['Pl', 'Pr', 'L', 'M']
154 self.logger.debug(f'Starting heterogeneous setup. Moving {CPU_only} and copying {both} to CPU')
156 if self.useGPU:
157 for key in CPU_only:
158 self.logger.debug(f'Moving {key} to CPU')
159 setattr(self.spectral, key, getattr(self.spectral, key).get())
161 for key in both:
162 self.logger.debug(f'Copying {key} to CPU')
163 setattr(self, f'{key}_CPU', getattr(self, key).get())
164 else:
165 for key in both:
166 setattr(self, f'{key}_CPU', getattr(self, key))
168 self.logger.debug('Done with heterogeneous setup')
169 self.__heterogeneous_setup = True
171 def __getattr__(self, name):
172 """
173 Pass requests on to the helper if they are not directly attributes of this class for convenience.
175 Args:
176 name (str): Name of the attribute you want
178 Returns:
179 request
180 """
181 # While unpickling, `__dict__` is still empty, so looking up `spectral` here would recurse forever
182 if name == 'spectral' or name.startswith('__'):
183 raise AttributeError(name)
184 return getattr(self.spectral, name)
186 def _setup_operator(self, LHS):
187 """
188 Setup a sparse linear operator by adding relationships. See documentation for ``GenericSpectralLinear.setup_L`` to learn more.
190 Args:
191 LHS (dict): Equations to be added to the operator
193 Returns:
194 sparse linear operator
195 """
196 operator = self.spectral.get_empty_operator_matrix()
197 for line, equation in LHS.items():
198 self.spectral.add_equation_lhs(operator, line, equation)
199 return self.spectral.convert_operator_matrix_to_operator(operator)
201 def setup_L(self, LHS):
202 """
203 Setup the left hand side of the linear operator L and store it in ``self.L``.
205 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:
207 .. code-block:: python
209 Dx = self.get_differentiation_matrix(axes=(0,))
210 I = self.get_Id()
211 LHS = {'ux': {'u': Dx, 'ux': -I}}
212 self.setup_L(LHS)
214 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.
216 Args:
217 LHS (dict): Dictionary containing the equations.
218 """
219 self.L = self._setup_operator(LHS)
220 self.logger.debug('Set up L matrix')
222 def setup_M(self, LHS):
223 '''
224 Setup mass matrix, see documentation of ``GenericSpectralLinear.setup_L``.
225 '''
226 diff_index = list(LHS.keys())
227 self.diff_mask = [me in diff_index for me in self.components]
228 self.M = self._setup_operator(LHS)
229 self.logger.debug('Set up M matrix')
231 def setup_preconditioner(self, Dirichlet_recombination=True, left_preconditioner=True):
232 """
233 Get left and right preconditioners.
235 Args:
236 Dirichlet_recombination (bool): Basis conversion for right preconditioner. Useful for Chebychev and Ultraspherical methods. 10/10 would recommend.
237 left_preconditioner (bool): If True, it will interleave the variables and reverse the Kronecker product
238 """
239 N = np.prod(self.init[0][1:])
240 if self.useGPU:
241 from cupy_backends.cuda.libs.cusparse import CuSparseError as MemError
242 else:
243 MemError = MemoryError
245 if left_preconditioner:
246 self.logger.debug(f'Setting up left preconditioner with {N} local degrees of freedom')
248 # reverse Kronecker product
249 nc = self.ncomponents
251 rows = self.xp.arange(N * nc)
252 cols = (rows % nc) * N + rows // nc
254 self.Pl = self.spectral.sparse_lib.csc_matrix(
255 (self.xp.ones(N * nc, dtype=complex), (rows, cols)), shape=(N * nc, N * nc)
256 )
258 self.logger.debug('Finished setup of left preconditioner')
259 else:
260 Id = self.spectral.sparse_lib.eye(N)
261 Pl_lhs = {comp: {comp: Id} for comp in self.components}
262 self.Pl = self._setup_operator(Pl_lhs)
264 if Dirichlet_recombination and type(self.axes[-1]).__name__ in ['ChebychevHelper', 'UltrasphericalHelper']:
265 self.logger.debug('Using Dirichlet recombination as right preconditioner')
266 _Pr = self.spectral.get_Dirichlet_recombination_matrix(axis=-1)
267 else:
268 _Pr = self.spectral.sparse_lib.eye(N)
270 Pr_lhs = {comp: {comp: _Pr} for comp in self.components}
271 operator = self._setup_operator(Pr_lhs)
273 try:
274 self.Pr = operator @ self.Pl.T
275 except MemError:
276 self.logger.debug('Setting up right preconditioner on CPU due to memory error')
277 self.Pr = self.spectral.sparse_lib.csc_matrix(operator.get() @ self.Pl.T.get())
279 def use_restricted_boundary_values(self, u, t):
280 """
281 Impose the boundary values of the restricted solution `u` in the solves at the times `t` instead of the data,
282 see :meth:`pySDC.core.problem.Problem.use_restricted_boundary_values`. The values are the boundary condition
283 operator applied to the coefficients, so this works for any boundary condition and for spatial coarsening, where
284 the boundary values of the restricted coefficients need not be the restricted boundary values.
286 Args:
287 u (list): Restricted solution at the coarse nodes
288 t (list): Times of the coarse nodes
289 """
290 self.heterogeneous_setup()
291 on_CPU = self.heterogeneous and self.useGPU # then the BC operator lives on the CPU
293 self._restricted_BCs_hat = {}
294 for _u, _t in zip(u, t, strict=True):
295 u_hat = (_u if self.spectral_space else self.spectral.transform(_u)).flatten()
296 BCs_hat = self.spectral.BCs @ (u_hat.get() if on_CPU else u_hat)
297 self._restricted_BCs_hat[_t] = self.xp.asarray(BCs_hat).reshape(self.spectral.rhs_BCs_hat.shape)
299 def solve_system(self, rhs, dt, u0=None, t=None, *args, skip_itransform=False, **kwargs):
300 """
301 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
302 ``GenericSpectralLinear.setup_L`` and ``GenericSpectralLinear.setup_M``.
304 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.
305 This means solving with dt=0 to mimic explicit methods does not work for all problems, in particular simple DAEs.
307 Note that by putting M rhs on the right hand side, this function can only solve algebraic conditions equal to
308 zero. If you want something else, it should be easy to overload this function.
309 """
311 sp = self.spectral.sparse_lib
313 self.heterogeneous_setup()
315 if self.spectral_space:
316 rhs_hat = rhs.copy()
317 if u0 is not None:
318 u0_hat = u0.copy().flatten()
319 else:
320 u0_hat = None
321 else:
322 rhs_hat = self.spectral.transform(rhs)
323 if u0 is not None:
324 u0_hat = self.spectral.transform(u0).flatten()
325 else:
326 u0_hat = None
328 # apply inverse right preconditioner to initial guess
329 if u0_hat is not None and 'direct' not in self.solver_type:
330 if not hasattr(self, '_Pr_inv'):
331 self._Pr_inv = self.linalg.splu(self.Pr.astype(complex)).solve
332 u0_hat[...] = self._Pr_inv(u0_hat)
334 rhs_hat = (self.M @ rhs_hat.flatten()).reshape(rhs_hat.shape)
335 rhs_hat = self.spectral.put_BCs_in_rhs_hat(rhs_hat, self._restricted_BCs_hat.get(t))
336 rhs_hat = self.Pl @ rhs_hat.flatten()
338 if dt not in self.cached_factorizations.keys() or not self.solver_type.lower() == 'cached_direct':
339 if self.heterogeneous:
340 M = self.M_CPU
341 L = self.L_CPU
342 Pl = self.Pl_CPU
343 Pr = self.Pr_CPU
344 else:
345 M = self.M
346 L = self.L
347 Pl = self.Pl
348 Pr = self.Pr
350 A = M + dt * L
351 A = Pl @ self.spectral.put_BCs_in_matrix(A) @ Pr
353 # if A.shape[0] < 200e20:
354 # import matplotlib.pyplot as plt
356 # # M = self.spectral.put_BCs_in_matrix(self.L.copy())
357 # M = A # self.L
358 # im = plt.spy(M)
359 # plt.show()
361 if 'ilu' in self.solver_type.lower():
362 if dt not in self.cached_factorizations.keys():
363 if len(self.cached_factorizations) >= self.max_cached_factorizations:
364 to_evict = list(self.cached_factorizations.keys())[0]
365 self.cached_factorizations.pop(to_evict)
366 self.logger.debug(f'Evicted matrix factorization for {to_evict=:.6f} from cache')
367 iLU = self.linalg.spilu(
368 A, **{**self.preconditioner_args, 'drop_tol': dt * self.preconditioner_args['drop_tol']}
369 )
370 self.cached_factorizations[dt] = self.linalg.LinearOperator(A.shape, iLU.solve)
371 self.logger.debug(f'Cached incomplete LU factorization for {dt=:.6f}')
372 self.work_counters['factorizations']()
373 M = self.cached_factorizations[dt]
374 else:
375 M = None
376 info = 0
378 if self.solver_type.lower() == 'cached_direct':
379 if dt not in self.cached_factorizations.keys():
380 if len(self.cached_factorizations) >= self.max_cached_factorizations:
381 self.cached_factorizations.pop(list(self.cached_factorizations.keys())[0])
382 self.logger.debug(f'Evicted matrix factorization for {dt=:.6f} from cache')
384 if self.heterogeneous:
385 import scipy.sparse as sp
387 cpu_decomp = sp.linalg.splu(A)
389 if self.useGPU:
390 from cupyx.scipy.sparse.linalg import SuperLU
392 solver = SuperLU(cpu_decomp).solve
393 else:
394 solver = cpu_decomp.solve
395 else:
396 solver = self.spectral.linalg.factorized(A)
398 self.cached_factorizations[dt] = solver
399 self.logger.debug(f'Cached matrix factorization for {dt=:.6f}')
400 self.work_counters['factorizations']()
402 _sol_hat = self.cached_factorizations[dt](rhs_hat)
403 self.work_counters[self.solver_type]()
404 self.logger.debug(f'Used cached matrix factorization for {dt=:.6f}')
406 elif self.solver_type.lower() == 'direct':
407 _sol_hat = sp.linalg.spsolve(A, rhs_hat)
408 elif 'gmres' in self.solver_type.lower():
409 _sol_hat, _ = sp.linalg.gmres(
410 A,
411 rhs_hat,
412 x0=u0_hat,
413 **self.solver_args,
414 callback=self.work_counters[self.solver_type],
415 callback_type='pr_norm',
416 M=M,
417 )
418 elif self.solver_type.lower() == 'cg':
419 _sol_hat, info = sp.linalg.cg(
420 A, rhs_hat, x0=u0_hat, **self.solver_args, callback=self.work_counters[self.solver_type]
421 )
422 elif 'bicgstab' in self.solver_type.lower():
423 _sol_hat, info = self.linalg.bicgstab(
424 A,
425 rhs_hat,
426 x0=u0_hat,
427 **self.solver_args,
428 callback=self.work_counters[self.solver_type],
429 M=M,
430 )
431 else:
432 raise NotImplementedError(f'Solver {self.solver_type=} not implemented in {type(self).__name__}!')
434 if info != 0:
435 self.logger.warn(f'{self.solver_type} not converged! {info=}')
437 sol_hat = self.spectral.u_init_forward
438 sol_hat[...] = (self.Pr @ _sol_hat).reshape(sol_hat.shape)
440 if self.spectral_space:
441 return sol_hat
442 else:
443 sol = self.spectral.u_init
444 sol[:] = self.spectral.itransform(sol_hat).real
446 if self.spectral.debug:
447 self.spectral.check_BCs(sol)
449 return sol
451 def setUpFieldsIO(self):
452 """
453 Set up the MPI mode of `Rectilinear` output files with the local slice of this problem in physical space.
455 Returns:
456 None
457 """
458 Rectilinear.setupMPI(
459 comm=self.comm.commMPI if self.useGPU else self.comm,
460 iLoc=[me.start for me in self.local_slice(False)],
461 nLoc=[me.stop - me.start for me in self.local_slice(False)],
462 )
464 def getOutputFile(self, fileName):
465 """
466 Set up a `Rectilinear` output file on the grid of this problem, with one variable per component.
468 Args:
469 fileName (str): Name of the file
471 Returns:
472 pySDC.helpers.fieldsIO.Rectilinear: The initialized output file
473 """
474 self.setUpFieldsIO()
476 coords = [me.get_1dgrid() for me in self.spectral.axes]
477 if self.spectral.useGPU:
478 coords = [me.get() for me in coords]
480 assert np.allclose([len(me) for me in coords], self.spectral.global_shape[1:])
482 fOut = Rectilinear(np.float64, fileName=fileName)
483 fOut.setHeader(nVar=len(self.components), coords=coords)
484 fOut.initialize()
485 return fOut
487 def processSolutionForOutput(self, u):
488 """
489 Prepare a solution for output: its real part in physical space, moved to the CPU if running on GPU.
491 Args:
492 u (dtype_u): Solution, in spectral space if `spectral_space` is set, else in physical space
494 Returns:
495 numpy.ndarray: The solution as contiguous array
496 """
497 if self.spectral_space:
498 u = self.itransform(u).real
499 else:
500 u = u.real
502 if self.spectral.useGPU:
503 u = u.get()
505 return np.ascontiguousarray(u.view(np.ndarray))
508def compute_residual_DAE(self, stage=''):
509 """
510 Computation of the residual that does not add u_0 - u_m in algebraic equations.
512 Args:
513 stage (str): The current stage of the step the level belongs to
514 """
516 # get current level and problem description
517 L = self.level
519 # Check if we want to skip the residual computation to gain performance
520 # Keep in mind that skipping any residual computation is likely to give incorrect outputs of the residual!
521 if stage in self.params.skip_residual_computation:
522 L.status.residual = 0.0 if L.status.residual is None else L.status.residual
523 return None
525 # check if there are new values (e.g. from a sweep)
526 # assert L.status.updated
528 # compute the residual for each node
530 # build QF(u)
531 res_norm = []
532 res = self.integrate()
533 mask = L.prob.diff_mask
534 for m in range(self.coll.num_nodes):
535 res[m][mask] += L.u[0][mask] - L.u[m + 1][mask]
536 # add tau if associated
537 if L.tau[m] is not None:
538 res[m] += L.tau[m]
539 # use abs function from data type here
540 res_norm.append(abs(res[m]))
542 # find maximal residual over the nodes
543 if L.params.residual_type == 'full_abs':
544 L.status.residual = max(res_norm)
545 elif L.params.residual_type == 'last_abs':
546 L.status.residual = res_norm[-1]
547 elif L.params.residual_type == 'full_rel':
548 L.status.residual = max(res_norm) / abs(L.u[0])
549 elif L.params.residual_type == 'last_rel':
550 L.status.residual = res_norm[-1] / abs(L.u[0])
551 else:
552 raise ParameterError(
553 f'residual_type = {L.params.residual_type} not implemented, choose '
554 f'full_abs, last_abs, full_rel or last_rel instead'
555 )
557 # indicate that the residual has seen the new values
558 L.status.updated = False
560 return None
563def compute_residual_DAE_MPI(self, stage=None):
564 """
565 Computation of the residual using the collocation matrix Q
567 Args:
568 stage (str): The current stage of the step the level belongs to
569 """
570 from mpi4py import MPI
572 L = self.level
574 # Check if we want to skip the residual computation to gain performance
575 # Keep in mind that skipping any residual computation is likely to give incorrect outputs of the residual!
576 if stage in self.params.skip_residual_computation:
577 L.status.residual = 0.0 if L.status.residual is None else L.status.residual
578 return None
580 # compute the residual for each node
582 # build QF(u)
583 res = self.integrate(last_only=L.params.residual_type[:4] == 'last')
584 mask = L.prob.diff_mask
585 res[mask] += L.u[0][mask] - L.u[self.rank + 1][mask]
586 # add tau if associated
587 if L.tau[self.rank] is not None:
588 res += L.tau[self.rank]
589 # use abs function from data type here
590 res_norm = abs(res)
592 # find maximal residual over the nodes
593 if L.params.residual_type == 'full_abs':
594 L.status.residual = self.comm.allreduce(res_norm, op=MPI.MAX)
595 elif L.params.residual_type == 'last_abs':
596 L.status.residual = self.comm.bcast(res_norm, root=self.comm.size - 1)
597 elif L.params.residual_type == 'full_rel':
598 L.status.residual = self.comm.allreduce(res_norm / abs(L.u[0]), op=MPI.MAX)
599 elif L.params.residual_type == 'last_rel':
600 L.status.residual = self.comm.bcast(res_norm / abs(L.u[0]), root=self.comm.size - 1)
601 else:
602 raise NotImplementedError(f'residual type \"{L.params.residual_type}\" not implemented!')
604 # indicate that the residual has seen the new values
605 L.status.updated = False
607 return None