# Coverage for pySDC/implementations/problem_classes/NonlinearSchroedinger_MPIFFT.py: 93%

## 60 statements

, created at 2024-09-20 17:10 +0000

1import numpy as np

2from scipy.optimize import newton_krylov

3from scipy.optimize import NoConvergence

5from pySDC.core.errors import ProblemError

6from pySDC.core.problem import WorkCounter

7from pySDC.implementations.problem_classes.generic_MPIFFT_Laplacian import IMEX_Laplacian_MPIFFT

8from pySDC.implementations.datatype_classes.mesh import mesh

11class nonlinearschroedinger_imex(IMEX_Laplacian_MPIFFT):

12 r"""

13 Example implementing the :math:N-dimensional nonlinear Schrödinger equation with periodic boundary conditions

15 .. math::

16 \frac{\partial u}{\partial t} = -i \Delta u + 2 c i |u|^2 u

18 for fixed parameter :math:c and :math:N=2, 3. The linear parts of the problem will be solved using

19 mpi4py-fft [1]_. *Semi-explicit* time-stepping is used here to solve the problem in the temporal dimension, i.e., the

20 Laplacian will be handled implicitly.

22 Parameters

23 ----------

24 nvars : tuple, optional

25 Spatial resolution

26 spectral : bool, optional

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

28 L : float, optional

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

30 c : float, optional

31 Nonlinearity parameter.

32 comm : MPI.COMM_World

33 Communicator for parallelisation.

35 Attributes

36 ----------

37 fft : PFFT

38 Object for parallel FFT transforms.

39 X : mesh-grid

40 Grid coordinates in real space.

41 K2 : matrix

42 Laplace operator in spectral space.

44 References

45 ----------

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

47 Journal of Parallel and Distributed Computing (2019).

48 """

50 def __init__(self, c=1.0, **kwargs):

51 """Initialization routine"""

52 super().__init__(L=2 * np.pi, alpha=1j, dtype='D', **kwargs)

54 if not (c == 0.0 or c == 1.0):

55 raise ProblemError(f'Setup not implemented, c has to be 0 or 1, got {c}')

56 self._makeAttributeAndRegister('c', localVars=locals(), readOnly=True)

58 def _eval_explicit_part(self, u, t, f_expl):

59 f_expl[:] = self.ndim * self.c * 2j * self.xp.absolute(u) ** 2 * u

60 return f_expl

62 def u_exact(self, t, **kwargs):

63 r"""

64 Routine to compute the exact solution at time :math:t, see (39) from https://doi.org/10.1007/BF01017105 for details

66 Parameters

67 ----------

68 t : float

69 Time of the exact solution.

71 Returns

72 -------

73 u : dtype_u

74 The exact solution.

75 """

76 if 'u_init' in kwargs.keys() or 't_init' in kwargs.keys():

77 self.logger.warning(

78 f'{type(self).__name__} uses an analytic exact solution from t=0. If you try to compute the local error, you will get the global error instead!'

79 )

81 def nls_exact_1D(t, x, c):

82 ae = 1.0 / np.sqrt(2.0) * np.exp(1j * t)

83 if c != 0:

84 u = ae * ((np.cosh(t) + 1j * np.sinh(t)) / (np.cosh(t) - 1.0 / np.sqrt(2.0) * self.xp.cos(x)) - 1.0)

85 else:

86 u = self.xp.sin(x) * np.exp(-t * 1j)

88 return u

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

92 if self.spectral:

93 tmp = nls_exact_1D(self.ndim * t, sum(self.X), self.c)

94 me[:] = self.fft.forward(tmp)

95 else:

96 me[:] = nls_exact_1D(self.ndim * t, sum(self.X), self.c)

98 return me

101class nonlinearschroedinger_fully_implicit(nonlinearschroedinger_imex):

102 r"""

103 Example implementing the :math:N-dimensional nonlinear Schrödinger equation with periodic boundary conditions

105 .. math::

106 \frac{\partial u}{\partial t} = -i \Delta u + 2 c i |u|^2 u

108 for fixed parameter :math:c and :math:N=2, 3. The linear parts of the problem will be discretized using

109 mpi4py-fft [1]_. For time-stepping, the problem will be solved *fully-implicitly*, i.e., the nonlinear system containing

110 the full right-hand side is solved by GMRES method.

112 References

113 ----------

114 .. [1] https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.newton_krylov.html

115 """

117 dtype_u = mesh

118 dtype_f = mesh

120 def __init__(self, lintol=1e-9, liniter=99, **kwargs):

121 assert kwargs.get('useGPU', False) is False

123 super().__init__(**kwargs)

124 self._makeAttributeAndRegister('liniter', 'lintol', localVars=locals(), readOnly=False)

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

128 def eval_f(self, u, t):

129 """

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

132 Parameters

133 ----------

134 u : dtype_u

135 Current values of the numerical solution.

136 t : float

137 Current time at which the numerical solution is computed.

139 Returns

140 -------

141 f : dtype_f

142 The right-hand side of the problem.

143 """

145 f = self.dtype_f(self.init)

147 if self.spectral:

148 tmp = self.fft.backward(u)

149 tmpf = self.ndim * self.c * 2j * self.xp.absolute(tmp) ** 2 * tmp

150 f[:] = -self.K2 * 1j * u + self.fft.forward(tmpf)

152 else:

153 u_hat = self.fft.forward(u)

154 lap_u_hat = -self.K2 * 1j * u_hat

155 f[:] = self.fft.backward(lap_u_hat) + self.ndim * self.c * 2j * self.xp.absolute(u) ** 2 * u

157 self.work_counters['rhs']()

158 return f

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

161 r"""

162 Solve the nonlinear system :math:(1 - factor \cdot f)(\vec{u}) = \vec{rhs} using a SciPy Newton-Krylov

163 solver. See page [1]_ for details on the solver.

165 Parameters

166 ----------

167 rhs : dtype_f

168 Right-hand side for the linear system.

169 factor : float

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

171 u0 : dtype_u

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

173 t : float

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

176 Returns

177 -------

178 me : dtype_u

179 The solution as mesh.

180 """

181 me = self.dtype_u(self.init)

183 # assemble the nonlinear function F for the solver

184 def F(x):

185 r"""

186 Nonlinear function for the SciPy solver.

188 Args:

189 x : dtype_u

190 Current solution

191 """

192 self.work_counters['rhs'].decrement()

193 return x - factor * self.eval_f(u=x.reshape(self.init[0]), t=t).reshape(x.shape) - rhs.reshape(x.shape)

195 try:

196 sol = newton_krylov(

197 F=F,

198 xin=u0.copy(),

199 maxiter=self.liniter,

200 x_tol=self.lintol,

201 callback=self.work_counters['newton'],

202 method='gmres',

203 )

204 except NoConvergence as e:

205 sol = e.args[0]

207 me[:] = sol

208 return me