Coverage for pySDC/implementations/problem_classes/AllenCahn_Temp_MPIFFT.py: 81%

123 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2024-12-20 14:51 +0000

1import numpy as np 

2from mpi4py_fft import PFFT 

3 

4from pySDC.core.errors import ProblemError 

5from pySDC.core.problem import Problem 

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

7 

8from mpi4py_fft import newDistArray 

9 

10 

11class allencahn_temp_imex(Problem): 

12 r""" 

13 This class implements the :math:`N`-dimensional Allen-Cahn equation with periodic boundary conditions :math:`u \in [0, 1]^2` 

14 

15 .. math:: 

16 \frac{\partial u}{\partial t} = D \Delta u - \frac{2}{\varepsilon^2} u (1 - u) (1 - 2u) 

17 - 6 d_w \frac{u - T_M}{T_M}u (1 - u) 

18 

19 on a spatial domain :math:`[-\frac{L}{2}, \frac{L}{2}]^2`, with driving force :math:`d_w`, and :math:`N=2,3`. :math:`D` and 

20 :math:`T_M` are fixed parameters. Different initial conditions can be used, for example, circles of the form 

21 

22 .. math:: 

23 u({\bf x}, 0) = \tanh\left(\frac{r - \sqrt{(x_i-0.5)^2 + (y_j-0.5)^2}}{\sqrt{2}\varepsilon}\right), 

24 

25 for :math:`i, j=0,..,N-1`, where :math:`N` is the number of spatial grid points. For time-stepping, the problem is treated 

26 *semi-implicitly*, i.e., the nonlinear system is solved by Fast-Fourier Tranform (FFT) and the linear parts in the right-hand 

27 side will be treated explicitly using ``mpi4py-fft`` [1]_ to solve them. 

28 

29 Attributes 

30 ---------- 

31 nvars : List of int tuples, optional 

32 Number of unknowns in the problem, e.g. ``nvars=[(128, 128), (64, 64)]``. 

33 eps : float, optional 

34 Scaling parameter :math:`\varepsilon`. 

35 radius : float, optional 

36 Radius of the circles. 

37 spectral : bool, optional 

38 Indicates if spectral initial condition is used. 

39 TM : float, optional 

40 Problem parameter :math:`T_M`. 

41 D : float, optional 

42 Problem parameter :math:`D`. 

43 dw : float, optional 

44 Driving force. 

45 L : float, optional 

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

47 init_type : str, optional 

48 Initialises type of initial state. 

49 comm : bool, optional 

50 Communicator. 

51 

52 Attributes 

53 ---------- 

54 fft : fft object 

55 Object for FFT. 

56 X : np.ogrid 

57 Grid coordinates in real space. 

58 K2 : np.ndarray 

59 Laplace operator in spectral space. 

60 dx : float 

61 Mesh width in x direction. 

62 dy : float 

63 Mesh width in y direction. 

64 

65 References 

66 ---------- 

67 .. [1] https://mpi4py-fft.readthedocs.io/en/latest/ 

68 """ 

69 

70 dtype_u = mesh 

71 dtype_f = imex_mesh 

72 

73 def __init__( 

74 self, 

75 nvars=None, 

76 eps=0.04, 

77 radius=0.25, 

78 spectral=None, 

79 TM=1.0, 

80 D=10.0, 

81 dw=0.0, 

82 L=1.0, 

83 init_type='circle', 

84 comm=None, 

85 ): 

86 """Initialization routine""" 

87 

88 if nvars is None: 

89 nvars = [(128, 128)] 

90 

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

92 raise ProblemError('Need at least two dimensions') 

93 

94 # creating FFT structure 

95 ndim = len(nvars) 

96 axes = tuple(range(ndim)) 

97 self.fft = PFFT(comm, list(nvars), axes=axes, dtype=np.float64, collapse=True) 

98 

99 # get test data to figure out type and dimensions 

100 tmp_u = newDistArray(self.fft, spectral) 

101 

102 # add two components to contain field and temperature 

103 self.ncomp = 2 

104 sizes = tmp_u.shape + (self.ncomp,) 

105 

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

107 super().__init__(init=(sizes, comm, tmp_u.dtype)) 

108 self._makeAttributeAndRegister( 

109 'nvars', 

110 'eps', 

111 'radius', 

112 'spectral', 

113 'TM', 

114 'D', 

115 'dw', 

116 'L', 

117 'init_type', 

118 'comm', 

119 localVars=locals(), 

120 readOnly=True, 

121 ) 

122 

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

124 

125 # get local mesh 

126 X = list(np.ogrid[self.fft.local_slice(False)]) 

127 N = self.fft.global_shape() 

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

129 X[i] = X[i] * L[i] / N[i] 

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

131 

132 # get local wavenumbers and Laplace operator 

133 s = self.fft.local_slice() 

134 N = self.fft.global_shape() 

135 k = [np.fft.fftfreq(n, 1.0 / n).astype(int) for n in N[:-1]] 

136 k.append(np.fft.rfftfreq(N[-1], 1.0 / N[-1]).astype(int)) 

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

138 Ks = list(np.meshgrid(*K, indexing='ij', sparse=True)) 

139 Lp = 2 * np.pi / L 

140 for i in range(ndim): 

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

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

143 K = np.array(K).astype(float) 

144 self.K2 = np.sum(K * K, 0, dtype=float) 

145 

146 # Need this for diagnostics 

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

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

149 

150 def eval_f(self, u, t): 

151 """ 

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

153 

154 Parameters 

155 ---------- 

156 u : dtype_u 

157 Current values of the numerical solution. 

158 t : float 

159 Current time of the numerical solution is computed. 

160 

161 Returns 

162 ------- 

163 f : dtype_f 

164 The right-hand side of the problem. 

165 """ 

166 

167 f = self.dtype_f(self.init) 

168 

169 if self.spectral: 

170 f.impl[..., 0] = -self.K2 * u[..., 0] 

171 

172 if self.eps > 0: 

173 tmp_u = newDistArray(self.fft, False) 

174 tmp_T = newDistArray(self.fft, False) 

175 tmp_u = self.fft.backward(u[..., 0], tmp_u) 

176 tmp_T = self.fft.backward(u[..., 1], tmp_T) 

177 tmpf = -2.0 / self.eps**2 * tmp_u * (1.0 - tmp_u) * (1.0 - 2.0 * tmp_u) - 6.0 * self.dw * ( 

178 tmp_T - self.TM 

179 ) / self.TM * tmp_u * (1.0 - tmp_u) 

180 f.expl[..., 0] = self.fft.forward(tmpf) 

181 

182 f.impl[..., 1] = -self.D * self.K2 * u[..., 1] 

183 f.expl[..., 1] = f.impl[..., 0] + f.expl[..., 0] 

184 

185 else: 

186 u_hat = self.fft.forward(u[..., 0]) 

187 lap_u_hat = -self.K2 * u_hat 

188 f.impl[..., 0] = self.fft.backward(lap_u_hat, f.impl[..., 0]) 

189 

190 if self.eps > 0: 

191 f.expl[..., 0] = -2.0 / self.eps**2 * u[..., 0] * (1.0 - u[..., 0]) * (1.0 - 2.0 * u[..., 0]) 

192 f.expl[..., 0] -= 6.0 * self.dw * (u[..., 1] - self.TM) / self.TM * u[..., 0] * (1.0 - u[..., 0]) 

193 

194 u_hat = self.fft.forward(u[..., 1]) 

195 lap_u_hat = -self.D * self.K2 * u_hat 

196 f.impl[..., 1] = self.fft.backward(lap_u_hat, f.impl[..., 1]) 

197 f.expl[..., 1] = f.impl[..., 0] + f.expl[..., 0] 

198 

199 return f 

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 

222 if self.spectral: 

223 me = self.dtype_u(self.init) 

224 me[..., 0] = rhs[..., 0] / (1.0 + factor * self.K2) 

225 me[..., 1] = rhs[..., 1] / (1.0 + factor * self.D * self.K2) 

226 

227 else: 

228 me = self.dtype_u(self.init) 

229 rhs_hat = self.fft.forward(rhs[..., 0]) 

230 rhs_hat /= 1.0 + factor * self.K2 

231 me[..., 0] = self.fft.backward(rhs_hat) 

232 rhs_hat = self.fft.forward(rhs[..., 1]) 

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

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

235 

236 return me 

237 

238 def u_exact(self, t): 

239 r""" 

240 Routine to compute the exact solution at time :math:`t`. 

241 

242 Parameters 

243 ---------- 

244 t : float 

245 Time of the exact solution. 

246 

247 Returns 

248 ------- 

249 me : dtype_u 

250 The exact solution. 

251 """ 

252 

253 def circle(): 

254 tmp_me = newDistArray(self.fft, self.spectral) 

255 r2 = (self.X[0] - 0.5) ** 2 + (self.X[1] - 0.5) ** 2 

256 if self.spectral: 

257 tmp = 0.5 * (1.0 + np.tanh((self.radius - np.sqrt(r2)) / (np.sqrt(2) * self.eps))) 

258 tmp_me[:] = self.fft.forward(tmp) 

259 else: 

260 tmp_me[:] = 0.5 * (1.0 + np.tanh((self.radius - np.sqrt(r2)) / (np.sqrt(2) * self.eps))) 

261 return tmp_me 

262 

263 def circle_rand(): 

264 tmp_me = newDistArray(self.fft, self.spectral) 

265 ndim = len(tmp_me.shape) 

266 L = int(self.L) 

267 # get random radii for circles/spheres 

268 np.random.seed(1) 

269 lbound = 3.0 * self.eps 

270 ubound = 0.5 - self.eps 

271 rand_radii = (ubound - lbound) * np.random.random_sample(size=tuple([L] * ndim)) + lbound 

272 # distribute circles/spheres 

273 tmp = newDistArray(self.fft, False) 

274 if ndim == 2: 

275 for i in range(0, L): 

276 for j in range(0, L): 

277 # build radius 

278 r2 = (self.X[0] + i - L + 0.5) ** 2 + (self.X[1] + j - L + 0.5) ** 2 

279 # add this blob, shifted by 1 to avoid issues with adding up negative contributions 

280 tmp += np.tanh((rand_radii[i, j] - np.sqrt(r2)) / (np.sqrt(2) * self.eps)) + 1 

281 # normalize to [0,1] 

282 tmp *= 0.5 

283 assert np.all(tmp <= 1.0) 

284 if self.spectral: 

285 tmp_me[:] = self.fft.forward(tmp) 

286 else: 

287 tmp_me[:] = tmp[:] 

288 return tmp_me 

289 

290 def sine(): 

291 tmp_me = newDistArray(self.fft, self.spectral) 

292 if self.spectral: 

293 tmp = 0.5 * (2.0 + 0.0 * np.sin(2 * np.pi * self.X[0]) * np.sin(2 * np.pi * self.X[1])) 

294 tmp_me[:] = self.fft.forward(tmp) 

295 else: 

296 tmp_me[:] = 0.5 * (2.0 + 0.0 * np.sin(2 * np.pi * self.X[0]) * np.sin(2 * np.pi * self.X[1])) 

297 return tmp_me 

298 

299 assert t == 0, 'ERROR: u_exact only valid for t=0' 

300 me = self.dtype_u(self.init, val=0.0) 

301 if self.init_type == 'circle': 

302 me[..., 0] = circle() 

303 me[..., 1] = sine() 

304 elif self.init_type == 'circle_rand': 

305 me[..., 0] = circle_rand() 

306 me[..., 1] = sine() 

307 else: 

308 raise NotImplementedError('type of initial value not implemented, got %s' % self.init_type) 

309 

310 return me 

311 

312 

313# class allencahn_temp_imex_timeforcing(allencahn_temp_imex): 

314# """ 

315# Example implementing Allen-Cahn equation in 2-3D using mpi4py-fft for solving linear parts, IMEX time-stepping, 

316# time-dependent forcing 

317# """ 

318# 

319# def eval_f(self, u, t): 

320# """ 

321# Routine to evaluate the RHS 

322# 

323# Args: 

324# u (dtype_u): current values 

325# t (float): current time 

326# 

327# Returns: 

328# dtype_f: the RHS 

329# """ 

330# 

331# f = self.dtype_f(self.init) 

332# 

333# if self.spectral: 

334# 

335# f.impl = -self.K2 * u 

336# 

337# tmp = newDistArray(self.fft, False) 

338# tmp[:] = self.fft.backward(u, tmp) 

339# 

340# if self.eps > 0: 

341# tmpf = -2.0 / self.eps ** 2 * tmp * (1.0 - tmp) * (1.0 - 2.0 * tmp) 

342# else: 

343# tmpf = self.dtype_f(self.init, val=0.0) 

344# 

345# # build sum over RHS without driving force 

346# Rt_local = float(np.sum(self.fft.backward(f.impl) + tmpf)) 

347# if self.comm is not None: 

348# Rt_global = self.comm.allreduce(sendobj=Rt_local, op=MPI.SUM) 

349# else: 

350# Rt_global = Rt_local 

351# 

352# # build sum over driving force term 

353# Ht_local = float(np.sum(6.0 * tmp * (1.0 - tmp))) 

354# if self.comm is not None: 

355# Ht_global = self.comm.allreduce(sendobj=Ht_local, op=MPI.SUM) 

356# else: 

357# Ht_global = Rt_local 

358# 

359# # add/substract time-dependent driving force 

360# if Ht_global != 0.0: 

361# dw = Rt_global / Ht_global 

362# else: 

363# dw = 0.0 

364# 

365# tmpf -= 6.0 * dw * tmp * (1.0 - tmp) 

366# f.expl[:] = self.fft.forward(tmpf) 

367# 

368# else: 

369# 

370# u_hat = self.fft.forward(u) 

371# lap_u_hat = -self.K2 * u_hat 

372# f.impl[:] = self.fft.backward(lap_u_hat, f.impl) 

373# 

374# if self.eps > 0: 

375# f.expl = -2.0 / self.eps ** 2 * u * (1.0 - u) * (1.0 - 2.0 * u) 

376# 

377# # build sum over RHS without driving force 

378# Rt_local = float(np.sum(f.impl + f.expl)) 

379# if self.comm is not None: 

380# Rt_global = self.comm.allreduce(sendobj=Rt_local, op=MPI.SUM) 

381# else: 

382# Rt_global = Rt_local 

383# 

384# # build sum over driving force term 

385# Ht_local = float(np.sum(6.0 * u * (1.0 - u))) 

386# if self.comm is not None: 

387# Ht_global = self.comm.allreduce(sendobj=Ht_local, op=MPI.SUM) 

388# else: 

389# Ht_global = Rt_local 

390# 

391# # add/substract time-dependent driving force 

392# if Ht_global != 0.0: 

393# dw = Rt_global / Ht_global 

394# else: 

395# dw = 0.0 

396# 

397# f.expl -= 6.0 * dw * u * (1.0 - u) 

398# 

399# return f