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

1import numpy as np 

2from scipy.optimize import newton_krylov 

3from scipy.optimize 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 (39) from https://doi.org/10.1007/BF01017105 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 assert kwargs.get('useGPU', False) is False 

122 

123 super().__init__(**kwargs) 

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

125 

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

127 

128 def eval_f(self, u, t): 

129 """ 

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

131 

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. 

138 

139 Returns 

140 ------- 

141 f : dtype_f 

142 The right-hand side of the problem. 

143 """ 

144 

145 f = self.dtype_f(self.init) 

146 

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) 

151 

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 

156 

157 self.work_counters['rhs']() 

158 return f 

159 

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. 

164 

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

175 

176 Returns 

177 ------- 

178 me : dtype_u 

179 The solution as mesh. 

180 """ 

181 me = self.dtype_u(self.init) 

182 

183 # assemble the nonlinear function F for the solver 

184 def F(x): 

185 r""" 

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

187 

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) 

194 

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] 

206 

207 me[:] = sol 

208 return me