Coverage for pySDC/implementations/problem_classes/generic_spectral.py: 70%
230 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-18 07:07 +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` class 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 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
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)
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)
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')
117 for base in bases:
118 self.spectral.add_axis(**base)
119 self.spectral.add_component(components)
121 self.spectral.setup_fft(real_spectral_coefficients)
123 super().__init__(init=self.spectral.init_forward if spectral_space else self.spectral.init)
125 self.work_counters[solver_type] = WorkCounter()
126 self.work_counters['factorizations'] = WorkCounter()
128 self.setup_preconditioner(Dirichlet_recombination, left_preconditioner)
130 self.cached_factorizations = {}
132 if self.heterogeneous:
133 self.__heterogeneous_setup = False
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']
140 if self.useGPU:
141 for key in CPU_only:
142 setattr(self.spectral, key, getattr(self.spectral, key).get())
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))
150 self.__heterogeneous_setup = True
152 def __getattr__(self, name):
153 """
154 Pass requests on to the helper if they are not directly attributes of this class for convenience.
156 Args:
157 name (str): Name of the attribute you want
159 Returns:
160 request
161 """
162 return getattr(self.spectral, name)
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.
168 Args:
169 LHS (dict): Equations to be added to the operator
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)
179 def setup_L(self, LHS):
180 """
181 Setup the left hand side of the linear operator L and store it in ``self.L``.
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:
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 ```
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.
194 Args:
195 LHS (dict): Dictionary containing the equations.
196 """
197 self.L = self._setup_operator(LHS)
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)
207 def setup_preconditioner(self, Dirichlet_recombination=True, left_preconditioner=True):
208 """
209 Get left and right preconditioners.
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:])
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
224 R = sp.lil_matrix((self.ncomponents * N,) * 2, dtype=int)
226 for j in range(self.ncomponents):
227 for i in range(N):
228 R[i * self.ncomponents + j, j * N + i] = 1
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)
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)
241 Pr_lhs = {comp: {comp: _Pr} for comp in self.components}
242 self.Pr = self._setup_operator(Pr_lhs) @ self.Pl.T
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``.
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.
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 """
256 sp = self.spectral.sparse_lib
258 self.heterogeneous_setup()
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
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)
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()
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
295 A = M + dt * L
296 A = Pl @ self.spectral.put_BCs_in_matrix(A) @ Pr
298 # if A.shape[0] < 200e20:
299 # import matplotlib.pyplot as plt
301 # # M = self.spectral.put_BCs_in_matrix(self.L.copy())
302 # M = A # self.L
303 # im = plt.spy(M)
304 # plt.show()
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
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')
329 if self.heterogeneous:
330 import scipy.sparse as sp
332 cpu_decomp = sp.linalg.splu(A)
333 if self.useGPU:
334 from cupyx.scipy.sparse.linalg import SuperLU
336 solver = SuperLU(cpu_decomp).solve
337 else:
338 solver = cpu_decomp.solve
339 else:
340 solver = self.spectral.linalg.factorized(A)
342 self.cached_factorizations[dt] = solver
343 self.logger.debug(f'Cached matrix factorization for {dt=:.6f}')
344 self.work_counters['factorizations']()
346 _sol_hat = self.cached_factorizations[dt](rhs_hat)
347 self.logger.debug(f'Used cached matrix factorization for {dt=:.6f}')
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__}!')
377 if info != 0:
378 self.logger.warn(f'{self.solver_type} not converged! {info=}')
380 sol_hat = self.spectral.u_init_forward
381 sol_hat[...] = (self.Pr @ _sol_hat).reshape(sol_hat.shape)
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
389 if self.spectral.debug:
390 self.spectral.check_BCs(sol)
392 return sol
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 )
401 def getOutputFile(self, fileName):
402 self.setUpFieldsIO()
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:])
407 fOut = Rectilinear(np.float64, fileName=fileName)
408 fOut.setHeader(nVar=len(self.components), coords=coords)
409 fOut.initialize()
410 return fOut
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)
419def compute_residual_DAE(self, stage=''):
420 """
421 Computation of the residual that does not add u_0 - u_m in algebraic equations.
423 Args:
424 stage (str): The current stage of the step the level belongs to
425 """
427 # get current level and problem description
428 L = self.level
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
436 # check if there are new values (e.g. from a sweep)
437 # assert L.status.updated
439 # compute the residual for each node
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]))
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 )
468 # indicate that the residual has seen the new values
469 L.status.updated = False
471 return None
474def compute_residual_DAE_MPI(self, stage=None):
475 """
476 Computation of the residual using the collocation matrix Q
478 Args:
479 stage (str): The current stage of the step the level belongs to
480 """
481 from mpi4py import MPI
483 L = self.level
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
491 # compute the residual for each node
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)
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!')
515 # indicate that the residual has seen the new values
516 L.status.updated = False
518 return None
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`.
526 Args:
527 S (pySDC.Step): The current step
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