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

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 u 

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] * 30 

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 assert self.ndim == 2, 'The initial conditions are 2D for now..' 

240 

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

242 

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

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

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

246 

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

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

249 -80.0 

250 * ( 

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

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

253 ) 

254 ) 

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

256 -80.0 

257 * ( 

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

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

260 ) 

261 ) 

262 

263 _u += 1 

264 else: 

265 raise NotImplementedError 

266 

267 u = self.u_init 

268 if self.spectral: 

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

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

271 else: 

272 u[0, ...] = _u 

273 u[1, ...] = _v 

274 

275 return u 

276 

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

278 """ 

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

280 

281 Args: 

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

283 

284 Returns 

285 ------- 

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

287 """ 

288 import matplotlib.pyplot as plt 

289 from mpl_toolkits.axes_grid1 import make_axes_locatable 

290 

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

292 

293 if n_comps == 2: 

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

295 divider = make_axes_locatable(axs[1]) 

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

297 else: 

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

299 divider = make_axes_locatable(ax) 

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

301 return self.fig 

302 

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

304 r""" 

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

306 

307 Parameters 

308 ---------- 

309 u : dtype_u 

310 Solution to be plotted 

311 t : float 

312 Time to display at the top of the figure 

313 fig : matplotlib.pyplot.figure.Figure 

314 Figure with the correct structure 

315 

316 Returns 

317 ------- 

318 None 

319 """ 

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

321 axs = fig.axes 

322 

323 vmin = u.min() 

324 vmax = u.max() 

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

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

327 axs[i].set_aspect(1) 

328 axs[i].set_title(label) 

329 

330 if t is not None: 

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

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

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

334 fig.colorbar(im, self.cax) 

335 

336 

337class grayscott_imex_linear(grayscott_imex_diffusion): 

338 r""" 

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

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

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

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

343 

344 .. math:: 

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

346 

347 .. math:: 

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

349 

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

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

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

353 

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

355 part is computed in an explicit way). 

356 """ 

357 

358 def __init__(self, **kwargs): 

359 super().__init__(**kwargs) 

360 self.Ku -= self.A 

361 self.Kv -= self.B 

362 

363 def eval_f(self, u, t): 

364 """ 

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

366 

367 Parameters 

368 ---------- 

369 u : dtype_u 

370 Current values of the numerical solution. 

371 t : float 

372 Current time of the numerical solution is computed. 

373 

374 Returns 

375 ------- 

376 f : dtype_f 

377 The right-hand side of the problem. 

378 """ 

379 

380 f = self.dtype_f(self.init) 

381 

382 if self.spectral: 

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

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

385 tmpu = newDistArray(self.fft, False) 

386 tmpv = newDistArray(self.fft, False) 

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

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

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

390 tmpfv = tmpu * tmpv**2 

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

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

393 

394 else: 

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

396 lap_u_hat = self.Ku * u_hat 

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

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

399 lap_u_hat = self.Kv * u_hat 

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

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

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

403 

404 self.work_counters['rhs']() 

405 return f 

406 

407 

408class grayscott_mi_diffusion(grayscott_imex_diffusion): 

409 r""" 

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

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

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

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

414 

415 .. math:: 

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

417 

418 .. math:: 

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

420 

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

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

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

424 

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

426 implicitly. 

427 

428 Parameters 

429 ---------- 

430 nvars : tuple of int, optional 

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

432 Du : float, optional 

433 Diffusion rate for :math:`u`. 

434 Dv: float, optional 

435 Diffusion rate for :math:`v`. 

436 A : float, optional 

437 Feed rate for :math:`v`. 

438 B : float, optional 

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

440 spectral : bool, optional 

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

442 L : int, optional 

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

444 comm : COMM_WORLD, optional 

445 Communicator for ``mpi4py-fft``. 

446 

447 Attributes 

448 ---------- 

449 fft : PFFT 

450 Object for parallel FFT transforms. 

451 X : mesh-grid 

452 Grid coordinates in real space. 

453 ndim : int 

454 Number of spatial dimensions. 

455 Ku : matrix 

456 Laplace operator in spectral space (u component). 

457 Kv : matrix 

458 Laplace operator in spectral space (v component). 

459 

460 References 

461 ---------- 

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

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

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

465 Journal of Parallel and Distributed Computing (2019). 

466 """ 

467 

468 dtype_f = comp2_mesh 

469 

470 def __init__( 

471 self, 

472 newton_maxiter=100, 

473 newton_tol=1e-12, 

474 **kwargs, 

475 ): 

476 """Initialization routine""" 

477 super().__init__(**kwargs) 

478 # This may not run in parallel yet.. 

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

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

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

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

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

484 

485 def eval_f(self, u, t): 

486 """ 

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

488 

489 Parameters 

490 ---------- 

491 u : dtype_u 

492 Current values of the numerical solution. 

493 t : float 

494 Current time of the numerical solution is computed. 

495 

496 Returns 

497 ------- 

498 f : dtype_f 

499 The right-hand side of the problem. 

500 """ 

501 

502 f = self.dtype_f(self.init) 

503 

504 if self.spectral: 

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

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

507 tmpu = newDistArray(self.fft, False) 

508 tmpv = newDistArray(self.fft, False) 

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

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

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

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

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

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

515 

516 else: 

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

518 lap_u_hat = self.Ku * u_hat 

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

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

521 lap_u_hat = self.Kv * u_hat 

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

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

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

525 

526 self.work_counters['rhs']() 

527 return f 

528 

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

530 """ 

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

532 

533 Parameters 

534 ---------- 

535 rhs : dtype_f 

536 Right-hand side for the linear system. 

537 factor : float 

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

539 u0 : dtype_u 

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

541 t : float 

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

543 

544 Returns 

545 ------- 

546 me : dtype_u 

547 The solution as mesh. 

548 """ 

549 

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

551 return me 

552 

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

554 """ 

555 Newton-Solver for the second component. 

556 

557 Parameters 

558 ---------- 

559 rhs : dtype_f 

560 Right-hand side for the linear system. 

561 factor float 

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

563 u0 : dtype_u 

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

565 t : float 

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

567 

568 Returns 

569 ------- 

570 me : dtype_u 

571 The solution as mesh. 

572 """ 

573 u = self.dtype_u(u0) 

574 

575 if self.spectral: 

576 tmpu = newDistArray(self.fft, False) 

577 tmpv = newDistArray(self.fft, False) 

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

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

580 tmprhsu = newDistArray(self.fft, False) 

581 tmprhsv = newDistArray(self.fft, False) 

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

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

584 

585 else: 

586 tmpu = u[0, ...] 

587 tmpv = u[1, ...] 

588 tmprhsu = rhs[0, ...] 

589 tmprhsv = rhs[1, ...] 

590 

591 # start newton iteration 

592 n = 0 

593 res = 99 

594 while n < self.newton_maxiter: 

595 # print(n, res) 

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

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

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

599 

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

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

602 if res < self.newton_tol: 

603 break 

604 

605 # assemble dg 

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

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

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

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

610 

611 # interleave and unravel to put into sparse matrix 

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

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

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

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

616 

617 # put into sparse matrix 

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

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

620 

621 # interleave g terms to apply inverse to it 

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

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

624 ) 

625 # invert dg matrix 

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

627 # update real space vectors 

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

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

630 

631 # increase iteration count 

632 n += 1 

633 self.work_counters['newton']() 

634 

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

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

637 elif self.xp.isnan(res): 

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

639 

640 if n == self.newton_maxiter: 

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

642 

643 me = self.dtype_u(self.init) 

644 if self.spectral: 

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

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

647 else: 

648 me[0, ...] = tmpu 

649 me[1, ...] = tmpv 

650 return me 

651 

652 

653class grayscott_mi_linear(grayscott_imex_linear): 

654 r""" 

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

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

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

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

659 

660 .. math:: 

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

662 

663 .. math:: 

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

665 

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

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

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

669 

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

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

672 """ 

673 

674 dtype_f = comp2_mesh 

675 

676 def __init__( 

677 self, 

678 newton_maxiter=100, 

679 newton_tol=1e-12, 

680 **kwargs, 

681 ): 

682 """Initialization routine""" 

683 super().__init__(**kwargs) 

684 # This may not run in parallel yet.. 

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

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

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

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

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

690 

691 def eval_f(self, u, t): 

692 """ 

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

694 

695 Parameters 

696 ---------- 

697 u : dtype_u 

698 Current values of the numerical solution. 

699 t : float 

700 Current time of the numerical solution is computed. 

701 

702 Returns 

703 ------- 

704 f : dtype_f 

705 The right-hand side of the problem. 

706 """ 

707 

708 f = self.dtype_f(self.init) 

709 

710 if self.spectral: 

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

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

713 tmpu = newDistArray(self.fft, False) 

714 tmpv = newDistArray(self.fft, False) 

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

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

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

718 tmpfv = tmpu * tmpv**2 

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

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

721 

722 else: 

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

724 lap_u_hat = self.Ku * u_hat 

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

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

727 lap_u_hat = self.Kv * u_hat 

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

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

730 f.comp2[1, ...] = u[0, ...] * u[1, ...] ** 2 

731 

732 self.work_counters['rhs']() 

733 return f 

734 

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

736 """ 

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

738 

739 Parameters 

740 ---------- 

741 rhs : dtype_f 

742 Right-hand side for the linear system. 

743 factor : float 

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

745 u0 : dtype_u 

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

747 t : float 

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

749 

750 Returns 

751 ------- 

752 me : dtype_u 

753 The solution as mesh. 

754 """ 

755 

756 me = super(grayscott_mi_linear, self).solve_system(rhs, factor, u0, t) 

757 return me 

758 

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

760 """ 

761 Newton-Solver for the second component. 

762 

763 Parameters 

764 ---------- 

765 rhs : dtype_f 

766 Right-hand side for the linear system. 

767 factor : float 

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

769 u0 : dtype_u