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

1import numpy as np 

2from scipy.optimize import newton_krylov 

3from scipy.optimize.nonlin import NoConvergence 

4 

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 

9 

10 

11class nonlinearschroedinger_imex(IMEX_Laplacian_MPIFFT): 

12 r""" 

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

14 

15 .. math:: 

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

17 

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. 

21 

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. 

34 

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. 

43 

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

49 

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

51 """Initialization routine""" 

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

53 

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) 

57 

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 

61 

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

63 r""" 

64 Routine to compute the exact solution at time :math:`t`, see (1.3) https://arxiv.org/pdf/nlin/0702010.pdf for details 

65 

66 Parameters 

67 ---------- 

68 t : float 

69 Time of the exact solution. 

70 

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 ) 

80 

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) 

87 

88 return u 

89 

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

91 

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) 

97 

98 return me 

99 

100 

101class nonlinearschroedinger_fully_implicit(nonlinearschroedinger_imex): 

102 r""" 

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

104 

105 .. math:: 

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

107 

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. 

111 

112 References 

113 ---------- 

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

115 """ 

116 

117 dtype_u = mesh 

118 dtype_f = mesh 

119 

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

121 super().__init__(**kwargs) 

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

123 

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

125 

126 def eval_f(self, u, t): 

127 """ 

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

129 

130 Parameters 

131 ---------- 

132 u : dtype_u 

133 Current values of the numerical solution. 

134 t : float 

135 Current time at which the numerical solution is computed. 

136 

137 Returns 

138 ------- 

139 f : dtype_f 

140 The right-hand side of the problem. 

141 """ 

142 

143 f = self.dtype_f(self.init) 

144 

145 if self.spectral: 

146 tmp = self.fft.backward(u) 

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

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

149 

150 else: 

151 u_hat = self.fft.forward(u) 

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

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

154 

155 self.work_counters['rhs']() 

156 return f 

157 

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

159 r""" 

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

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

162 

163 Parameters 

164 ---------- 

165 rhs : dtype_f 

166 Right-hand side for the linear system. 

167 factor : float 

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

169 u0 : dtype_u 

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

171 t : float 

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

173 

174 Returns 

175 ------- 

176 me : dtype_u 

177 The solution as mesh. 

178 """ 

179 me = self.dtype_u(self.init) 

180 

181 # assemble the nonlinear function F for the solver 

182 def F(x): 

183 r""" 

184 Nonlinear function for the ``SciPy`` solver. 

185 

186 Args: 

187 x : dtype_u 

188 Current solution 

189 """ 

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

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

192 

193 try: 

194 sol = newton_krylov( 

195 F=F, 

196 xin=u0.copy(), 

197 maxiter=self.liniter, 

198 x_tol=self.lintol, 

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

200 method='gmres', 

201 ) 

202 except NoConvergence as e: 

203 sol = e.args[0] 

204 

205 me[:] = sol 

206 return me