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

1import numpy as np 

2from mpi4py import MPI 

3from mpi4py_fft import newDistArray 

4 

5from pySDC.helpers.fft_helper import PFFT 

6 

7from pySDC.core.errors import ProblemError 

8from pySDC.core.problem import Problem, WorkCounter 

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

10 

11 

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]_. 

16 Works in two and three dimensions. 

17 

18 Parameters 

19 ---------- 

20 nvars : tuple, optional 

21 Spatial resolution 

22 spectral : bool, optional 

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

24 L : float, optional 

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

26 alpha : float, optional 

27 Multiplicative factor before the Laplacian 

28 comm : MPI.COMM_World 

29 Communicator for parallelisation. 

30 dtype : numpy.dtype or str, optional 

31 Datatype of the solution in real space, e.g. ``'d'`` for real or ``'D'`` for complex problems. 

32 useGPU : bool, optional 

33 Run on the GPU with CuPy and NCCL instead of on the CPU with NumPy and MPI. 

34 x0 : float, optional 

35 Coordinate of the left boundary of the domain in each direction. 

36 

37 Attributes 

38 ---------- 

39 fft : PFFT 

40 Object for parallel FFT transforms. 

41 X : mesh-grid 

42 Grid coordinates in real space. 

43 K2 : matrix 

44 Laplace operator in spectral space. 

45 

46 References 

47 ---------- 

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

49 Journal of Parallel and Distributed Computing (2019). 

50 """ 

51 

52 dtype_u = mesh 

53 dtype_f = imex_mesh 

54 

55 xp = np 

56 fft_backend = 'fftw' 

57 fft_comm_backend = 'MPI' 

58 

59 def setup_GPU(self): 

60 """switch to GPU modules""" 

61 import cupy as cp 

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

63 

64 self.xp = cp 

65 

66 self.dtype_u = cupy_mesh 

67 self.dtype_f = imex_cupy_mesh 

68 

69 self.fft_backend = 'cupy' 

70 self.fft_comm_backend = 'NCCL' 

71 

72 def __init__( 

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

74 ): 

75 """Initialization routine""" 

76 if useGPU: 

77 self.setup_GPU() 

78 

79 if nvars is None: 

80 nvars = (128, 128) 

81 

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

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

84 

85 # Creating FFT structure 

86 self.ndim = len(nvars) 

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

88 self.fft = PFFT( 

89 comm, 

90 list(nvars), 

91 axes=axes, 

92 dtype=dtype, 

93 collapse=True, 

94 backend=self.fft_backend, 

95 comm_backend=self.fft_comm_backend, 

96 grid=(-1,), 

97 ) 

98 

99 # get test data to figure out type and dimensions 

100 tmp_u = newDistArray(self.fft, spectral) 

101 

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

103 

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

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

106 self._makeAttributeAndRegister( 

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

108 ) 

109 

110 self.getLocalGrid() 

111 self.getLaplacian() 

112 

113 # Need this for diagnostics 

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

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

116 

117 # work counters 

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

119 

120 def getLocalGrid(self): 

121 """ 

122 Compute the coordinates of the local part of the grid in real space and store them in ``self.X``, one array per 

123 dimension broadcast to the local shape, with the points :math:`x_0 + i L / N`. 

124 

125 Returns 

126 ------- 

127 None 

128 """ 

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

130 N = self.fft.global_shape() 

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

132 X[i] = self.x0 + (X[i] * self.L[i] / N[i]) 

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

134 

135 def getLaplacian(self): 

136 r""" 

137 Compute the squared wave numbers :math:`|k|^2` of the local part of the spectral grid, with the wave numbers 

138 scaled by :math:`2\pi / L`, and store them in ``self.K2``. The Laplacian in spectral space is :math:`-|k|^2`. 

139 

140 Returns 

141 ------- 

142 None 

143 """ 

144 s = self.fft.local_slice() 

145 N = self.fft.global_shape() 

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

147 K = [ki[si] for ki, si in zip(k, s, strict=True)] 

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

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

150 for i in range(self.ndim): 

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

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

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

154 self.K2 = self.xp.sum(K * K, 0, dtype=float) 

155 

156 def eval_f(self, u, t): 

157 """ 

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

159 

160 Parameters 

161 ---------- 

162 u : dtype_u 

163 Current values of the numerical solution. 

164 t : float 

165 Current time at which the numerical solution is computed. 

166 

167 Returns 

168 ------- 

169 f : dtype_f 

170 The right-hand side of the problem. 

171 """ 

172 

173 f = self.dtype_f(self.init) 

174 

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

176 

177 if self.spectral: 

178 tmp = self.fft.backward(u) 

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

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

181 

182 else: 

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

184 

185 self.work_counters['rhs']() 

186 return f 

187 

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

189 alpha = alpha if alpha else self.alpha 

190 if self.spectral: 

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

192 else: 

193 u_hat = self.fft.forward(u) 

194 lap_u_hat = -alpha * self.K2 * u_hat 

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

196 return f_impl 

197 

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

199 return f_expl 

200 

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

202 """ 

203 Simple FFT solver for the diffusion part. 

204 

205 Parameters 

206 ---------- 

207 rhs : dtype_f 

208 Right-hand side for the linear system. 

209 factor : float 

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

211 u0 : dtype_u 

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

213 t : float 

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

215 

216 Returns 

217 ------- 

218 me : dtype_u 

219 The solution as mesh. 

220 """ 

221 me = self.dtype_u(self.init) 

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

223 

224 return me 

225 

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

227 alpha = alpha if alpha else self.alpha 

228 if self.spectral: 

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

230 

231 else: 

232 rhs_hat = self.fft.forward(rhs) 

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

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

235 return me