Coverage for pySDC/implementations/problem_classes/generic_ND

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

9import scipy.sparse as sp

10import scipy.sparse.linalg as spla

12from pySDC.core.errors import ProblemError

13from pySDC.core.problem import Problem, WorkCounter

14from pySDC.helpers import problem_helper

15from pySDC.implementations.datatype_classes.mesh import mesh

18class GenericNDimFinDiff(Problem):

19 r"""

20 Base class for finite difference spatial discretisation in :math:`N` dimensions

22 .. math::

23 \frac{d u}{dt} = A u,

25 where :math:`A \in \mathbb{R}^{nN \times nN}` is a matrix arising from finite difference discretisation of spatial

26 derivatives with :math:`n` degrees of freedom per dimension and :math:`N` dimensions. This generic class follows the MOL

27 (method-of-lines) approach and can be used to discretize partial differential equations such as the advection

28 equation and the heat equation.

30 Parameters

31 ----------

32 nvars : int, optional

33 Spatial resolution for the ND problem. For :math:`N = 2`,

34 set ``nvars=(16, 16)``.

35 coeff : float, optional

36 Factor for finite difference matrix :math:`A`.

37 derivative : int, optional

38 Order of the spatial derivative.

39 freq : tuple of int, optional

40 Spatial frequency, can be a tuple.

41 stencil_type : str, optional

42 Stencil type for finite differences.

43 order : int, optional

44 Order of accuracy of the finite difference discretization.

45 lintol : float, optional

46 Tolerance for spatial solver.

47 liniter : int, optional

48 Maximum number of iterations for linear solver.

49 dtype : dtype-like, optional

50 Precision the state is stored at. ``float64`` by default, which is what every caller got

51 before this existed. The operators follow at ``promote_types(dtype, float32)``, since SciPy

52 has no half-precision sparse matrix and hardware that stores half precision computes in

53 single anyway -- so ``float16`` here means genuinely half-precision *storage* with

54 single-precision arithmetic, which is the arrangement it has on real hardware too.

55 solver_type : str, optional

56 Type of solver. Can be ``'direct'``, ``'GMRES'``, ``'CG'`` or, on periodic grids, ``'FFT'``:

57 the operator is then circulant, so the FFT diagonalises it and a solve is two transforms

58 and a division -- exact, and the cheapest solve there is on a GPU. ``'direct'`` keeps the

59 factorizations of the last ``max_cached_factorizations`` factors and reuses them.

60 solve_precision : dtype-like or None, optional

61 Precision the implicit solve runs at, independently of the state's ``dtype``. ``None``

62 (default) solves at the operator's precision, exactly as without it. Otherwise the solve

63 genuinely works in that precision: the sparse solvers on a copy of the operator held at it,

64 ``'FFT'`` with transforms at it, and the result is widened back to ``dtype``. Sparse

65 matrices stop at single precision in SciPy and CuPy alike, so ``float16`` needs ``'FFT'``.

66 On a GPU those are cuFFT's half-precision transforms (power-of-two grids), with every value

67 stored and every product rounded in half precision. A CPU has no half-precision FFT, so

68 there ``float16`` is *emulated*, by rounding through it around a single-precision

69 transform -- which is optimistic, since the transform's own arithmetic is then single.

70 For a complex state the solve is complex at the same precision.

71 bc : str or tuple of 2 string, optional

72 Type of boundary conditions. Default is ``'periodic'``.

73 To define two different types of boundary condition for each side,

74 you can use a tuple, for instance ``bc=("dirichlet", "neumann")``

75 uses Dirichlet BC on the left side, and Neumann BC on the right side.

76 bcParams : dict, optional

77 Parameters for boundary conditions, that can contains those keys :

79 - **val** : value for the boundary value (Dirichlet) or derivative

80 (Neumann), default to 0

81 - **reduce** : if true, reduce the order of the A matrix close to the

82 boundary. If false (default), use shifted stencils close to the

83 boundary.

84 - **neumann_bc_order** : finite difference order that should be used

85 for the neumann BC derivative. If None (default), uses the same

86 order as the discretization for A.

88 Default is None, which takes the default values for each parameters.

89 You can also define a tuple to set different parameters for each

90 side.

91 useGPU : bool, optional

92 Run on the GPU with CuPy instead of on the CPU with NumPy.

94 Attributes

95 ----------

96 A : sparse matrix (CSC)

97 FD discretization matrix of the ND operator.

98 Id : sparse matrix (CSC)

99 Identity matrix of the same dimension as A.

100 xvalues : np.1darray

101 Values of spatial grid.

102 """

103

104 dtype_u = mesh

105 dtype_f = mesh

106 xp = np

107 xsp = sp

108 linalg = spla

109 # direct solves keep this many factorizations of I - factor * A, see `_direct_solve`

110 max_cached_factorizations = 12

111

112 def setup_GPU(self):

113 """

114 Switch to GPU modules

115 """

116 import cupy as cp

117 import cupyx.scipy.sparse as csp

118 import cupyx.scipy.sparse.linalg as cspla

119

120 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh

121

122 self.xp = cp

123 self.xsp = csp

124 self.linalg = cspla

125 self.dtype_u = cupy_mesh

126 self.dtype_f = cupy_mesh

127

128 def __init__(

129 self,

130 nvars=512,

131 coeff=1.0,

132 derivative=1,

133 freq=2,

134 stencil_type='center',

135 order=2,

136 lintol=1e-12,

137 liniter=10000,

138 solver_type='direct',

139 bc='periodic',

140 bcParams=None,

141 dtype='float64',

142 useGPU=False,

143 solve_precision=None,

144 ):

145 """Initialization routine"""

146 if useGPU:

147 self.setup_GPU()

148

149 # make sure parameters have the correct types

150 if type(nvars) not in [int, tuple]:

151 raise ProblemError('nvars should be either tuple or int')

152 if type(freq) not in [int, tuple]:

153 raise ProblemError('freq should be either tuple or int')

154

155 # transforms nvars into a tuple

156 if type(nvars) is int:

157 nvars = (nvars,)

158

159 # automatically determine ndim from nvars

160 ndim = len(nvars)

161 if ndim > 3:

162 raise ProblemError(f'can work with up to three dimensions, got {ndim}')

163

164 # eventually extend freq to other dimension

165 if type(freq) is int:

166 freq = (freq,) * ndim

167 if len(freq) != ndim:

168 raise ProblemError(f'len(freq)={len(freq)}, different to ndim={ndim}')

169

170 # check values for freq and nvars

171 for f in freq:

172 if ndim == 1 and f == -1:

173 # use Gaussian initial solution in 1D

174 bc = 'periodic'

175 break

176 if f % 2 != 0 and bc == 'periodic':

177 raise ProblemError('need even number of frequencies due to periodic BCs')

178 for nvar in nvars:

179 if nvar % 2 != 0 and bc == 'periodic':

180 raise ProblemError('the setup requires nvars = 2^p per dimension')

181 if (nvar + 1) % 2 != 0 and bc == 'dirichlet-zero':

182 raise ProblemError('setup requires nvars = 2^p - 1')

183 if ndim > 1 and nvars[1:] != nvars[:-1]:

184 raise ProblemError('need a square domain, got %s' % nvars)

185

186 # invoke super init, passing number of dofs and the precision to store them at

187 dtype = np.dtype(dtype)

188

189 # SciPy holds no float16 sparse matrix, and hardware that stores half precision computes in

190 # single anyway, so the operators sit at the smallest single-or-wider type that holds `dtype`

191 operator_dtype = np.promote_types(dtype, np.float32)

192

193 super().__init__(init=(nvars[0] if ndim == 1 else nvars, None, dtype))

194

195 dx, xvalues = problem_helper.get_1d_grid(size=nvars[0], bc=bc, left_boundary=0.0, right_boundary=1.0)

196

197 self.A, _ = problem_helper.get_finite_difference_matrix(

198 derivative=derivative,

199 order=order,

200 stencil_type=stencil_type,

201 dx=dx,

202 size=nvars[0],

203 dim=ndim,

204 bc=bc,

205 cupy=useGPU,

206 )

207 self.A *= coeff

208

209 self.A = self.A.astype(operator_dtype)

210

211 # SciPy's sparse direct solver wants CSC and CuPy's wants CSR. Whichever one is handed the

212 # wrong layout converts the whole matrix on every call -- that is what cupyx's

213 # `SparseEfficiencyWarning: CSR format is required` is reporting, once per solve. CSR is

214 # also the better layout for the matrix-vector product in `eval_f`.

215 self.A = self.A.tocsr() if useGPU else self.A.tocsc()

216

217 # the grid feeds every `u_exact`, so it has to live where the solution does

218 self.xvalues = self.xp.asarray(xvalues)

219 self.Id = self.xsp.eye(np.prod(nvars), format='csr' if useGPU else 'csc', dtype=operator_dtype)

220

221 # store attribute and register them as parameters

222 self._makeAttributeAndRegister('nvars', 'stencil_type', 'order', 'bc', localVars=locals(), readOnly=True)

223 self._makeAttributeAndRegister('freq', 'lintol', 'liniter', 'solver_type', localVars=locals())

224 self.dtype = dtype

225 self.operator_dtype = operator_dtype

226

227 if self.solver_type != 'direct':

228 self.work_counters[self.solver_type] = WorkCounter()

229

230 self._factorizations = {}

231

232 solve_precision = None if solve_precision is None else np.dtype(solve_precision)

233 self._makeAttributeAndRegister('solve_precision', localVars=locals(), readOnly=True)

234 self._setup_solve(solver_type, solve_precision, bc)

235

236 def _setup_solve(self, solver_type, solve_precision, bc):

237 """Prepare the reduced-precision operator, or the eigenvalues the FFT solve divides by."""

238 # half precision is the one format that is not a complex dtype for a complex state

239 self._half = solve_precision == np.float16

240 precision = self.operator_dtype if solve_precision is None or self._half else solve_precision

241 if self.dtype.kind == 'c':

242 precision = np.promote_types(precision, np.complex64)

243 self._solve_dtype = precision

244

245 if solver_type == 'FFT':

246 if bc != 'periodic':

247 raise ProblemError(f'the FFT solve diagonalises a circulant operator and needs periodic BCs, got {bc}')

248 if self._half and self.xp is not np and any(n & (n - 1) for n in self.nvars):

249 raise ProblemError(f'cuFFT computes in half precision on power-of-two grids only, got {self.nvars}')

250 self._fft = scipy.fft if self.xp is np else self.xp.fft

251 # the first column of a circulant transforms to its eigenvalues

252 first_column = self.xp.zeros(self.A.shape[0], dtype=self.operator_dtype)

253 first_column[0] = 1.0

254 self._eigenvalues = self._fft.fftn((self.A @ first_column).reshape(self.nvars))

255 self._half_plan = None

256 elif solve_precision is not None:

257 if self._half:

258 raise ProblemError('sparse matrices stop at single precision; float16 needs solver_type="FFT"')

259 self._A_solve = self.A.astype(precision)

260 self._Id_solve = self.Id.astype(precision)

261

262 @property

263 def ndim(self):

264 """Number of dimensions of the spatial problem"""

265 return len(self.nvars)

266

267 @property

268 def dx(self):

269 """Size of the mesh (in all dimensions)"""

270 return self.xvalues[1] - self.xvalues[0]

271

272 @property

273 def grids(self):

274 """ND grids associated to the problem"""

275 x = self.xvalues

276 if self.ndim == 1:

277 return x

278 if self.ndim == 2:

279 return x[None, :], x[:, None]

280 if self.ndim == 3:

281 return x[None, :, None], x[:, None, None], x[None, None, :]

282

283 @classmethod

284 def get_default_sweeper_class(cls):

285 """

286 Default sweeper for these problems, which are treated fully implicitly.

287

288 Returns

289 -------

290 type

291 The sweeper class ``generic_implicit``.

292 """

293 from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit

294

295 return generic_implicit

296

297 def eval_f(self, u, t):

298 """

299 Routine to evaluate the right-hand side of the problem.

300

301 Parameters

302 ----------

303 u : dtype_u

304 Current values.

305 t : float

306 Current time.

307

308 Returns

309 -------

310 f : dtype_f

311 Values of the right-hand side of the problem.

312 """

313 f = self.f_init

314 f[:] = self.A.dot(u.flatten()).reshape(self.nvars)

315 return f

316

317 def eval_f_increment(self, base, delta, t):

318 r"""

319 Evaluate :math:`f(w + \delta) - f(w) = A\delta`, which carries an explicit factor

320 :math:`\delta`.

321

322 The operator is linear, so the increment is the operator applied to the correction and the

323 base state does not enter. Supplying it means a sweeper never has to form the increment by

324 subtracting two stored right-hand sides, whose cancellation error carries

325 :math:`\varepsilon\|A\|` -- see :class:`pySDC.core.problem.Problem`.

326

327 Parameters

328 ----------

329 base : dtype_u

330 The base state, unused for a linear operator.

331 delta : dtype_u

332 The correction.

333 t : float

334 Current time, accepted for interface compatibility.

335

336 Returns

337 -------

338 f : dtype_f

339 The increment.

340 """

341 f = self.dtype_f(self.init)

342 f[:] = self.A.dot(delta.flatten()).reshape(self.nvars)

343 return f

344

345 def solve_system(self, rhs, factor, u0, t):

346 r"""

347 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.

348

349 Parameters

350 ----------

351 rhs : dtype_f

352 Right-hand side for the linear system.

353 factor : float

354 Abbrev. for the local stepsize (or any other factor required).

355 u0 : dtype_u

356 Initial guess for the iterative solver.

357 t : float

358 Current time (e.g. for time-dependent BCs).

359

360 Returns

361 -------

362 sol : dtype_u

363 The solution of the linear solver.

364 """

365 if self.solver_type == 'FFT' or self.solve_precision is not None:

366 return self._solve_at_precision(rhs, factor, u0)

367

368 solver_type, Id, A, nvars, lintol, liniter, sol = (

369 self.solver_type,

370 self.Id,

371 self.A,

372 self.nvars,

373 self.lintol,

374 self.liniter,

375 self.u_init,

376 )

377

378 if solver_type == 'direct':

379 sol[:] = self._direct_solve(factor, lambda: Id - factor * A, rhs.flatten()).reshape(nvars)

380 elif solver_type == 'GMRES':

381 sol[:] = self.linalg.gmres(

382 Id - factor * A,

383 rhs.flatten(),

384 x0=u0.flatten(),

385 rtol=lintol,

386 maxiter=liniter,

387 atol=0,

388 callback=self.work_counters[solver_type],

389 callback_type='legacy',

390 )[0].reshape(nvars)

391 elif solver_type == 'CG':

392 sol[:] = self.linalg.cg(

393 Id - factor * A,

394 rhs.flatten(),

395 x0=u0.flatten(),

396 rtol=lintol,

397 maxiter=liniter,

398 atol=0,

399 callback=self.work_counters[solver_type],

400 )[0].reshape(nvars)

401 else:

402 raise ValueError(f'solver type "{solver_type}" not known in generic advection-diffusion implementation!')

403

404 return sol

405

406 def _direct_solve(self, factor, matrix, b):

407 r"""

408 Solve ``matrix() @ x = b`` directly, factorizing :math:`I - factor\cdot A` once per factor.

409

410 A sweep repeats the same few factors (one per node and level at a fixed step size), so

411 refactorizing on every call, as ``spsolve`` does, repeats identical work: the cached

412 factorization gives the same solution bit for bit, 14 to 40 times faster per solve in one

413 and two dimensions. Like ``spsolve``, the matrix is promoted to the right-hand side's dtype,

414 and the cache is keyed by that dtype too: a single-precision problem can be handed a

415 double-precision right-hand side, since a NumPy float64 scalar such as a quadrature weight

416 widens the array it multiplies, and a factorization's solve refuses to cast that down. The

417 oldest factorization is evicted beyond ``max_cached_factorizations``, as in

418 ``GenericSpectralLinear``, which bounds memory when an adaptive step size keeps changing the

419 factor. On a GPU the solve is left to CuPy's ``spsolve``.

420

421 Parameters

422 ----------

423 factor : float

424 The factor in front of :math:`A`, which identifies the matrix.

425 matrix : callable

426 Builds the matrix, only called when it is not factorized yet.

427 b : array

428 Flattened right-hand side.

429

430 Returns

431 -------

432 array

433 The flattened solution.

434 """

435 if self.xp is not np:

436 return self.linalg.spsolve(matrix(), b)

437 key = (factor, b.dtype)

438 if key not in self._factorizations:

439 if len(self._factorizations) >= self.max_cached_factorizations:

440 self._factorizations.pop(next(iter(self._factorizations)))

441 M = matrix()

442 self._factorizations[key] = self.linalg.factorized(M.astype(np.promote_types(M.dtype, b.dtype)))

443 return self._factorizations[key](b)

444

445 def _solve_at_precision(self, rhs, factor, u0):

446 r"""

447 Solve :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}` at ``solve_precision``, or by FFT.

448

449 Returns

450 -------

451 sol : dtype_u

452 The solution, widened to the state's precision.

453 """

454 sol = self.dtype_u(self.init)

455 if self.solver_type == 'FFT':

456 sol[:] = self._solve_fft(rhs.view(self.xp.ndarray).reshape(self.nvars), factor)

457 self.work_counters['FFT']()

458 return sol

459

460 precision = self._solve_dtype

461

462 def matrix():

463 # the factor at the solve's own precision: a float64 scalar would widen the operator back

464 return self._Id_solve - precision.type(factor) * self._A_solve

465

466 b = rhs.view(self.xp.ndarray).flatten().astype(precision)

467 if self.solver_type == 'direct':

468 x = self._direct_solve(factor, matrix, b)

469 elif self.solver_type in ['GMRES', 'CG']:

470 x0 = u0.view(self.xp.ndarray).flatten().astype(precision)

471 options = {'callback_type': 'legacy'} if self.solver_type == 'GMRES' else {}

472 # A Krylov solve cannot meet a tolerance below its own precision. Asked to, it iterates to

473 # `liniter` and, in CG, divides by a dot product that has rounded to zero: NaN. A hundred

474 # times the precision's epsilon is what single precision still reaches.

475 rtol = max(self.lintol, 100 * float(np.finfo(precision).eps))

476 x = getattr(self.linalg, self.solver_type.lower())(

477 matrix(),

478 b,

479 x0=x0,

480 rtol=rtol,

481 maxiter=self.liniter,

482 atol=0,

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

484 **options,

485 )[0]

486 else:

487 raise ValueError(

488 f'solver type "{self.solver_type}" not known in generic advection-diffusion implementation!'

489 )

490 sol[:] = x.reshape(self.nvars)

491 return sol

492

493 def _solve_fft(self, b, factor):

494 """

495 ``ifft(fft(b) / (1 - factor * eigenvalues))`` at the solve precision, at the state's precision.

496

497 The right-hand side is scaled to a maximum of one first and the result scaled back, which is

498 exact for a linear solve and keeps a half-precision solve clear of underflow: a correction of

499 1e-10 is below ``float16``'s smallest subnormal.

500 """

501 scale = float(abs(b).max())

502 if scale == 0.0:

503 return self.xp.zeros_like(b)

504 inverse = 1.0 / (1.0 - factor * self._eigenvalues)

505 if self._half:

506 x = self._solve_fft_half(b / scale, inverse)

507 else:

508 precision = np.promote_types(self._solve_dtype, np.complex64)

509 x = self._fft.ifftn(self._fft.fftn((b / scale).astype(precision)) * inverse.astype(precision))

510 # widen first, scale second: scaling at half precision would underflow on the way out

511 x = x.astype(np.promote_types(self.dtype, np.complex64))

512 x = x if self.dtype.kind == 'c' else x.real

513 return x * scale

514

515 def _solve_fft_half(self, y, inverse):

516 """

517 ``ifft(fft(y) * inverse)`` in half precision for ``|y| <= 1``, returned in single precision.

518

519 Both transforms are scaled by ``1 / sqrt(N)``, so the forward one is bounded by ``sqrt(N)``

520 and the inverse one by ``sqrt(N)`` times the solution: 1024 on a 1024 x 1024 grid, where an

521 unscaled transform would reach 1e6 and ``float16`` stops at 65504.

522 """

523 unit = 1.0 / np.sqrt(y.size)

524 if self.xp is np:

525 # an emulation: rounded through float16 at every stage, computed in single precision

526

527 def rounded(z):

528 return z.real.astype(np.float16).astype(np.float32) + 1j * z.imag.astype(np.float16).astype(np.float32)

529

530 spectrum = rounded(self._fft.fftn(rounded((y * unit).astype(np.complex64))))

531 spectrum = rounded(spectrum * rounded(inverse.astype(np.complex64)))

532 return rounded(self._fft.ifftn(spectrum, norm='forward')) * unit

533

534 import cupy as cp

535 from cupy.cuda import cufft

536

537 shape, n = tuple(self.nvars), y.size

538 if self._half_plan is None:

539 # complex32 in and out, stored as float16 (re, im) pairs along the last axis

540 self._half_plan = cufft.XtPlanNd(

541 shape, shape, 1, n, 'E', shape, 1, n, 'E', 1, 'E', order='C', last_axis=-1, last_size=None

542 )

543 data = cp.zeros(shape[:-1] + (2 * shape[-1],), dtype=cp.float16)

544 data[..., 0::2] = (y * unit).real

545 if y.dtype.kind == 'c':

546 data[..., 1::2] = (y * unit).imag

547 spectrum = cp.empty_like(data)

548 self._half_plan.fft(data, spectrum, cufft.CUFFT_FORWARD)

549 re, im = spectrum[..., 0::2].copy(), spectrum[..., 1::2].copy()

550 ir, ii = inverse.real.astype(cp.float16), inverse.imag.astype(cp.float16)

551 spectrum[..., 0::2] = re * ir - im * ii

552 spectrum[..., 1::2] = re * ii + im * ir

553 self._half_plan.fft(spectrum, data, cufft.CUFFT_INVERSE)

554 return (data[..., 0::2].astype(cp.float32) + 1j * data[..., 1::2].astype(cp.float32)) * unit

Coverage for pySDC/implementations/problem_classes/generic_ND_FD.py: 99%

188 statements