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

86 statements

, created at 2024-09-20 17:10 +0000

1import numpy as np

2from mpi4py import MPI

3from mpi4py_fft import PFFT

5from pySDC.core.errors import ProblemError

6from pySDC.core.problem import Problem, WorkCounter

7from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh

9from mpi4py_fft import newDistArray

12class IMEX_Laplacian_MPIFFT(Problem):

13 r"""

14 Generic base class for IMEX problems using a spectral method to solve the Laplacian implicitly and a possible rest

15 explicitly. The FFTs are done with``mpi4py-fft`` [1]_.

17 Parameters

18 ----------

19 nvars : tuple, optional

20 Spatial resolution

21 spectral : bool, optional

22 If True, the solution is computed in spectral space.

23 L : float, optional

24 Denotes the period of the function to be approximated for the Fourier transform.

25 alpha : float, optional

26 Multiplicative factor before the Laplacian

27 comm : MPI.COMM_World

28 Communicator for parallelisation.

30 Attributes

31 ----------

32 fft : PFFT

33 Object for parallel FFT transforms.

34 X : mesh-grid

35 Grid coordinates in real space.

36 K2 : matrix

37 Laplace operator in spectral space.

39 References

40 ----------

41 .. [1] Lisandro Dalcin, Mikael Mortensen, David E. Keyes. Fast parallel multidimensional FFT using advanced MPI.

42 Journal of Parallel and Distributed Computing (2019).

43 """

45 dtype_u = mesh

46 dtype_f = imex_mesh

48 xp = np

49 fft_backend = 'fftw'

50 fft_comm_backend = 'MPI'

52 @classmethod

53 def setup_GPU(cls):

54 """switch to GPU modules"""

55 import cupy as cp

56 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh, imex_cupy_mesh

58 cls.xp = cp

60 cls.dtype_u = cupy_mesh

61 cls.dtype_f = imex_cupy_mesh

63 cls.fft_backend = 'cupy'

64 cls.fft_comm_backend = 'NCCL'

66 def __init__(

67 self, nvars=None, spectral=False, L=2 * np.pi, alpha=1.0, comm=MPI.COMM_WORLD, dtype='d', useGPU=False, x0=0.0

68 ):

69 """Initialization routine"""

71 if useGPU:

72 self.setup_GPU()

74 if nvars is None:

75 nvars = (128, 128)

77 if not (isinstance(nvars, tuple) and len(nvars) > 1):

78 raise ProblemError('Need at least two dimensions for distributed FFTs')

80 # Creating FFT structure

81 self.ndim = len(nvars)

82 axes = tuple(range(self.ndim))

83 self.fft = PFFT(

84 comm,

85 list(nvars),

86 axes=axes,

87 dtype=dtype,

88 collapse=True,

89 backend=self.fft_backend,

90 comm_backend=self.fft_comm_backend,

91 )

93 # get test data to figure out type and dimensions

94 tmp_u = newDistArray(self.fft, spectral)

96 L = np.array([L] * self.ndim, dtype=float)

98 # invoke super init, passing the communicator and the local dimensions as init

99 super().__init__(init=(tmp_u.shape, comm, tmp_u.dtype))

100 self._makeAttributeAndRegister(

101 'nvars', 'spectral', 'L', 'alpha', 'comm', 'x0', localVars=locals(), readOnly=True

102 )

104 # get local mesh

105 X = list(self.xp.ogrid[self.fft.local_slice(False)])

106 N = self.fft.global_shape()

107 for i in range(len(N)):

108 X[i] = x0 + (X[i] * L[i] / N[i])

109 self.X = [self.xp.broadcast_to(x, self.fft.shape(False)) for x in X]

111 # get local wavenumbers and Laplace operator

112 s = self.fft.local_slice()

113 N = self.fft.global_shape()

114 k = [self.xp.fft.fftfreq(n, 1.0 / n).astype(int) for n in N]

115 K = [ki[si] for ki, si in zip(k, s)]

116 Ks = list(self.xp.meshgrid(*K, indexing='ij', sparse=True))

117 Lp = 2 * np.pi / self.L

118 for i in range(self.ndim):

119 Ks[i] = (Ks[i] * Lp[i]).astype(float)

120 K = [self.xp.broadcast_to(k, self.fft.shape(True)) for k in Ks]

121 K = self.xp.array(K).astype(float)

122 self.K2 = self.xp.sum(K * K, 0, dtype=float) # Laplacian in spectral space

124 # Need this for diagnostics

125 self.dx = self.L[0] / nvars[0]

126 self.dy = self.L[1] / nvars[1]

128 # work counters

129 self.work_counters['rhs'] = WorkCounter()

131 def eval_f(self, u, t):

132 """

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

135 Parameters

136 ----------

137 u : dtype_u

138 Current values of the numerical solution.

139 t : float

140 Current time at which the numerical solution is computed.

142 Returns

143 -------

144 f : dtype_f

145 The right-hand side of the problem.

146 """

148 f = self.dtype_f(self.init)

150 f.impl[:] = self._eval_Laplacian(u, f.impl)

152 if self.spectral:

153 tmp = self.fft.backward(u)

154 tmp[:] = self._eval_explicit_part(tmp, t, tmp)

155 f.expl[:] = self.fft.forward(tmp)

157 else:

158 f.expl[:] = self._eval_explicit_part(u, t, f.expl)

160 self.work_counters['rhs']()

161 return f

163 def _eval_Laplacian(self, u, f_impl, alpha=None):

164 alpha = alpha if alpha else self.alpha

165 if self.spectral:

166 f_impl[:] = -alpha * self.K2 * u

167 else:

168 u_hat = self.fft.forward(u)

169 lap_u_hat = -alpha * self.K2 * u_hat

170 f_impl[:] = self.fft.backward(lap_u_hat, f_impl)

171 return f_impl

173 def _eval_explicit_part(self, u, t, f_expl):

174 return f_expl

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

177 """

178 Simple FFT solver for the diffusion part.

180 Parameters

181 ----------

182 rhs : dtype_f

183 Right-hand side for the linear system.

184 factor : float

185 Abbrev. for the node-to-node stepsize (or any other factor required).

186 u0 : dtype_u

187 Initial guess for the iterative solver (not used here so far).

188 t : float

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

191 Returns

192 -------

193 me : dtype_u

194 The solution as mesh.

195 """

196 me = self.dtype_u(self.init)

197 me[:] = self._invert_Laplacian(me, factor, rhs)

199 return me

201 def _invert_Laplacian(self, me, factor, rhs, alpha=None):

202 alpha = alpha if alpha else self.alpha

203 if self.spectral:

204 me[:] = rhs / (1.0 + factor * alpha * self.K2)

206 else:

207 rhs_hat = self.fft.forward(rhs)

208 rhs_hat /= 1.0 + factor * alpha * self.K2

209 me[:] = self.fft.backward(rhs_hat)

210 return me