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

89 statements  

« prev     ^ index     » next       coverage.py v7.6.12, created at 2025-03-04 07:15 +0000

1import numpy as np 

2from mpi4py import MPI 

3from mpi4py_fft import PFFT, newDistArray 

4 

5from pySDC.core.errors import ProblemError 

6from pySDC.core.problem import Problem, WorkCounter 

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

8 

9 

10class IMEX_Laplacian_MPIFFT(Problem): 

11 r""" 

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

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

14 Works in two and three dimensions. 

15 

16 Parameters 

17 ---------- 

18 nvars : tuple, optional 

19 Spatial resolution 

20 spectral : bool, optional 

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

22 L : float, optional 

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

24 alpha : float, optional 

25 Multiplicative factor before the Laplacian 

26 comm : MPI.COMM_World 

27 Communicator for parallelisation. 

28 

29 Attributes 

30 ---------- 

31 fft : PFFT 

32 Object for parallel FFT transforms. 

33 X : mesh-grid 

34 Grid coordinates in real space. 

35 K2 : matrix 

36 Laplace operator in spectral space. 

37 

38 References 

39 ---------- 

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

41 Journal of Parallel and Distributed Computing (2019). 

42 """ 

43 

44 dtype_u = mesh 

45 dtype_f = imex_mesh 

46 

47 xp = np 

48 fft_backend = 'fftw' 

49 fft_comm_backend = 'MPI' 

50 

51 @classmethod 

52 def setup_GPU(cls): 

53 """switch to GPU modules""" 

54 import cupy as cp 

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

56 

57 cls.xp = cp 

58 

59 cls.dtype_u = cupy_mesh 

60 cls.dtype_f = imex_cupy_mesh 

61 

62 cls.fft_backend = 'cupy' 

63 cls.fft_comm_backend = 'NCCL' 

64 

65 def __init__( 

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

67 ): 

68 if useGPU: 

69 self.setup_GPU() 

70 

71 if nvars is None: 

72 nvars = (128, 128) 

73 

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

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

76 

77 # Creating FFT structure 

78 self.ndim = len(nvars) 

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

80 self.fft = PFFT( 

81 comm, 

82 list(nvars), 

83 axes=axes, 

84 dtype=dtype, 

85 collapse=True, 

86 backend=self.fft_backend, 

87 comm_backend=self.fft_comm_backend, 

88 grid=(-1,), 

89 ) 

90 

91 # get test data to figure out type and dimensions 

92 tmp_u = newDistArray(self.fft, spectral) 

93 

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

95 

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

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

98 self._makeAttributeAndRegister( 

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

100 ) 

101 

102 self.getLocalGrid() 

103 self.getLaplacian() 

104 

105 # Need this for diagnostics 

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

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

108 

109 # work counters 

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

111 

112 def getLocalGrid(self): 

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

114 N = self.fft.global_shape() 

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

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

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

118 

119 def getLaplacian(self): 

120 s = self.fft.local_slice() 

121 N = self.fft.global_shape() 

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

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

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

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

126 for i in range(self.ndim): 

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

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

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

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

131 

132 def eval_f(self, u, t): 

133 """ 

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

135 

136 Parameters 

137 ---------- 

138 u : dtype_u 

139 Current values of the numerical solution. 

140 t : float 

141 Current time at which the numerical solution is computed. 

142 

143 Returns 

144 ------- 

145 f : dtype_f 

146 The right-hand side of the problem. 

147 """ 

148 

149 f = self.dtype_f(self.init) 

150 

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

152 

153 if self.spectral: 

154 tmp = self.fft.backward(u) 

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

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

157 

158 else: 

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

160 

161 self.work_counters['rhs']() 

162 return f 

163 

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

165 alpha = alpha if alpha else self.alpha 

166 if self.spectral: 

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

168 else: 

169 u_hat = self.fft.forward(u) 

170 lap_u_hat = -alpha * self.K2 * u_hat 

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

172 return f_impl 

173 

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

175 return f_expl 

176 

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

178 """ 

179 Simple FFT solver for the diffusion part. 

180 

181 Parameters 

182 ---------- 

183 rhs : dtype_f 

184 Right-hand side for the linear system. 

185 factor : float 

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

187 u0 : dtype_u 

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

189 t : float 

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

191 

192 Returns 

193 ------- 

194 me : dtype_u 

195 The solution as mesh. 

196 """ 

197 me = self.dtype_u(self.init) 

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

199 

200 return me 

201 

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

203 alpha = alpha if alpha else self.alpha 

204 if self.spectral: 

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

206 

207 else: 

208 rhs_hat = self.fft.forward(rhs) 

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

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

211 return me