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

1import numpy as np 

2from mpi4py import MPI 

3from mpi4py_fft import PFFT 

4 

5from pySDC.core.Errors import ProblemError 

6from pySDC.core.Problem import ptype, WorkCounter 

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

8 

9from mpi4py_fft import newDistArray 

10 

11 

12class IMEX_Laplacian_MPIFFT(ptype): 

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 

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. 

29 

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. 

38 

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 """ 

44 

45 dtype_u = mesh 

46 dtype_f = imex_mesh 

47 

48 xp = np 

49 

50 def __init__(self, nvars=None, spectral=False, L=2 * np.pi, alpha=1.0, comm=MPI.COMM_WORLD, dtype='d', x0=0.0): 

51 """Initialization routine""" 

52 

53 if nvars is None: 

54 nvars = (128, 128) 

55 

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

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

58 

59 # Creating FFT structure 

60 self.ndim = len(nvars) 

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

62 self.fft = PFFT(comm, list(nvars), axes=axes, dtype=dtype, collapse=True) 

63 

64 # get test data to figure out type and dimensions 

65 tmp_u = newDistArray(self.fft, spectral) 

66 

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

68 

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

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

71 self._makeAttributeAndRegister( 

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

73 ) 

74 

75 # get local mesh 

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

77 N = self.fft.global_shape() 

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

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

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

81 

82 # get local wavenumbers and Laplace operator 

83 s = self.fft.local_slice() 

84 N = self.fft.global_shape() 

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

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

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

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

89 for i in range(self.ndim): 

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

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

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

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

94 

95 # Need this for diagnostics 

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

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

98 

99 # work counters 

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

101 

102 def eval_f(self, u, t): 

103 """ 

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

105 

106 Parameters 

107 ---------- 

108 u : dtype_u 

109 Current values of the numerical solution. 

110 t : float 

111 Current time at which the numerical solution is computed. 

112 

113 Returns 

114 ------- 

115 f : dtype_f 

116 The right-hand side of the problem. 

117 """ 

118 

119 f = self.dtype_f(self.init) 

120 

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

122 

123 if self.spectral: 

124 tmp = self.fft.backward(u) 

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

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

127 

128 else: 

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

130 

131 self.work_counters['rhs']() 

132 return f 

133 

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

135 alpha = alpha if alpha else self.alpha 

136 if self.spectral: 

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

138 else: 

139 u_hat = self.fft.forward(u) 

140 lap_u_hat = -alpha * self.K2 * u_hat 

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

142 return f_impl 

143 

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

145 return f_expl 

146 

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

148 """ 

149 Simple FFT solver for the diffusion part. 

150 

151 Parameters 

152 ---------- 

153 rhs : dtype_f 

154 Right-hand side for the linear system. 

155 factor : float 

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

157 u0 : dtype_u 

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

159 t : float 

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

161 

162 Returns 

163 ------- 

164 me : dtype_u 

165 The solution as mesh. 

166 """ 

167 me = self.dtype_u(self.init) 

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

169 

170 return me 

171 

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

173 alpha = alpha if alpha else self.alpha 

174 if self.spectral: 

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

176 

177 else: 

178 rhs_hat = self.fft.forward(rhs) 

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

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

181 return me