Coverage for pySDC/implementations/problem_classes/generic_ND_FD.py: 98%
96 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-25 20:28 +0000
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3"""
4Created on Sat Feb 11 22:39:30 2023
5"""
7import numpy as np
8import scipy.sparse as sp
9import scipy.sparse.linalg as spla
11from pySDC.core.errors import ProblemError
12from pySDC.core.problem import Problem, WorkCounter
13from pySDC.helpers import problem_helper
14from pySDC.implementations.datatype_classes.mesh import mesh
17class GenericNDimFinDiff(Problem):
18 r"""
19 Base class for finite difference spatial discretisation in :math:`N` dimensions
21 .. math::
22 \frac{d u}{dt} = A u,
24 where :math:`A \in \mathbb{R}^{nN \times nN}` is a matrix arising from finite difference discretisation of spatial
25 derivatives with :math:`n` degrees of freedom per dimension and :math:`N` dimensions. This generic class follows the MOL
26 (method-of-lines) approach and can be used to discretize partial differential equations such as the advection
27 equation and the heat equation.
29 Parameters
30 ----------
31 nvars : int, optional
32 Spatial resolution for the ND problem. For :math:`N = 2`,
33 set ``nvars=(16, 16)``.
34 coeff : float, optional
35 Factor for finite difference matrix :math:`A`.
36 derivative : int, optional
37 Order of the spatial derivative.
38 freq : tuple of int, optional
39 Spatial frequency, can be a tuple.
40 stencil_type : str, optional
41 Stencil type for finite differences.
42 order : int, optional
43 Order of accuracy of the finite difference discretization.
44 lintol : float, optional
45 Tolerance for spatial solver.
46 liniter : int, optional
47 Maximum number of iterations for linear solver.
48 dtype : dtype-like, optional
49 Precision the state is stored at. ``float64`` by default, which is what every caller got
50 before this existed. The operators follow at ``promote_types(dtype, float32)``, since SciPy
51 has no half-precision sparse matrix and hardware that stores half precision computes in
52 single anyway -- so ``float16`` here means genuinely half-precision *storage* with
53 single-precision arithmetic, which is the arrangement it has on real hardware too.
54 solver_type : str, optional
55 Type of solver. Can be ``'direct'``, ``'GMRES'`` or ``'CG'``.
56 bc : str or tuple of 2 string, optional
57 Type of boundary conditions. Default is ``'periodic'``.
58 To define two different types of boundary condition for each side,
59 you can use a tuple, for instance ``bc=("dirichlet", "neumann")``
60 uses Dirichlet BC on the left side, and Neumann BC on the right side.
61 bcParams : dict, optional
62 Parameters for boundary conditions, that can contains those keys :
64 - **val** : value for the boundary value (Dirichlet) or derivative
65 (Neumann), default to 0
66 - **reduce** : if true, reduce the order of the A matrix close to the
67 boundary. If false (default), use shifted stencils close to the
68 boundary.
69 - **neumann_bc_order** : finite difference order that should be used
70 for the neumann BC derivative. If None (default), uses the same
71 order as the discretization for A.
73 Default is None, which takes the default values for each parameters.
74 You can also define a tuple to set different parameters for each
75 side.
77 Attributes
78 ----------
79 A : sparse matrix (CSC)
80 FD discretization matrix of the ND operator.
81 Id : sparse matrix (CSC)
82 Identity matrix of the same dimension as A.
83 xvalues : np.1darray
84 Values of spatial grid.
85 """
87 dtype_u = mesh
88 dtype_f = mesh
89 xp = np
90 xsp = sp
91 linalg = spla
93 def setup_GPU(self):
94 """
95 Switch to GPU modules
96 """
97 import cupy as cp
98 import cupyx.scipy.sparse as csp
99 import cupyx.scipy.sparse.linalg as cspla
101 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh
103 self.xp = cp
104 self.xsp = csp
105 self.linalg = cspla
106 self.dtype_u = cupy_mesh
107 self.dtype_f = cupy_mesh
109 def __init__(
110 self,
111 nvars=512,
112 coeff=1.0,
113 derivative=1,
114 freq=2,
115 stencil_type='center',
116 order=2,
117 lintol=1e-12,
118 liniter=10000,
119 solver_type='direct',
120 bc='periodic',
121 bcParams=None,
122 dtype='float64',
123 useGPU=False,
124 ):
125 if useGPU:
126 self.setup_GPU()
128 # make sure parameters have the correct types
129 if type(nvars) not in [int, tuple]:
130 raise ProblemError('nvars should be either tuple or int')
131 if type(freq) not in [int, tuple]:
132 raise ProblemError('freq should be either tuple or int')
134 # transforms nvars into a tuple
135 if type(nvars) is int:
136 nvars = (nvars,)
138 # automatically determine ndim from nvars
139 ndim = len(nvars)
140 if ndim > 3:
141 raise ProblemError(f'can work with up to three dimensions, got {ndim}')
143 # eventually extend freq to other dimension
144 if type(freq) is int:
145 freq = (freq,) * ndim
146 if len(freq) != ndim:
147 raise ProblemError(f'len(freq)={len(freq)}, different to ndim={ndim}')
149 # check values for freq and nvars
150 for f in freq:
151 if ndim == 1 and f == -1:
152 # use Gaussian initial solution in 1D
153 bc = 'periodic'
154 break
155 if f % 2 != 0 and bc == 'periodic':
156 raise ProblemError('need even number of frequencies due to periodic BCs')
157 for nvar in nvars:
158 if nvar % 2 != 0 and bc == 'periodic':
159 raise ProblemError('the setup requires nvars = 2^p per dimension')
160 if (nvar + 1) % 2 != 0 and bc == 'dirichlet-zero':
161 raise ProblemError('setup requires nvars = 2^p - 1')
162 if ndim > 1 and nvars[1:] != nvars[:-1]:
163 raise ProblemError('need a square domain, got %s' % nvars)
165 # invoke super init, passing number of dofs and the precision to store them at
166 dtype = np.dtype(dtype)
168 # SciPy holds no float16 sparse matrix, and hardware that stores half precision computes in
169 # single anyway, so the operators sit at the smallest single-or-wider type that holds `dtype`
170 operator_dtype = np.promote_types(dtype, np.float32)
172 super().__init__(init=(nvars[0] if ndim == 1 else nvars, None, dtype))
174 dx, xvalues = problem_helper.get_1d_grid(size=nvars[0], bc=bc, left_boundary=0.0, right_boundary=1.0)
176 self.A, _ = problem_helper.get_finite_difference_matrix(
177 derivative=derivative,
178 order=order,
179 stencil_type=stencil_type,
180 dx=dx,
181 size=nvars[0],
182 dim=ndim,
183 bc=bc,
184 cupy=useGPU,
185 )
186 self.A *= coeff
188 self.A = self.A.astype(operator_dtype)
190 # SciPy's sparse direct solver wants CSC and CuPy's wants CSR. Whichever one is handed the
191 # wrong layout converts the whole matrix on every call -- that is what cupyx's
192 # `SparseEfficiencyWarning: CSR format is required` is reporting, once per solve. CSR is
193 # also the better layout for the matrix-vector product in `eval_f`.
194 self.A = self.A.tocsr() if useGPU else self.A.tocsc()
196 # the grid feeds every `u_exact`, so it has to live where the solution does
197 self.xvalues = self.xp.asarray(xvalues)
198 self.Id = self.xsp.eye(np.prod(nvars), format='csr' if useGPU else 'csc', dtype=operator_dtype)
200 # store attribute and register them as parameters
201 self._makeAttributeAndRegister('nvars', 'stencil_type', 'order', 'bc', localVars=locals(), readOnly=True)
202 self._makeAttributeAndRegister('freq', 'lintol', 'liniter', 'solver_type', localVars=locals())
203 self.dtype = dtype
204 self.operator_dtype = operator_dtype
206 if self.solver_type != 'direct':
207 self.work_counters[self.solver_type] = WorkCounter()
209 @property
210 def ndim(self):
211 """Number of dimensions of the spatial problem"""
212 return len(self.nvars)
214 @property
215 def dx(self):
216 """Size of the mesh (in all dimensions)"""
217 return self.xvalues[1] - self.xvalues[0]
219 @property
220 def grids(self):
221 """ND grids associated to the problem"""
222 x = self.xvalues
223 if self.ndim == 1:
224 return x
225 if self.ndim == 2:
226 return x[None, :], x[:, None]
227 if self.ndim == 3:
228 return x[None, :, None], x[:, None, None], x[None, None, :]
230 @classmethod
231 def get_default_sweeper_class(cls):
232 from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
234 return generic_implicit
236 def eval_f(self, u, t):
237 """
238 Routine to evaluate the right-hand side of the problem.
240 Parameters
241 ----------
242 u : dtype_u
243 Current values.
244 t : float
245 Current time.
247 Returns
248 -------
249 f : dtype_f
250 Values of the right-hand side of the problem.
251 """
252 f = self.f_init
253 f[:] = self.A.dot(u.flatten()).reshape(self.nvars)
254 return f
256 def eval_f_increment(self, base, delta, t):
257 r"""
258 Evaluate :math:`f(w + \delta) - f(w) = A\delta`, which carries an explicit factor
259 :math:`\delta`.
261 The operator is linear, so the increment is the operator applied to the correction and the
262 base state does not enter. Supplying it means a sweeper never has to form the increment by
263 subtracting two stored right-hand sides, whose cancellation error carries
264 :math:`\varepsilon\|A\|` -- see :class:`pySDC.core.problem.Problem`.
266 Parameters
267 ----------
268 base : dtype_u
269 The base state, unused for a linear operator.
270 delta : dtype_u
271 The correction.
272 t : float
273 Current time, accepted for interface compatibility.
275 Returns
276 -------
277 f : dtype_f
278 The increment.
279 """
280 f = self.dtype_f(self.init)
281 f[:] = self.A.dot(delta.flatten()).reshape(self.nvars)
282 return f
284 def solve_system(self, rhs, factor, u0, t):
285 r"""
286 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.
288 Parameters
289 ----------
290 rhs : dtype_f
291 Right-hand side for the linear system.
292 factor : float
293 Abbrev. for the local stepsize (or any other factor required).
294 u0 : dtype_u
295 Initial guess for the iterative solver.
296 t : float
297 Current time (e.g. for time-dependent BCs).
299 Returns
300 -------
301 sol : dtype_u
302 The solution of the linear solver.
303 """
304 solver_type, Id, A, nvars, lintol, liniter, sol = (
305 self.solver_type,
306 self.Id,
307 self.A,
308 self.nvars,
309 self.lintol,
310 self.liniter,
311 self.u_init,
312 )
314 if solver_type == 'direct':
315 sol[:] = self.linalg.spsolve(Id - factor * A, rhs.flatten()).reshape(nvars)
316 elif solver_type == 'GMRES':
317 sol[:] = self.linalg.gmres(
318 Id - factor * A,
319 rhs.flatten(),
320 x0=u0.flatten(),
321 rtol=lintol,
322 maxiter=liniter,
323 atol=0,
324 callback=self.work_counters[solver_type],
325 callback_type='legacy',
326 )[0].reshape(nvars)
327 elif solver_type == 'CG':
328 sol[:] = self.linalg.cg(
329 Id - factor * A,
330 rhs.flatten(),
331 x0=u0.flatten(),
332 rtol=lintol,
333 maxiter=liniter,
334 atol=0,
335 callback=self.work_counters[solver_type],
336 )[0].reshape(nvars)
337 else:
338 raise ValueError(f'solver type "{solver_type}" not known in generic advection-diffusion implementation!')
340 return sol