Coverage for pySDC/implementations/problem_classes/TestEquation_0D.py: 100%
101 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1import numpy as np
2import scipy.sparse as nsp
4from pySDC.core.problem import Problem, WorkCounter
5from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh
6from pySDC.helpers.fieldsIO import Scalar
9class testequation0d(Problem):
10 r"""
11 Dahlquist test equation for many values of :math:`\lambda` at once, treated fully implicitly.
13 It is of the form
15 .. math::
16 \frac{d u(t)}{dt} = A u(t)
18 for :math:`A = diag(\lambda_1, .. ,\lambda_n)`.
20 Parameters
21 ----------
22 lambdas : sequence of array_like, optional
23 List of lambda parameters.
24 u0 : sequence of array_like, optional
25 Initial condition.
26 useGPU : bool, optional
27 Run on the GPU with CuPy instead of on the CPU with NumPy.
29 Attributes
30 ----------
31 A : scipy.sparse.csc_matrix
32 Diagonal matrix containing :math:`\lambda_1,..,\lambda_n`.
33 """
35 xp = np
36 xsp = nsp
37 dtype_u = mesh
38 dtype_f = mesh
40 def setup_GPU(self):
41 """
42 Switch to GPU modules
43 """
44 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh
45 import cupy as cp
46 import cupyx.scipy.sparse as csp
48 self.xp = cp
49 self.xsp = csp
50 self.dtype_u = cupy_mesh
51 self.dtype_f = cupy_mesh
53 def __init__(self, lambdas=None, u0=0.0, useGPU=False):
54 """Initialization routine"""
55 if useGPU:
56 self.setup_GPU()
58 if lambdas is None:
59 re = self.xp.linspace(-30, 19, 50)
60 im = self.xp.linspace(-50, 49, 50)
61 lambdas = self.xp.array([[complex(re[i], im[j]) for i in range(len(re))] for j in range(len(im))]).reshape(
62 (len(re) * len(im))
63 )
64 lambdas = self.xp.asarray(lambdas)
65 assert lambdas.ndim == 1, f'expect flat list here, got {lambdas}'
66 nvars = lambdas.size
67 assert nvars > 0, 'expect at least one lambda parameter here'
69 # invoke super init, passing number of dofs, dtype_u and dtype_f
70 super().__init__(init=(nvars, None, self.xp.dtype('complex128')))
72 lambdas = self.xp.array(lambdas)
73 self.A = self.xsp.diags(lambdas)
74 self._makeAttributeAndRegister('nvars', 'lambdas', 'u0', 'useGPU', localVars=locals(), readOnly=True)
75 self.work_counters['rhs'] = WorkCounter()
77 def eval_f(self, u, t):
78 """
79 Routine to evaluate the right-hand side of the problem.
81 Parameters
82 ----------
83 u : dtype_u
84 Current values of the numerical solution.
85 t : float
86 Current time of the numerical solution is computed.
88 Returns
89 -------
90 f : dtype_f
91 The right-hand side of the problem.
92 """
94 f = self.dtype_f(self.init)
95 f[:] = u
96 f *= self.lambdas
97 self.work_counters['rhs']()
98 return f
100 def solve_system(self, rhs, factor, u0, t):
101 r"""
102 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.
104 Parameters
105 ----------
106 rhs : dtype_f
107 Right-hand side for the linear system.
108 factor : float
109 Abbrev. for the local stepsize (or any other factor required).
110 u0 : dtype_u
111 Initial guess for the iterative solver.
112 t : float
113 Current time (e.g. for time-dependent BCs).
115 Returns
116 -------
117 me : dtype_u
118 The solution as mesh.
119 """
120 me = self.dtype_u(self.init)
121 L = 1 - factor * self.lambdas
122 L[L == 0] = 1 # to avoid potential divisions by zeros
123 me[:] = rhs
124 me /= L
125 return me
127 def u_exact(self, t, u_init=None, t_init=None):
128 """
129 Routine to compute the exact solution at time t.
131 Parameters
132 ----------
133 t : float
134 Time of the exact solution.
135 u_init : pySDC.problem.testequation0d.dtype_u
136 Initial solution.
137 t_init : float
138 The initial time.
140 Returns
141 -------
142 me : dtype_u
143 The exact solution.
144 """
146 u_init = (self.u0 if u_init is None else u_init) * 1.0
147 t_init = 0.0 if t_init is None else t_init * 1.0
149 me = self.dtype_u(self.init)
150 me[:] = u_init * self.xp.exp((t - t_init) * self.lambdas)
151 return me
153 def getOutputFile(self, fileName):
154 r"""
155 Set up a ``Scalar`` output file for complex values, with one variable per :math:`\lambda`.
157 Parameters
158 ----------
159 fileName : str
160 Name of the output file.
162 Returns
163 -------
164 fOut : pySDC.helpers.fieldsIO.Scalar
165 The initialized output file.
166 """
167 fOut = Scalar(np.complex128, fileName=fileName)
168 fOut.setHeader(self.lambdas.size)
169 fOut.initialize()
170 return fOut
172 def processSolutionForOutput(self, u):
173 """
174 Flatten the solution for output.
176 Parameters
177 ----------
178 u : dtype_u
179 Solution to be written.
181 Returns
182 -------
183 np.1darray
184 The flattened solution.
185 """
186 return u.flatten()
189class test_equation_IMEX(Problem):
190 r"""
191 Test equation with the right-hand side split into an implicit and an explicit part
193 .. math::
194 \frac{d u(t)}{dt} = A_I u(t) + A_E u(t)
196 for :math:`A_I = diag(\lambda_{I,1}, .. ,\lambda_{I,n})` and :math:`A_E = diag(\lambda_{E,1}, .. ,\lambda_{E,n})`.
197 For IMEX time-stepping, :math:`A_I u` is treated implicitly and :math:`A_E u` explicitly.
199 Parameters
200 ----------
201 lambdas_implicit : sequence of array_like, optional
202 Flat list of the parameters :math:`\lambda_{I,k}` of the implicit part. ``None`` means 2500 complex values
203 :math:`a + b i` with integers :math:`a \in [-30, 19]` and :math:`b \in [-50, 49]`.
204 lambdas_explicit : sequence of array_like, optional
205 Flat list of the parameters :math:`\lambda_{E,k}` of the explicit part, of the same shape as
206 ``lambdas_implicit``. ``None`` means a copy of ``lambdas_implicit``.
207 u0 : sequence of array_like, optional
208 Initial condition at :math:`t = 0`, either a scalar or one value per equation.
210 Attributes
211 ----------
212 A : scipy.sparse.dia_matrix
213 Diagonal matrix containing the implicit parameters :math:`\lambda_{I,1},..,\lambda_{I,n}`.
214 """
216 dtype_f = imex_mesh
217 dtype_u = mesh
218 xp = np
219 xsp = nsp
221 def __init__(self, lambdas_implicit=None, lambdas_explicit=None, u0=0.0):
222 """Initialization routine"""
224 if lambdas_implicit is None:
225 re = self.xp.linspace(-30, 19, 50)
226 im = self.xp.linspace(-50, 49, 50)
227 lambdas_implicit = self.xp.array(
228 [[complex(re[i], im[j]) for i in range(len(re))] for j in range(len(im))]
229 ).reshape((len(re) * len(im)))
230 if lambdas_explicit is None:
231 lambdas_explicit = lambdas_implicit.copy()
233 lambdas_implicit = self.xp.asarray(lambdas_implicit)
234 lambdas_explicit = self.xp.asarray(lambdas_explicit)
236 assert lambdas_implicit.ndim == 1, f'expect flat list here, got {lambdas_implicit}'
237 assert lambdas_explicit.shape == lambdas_implicit.shape
238 nvars = lambdas_implicit.size
239 assert nvars > 0, 'expect at least one lambda parameter here'
241 # invoke super init, passing number of dofs, dtype_u and dtype_f
242 super().__init__(init=(nvars, None, self.xp.dtype('complex128')))
244 self.A = self.xsp.diags(lambdas_implicit)
245 self._makeAttributeAndRegister(
246 'nvars', 'lambdas_implicit', 'lambdas_explicit', 'u0', localVars=locals(), readOnly=True
247 )
248 self.work_counters['rhs'] = WorkCounter()
250 def eval_f(self, u, t):
251 """
252 Routine to evaluate the right-hand side of the problem.
254 Parameters
255 ----------
256 u : dtype_u
257 Current values of the numerical solution.
258 t : float
259 Current time of the numerical solution is computed.
261 Returns
262 -------
263 f : dtype_f
264 The right-hand side of the problem.
265 """
267 f = self.dtype_f(self.init)
268 f.impl[:] = u * self.lambdas_implicit
269 f.expl[:] = u * self.lambdas_explicit
270 self.work_counters['rhs']()
271 return f
273 def solve_system(self, rhs, factor, u0, t):
274 r"""
275 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.
277 Parameters
278 ----------
279 rhs : dtype_f
280 Right-hand side for the linear system.
281 factor : float
282 Abbrev. for the local stepsize (or any other factor required).
283 u0 : dtype_u
284 Initial guess for the iterative solver.
285 t : float
286 Current time (e.g. for time-dependent BCs).
288 Returns
289 -------
290 me : dtype_u
291 The solution as mesh.
292 """
293 me = self.dtype_u(self.init)
294 L = 1 - factor * self.lambdas_implicit
295 L[L == 0] = 1 # to avoid potential divisions by zeros
296 me[:] = rhs
297 me /= L
298 return me
300 def u_exact(self, t, u_init=None, t_init=None):
301 """
302 Routine to compute the exact solution at time t.
304 Parameters
305 ----------
306 t : float
307 Time of the exact solution.
308 u_init : pySDC.problem.testequation0d.dtype_u
309 Initial solution.
310 t_init : float
311 The initial time.
313 Returns
314 -------
315 me : dtype_u
316 The exact solution.
317 """
319 u_init = (self.u0 if u_init is None else u_init) * 1.0
320 t_init = 0.0 if t_init is None else t_init * 1.0
322 me = self.dtype_u(self.init)
323 me[:] = u_init * self.xp.exp((t - t_init) * (self.lambdas_implicit + self.lambdas_explicit))
324 return me