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

267 statements

, created at 2024-09-09 14:59 +0000

1import scipy.sparse as sp

2from mpi4py import MPI

3from mpi4py_fft import PFFT

5from pySDC.core.errors import ProblemError

6from pySDC.core.problem import Problem, WorkCounter

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

8from pySDC.implementations.problem_classes.generic_MPIFFT_Laplacian import IMEX_Laplacian_MPIFFT

10from mpi4py_fft import newDistArray

13class grayscott_imex_diffusion(IMEX_Laplacian_MPIFFT):

14 r"""

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

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

17 whereas :math:v is produced with feed rate :math:A. :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

20 .. math::

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

23 .. math::

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

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 [2]_, see also

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

31 is computed in explicit fashion).

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:v.

43 B : float, optional

44 Overall decay rate for :math:u.

45 spectral : bool, optional

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

47 L : int, 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.

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).

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

74 def __init__(self, Du=1.0, Dv=0.01, A=0.09, B=0.086, **kwargs):

75 kwargs['L'] = 2.0

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

78 # prepare the array with two components

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

80 self.iU = 0

81 self.iV = 1

82 self.ncomp = 2 # needed for transfer class

83 self.init = (shape, self.comm, self.xp.dtype('float'))

85 self._makeAttributeAndRegister('Du', 'Dv', 'A', 'B', localVars=locals(), readOnly=True)

87 # prepare "Laplacians"

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

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

91 def eval_f(self, u, t):

92 """

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

95 Parameters

96 ----------

97 u : dtype_u

98 Current values of the numerical solution.

99 t : float

100 Current time of the numerical solution is computed.

102 Returns

103 -------

104 f : dtype_f

105 The right-hand side of the problem.

106 """

108 f = self.dtype_f(self.init)

110 if self.spectral:

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

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

113 tmpu = newDistArray(self.fft, False)

114 tmpv = newDistArray(self.fft, False)

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

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

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

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

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

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

122 else:

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

124 lap_u_hat = self.Ku * u_hat

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

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

127 lap_u_hat = self.Kv * u_hat

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

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

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

132 self.work_counters['rhs']()

133 return f

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

136 """

137 Simple FFT solver for the diffusion part.

139 Parameters

140 ----------

141 rhs : dtype_f

142 Right-hand side for the linear system.

143 factor : float

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

145 u0 : dtype_u

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

147 t : float

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

150 Returns

151 -------

152 me : dtype_u

153 Solution.

154 """

156 me = self.dtype_u(self.init)

157 if self.spectral:

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

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

161 else:

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

163 rhs_hat /= 1.0 - factor * self.Ku

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

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

166 rhs_hat /= 1.0 - factor * self.Kv

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

169 return me

171 def u_exact(self, t):

172 r"""

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

175 Parameters

176 ----------

177 t : float

178 Time of the exact solution.

180 Returns

181 -------

182 me : dtype_u

183 Exact solution.

184 """

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

186 assert self.ndim == 2, 'The initial conditions are 2D for now..'

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

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

191 if self.spectral:

192 tmp = 1.0 - self.xp.exp(-80.0 * ((self.X[0] + 0.05) ** 2 + (self.X[1] + 0.02) ** 2))

193 me[0, ...] = self.fft.forward(tmp)

194 tmp = self.xp.exp(-80.0 * ((self.X[0] - 0.05) ** 2 + (self.X[1] - 0.02) ** 2))

195 me[1, ...] = self.fft.forward(tmp)

196 else:

197 me[0, ...] = 1.0 - self.xp.exp(-80.0 * ((self.X[0] + 0.05) ** 2 + (self.X[1] + 0.02) ** 2))

198 me[1, ...] = self.xp.exp(-80.0 * ((self.X[0] - 0.05) ** 2 + (self.X[1] - 0.02) ** 2))

200 return me

203class grayscott_imex_linear(grayscott_imex_diffusion):

204 r"""

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

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

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

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

210 .. math::

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

213 .. math::

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

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

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

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

221 part is computed in an explicit way).

222 """

224 def __init__(self, **kwargs):

225 super().__init__(**kwargs)

226 self.Ku -= self.A

227 self.Kv -= self.B

229 def eval_f(self, u, t):

230 """

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

233 Parameters

234 ----------

235 u : dtype_u

236 Current values of the numerical solution.

237 t : float

238 Current time of the numerical solution is computed.

240 Returns

241 -------

242 f : dtype_f

243 The right-hand side of the problem.

244 """

246 f = self.dtype_f(self.init)

248 if self.spectral:

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

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

251 tmpu = newDistArray(self.fft, False)

252 tmpv = newDistArray(self.fft, False)

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

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

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

256 tmpfv = tmpu * tmpv**2

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

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

260 else:

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

262 lap_u_hat = self.Ku * u_hat

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

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

265 lap_u_hat = self.Kv * u_hat

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

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

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

270 self.work_counters['rhs']()

271 return f

274class grayscott_mi_diffusion(grayscott_imex_diffusion):

275 r"""

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

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

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

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

281 .. math::

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

284 .. math::

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

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

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

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

292 implicitly.

294 Parameters

295 ----------

296 nvars : tuple of int, optional

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

298 Du : float, optional

299 Diffusion rate for :math:u.

300 Dv: float, optional

301 Diffusion rate for :math:v.

302 A : float, optional

303 Feed rate for :math:v.

304 B : float, optional

305 Overall decay rate for :math:u.

306 spectral : bool, optional

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

308 L : int, optional

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

310 comm : COMM_WORLD, optional

311 Communicator for mpi4py-fft.

313 Attributes

314 ----------

315 fft : PFFT

316 Object for parallel FFT transforms.

317 X : mesh-grid

318 Grid coordinates in real space.

319 ndim : int

320 Number of spatial dimensions.

321 Ku : matrix

322 Laplace operator in spectral space (u component).

323 Kv : matrix

324 Laplace operator in spectral space (v component).

326 References

327 ----------

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

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

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

331 Journal of Parallel and Distributed Computing (2019).

332 """

334 dtype_f = comp2_mesh

336 def __init__(

337 self,

338 newton_maxiter=100,

339 newton_tol=1e-12,

340 **kwargs,

341 ):

342 """Initialization routine"""

343 super().__init__(**kwargs)

344 # This may not run in parallel yet..

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

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

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

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

351 def eval_f(self, u, t):

352 """

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

355 Parameters

356 ----------

357 u : dtype_u

358 Current values of the numerical solution.

359 t : float

360 Current time of the numerical solution is computed.

362 Returns

363 -------

364 f : dtype_f

365 The right-hand side of the problem.

366 """

368 f = self.dtype_f(self.init)

370 if self.spectral:

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

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

373 tmpu = newDistArray(self.fft, False)

374 tmpv = newDistArray(self.fft, False)

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

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

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

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

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

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

382 else:

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

384 lap_u_hat = self.Ku * u_hat

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

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

387 lap_u_hat = self.Kv * u_hat

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

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

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

392 self.work_counters['rhs']()

393 return f

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

396 """

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

399 Parameters

400 ----------

401 rhs : dtype_f

402 Right-hand side for the linear system.

403 factor : float

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

405 u0 : dtype_u

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

407 t : float

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

410 Returns

411 -------

412 me : dtype_u

413 The solution as mesh.

414 """

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

417 return me

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

420 """

421 Newton-Solver for the second component.

423 Parameters

424 ----------

425 rhs : dtype_f

426 Right-hand side for the linear system.

427 factor float

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

429 u0 : dtype_u

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

431 t : float

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

434 Returns

435 -------

436 me : dtype_u

437 The solution as mesh.

438 """

439 u = self.dtype_u(u0)

441 if self.spectral:

442 tmpu = newDistArray(self.fft, False)

443 tmpv = newDistArray(self.fft, False)

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

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

446 tmprhsu = newDistArray(self.fft, False)

447 tmprhsv = newDistArray(self.fft, False)

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

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

451 else:

452 tmpu = u[0, ...]

453 tmpv = u[1, ...]

454 tmprhsu = rhs[0, ...]

455 tmprhsv = rhs[1, ...]

457 # start newton iteration

458 n = 0

459 res = 99

460 while n < self.newton_maxiter:

461 # print(n, res)

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

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

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

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

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

468 if res < self.newton_tol:

469 break

471 # assemble dg

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

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

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

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

477 # interleave and unravel to put into sparse matrix

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

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

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

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

483 # put into sparse matrix

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

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

487 # interleave g terms to apply inverse to it

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

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

490 )

491 # invert dg matrix

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

493 # update real space vectors

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

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

497 # increase iteration count

498 n += 1

499 self.work_counters['newton']()

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

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

503 elif self.xp.isnan(res):

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

506 if n == self.newton_maxiter:

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

509 me = self.dtype_u(self.init)

510 if self.spectral:

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

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

513 else:

514 me[0, ...] = tmpu

515 me[1, ...] = tmpv

516 return me

519class grayscott_mi_linear(grayscott_imex_linear):

520 r"""

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

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

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

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

526 .. math::

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

529 .. math::

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

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

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

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

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

538 """

540 dtype_f = comp2_mesh

542 def __init__(

543 self,

544 newton_maxiter=100,

545 newton_tol=1e-12,

546 **kwargs,

547 ):

548 """Initialization routine"""

549 super().__init__(**kwargs)

550 # This may not run in parallel yet..

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

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

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

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

557 def eval_f(self, u, t):

558 """

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

561 Parameters

562 ----------

563 u : dtype_u

564 Current values of the numerical solution.

565 t : float

566 Current time of the numerical solution is computed.

568 Returns

569 -------

570 f : dtype_f

571 The right-hand side of the problem.

572 """

574 f = self.dtype_f(self.init)

576 if self.spectral:

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

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

579 tmpu = newDistArray(self.fft, False)

580 tmpv = newDistArray(self.fft, False)

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

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

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

584 tmpfv = tmpu * tmpv**2

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

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

588 else:

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

590 lap_u_hat = self.Ku * u_hat

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

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

593 lap_u_hat = self.Kv * u_hat

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

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

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

598 self.work_counters['rhs']()

599 return f

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

602 """

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

605 Parameters

606 ----------

607 rhs : dtype_f

608 Right-hand side for the linear system.

609 factor : float

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

611 u0 : dtype_u

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

613 t : float

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

616 Returns

617 -------

618 me : dtype_u

619 The solution as mesh.

620 """

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

623 return me

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

626 """

627 Newton-Solver for the second component.

629 Parameters

630 ----------

631 rhs : dtype_f

632 Right-hand side for the linear system.

633 factor : float

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

635 u0 : dtype_u

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

637 t : float

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

640 Returns

641 -------

642 u : dtype_u

643 The solution as mesh.

644 """

645 u = self.dtype_u(u0)

647 if self.spectral:

648 tmpu = newDistArray(self.fft, False)

649 tmpv = newDistArray(self.fft, False)

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

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

652 tmprhsu = newDistArray(self.fft, False)

653 tmprhsv = newDistArray(self.fft, False)

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

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

657 else:

658 tmpu = u[0, ...]

659 tmpv = u[1, ...]

660 tmprhsu = rhs[0, ...]

661 tmprhsv = rhs[1, ...]

663 # start newton iteration

664 n = 0

665 res = 99

666 while n < self.newton_maxiter:

667 # print(n, res)

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

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

670 tmpgv = tmpv - tmprhsv - factor * (tmpu * tmpv**2)

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

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

674 if res < self.newton_tol:

675 break

677 # assemble dg

678 dg00 = 1 - factor * (-(tmpv**2))

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

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

681 dg11 = 1 - factor * (2 * tmpu * tmpv)

683 # interleave and unravel to put into sparse matrix

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

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

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

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

689 # put into sparse matrix

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

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

693 # interleave g terms to apply inverse to it

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

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

696 )

697 # invert dg matrix

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

699 # update real-space vectors

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

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

703 # increase iteration count

704 n += 1

705 self.work_counters['newton']()

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

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

709 elif self.xp.isnan(res):

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

712 if n == self.newton_maxiter:

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

715 me = self.dtype_u(self.init)

716 if self.spectral:

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

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

719 else:

720 me[0, ...] = tmpu

721 me[1, ...] = tmpv

722 return me