Coverage for pySDC / implementations / problem_classes / generic_ND_FD.py: 94%
71 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-27 07:06 +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
9from scipy.sparse.linalg import gmres, spsolve, cg
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 solver_type : str, optional
49 Type of solver. Can be ``'direct'``, ``'GMRES'`` or ``'CG'``.
50 bc : str or tuple of 2 string, optional
51 Type of boundary conditions. Default is ``'periodic'``.
52 To define two different types of boundary condition for each side,
53 you can use a tuple, for instance ``bc=("dirichlet", "neumann")``
54 uses Dirichlet BC on the left side, and Neumann BC on the right side.
55 bcParams : dict, optional
56 Parameters for boundary conditions, that can contains those keys :
58 - **val** : value for the boundary value (Dirichlet) or derivative
59 (Neumann), default to 0
60 - **reduce** : if true, reduce the order of the A matrix close to the
61 boundary. If false (default), use shifted stencils close to the
62 boundary.
63 - **neumann_bc_order** : finite difference order that should be used
64 for the neumann BC derivative. If None (default), uses the same
65 order as the discretization for A.
67 Default is None, which takes the default values for each parameters.
68 You can also define a tuple to set different parameters for each
69 side.
71 Attributes
72 ----------
73 A : sparse matrix (CSC)
74 FD discretization matrix of the ND operator.
75 Id : sparse matrix (CSC)
76 Identity matrix of the same dimension as A.
77 xvalues : np.1darray
78 Values of spatial grid.
79 """
81 dtype_u = mesh
82 dtype_f = mesh
84 def __init__(
85 self,
86 nvars=512,
87 coeff=1.0,
88 derivative=1,
89 freq=2,
90 stencil_type='center',
91 order=2,
92 lintol=1e-12,
93 liniter=10000,
94 solver_type='direct',
95 bc='periodic',
96 bcParams=None,
97 ):
98 # make sure parameters have the correct types
99 if type(nvars) not in [int, tuple]:
100 raise ProblemError('nvars should be either tuple or int')
101 if type(freq) not in [int, tuple]:
102 raise ProblemError('freq should be either tuple or int')
104 # transforms nvars into a tuple
105 if type(nvars) is int:
106 nvars = (nvars,)
108 # automatically determine ndim from nvars
109 ndim = len(nvars)
110 if ndim > 3:
111 raise ProblemError(f'can work with up to three dimensions, got {ndim}')
113 # eventually extend freq to other dimension
114 if type(freq) is int:
115 freq = (freq,) * ndim
116 if len(freq) != ndim:
117 raise ProblemError(f'len(freq)={len(freq)}, different to ndim={ndim}')
119 # check values for freq and nvars
120 for f in freq:
121 if ndim == 1 and f == -1:
122 # use Gaussian initial solution in 1D
123 bc = 'periodic'
124 break
125 if f % 2 != 0 and bc == 'periodic':
126 raise ProblemError('need even number of frequencies due to periodic BCs')
127 for nvar in nvars:
128 if nvar % 2 != 0 and bc == 'periodic':
129 raise ProblemError('the setup requires nvars = 2^p per dimension')
130 if (nvar + 1) % 2 != 0 and bc == 'dirichlet-zero':
131 raise ProblemError('setup requires nvars = 2^p - 1')
132 if ndim > 1 and nvars[1:] != nvars[:-1]:
133 raise ProblemError('need a square domain, got %s' % nvars)
135 # invoke super init, passing number of dofs
136 super().__init__(init=(nvars[0] if ndim == 1 else nvars, None, np.dtype('float64')))
138 dx, xvalues = problem_helper.get_1d_grid(size=nvars[0], bc=bc, left_boundary=0.0, right_boundary=1.0)
140 self.A, _ = problem_helper.get_finite_difference_matrix(
141 derivative=derivative,
142 order=order,
143 stencil_type=stencil_type,
144 dx=dx,
145 size=nvars[0],
146 dim=ndim,
147 bc=bc,
148 )
149 self.A *= coeff
151 self.xvalues = xvalues
152 self.Id = sp.eye(np.prod(nvars), format='csc')
154 # store attribute and register them as parameters
155 self._makeAttributeAndRegister('nvars', 'stencil_type', 'order', 'bc', localVars=locals(), readOnly=True)
156 self._makeAttributeAndRegister('freq', 'lintol', 'liniter', 'solver_type', localVars=locals())
158 if self.solver_type != 'direct':
159 self.work_counters[self.solver_type] = WorkCounter()
161 @property
162 def ndim(self):
163 """Number of dimensions of the spatial problem"""
164 return len(self.nvars)
166 @property
167 def dx(self):
168 """Size of the mesh (in all dimensions)"""
169 return self.xvalues[1] - self.xvalues[0]
171 @property
172 def grids(self):
173 """ND grids associated to the problem"""
174 x = self.xvalues
175 if self.ndim == 1:
176 return x
177 if self.ndim == 2:
178 return x[None, :], x[:, None]
179 if self.ndim == 3:
180 return x[None, :, None], x[:, None, None], x[None, None, :]
182 @classmethod
183 def get_default_sweeper_class(cls):
184 from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit
186 return generic_implicit
188 def eval_f(self, u, t):
189 """
190 Routine to evaluate the right-hand side of the problem.
192 Parameters
193 ----------
194 u : dtype_u
195 Current values.
196 t : float
197 Current time.
199 Returns
200 -------
201 f : dtype_f
202 Values of the right-hand side of the problem.
203 """
204 f = self.f_init
205 f[:] = self.A.dot(u.flatten()).reshape(self.nvars)
206 return f
208 def solve_system(self, rhs, factor, u0, t):
209 r"""
210 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.
212 Parameters
213 ----------
214 rhs : dtype_f
215 Right-hand side for the linear system.
216 factor : float
217 Abbrev. for the local stepsize (or any other factor required).
218 u0 : dtype_u
219 Initial guess for the iterative solver.
220 t : float
221 Current time (e.g. for time-dependent BCs).
223 Returns
224 -------
225 sol : dtype_u
226 The solution of the linear solver.
227 """
228 solver_type, Id, A, nvars, lintol, liniter, sol = (
229 self.solver_type,
230 self.Id,
231 self.A,
232 self.nvars,
233 self.lintol,
234 self.liniter,
235 self.u_init,
236 )
238 if solver_type == 'direct':
239 sol[:] = spsolve(Id - factor * A, rhs.flatten()).reshape(nvars)
240 elif solver_type == 'GMRES':
241 sol[:] = gmres(
242 Id - factor * A,
243 rhs.flatten(),
244 x0=u0.flatten(),
245 rtol=lintol,
246 maxiter=liniter,
247 atol=0,
248 callback=self.work_counters[solver_type],
249 callback_type='legacy',
250 )[0].reshape(nvars)
251 elif solver_type == 'CG':
252 sol[:] = cg(
253 Id - factor * A,
254 rhs.flatten(),
255 x0=u0.flatten(),
256 rtol=lintol,
257 maxiter=liniter,
258 atol=0,
259 callback=self.work_counters[solver_type],
260 )[0].reshape(nvars)
261 else:
262 raise ValueError(f'solver type "{solver_type}" not known in generic advection-diffusion implementation!')
264 return sol