Coverage for pySDC/implementations/problem_classes/GrayScott_MPIFFT.py: 98%

1import scipy.sparse as sp 

2 

3from pySDC.core.errors import ProblemError 

4from pySDC.core.problem import WorkCounter 

5from pySDC.implementations.datatype_classes.mesh import comp2_mesh 

6from pySDC.implementations.problem_classes.generic_MPIFFT_Laplacian import IMEX_Laplacian_MPIFFT 

7 

8from mpi4py_fft import newDistArray 

9 

10 

11class grayscott_imex_diffusion(IMEX_Laplacian_MPIFFT): 

12 r""" 

13 Periodic Gray-Scott system with mpi4py-fft FFTs, IMEX with diffusion implicit and the reaction explicit. 

14 

15 The Gray-Scott system [#]_ describes a reaction-diffusion process of two substances :math:`u` and :math:`v`, 

16 where they diffuse over time. :math:`u` is fed with rate :math:`A` and used up by the reaction, which produces 

17 :math:`v`, and :math:`v` is removed with rate :math:`B`. :math:`D_u,\, D_v` are the diffusion rates for 

18 :math:`u,\, v`. Here, the process is described by the :math:`N`-dimensional model 

19 

20 .. math:: 

21 \frac{\partial u}{\partial t} = D_u \Delta u - u v^2 + A (1 - u), 

22 

23 .. math:: 

24 \frac{\partial v}{\partial t} = D_v \Delta v + u v^2 - B v 

25 

26 in :math:`x \in \Omega:=[-L/2, L/2]^N` with :math:`N=2,3`. Spatial discretization is done by using 

27 Fast Fourier transformation for solving the linear parts provided by ``mpi4py-fft`` [#]_, see also 

28 https://mpi4py-fft.readthedocs.io/en/latest/. 

29 

30 This class implements the problem for *semi-explicit* time-stepping (diffusion is treated implicitly, and reaction 

31 is computed in explicit fashion). 

32 

33 Parameters 

34 ---------- 

35 nvars : tuple of int, optional 

36 Spatial resolution, i.e., number of degrees of freedom in space. Should be a tuple, e.g. ``nvars=(127, 127)``. 

37 Du : float, optional 

38 Diffusion rate for :math:`u`. 

39 Dv : float, optional 

40 Diffusion rate for :math:`v`. 

41 A : float, optional 

42 Feed rate for :math:`u`. 

43 B : float, optional 

44 Removal rate for :math:`v`. 

45 spectral : bool, optional 

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

47 L : float, optional 

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

49 comm : COMM_WORLD, optional 

50 Communicator for ``mpi4py-fft``. 

51 num_blobs : int, optional 

52 Number of blobs in the initial conditions. Negative values give rectangles. 

53 

54 Attributes 

55 ---------- 

56 fft : PFFT 

57 Object for parallel FFT transforms. 

58 X : mesh-grid 

59 Grid coordinates in real space. 

60 ndim : int 

61 Number of spatial dimensions. 

62 Ku : matrix 

63 Laplace operator in spectral space (u component). 

64 Kv : matrix 

65 Laplace operator in spectral space (v component). 

66 

67 References 

68 ---------- 

69 .. [#] Autocatalytic reactions in the isothermal, continuous stirred tank reactor: Isolas and other forms 

70 of multistability. P. Gray, S. K. Scott. Chem. Eng. Sci. 38, 1 (1983). 

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

72 Journal of Parallel and Distributed Computing (2019). 

73 """ 

74 

75 def __init__( 

76 self, 

77 Du=1.0, 

78 Dv=0.01, 

79 A=0.09, 

80 B=0.086, 

81 L=2.0, 

82 num_blobs=1, 

83 **kwargs, 

84 ): 

85 """Initialization routine. ``kwargs`` are passed on to ``IMEX_Laplacian_MPIFFT``.""" 

86 super().__init__(dtype='d', alpha=1.0, x0=-L / 2.0, L=L, **kwargs) 

87 

88 # prepare the array with two components 

89 shape = (2,) + (self.init[0]) 

90 self.iU = 0 

91 self.iV = 1 

92 self.ncomp = 2 # needed for transfer class 

93 

94 self.init = (shape, self.comm, self.xp.dtype('complex') if self.spectral else self.xp.dtype('float')) 

95 

96 self._makeAttributeAndRegister( 

97 'Du', 

98 'Dv', 

99 'A', 

100 'B', 

101 'num_blobs', 

102 localVars=locals(), 

103 readOnly=True, 

104 ) 

105 

106 # prepare "Laplacians" 

107 self.Ku = -self.Du * self.K2 

108 self.Kv = -self.Dv * self.K2 

109 

110 def eval_f(self, u, t): 

111 """ 

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

113 

114 Parameters 

115 ---------- 

116 u : dtype_u 

117 Current values of the numerical solution. 

118 t : float 

119 Current time of the numerical solution is computed. 

120 

121 Returns 

122 ------- 

123 f : dtype_f 

124 The right-hand side of the problem. 

125 """ 

126 

127 f = self.dtype_f(self.init) 

128 

129 if self.spectral: 

130 f.impl[0, ...] = self.Ku * u[0, ...] 

131 f.impl[1, ...] = self.Kv * u[1, ...] 

132 tmpu = newDistArray(self.fft, False) 

133 tmpv = newDistArray(self.fft, False) 

134 tmpu[:] = self.fft.backward(u[0, ...], tmpu) 

135 tmpv[:] = self.fft.backward(u[1, ...], tmpv) 

136 tmpfu = -tmpu * tmpv**2 + self.A * (1 - tmpu) 

137 tmpfv = tmpu * tmpv**2 - self.B * tmpv 

138 f.expl[0, ...] = self.fft.forward(tmpfu) 

139 f.expl[1, ...] = self.fft.forward(tmpfv) 

140 

141 else: 

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

143 lap_u_hat = self.Ku * u_hat 

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

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

146 lap_u_hat = self.Kv * u_hat 

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

148 f.expl[0, ...] = -u[0, ...] * u[1, ...] ** 2 + self.A * (1 - u[0, ...]) 

149 f.expl[1, ...] = u[0, ...] * u[1, ...] ** 2 - self.B * u[1, ...] 

150 

151 self.work_counters['rhs']() 

152 return f 

153 

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

155 """ 

156 Simple FFT solver for the diffusion part. 

157 

158 Parameters 

159 ---------- 

160 rhs : dtype_f 

161 Right-hand side for the linear system. 

162 factor : float 

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

164 u0 : dtype_u 

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

166 t : float 

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

168 

169 Returns 

170 ------- 

171 me : dtype_u 

172 Solution. 

173 """ 

174 

175 me = self.dtype_u(self.init) 

176 if self.spectral: 

177 me[0, ...] = rhs[0, ...] / (1.0 - factor * self.Ku) 

178 me[1, ...] = rhs[1, ...] / (1.0 - factor * self.Kv) 

179 

180 else: 

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

182 rhs_hat /= 1.0 - factor * self.Ku 

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

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

185 rhs_hat /= 1.0 - factor * self.Kv 

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

187 

188 return me 

189 

190 def u_exact(self, t, seed=10700000): 

191 r""" 

192 Routine to compute the exact solution at time :math:`t = 0`, see 

193 `Chebfun's Gray-Scott example <https://www.chebfun.org/examples/pde/GrayScott.html>`__. 

194 

195 Parameters 

196 ---------- 

197 t : float 

198 Time of the exact solution. 

199 seed : int, optional 

200 Seed for the random number generator that places the blobs or rectangles. 

201 

202 Returns 

203 ------- 

204 me : dtype_u 

205 Exact solution. 

206 """ 

207 assert t == 0.0, 'Exact solution only valid as initial condition' 

208 

209 xp = self.xp 

210 

211 _u = xp.zeros_like(self.X[0]) 

212 _v = xp.zeros_like(self.X[0]) 

213 

214 rng = xp.random.default_rng(seed) 

215 

216 if self.num_blobs < 0: 

217 """ 

218 Rectangles with stationary background, see arXiv:1501.01990 

219 """ 

220 F, k = self.A, self.B - self.A 

221 A = xp.sqrt(F) / (F + k) 

222 

223 # set stable background state from Equation 2 

224 assert 2 * k < xp.sqrt(F) - 2 * F, 'Kill rate is too large to facilitate stable background' 

225 _u[...] = (A - xp.sqrt(A**2 - 4)) / (2 * A) 

226 _v[...] = xp.sqrt(F) * (A + xp.sqrt(A**2 - 4)) / 2 

227 

228 for _ in range(-self.num_blobs): 

229 x0 = rng.random(size=self.ndim) * self.L[0] - self.L[0] / 2 

230 l = rng.random(size=self.ndim) * self.L[0] / self.nvars[0] * 80 

231 

232 masks = [xp.logical_and(self.X[i] > x0[i], self.X[i] < x0[i] + l[i]) for i in range(self.ndim)] 

233 mask = masks[0] 

234 for m in masks[1:]: 

235 mask = xp.logical_and(mask, m) 

236 

237 _u[mask] = rng.random() 

238 _v[mask] = rng.random() 

239 

240 elif self.num_blobs > 0: 

241 """ 

242 Blobs as in https://www.chebfun.org/examples/pde/GrayScott.html 

243 """ 

244 inc = self.L[0] / (self.num_blobs + 1) 

245 

246 for i in range(1, self.num_blobs + 1): 

247 for j in range(1, self.num_blobs + 1): 

248 signs = (-1) ** rng.integers(low=0, high=2, size=self.ndim) 

249 

250 if self.ndim == 2: 

251 # This assumes that the box is [-L/2, L/2]^2 

252 _u[...] += -xp.exp( 

253 -80.0 

254 * ( 

255 (self.X[0] + self.x0 + inc * i + signs[0] * 0.05) ** 2 

256 + (self.X[1] + self.x0 + inc * j + signs[1] * 0.02) ** 2 

257 ) 

258 ) 

259 _v[...] += xp.exp( 

260 -80.0 

261 * ( 

262 (self.X[0] + self.x0 + inc * i - signs[0] * 0.05) ** 2 

263 + (self.X[1] + self.x0 + inc * j - signs[1] * 0.02) ** 2 

264 ) 

265 ) 

266 elif self.ndim == 3: 

267 z_pos = self.x0 + rng.random() * self.L[2] 

268 # This assumes that the box is [-L/2, L/2]^3 

269 _u[...] += -xp.exp( 

270 -80.0 

271 * ( 

272 (self.X[0] + self.x0 + inc * i + signs[0] * 0.05) ** 2 

273 + (self.X[1] + self.x0 + inc * j + signs[1] * 0.02) ** 2 

274 + (self.X[2] - z_pos + signs[2] * 0.035) ** 2 

275 ) 

276 ) 

277 _v[...] += xp.exp( 

278 -80.0 

279 * ( 

280 (self.X[0] + self.x0 + inc * i - signs[0] * 0.05) ** 2 

281 + (self.X[1] + self.x0 + inc * j - signs[1] * 0.02) ** 2 

282 + (self.X[2] - z_pos - signs[2] * 0.035) ** 2 

283 ) 

284 ) 

285 else: 

286 raise NotImplementedError 

287 

288 _u += 1 

289 else: 

290 _u[...] = rng.random(_u.shape) 

291 _v[...] = rng.random(_v.shape) 

292 

293 u = self.u_init 

294 if self.spectral: 

295 u[0, ...] = self.fft.forward(_u) 

296 u[1, ...] = self.fft.forward(_v) 

297 else: 

298 u[0, ...] = _u 

299 u[1, ...] = _v 

300 

301 return u 

302 

303 def get_fig(self, n_comps=2): # pragma: no cover 

304 """ 

305 Get a figure suitable to plot the solution of this problem 

306 

307 Args: 

308 n_comps (int): Number of components that fit in the solution 

309 

310 Returns 

311 ------- 

312 self.fig : matplotlib.pyplot.figure.Figure 

313 """ 

314 import matplotlib.pyplot as plt 

315 from mpl_toolkits.axes_grid1 import make_axes_locatable 

316 

317 if n_comps == 2: 

318 self.fig, axs = plt.subplots(1, 2, sharex=True, sharey=True, figsize=((6, 3)), constrained_layout=True) 

319 divider = make_axes_locatable(axs[1]) 

320 self.cax = divider.append_axes('right', size='3%', pad=0.03) 

321 else: 

322 self.fig, ax = plt.subplots(1, 1, figsize=((6, 5)), constrained_layout=True) 

323 divider = make_axes_locatable(ax) 

324 self.cax = divider.append_axes('right', size='3%', pad=0.03) 

325 return self.fig 

326 

327 def plot(self, u, t=None, fig=None): # pragma: no cover 

328 r""" 

329 Plot the solution. Please supply a figure with the same structure as returned by ``self.get_fig``. 

330 

331 Parameters 

332 ---------- 

333 u : dtype_u 

334 Solution to be plotted 

335 t : float 

336 Time to display at the top of the figure 

337 fig : matplotlib.pyplot.figure.Figure 

338 Figure with the correct structure 

339 

340 Returns 

341 ------- 

342 None 

343 """ 

344 fig = self.get_fig(n_comps=2) if fig is None else fig 

345 axs = fig.axes 

346 

347 vmin = u.min() 

348 vmax = u.max() 

349 for i, label in zip([self.iU, self.iV], [r'$u$', r'$v$'], strict=True): 

350 im = axs[i].pcolormesh(self.X[0], self.X[1], u[i], vmin=vmin, vmax=vmax) 

351 axs[i].set_aspect(1) 

352 axs[i].set_title(label) 

353 

354 if t is not None: 

355 fig.suptitle(f't = {t:.2e}') 

356 axs[0].set_xlabel(r'$x$') 

357 axs[0].set_ylabel(r'$y$') 

358 fig.colorbar(im, self.cax) 

359 

360 

361class grayscott_imex_linear(grayscott_imex_diffusion): 

362 r""" 

363 Periodic Gray-Scott system with mpi4py-fft, IMEX with diffusion and linear reaction terms implicit, rest explicit. 

364 

365 The Gray-Scott system [#]_ describes a reaction-diffusion process of two substances :math:`u` and :math:`v`, 

366 where they diffuse over time. :math:`u` is fed with rate :math:`A` and used up by the reaction, which produces 

367 :math:`v`, and :math:`v` is removed with rate :math:`B`. :math:`D_u,\, D_v` are the diffusion rates for 

368 :math:`u,\, v`. This is the same model as in ``grayscott_imex_diffusion``, with the linear reaction terms grouped with the 

369 diffusion: 

370 

371 .. math:: 

372 \frac{d u}{d t} = (D_u \Delta - A) u - u v^2 + A, 

373 

374 .. math:: 

375 \frac{d v}{d t} = (D_v \Delta - B) v + u v^2 

376 

377 in :math:`x \in \Omega:=[-L/2, L/2]^N` with :math:`N=2,3`. Spatial discretization is done by using 

378 Fast Fourier transformation for solving the linear parts provided by ``mpi4py-fft`` [#]_, see also 

379 https://mpi4py-fft.readthedocs.io/en/latest/. 

380 

381 This class implements the problem for *semi-explicit* time-stepping (diffusion and the linear reaction terms are 

382 treated implicitly, the rest of the reaction explicitly). 

383 

384 References 

385 ---------- 

386 .. [#] Autocatalytic reactions in the isothermal, continuous stirred tank reactor: Isolas and other forms 

387 of multistability. P. Gray, S. K. Scott. Chem. Eng. Sci. 38, 1 (1983). 

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

389 Journal of Parallel and Distributed Computing (2019). 

390 """ 

391 

392 def __init__(self, **kwargs): 

393 """ 

394 Initialization routine 

395 

396 Parameters 

397 ---------- 

398 **kwargs 

399 Passed on to ``grayscott_imex_diffusion``, see there. The linear reaction terms :math:`-A u` and 

400 :math:`-B v` are then added to the implicit operators ``Ku`` and ``Kv``. 

401 """ 

402 super().__init__(**kwargs) 

403 self.Ku -= self.A 

404 self.Kv -= self.B 

405 

406 def eval_f(self, u, t): 

407 """ 

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

409 

410 Parameters 

411 ---------- 

412 u : dtype_u 

413 Current values of the numerical solution. 

414 t : float 

415 Current time of the numerical solution is computed. 

416 

417 Returns 

418 ------- 

419 f : dtype_f 

420 The right-hand side of the problem. 

421 """ 

422 

423 f = self.dtype_f(self.init) 

424 

425 if self.spectral: 

426 f.impl[0, ...] = self.Ku * u[0, ...] 

427 f.impl[1, ...] = self.Kv * u[1, ...] 

428 tmpu = newDistArray(self.fft, False) 

429 tmpv = newDistArray(self.fft, False) 

430 tmpu[:] = self.fft.backward(u[0, ...], tmpu) 

431 tmpv[:] = self.fft.backward(u[1, ...], tmpv) 

432 tmpfu = -tmpu * tmpv**2 + self.A 

433 tmpfv = tmpu * tmpv**2 

434 f.expl[0, ...] = self.fft.forward(tmpfu) 

435 f.expl[1, ...] = self.fft.forward(tmpfv) 

436 

437 else: 

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

439 lap_u_hat = self.Ku * u_hat 

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

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

442 lap_u_hat = self.Kv * u_hat 

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

444 f.expl[0, ...] = -u[0, ...] * u[1, ...] ** 2 + self.A 

445 f.expl[1, ...] = u[0, ...] * u[1, ...] ** 2 

446 

447 self.work_counters['rhs']() 

448 return f 

449 

450 

451class grayscott_mi_diffusion(grayscott_imex_diffusion): 

452 r""" 

453 Periodic Gray-Scott system with mpi4py-fft, multi-implicit: diffusion by FFT, reaction by Newton. 

454 

455 The Gray-Scott system [#]_ describes a reaction-diffusion process of two substances :math:`u` and :math:`v`, 

456 where they diffuse over time. :math:`u` is fed with rate :math:`A` and used up by the reaction, which produces 

457 :math:`v`, and :math:`v` is removed with rate :math:`B`. :math:`D_u,\, D_v` are the diffusion rates for 

458 :math:`u,\, v`. Here, the process is described by the :math:`N`-dimensional model 

459 

460 .. math:: 

461 \frac{\partial u}{\partial t} = D_u \Delta u - u v^2 + A (1 - u), 

462 

463 .. math:: 

464 \frac{\partial v}{\partial t} = D_v \Delta v + u v^2 - B v 

465 

466 in :math:`x \in \Omega:=[-L/2, L/2]^N` with :math:`N=2,3`. Spatial discretization is done by using 

467 Fast Fourier transformation for solving the linear parts provided by ``mpi4py-fft`` [#]_, see also 

468 https://mpi4py-fft.readthedocs.io/en/latest/. 

469 

470 This class implements the problem for *multi-implicit* time-stepping, i.e., both diffusion and reaction part will be treated 

471 implicitly. 

472 

473 Parameters 

474 ---------- 

475 nvars : tuple of int, optional 

476 Spatial resolution, i.e., number of degrees of freedom in space. Should be a tuple, e.g. ``nvars=(127, 127)``. 

477 Du : float, optional 

478 Diffusion rate for :math:`u`. 

479 Dv : float, optional 

480 Diffusion rate for :math:`v`. 

481 A : float, optional 

482 Feed rate for :math:`u`. 

483 B : float, optional 

484 Removal rate for :math:`v`. 

485 spectral : bool, optional 

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

487 L : float, optional 

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

489 comm : COMM_WORLD, optional 

490 Communicator for ``mpi4py-fft``. 

491 newton_maxiter : int, optional 

492 Maximum number of iterations for the Newton solver. 

493 newton_tol : float, optional 

494 Tolerance for Newton's method to terminate. 

495 stop_at_nan : bool, optional 

496 Raise a ``ProblemError`` if Newton's method produces ``nan``, instead of only logging a warning. 

497 

498 Attributes 

499 ---------- 

500 fft : PFFT 

501 Object for parallel FFT transforms. 

502 X : mesh-grid 

503 Grid coordinates in real space. 

504 ndim : int 

505 Number of spatial dimensions. 

506 Ku : matrix 

507 Laplace operator in spectral space (u component). 

508 Kv : matrix 

509 Laplace operator in spectral space (v component). 

510 

511 References 

512 ---------- 

513 .. [#] Autocatalytic reactions in the isothermal, continuous stirred tank reactor: Isolas and other forms 

514 of multistability. P. Gray, S. K. Scott. Chem. Eng. Sci. 38, 1 (1983). 

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

516 Journal of Parallel and Distributed Computing (2019). 

517 """ 

518 

519 dtype_f = comp2_mesh 

520 

521 def __init__( 

522 self, 

523 newton_maxiter=100, 

524 newton_tol=1e-12, 

525 stop_at_nan=True, 

526 **kwargs, 

527 ): 

528 """Initialization routine""" 

529 super().__init__(**kwargs) 

530 # This may not run in parallel yet.. 

531 assert self.comm.Get_size() == 1 

532 self.work_counters['newton'] = WorkCounter() 

533 self.Ku = -self.Du * self.K2 

534 self.Kv = -self.Dv * self.K2 

535 self._makeAttributeAndRegister( 

536 'newton_maxiter', 'newton_tol', 'stop_at_nan', localVars=locals(), readOnly=False 

537 ) 

538 

539 def eval_f(self, u, t): 

540 """ 

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

542 

543 Parameters 

544 ---------- 

545 u : dtype_u 

546 Current values of the numerical solution. 

547 t : float 

548 Current time of the numerical solution is computed. 

549 

550 Returns 

551 ------- 

552 f : dtype_f 

553 The right-hand side of the problem. 

554 """ 

555 

556 f = self.dtype_f(self.init) 

557 

558 if self.spectral: 

559 f.comp1[0, ...] = self.Ku * u[0, ...] 

560 f.comp1[1, ...] = self.Kv * u[1, ...] 

561 tmpu = newDistArray(self.fft, False) 

562 tmpv = newDistArray(self.fft, False) 

563 tmpu[:] = self.fft.backward(u[0, ...], tmpu) 

564 tmpv[:] = self.fft.backward(u[1, ...], tmpv) 

565 tmpfu = -tmpu * tmpv**2 + self.A * (1 - tmpu) 

566 tmpfv = tmpu * tmpv**2 - self.B * tmpv 

567 f.comp2[0, ...] = self.fft.forward(tmpfu) 

568 f.comp2[1, ...] = self.fft.forward(tmpfv) 

569 

570 else: 

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

572 lap_u_hat = self.Ku * u_hat 

573 f.comp1[0, ...] = self.fft.backward(lap_u_hat, f.comp1[0, ...]) 

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

575 lap_u_hat = self.Kv * u_hat 

576 f.comp1[1, ...] = self.fft.backward(lap_u_hat, f.comp1[1, ...]) 

577 f.comp2[0, ...] = -u[0, ...] * u[1, ...] ** 2 + self.A * (1 - u[0, ...]) 

578 f.comp2[1, ...] = u[0, ...] * u[1, ...] ** 2 - self.B * u[1, ...] 

579 

580 self.work_counters['rhs']() 

581 return f 

582 

583 def solve_system_1(self, rhs, factor, u0, t): 

584 """ 

585 Solver for the first component, can just call the super function. 

586 

587 Parameters 

588 ---------- 

589 rhs : dtype_f 

590 Right-hand side for the linear system. 

591 factor : float 

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

593 u0 : dtype_u 

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

595 t : float 

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

597 

598 Returns 

599 ------- 

600 me : dtype_u 

601 The solution as mesh. 

602 """ 

603 

604 me = super().solve_system(rhs, factor, u0, t) 

605 return me 

606 

607 def solve_system_2(self, rhs, factor, u0, t): 

608 """ 

609 Newton-Solver for the second component. 

610 

611 Parameters 

612 ---------- 

613 rhs : dtype_f 

614 Right-hand side for the linear system. 

615 factor : float 

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

617 u0 : dtype_u 

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

619 t : float 

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

621 

622 Returns 

623 ------- 

624 me : dtype_u 

625 The solution as mesh. 

626 """ 

627 u = self.dtype_u(u0) 

628 

629 if self.spectral: 

630 tmpu = newDistArray(self.fft, False) 

631 tmpv = newDistArray(self.fft, False) 

632 tmpu[:] = self.fft.backward(u[0, ...], tmpu) 

633 tmpv[:] = self.fft.backward(u[1, ...], tmpv) 

634 tmprhsu = newDistArray(self.fft, False) 

635 tmprhsv = newDistArray(self.fft, False) 

636 tmprhsu[:] = self.fft.backward(rhs[0, ...], tmprhsu) 

637 tmprhsv[:] = self.fft.backward(rhs[1, ...], tmprhsv) 

638 

639 else: 

640 tmpu = u[0, ...] 

641 tmpv = u[1, ...] 

642 tmprhsu = rhs[0, ...] 

643 tmprhsv = rhs[1, ...] 

644 

645 # start newton iteration 

646 n = 0 

647 res = 99 

648 while n < self.newton_maxiter: 

649 # print(n, res) 

650 # form the function g with g(u) = 0 

651 tmpgu = tmpu - tmprhsu - factor * (-tmpu * tmpv**2 + self.A * (1 - tmpu)) 

652 tmpgv = tmpv - tmprhsv - factor * (tmpu * tmpv**2 - self.B * tmpv) 

653 

654 # if g is close to 0, then we are done 

655 res = max(self.xp.linalg.norm(tmpgu, self.xp.inf), self.xp.linalg.norm(tmpgv, self.xp.inf)) 

656 if res < self.newton_tol: 

657 break 

658 

659 # assemble dg 

660 dg00 = 1 - factor * (-(tmpv**2) - self.A) 

661 dg01 = -factor * (-2 * tmpu * tmpv) 

662 dg10 = -factor * (tmpv**2) 

663 dg11 = 1 - factor * (2 * tmpu * tmpv - self.B) 

664 

665 # interleave and unravel to put into sparse matrix 

666 dg00I = self.xp.ravel(self.xp.kron(dg00, self.xp.array([1, 0]))) 

667 dg01I = self.xp.ravel(self.xp.kron(dg01, self.xp.array([1, 0]))) 

668 dg10I = self.xp.ravel(self.xp.kron(dg10, self.xp.array([1, 0]))) 

669 dg11I = self.xp.ravel(self.xp.kron(dg11, self.xp.array([0, 1]))) 

670 

671 # put into sparse matrix 

672 dg = sp.diags(dg00I, offsets=0) + sp.diags(dg11I, offsets=0) 

673 dg += sp.diags(dg01I, offsets=1, shape=dg.shape) + sp.diags(dg10I, offsets=-1, shape=dg.shape) 

674 

675 # interleave g terms to apply inverse to it 

676 g = self.xp.kron(tmpgu.flatten(), self.xp.array([1, 0])) + self.xp.kron( 

677 tmpgv.flatten(), self.xp.array([0, 1]) 

678 ) 

679 # invert dg matrix 

680 b = sp.linalg.spsolve(dg, g) 

681 # update real space vectors 

682 tmpu[:] -= b[::2].reshape(self.nvars) 

683 tmpv[:] -= b[1::2].reshape(self.nvars) 

684 

685 # increase iteration count 

686 n += 1 

687 self.work_counters['newton']() 

688 

689 if self.xp.isnan(res) and self.stop_at_nan: 

690 raise ProblemError('Newton got nan after %i iterations, aborting...' % n) 

691 elif self.xp.isnan(res): 

692 self.logger.warning('Newton got nan after %i iterations...' % n) 

693 

694 if n == self.newton_maxiter: 

695 self.logger.warning('Newton did not converge after %i iterations, error is %s' % (n, res)) 

696 

697 me = self.dtype_u(self.init) 

698 if self.spectral: 

699 me[0, ...] = self.fft.forward(tmpu) 

700 me[1, ...] = self.fft.forward(tmpv) 

701 else: 

702 me[0, ...] = tmpu 

703 me[1, ...] = tmpv 

704 return me 

705 

706 

707class grayscott_mi_linear(grayscott_imex_linear): 

708 r""" 

709 Periodic Gray-Scott system with mpi4py-fft, multi-implicit: diffusion and linear terms by FFT, the rest by Newton. 

710 

711 The original Gray-Scott system [#]_ describes a reaction-diffusion process of two substances :math:`u` and :math:`v`, 

712 where they diffuse over time. :math:`u` is fed with rate :math:`A` and used up by the reaction, which produces 

713 :math:`v`, and :math:`v` is removed with rate :math:`B`. :math:`D_u,\, D_v` are the diffusion rates for 

714 :math:`u,\, v`. This is the same model as in ``grayscott_imex_diffusion``, with the linear reaction terms grouped with the 

715 diffusion: 

716 

717 .. math:: 

718 \frac{\partial u}{\partial t} = (D_u \Delta - A) u - u v^2 + A, 

719 

720 .. math:: 

721 \frac{\partial v}{\partial t} = (D_v \Delta - B) v + u v^2 

722 

723 in :math:`x \in \Omega:=[-L/2, L/2]^N` with :math:`N=2,3`. Spatial discretization is done by using 

724 Fast Fourier transformation for solving the linear parts provided by ``mpi4py-fft`` [#]_, see also 

725 https://mpi4py-fft.readthedocs.io/en/latest/. 

726 

727 The problem in this class will be treated in a *multi-implicit* way for time-stepping, i.e., the system for the 

728 diffusion and the linear reaction terms is solved by FFT, and the one for the rest of the reaction by a Newton 

729 solver. 

730 

731 Parameters 

732 ---------- 

733 newton_maxiter : int, optional 

734 Maximum number of iterations for Newton's method in ``solve_system_2``. 

735 newton_tol : float, optional 

736 Absolute tolerance for Newton's method to terminate, applied to the maximum norm of the residuals of both 

737 components in real space. 

738 stop_at_nan : bool, optional 

739 Raise a ``ProblemError`` if Newton's method produces ``nan``, instead of only logging a warning. 

740 **kwargs 

741 Passed on to ``grayscott_imex_diffusion``, see there: ``Du``, ``Dv``, ``A``, ``B``, ``L``, ``num_blobs``, 

742 ``nvars``, ``spectral``, ``comm`` and ``useGPU``. The Newton solver runs on a single process only, so 

743 ``comm`` has to have size one. 

744 

745 References 

746 ---------- 

747 .. [#] Autocatalytic reactions in the isothermal, continuous stirred tank reactor: Isolas and other forms 

748 of multistability. P. Gray, S. K. Scott. Chem. Eng. Sci. 38, 1 (1983). 

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

750 Journal of Parallel and Distributed Computing (2019). 

751 """ 

752 

753 dtype_f = comp2_mesh 

754 

755 def __init__( 

756 self, 

757 newton_maxiter=100, 

758 newton_tol=1e-12, 

759 stop_at_nan=True, 

760 **kwargs, 

761 ): 

762 """Initialization routine""" 

763 super().__init__(**kwargs) 

764 # This may not run in parallel yet.. 

765 assert self.comm.Get_size() == 1 

766 self.work_counters['newton'] = WorkCounter() 

767 self.Ku = -self.Du * self.K2 - self.A 

768 self.Kv = -self.Dv * self.K2 - self.B 

769 self._makeAttributeAndRegister( 

770 'newton_maxiter', 'newton_tol', 'stop_at_nan', localVars=locals(), readOnly=False 

771 ) 

772 

773 def eval_f(self, u, t): 

774 """ 

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

776 

777 Parameters 

778 ---------- 

779 u : dtype_u 

780 Current values of the numerical solution. 

781 t : float 

782 Current time of the numerical solution is computed. 

783 

784 Returns 

785 ------- 

786 f : dtype_f 

787 The right-hand side of the problem. 

788 """ 

789 

790 f = self.dtype_f(self.init)