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

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 The Gray-Scott system [1]_ describes a reaction-diffusion process of two substances :math:`u` and :math:`v`, 

14 where they diffuse over time. During the reaction :math:`u` is used up with overall decay rate :math:`B`, 

15 whereas :math:`v` is produced with feed rate :math:`A`. :math:`D_u,\, D_v` are the diffusion rates for 

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

17 

18 .. math:: 

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

20 

21 .. math:: 

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

23 

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

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

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

27 

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

29 is computed in explicit fashion). 

30 

31 Parameters 

32 ---------- 

33 nvars : tuple of int, optional 

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

35 Du : float, optional 

36 Diffusion rate for :math:`u`. 

37 Dv: float, optional 

38 Diffusion rate for :math:`v`. 

39 A : float, optional 

40 Feed rate for :math:`v`. 

41 B : float, optional 

42 Overall decay rate for :math:`u`. 

43 spectral : bool, optional 

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

45 L : int, optional 

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

47 comm : COMM_WORLD, optional 

48 Communicator for ``mpi4py-fft``. 

49 num_blobs : int, optional 

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

51 

52 Attributes 

53 ---------- 

54 fft : PFFT 

55 Object for parallel FFT transforms. 

56 X : mesh-grid 

57 Grid coordinates in real space. 

58 ndim : int 

59 Number of spatial dimensions. 

60 Ku : matrix 

61 Laplace operator in spectral space (u component). 

62 Kv : matrix 

63 Laplace operator in spectral space (v component). 

64 

65 References 

66 ---------- 

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

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

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

70 Journal of Parallel and Distributed Computing (2019). 

71 .. [3] https://www.chebfun.org/examples/pde/GrayScott.html 

72 """ 

73 

74 def __init__( 

75 self, 

76 Du=1.0, 

77 Dv=0.01, 

78 A=0.09, 

79 B=0.086, 

80 L=2.0, 

81 num_blobs=1, 

82 **kwargs, 

83 ): 

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

85 

86 # prepare the array with two components 

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

88 self.iU = 0 

89 self.iV = 1 

90 self.ncomp = 2 # needed for transfer class 

91 

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

93 

94 self._makeAttributeAndRegister( 

95 'Du', 

96 'Dv', 

97 'A', 

98 'B', 

99 'num_blobs', 

100 localVars=locals(), 

101 readOnly=True, 

102 ) 

103 

104 # prepare "Laplacians" 

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

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

107 

108 def eval_f(self, u, t): 

109 """ 

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

111 

112 Parameters 

113 ---------- 

114 u : dtype_u 

115 Current values of the numerical solution. 

116 t : float 

117 Current time of the numerical solution is computed. 

118 

119 Returns 

120 ------- 

121 f : dtype_f 

122 The right-hand side of the problem. 

123 """ 

124 

125 f = self.dtype_f(self.init) 

126 

127 if self.spectral: 

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

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

130 tmpu = newDistArray(self.fft, False) 

131 tmpv = newDistArray(self.fft, False) 

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

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

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

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

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

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

138 

139 else: 

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

141 lap_u_hat = self.Ku * u_hat 

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

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

144 lap_u_hat = self.Kv * u_hat 

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

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

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

148 

149 self.work_counters['rhs']() 

150 return f 

151 

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

153 """ 

154 Simple FFT solver for the diffusion part. 

155 

156 Parameters 

157 ---------- 

158 rhs : dtype_f 

159 Right-hand side for the linear system. 

160 factor : float 

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

162 u0 : dtype_u 

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

164 t : float 

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

166 

167 Returns 

168 ------- 

169 me : dtype_u 

170 Solution. 

171 """ 

172 

173 me = self.dtype_u(self.init) 

174 if self.spectral: 

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

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

177 

178 else: 

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

180 rhs_hat /= 1.0 - factor * self.Ku 

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

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

183 rhs_hat /= 1.0 - factor * self.Kv 

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

185 

186 return me 

187 

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

189 r""" 

190 Routine to compute the exact solution at time :math:`t = 0`, see [3]_. 

191 

192 Parameters 

193 ---------- 

194 t : float 

195 Time of the exact solution. 

196 

197 Returns 

198 ------- 

199 me : dtype_u 

200 Exact solution. 

201 """ 

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

203 

204 xp = self.xp 

205 

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

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

208 

209 rng = xp.random.default_rng(seed) 

210 

211 if self.num_blobs < 0: 

212 """ 

213 Rectangles with stationary background, see arXiv:1501.01990 

214 """ 

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

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

217 

218 # set stable background state from Equation 2 

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

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

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

222 

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

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

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

226 

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

228 mask = masks[0] 

229 for m in masks[1:]: 

230 mask = xp.logical_and(mask, m) 

231 

232 _u[mask] = rng.random() 

233 _v[mask] = rng.random() 

234 

235 elif self.num_blobs > 0: 

236 """ 

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

238 """ 

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

240 

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

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

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

244 

245 if self.ndim == 2: 

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

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

248 -80.0 

249 * ( 

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

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

252 ) 

253 ) 

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

255 -80.0 

256 * ( 

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

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

259 ) 

260 ) 

261 elif self.ndim == 3: 

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

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

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

265 -80.0 

266 * ( 

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

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

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

270 ) 

271 ) 

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

273 -80.0 

274 * ( 

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

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

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

278 ) 

279 ) 

280 else: 

281 raise NotImplementedError 

282 

283 _u += 1 

284 else: 

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

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

287 

288 u = self.u_init 

289 if self.spectral: 

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

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

292 else: 

293 u[0, ...] = _u 

294 u[1, ...] = _v 

295 

296 return u 

297 

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

299 """ 

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

301 

302 Args: 

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

304 

305 Returns 

306 ------- 

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

308 """ 

309 import matplotlib.pyplot as plt 

310 from mpl_toolkits.axes_grid1 import make_axes_locatable 

311 

312 plt.rcParams['figure.constrained_layout.use'] = True 

313 

314 if n_comps == 2: 

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

316 divider = make_axes_locatable(axs[1]) 

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

318 else: 

319 self.fig, ax = plt.subplots(1, 1, figsize=((6, 5))) 

320 divider = make_axes_locatable(ax) 

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

322 return self.fig 

323 

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

325 r""" 

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

327 

328 Parameters 

329 ---------- 

330 u : dtype_u 

331 Solution to be plotted 

332 t : float 

333 Time to display at the top of the figure 

334 fig : matplotlib.pyplot.figure.Figure 

335 Figure with the correct structure 

336 

337 Returns 

338 ------- 

339 None 

340 """ 

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

342 axs = fig.axes 

343 

344 vmin = u.min() 

345 vmax = u.max() 

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

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

348 axs[i].set_aspect(1) 

349 axs[i].set_title(label) 

350 

351 if t is not None: 

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

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

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

355 fig.colorbar(im, self.cax) 

356 

357 

358class grayscott_imex_linear(grayscott_imex_diffusion): 

359 r""" 

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

361 where they diffuse over time. During the reaction :math:`u` is used up with overall decay rate :math:`B`, 

362 whereas :math:`v` is produced with feed rate :math:`A`. :math:`D_u,\, D_v` are the diffusion rates for 

363 :math:`u,\, v`. The model with linear (reaction) part is described by the :math:`N`-dimensional model 

364 

365 .. math:: 

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

367 

368 .. math:: 

369 \frac{d v}{d t} = D_v \Delta v + u v^2 

370 

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

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

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

374 

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

376 part is computed in an explicit way). 

377 """ 

378 

379 def __init__(self, **kwargs): 

380 super().__init__(**kwargs) 

381 self.Ku -= self.A 

382 self.Kv -= self.B 

383 

384 def eval_f(self, u, t): 

385 """ 

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

387 

388 Parameters 

389 ---------- 

390 u : dtype_u 

391 Current values of the numerical solution. 

392 t : float 

393 Current time of the numerical solution is computed. 

394 

395 Returns 

396 ------- 

397 f : dtype_f 

398 The right-hand side of the problem. 

399 """ 

400 

401 f = self.dtype_f(self.init) 

402 

403 if self.spectral: 

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

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

406 tmpu = newDistArray(self.fft, False) 

407 tmpv = newDistArray(self.fft, False) 

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

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

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

411 tmpfv = tmpu * tmpv**2 

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

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

414 

415 else: 

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

417 lap_u_hat = self.Ku * u_hat 

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

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

420 lap_u_hat = self.Kv * u_hat 

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

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

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

424 

425 self.work_counters['rhs']() 

426 return f 

427 

428 

429class grayscott_mi_diffusion(grayscott_imex_diffusion): 

430 r""" 

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

432 where they diffuse over time. During the reaction :math:`u` is used up with overall decay rate :math:`B`, 

433 whereas :math:`v` is produced with feed rate :math:`A`. :math:`D_u,\, D_v` are the diffusion rates for 

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

435 

436 .. math:: 

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

438 

439 .. math:: 

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

441 

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

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

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

445 

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

447 implicitly. 

448 

449 Parameters 

450 ---------- 

451 nvars : tuple of int, optional 

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

453 Du : float, optional 

454 Diffusion rate for :math:`u`. 

455 Dv: float, optional 

456 Diffusion rate for :math:`v`. 

457 A : float, optional 

458 Feed rate for :math:`v`. 

459 B : float, optional 

460 Overall decay rate for :math:`u`. 

461 spectral : bool, optional 

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

463 L : int, optional 

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

465 comm : COMM_WORLD, optional 

466 Communicator for ``mpi4py-fft``. 

467 

468 Attributes 

469 ---------- 

470 fft : PFFT 

471 Object for parallel FFT transforms. 

472 X : mesh-grid 

473 Grid coordinates in real space. 

474 ndim : int 

475 Number of spatial dimensions. 

476 Ku : matrix 

477 Laplace operator in spectral space (u component). 

478 Kv : matrix 

479 Laplace operator in spectral space (v component). 

480 

481 References 

482 ---------- 

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

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

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

486 Journal of Parallel and Distributed Computing (2019). 

487 """ 

488 

489 dtype_f = comp2_mesh 

490 

491 def __init__( 

492 self, 

493 newton_maxiter=100, 

494 newton_tol=1e-12, 

495 **kwargs, 

496 ): 

497 """Initialization routine""" 

498 super().__init__(**kwargs) 

499 # This may not run in parallel yet.. 

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

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

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

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

504 self._makeAttributeAndRegister('newton_maxiter', 'newton_tol', localVars=locals(), readOnly=False) 

505 

506 def eval_f(self, u, t): 

507 """ 

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

509 

510 Parameters 

511 ---------- 

512 u : dtype_u 

513 Current values of the numerical solution. 

514 t : float 

515 Current time of the numerical solution is computed. 

516 

517 Returns 

518 ------- 

519 f : dtype_f 

520 The right-hand side of the problem. 

521 """ 

522 

523 f = self.dtype_f(self.init) 

524 

525 if self.spectral: 

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

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

528 tmpu = newDistArray(self.fft, False) 

529 tmpv = newDistArray(self.fft, False) 

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

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

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

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

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

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

536 

537 else: 

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

539 lap_u_hat = self.Ku * u_hat 

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

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

542 lap_u_hat = self.Kv * u_hat 

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

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

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

546 

547 self.work_counters['rhs']() 

548 return f 

549 

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

551 """ 

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

553 

554 Parameters 

555 ---------- 

556 rhs : dtype_f 

557 Right-hand side for the linear system. 

558 factor : float 

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

560 u0 : dtype_u 

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

562 t : float 

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

564 

565 Returns 

566 ------- 

567 me : dtype_u 

568 The solution as mesh. 

569 """ 

570 

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

572 return me 

573 

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

575 """ 

576 Newton-Solver for the second component. 

577 

578 Parameters 

579 ---------- 

580 rhs : dtype_f 

581 Right-hand side for the linear system. 

582 factor float 

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

584 u0 : dtype_u 

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

586 t : float 

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

588 

589 Returns 

590 ------- 

591 me : dtype_u 

592 The solution as mesh. 

593 """ 

594 u = self.dtype_u(u0) 

595 

596 if self.spectral: 

597 tmpu = newDistArray(self.fft, False) 

598 tmpv = newDistArray(self.fft, False) 

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

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

601 tmprhsu = newDistArray(self.fft, False) 

602 tmprhsv = newDistArray(self.fft, False) 

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

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

605 

606 else: 

607 tmpu = u[0, ...] 

608 tmpv = u[1, ...] 

609 tmprhsu = rhs[0, ...] 

610 tmprhsv = rhs[1, ...] 

611 

612 # start newton iteration 

613 n = 0 

614 res = 99 

615 while n < self.newton_maxiter: 

616 # print(n, res) 

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

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

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

620 

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

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

623 if res < self.newton_tol: 

624 break 

625 

626 # assemble dg 

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

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

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

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

631 

632 # interleave and unravel to put into sparse matrix 

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

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

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

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

637 

638 # put into sparse matrix 

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

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

641 

642 # interleave g terms to apply inverse to it 

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

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

645 ) 

646 # invert dg matrix 

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

648 # update real space vectors 

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

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

651 

652 # increase iteration count 

653 n += 1 

654 self.work_counters['newton']() 

655 

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

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

658 elif self.xp.isnan(res): 

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

660 

661 if n == self.newton_maxiter: 

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

663 

664 me = self.dtype_u(self.init) 

665 if self.spectral: 

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

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

668 else: 

669 me[0, ...] = tmpu 

670 me[1, ...] = tmpv 

671 return me 

672 

673 

674class grayscott_mi_linear(grayscott_imex_linear): 

675 r""" 

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

677 where they diffuse over time. During the reaction :math:`u` is used up with overall decay rate :math:`B`, 

678 whereas :math:`v` is produced with feed rate :math:`A`. :math:`D_u,\, D_v` are the diffusion rates for 

679 :math:`u,\, v`. The model with linear (reaction) part is described by the :math:`N`-dimensional model 

680 

681 .. math:: 

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

683 

684 .. math:: 

685 \frac{\partial v}{\partial t} = D_v \Delta v + u v^2 

686 

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

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

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

690 

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

692 the diffusion part will be solved by FFT, and for the linear part a Newton solver is used. 

693 """ 

694 

695 dtype_f = comp2_mesh 

696 

697 def __init__( 

698 self, 

699 newton_maxiter=100, 

700 newton_tol=1e-12, 

701 **kwargs, 

702 ): 

703 """Initialization routine""" 

704 super().__init__(**kwargs) 

705 # This may not run in parallel yet.. 

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

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

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

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

710 self._makeAttributeAndRegister('newton_maxiter', 'newton_tol', localVars=locals(), readOnly=False) 

711 

712 def eval_f(self, u, t): 

713 """ 

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

715 

716 Parameters 

717 ---------- 

718 u : dtype_u 

719 Current values of the numerical solution. 

720 t : float 

721 Current time of the numerical solution is computed. 

722 

723 Returns 

724 ------- 

725 f : dtype_f 

726 The right-hand side of the problem. 

727 """ 

728 

729 f = self.dtype_f(self.init) 

730 

731 if self.spectral: 

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

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

734 tmpu = newDistArray(self.fft, False) 

735 tmpv = newDistArray(self.fft, False) 

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

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

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

739 tmpfv = tmpu * tmpv**2 

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

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

742 

743 else: 

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

745 lap_u_hat = self.Ku * u_hat 

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

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

748 lap_u_hat = self.Kv * u_hat 

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

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