Coverage for pySDC/implementations/problem_classes/NonlinearSchroedinger_MPIFFT.py: 93%
60 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 12:08 +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 Periodic nonlinear Schrödinger equation with mpi4py-fft, IMEX with the Laplacian implicit, nonlinearity explicit.
15 Example implementing the :math:`N`-dimensional nonlinear Schrödinger equation with periodic boundary conditions
17 .. math::
18 \frac{\partial u}{\partial t} = i \Delta u + 2 N c i |u|^2 u
20 for fixed parameter :math:`c \in \{0, 1\}` and :math:`N=2, 3`, on :math:`[0, 2\pi]^N`. The factor :math:`N` makes
21 the one-dimensional breather solution (39) of https://doi.org/10.1007/BF01017105, taken along
22 :math:`x_1 + \dots + x_N` and at time :math:`N t`, an exact solution. The period :math:`2\pi` is fixed and cannot
23 be passed as ``L``. The linear parts of the problem will be solved using ``mpi4py-fft`` [#]_. *Semi-explicit*
24 time-stepping is used here to solve the problem in the temporal dimension, i.e., the Laplacian will be handled
25 implicitly.
27 Parameters
28 ----------
29 nvars : tuple, optional
30 Spatial resolution
31 spectral : bool, optional
32 If True, the solution is computed in spectral space.
33 c : float, optional
34 Nonlinearity parameter, either 0 or 1.
35 comm : MPI.COMM_World
36 Communicator for parallelisation.
38 Attributes
39 ----------
40 fft : PFFT
41 Object for parallel FFT transforms.
42 X : mesh-grid
43 Grid coordinates in real space.
44 K2 : matrix
45 Laplace operator in spectral space.
47 References
48 ----------
49 .. [#] Lisandro Dalcin, Mikael Mortensen, David E. Keyes. Fast parallel multidimensional FFT using advanced MPI.
50 Journal of Parallel and Distributed Computing (2019).
51 """
53 def __init__(self, c=1.0, **kwargs):
54 """Initialization routine"""
55 super().__init__(L=2 * np.pi, alpha=1j, dtype='D', **kwargs)
57 if not (c == 0.0 or c == 1.0):
58 raise ProblemError(f'Setup not implemented, c has to be 0 or 1, got {c}')
59 self._makeAttributeAndRegister('c', localVars=locals(), readOnly=True)
61 def _eval_explicit_part(self, u, t, f_expl):
62 f_expl[:] = self.ndim * self.c * 2j * self.xp.absolute(u) ** 2 * u
63 return f_expl
65 def u_exact(self, t, **kwargs):
66 r"""
67 Routine to compute the exact solution at time :math:`t`, see (39) from https://doi.org/10.1007/BF01017105 for details
69 Parameters
70 ----------
71 t : float
72 Time of the exact solution.
74 Returns
75 -------
76 u : dtype_u
77 The exact solution.
78 """
79 if 'u_init' in kwargs.keys() or 't_init' in kwargs.keys():
80 self.logger.warning(
81 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!'
82 )
84 def nls_exact_1D(t, x, c):
85 ae = 1.0 / np.sqrt(2.0) * np.exp(1j * t)
86 if c != 0:
87 u = ae * ((np.cosh(t) + 1j * np.sinh(t)) / (np.cosh(t) - 1.0 / np.sqrt(2.0) * self.xp.cos(x)) - 1.0)
88 else:
89 u = self.xp.sin(x) * np.exp(-t * 1j)
91 return u
93 me = self.dtype_u(self.init, val=0.0)
95 if self.spectral:
96 tmp = nls_exact_1D(self.ndim * t, sum(self.X), self.c)
97 me[:] = self.fft.forward(tmp)
98 else:
99 me[:] = nls_exact_1D(self.ndim * t, sum(self.X), self.c)
101 return me
104class nonlinearschroedinger_fully_implicit(nonlinearschroedinger_imex):
105 r"""
106 Periodic nonlinear Schrödinger equation with mpi4py-fft, fully implicit with SciPy's Newton-Krylov solver.
108 Example implementing the :math:`N`-dimensional nonlinear Schrödinger equation with periodic boundary conditions
110 .. math::
111 \frac{\partial u}{\partial t} = i \Delta u + 2 N c i |u|^2 u
113 for fixed parameter :math:`c \in \{0, 1\}` and :math:`N=2, 3`, on :math:`[0, 2\pi]^N`. The factor :math:`N` makes
114 the one-dimensional breather solution (39) of https://doi.org/10.1007/BF01017105, taken along
115 :math:`x_1 + \dots + x_N` and at time :math:`N t`, an exact solution. The period :math:`2\pi` is fixed and cannot
116 be passed as ``L``. The linear parts of the problem will be discretized using ``mpi4py-fft`` [#]_. For
117 time-stepping, the problem will be solved *fully-implicitly*, i.e., the nonlinear system containing the full
118 right-hand side is solved by a Newton-Krylov method [#]_.
120 Parameters
121 ----------
122 lintol : float, optional
123 Absolute tolerance for the ``SciPy`` Newton-Krylov solver, passed as ``x_tol``, i.e., applied to the
124 maximum norm of the Newton step. ``SciPy`` additionally requires the maximum norm of the residual to be below
125 its default ``f_tol`` of about ``6e-6``.
126 liniter : int, optional
127 Maximum number of Newton iterations of the Newton-Krylov solver, passed as ``maxiter``. If the solver does not
128 converge, the last iterate is used without a warning.
129 **kwargs
130 Passed on to ``nonlinearschroedinger_imex``, see there: ``c``, ``nvars``, ``spectral`` and ``comm``.
131 ``useGPU`` has to be False.
133 References
134 ----------
135 .. [#] Lisandro Dalcin, Mikael Mortensen, David E. Keyes. Fast parallel multidimensional FFT using advanced MPI.
136 Journal of Parallel and Distributed Computing (2019).
137 .. [#] https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.newton_krylov.html
138 """
140 dtype_u = mesh
141 dtype_f = mesh
143 def __init__(self, lintol=1e-9, liniter=99, **kwargs):
144 """Initialization routine"""
145 assert kwargs.get('useGPU', False) is False
147 super().__init__(**kwargs)
148 self._makeAttributeAndRegister('liniter', 'lintol', localVars=locals(), readOnly=False)
150 self.work_counters['newton'] = WorkCounter()
152 def eval_f(self, u, t):
153 """
154 Routine to evaluate the right-hand side of the problem.
156 Parameters
157 ----------
158 u : dtype_u
159 Current values of the numerical solution.
160 t : float
161 Current time at which the numerical solution is computed.
163 Returns
164 -------
165 f : dtype_f
166 The right-hand side of the problem.
167 """
169 f = self.dtype_f(self.init)
171 if self.spectral:
172 tmp = self.fft.backward(u)
173 tmpf = self.ndim * self.c * 2j * self.xp.absolute(tmp) ** 2 * tmp
174 f[:] = -self.K2 * 1j * u + self.fft.forward(tmpf)
176 else:
177 u_hat = self.fft.forward(u)
178 lap_u_hat = -self.K2 * 1j * u_hat
179 f[:] = self.fft.backward(lap_u_hat) + self.ndim * self.c * 2j * self.xp.absolute(u) ** 2 * u
181 self.work_counters['rhs']()
182 return f
184 def solve_system(self, rhs, factor, u0, t):
185 r"""
186 Solve the nonlinear system :math:`(1 - factor \cdot f)(\vec{u}) = \vec{rhs}` using a ``SciPy`` Newton-Krylov
187 solver, see `scipy.optimize.newton_krylov
188 <https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.newton_krylov.html>`__.
190 Parameters
191 ----------
192 rhs : dtype_f
193 Right-hand side for the linear system.
194 factor : float
195 Abbrev. for the node-to-node stepsize (or any other factor required).
196 u0 : dtype_u
197 Initial guess for the iterative solver (not used here so far).
198 t : float
199 Current time (e.g. for time-dependent BCs).
201 Returns
202 -------
203 me : dtype_u
204 The solution as mesh.
205 """
206 me = self.dtype_u(self.init)
208 # assemble the nonlinear function F for the solver
209 def F(x):
210 r"""
211 Nonlinear function for the ``SciPy`` solver.
213 Args:
214 x : dtype_u
215 Current solution
216 """
217 self.work_counters['rhs'].decrement()
218 return x - factor * self.eval_f(u=x.reshape(self.init[0]), t=t).reshape(x.shape) - rhs.reshape(x.shape)
220 try:
221 sol = newton_krylov(
222 F=F,
223 xin=u0.copy(),
224 maxiter=self.liniter,
225 x_tol=self.lintol,
226 callback=self.work_counters['newton'],
227 method='gmres',
228 )
229 except NoConvergence as e:
230 sol = e.args[0]
232 me[:] = sol
233 return me