Coverage for pySDC/implementations/problem_classes/TestEquation_0D.py: 100%
102 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-04 15:08 +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 This class implements the simple test equation of the form
13 .. math::
14 \frac{d u(t)}{dt} = A u(t)
16 for :math:`A = diag(\lambda_1, .. ,\lambda_n)`.
18 Parameters
19 ----------
20 lambdas : sequence of array_like, optional
21 List of lambda parameters.
22 u0 : sequence of array_like, optional
23 Initial condition.
25 Attributes
26 ----------
27 A : scipy.sparse.csc_matrix
28 Diagonal matrix containing :math:`\lambda_1,..,\lambda_n`.
29 """
31 xp = np
32 xsp = nsp
33 dtype_u = mesh
34 dtype_f = mesh
36 @classmethod
37 def setup_GPU(cls):
38 """
39 Switch to GPU modules
40 """
41 from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh
42 import cupy as cp
43 import cupyx.scipy.sparse as csp
45 cls.xp = cp
46 cls.xsp = csp
47 cls.dtype_u = cupy_mesh
48 cls.dtype_f = cupy_mesh
50 def __init__(self, lambdas=None, u0=0.0, useGPU=False):
51 """Initialization routine"""
52 if useGPU:
53 self.setup_GPU()
55 if lambdas is None:
56 re = self.xp.linspace(-30, 19, 50)
57 im = self.xp.linspace(-50, 49, 50)
58 lambdas = self.xp.array([[complex(re[i], im[j]) for i in range(len(re))] for j in range(len(im))]).reshape(
59 (len(re) * len(im))
60 )
61 lambdas = self.xp.asarray(lambdas)
62 assert lambdas.ndim == 1, f'expect flat list here, got {lambdas}'
63 nvars = lambdas.size
64 assert nvars > 0, 'expect at least one lambda parameter here'
66 # invoke super init, passing number of dofs, dtype_u and dtype_f
67 super().__init__(init=(nvars, None, self.xp.dtype('complex128')))
69 lambdas = self.xp.array(lambdas)
70 self.A = self.xsp.diags(lambdas)
71 self._makeAttributeAndRegister('nvars', 'lambdas', 'u0', 'useGPU', localVars=locals(), readOnly=True)
72 self.work_counters['rhs'] = WorkCounter()
74 def eval_f(self, u, t):
75 """
76 Routine to evaluate the right-hand side of the problem.
78 Parameters
79 ----------
80 u : dtype_u
81 Current values of the numerical solution.
82 t : float
83 Current time of the numerical solution is computed.
85 Returns
86 -------
87 f : dtype_f
88 The right-hand side of the problem.
89 """
91 f = self.dtype_f(self.init)
92 f[:] = u
93 f *= self.lambdas
94 self.work_counters['rhs']()
95 return f
97 def solve_system(self, rhs, factor, u0, t):
98 r"""
99 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.
101 Parameters
102 ----------
103 rhs : dtype_f
104 Right-hand side for the linear system.
105 factor : float
106 Abbrev. for the local stepsize (or any other factor required).
107 u0 : dtype_u
108 Initial guess for the iterative solver.
109 t : float
110 Current time (e.g. for time-dependent BCs).
112 Returns
113 -------
114 me : dtype_u
115 The solution as mesh.
116 """
117 me = self.dtype_u(self.init)
118 L = 1 - factor * self.lambdas
119 L[L == 0] = 1 # to avoid potential divisions by zeros
120 me[:] = rhs
121 me /= L
122 return me
124 def u_exact(self, t, u_init=None, t_init=None):
125 """
126 Routine to compute the exact solution at time t.
128 Parameters
129 ----------
130 t : float
131 Time of the exact solution.
132 u_init : pySDC.problem.testequation0d.dtype_u
133 Initial solution.
134 t_init : float
135 The initial time.
137 Returns
138 -------
139 me : dtype_u
140 The exact solution.
141 """
143 u_init = (self.u0 if u_init is None else u_init) * 1.0
144 t_init = 0.0 if t_init is None else t_init * 1.0
146 me = self.dtype_u(self.init)
147 me[:] = u_init * self.xp.exp((t - t_init) * self.lambdas)
148 return me
150 def getOutputFile(self, fileName):
151 fOut = Scalar(np.complex128, fileName=fileName)
152 fOut.setHeader(self.lambdas.size)
153 fOut.initialize()
154 return fOut
156 def processSolutionForOutput(self, u):
157 return u.flatten()
160class test_equation_IMEX(Problem):
161 dtype_f = imex_mesh
162 dtype_u = mesh
163 xp = np
164 xsp = nsp
166 def __init__(self, lambdas_implicit=None, lambdas_explicit=None, u0=0.0):
167 """Initialization routine"""
169 if lambdas_implicit is None:
170 re = self.xp.linspace(-30, 19, 50)
171 im = self.xp.linspace(-50, 49, 50)
172 lambdas_implicit = self.xp.array(
173 [[complex(re[i], im[j]) for i in range(len(re))] for j in range(len(im))]
174 ).reshape((len(re) * len(im)))
175 if lambdas_explicit is None:
176 lambdas_explicit = lambdas_implicit.copy()
178 lambdas_implicit = self.xp.asarray(lambdas_implicit)
179 lambdas_explicit = self.xp.asarray(lambdas_explicit)
181 assert lambdas_implicit.ndim == 1, f'expect flat list here, got {lambdas_implicit}'
182 assert lambdas_explicit.shape == lambdas_implicit.shape
183 nvars = lambdas_implicit.size
184 assert nvars > 0, 'expect at least one lambda parameter here'
186 # invoke super init, passing number of dofs, dtype_u and dtype_f
187 super().__init__(init=(nvars, None, self.xp.dtype('complex128')))
189 self.A = self.xsp.diags(lambdas_implicit)
190 self._makeAttributeAndRegister(
191 'nvars', 'lambdas_implicit', 'lambdas_explicit', 'u0', localVars=locals(), readOnly=True
192 )
193 self.work_counters['rhs'] = WorkCounter()
195 def eval_f(self, u, t):
196 """
197 Routine to evaluate the right-hand side of the problem.
199 Parameters
200 ----------
201 u : dtype_u
202 Current values of the numerical solution.
203 t : float
204 Current time of the numerical solution is computed.
206 Returns
207 -------
208 f : dtype_f
209 The right-hand side of the problem.
210 """
212 f = self.dtype_f(self.init)
213 f.impl[:] = u * self.lambdas_implicit
214 f.expl[:] = u * self.lambdas_explicit
215 self.work_counters['rhs']()
216 return f
218 def solve_system(self, rhs, factor, u0, t):
219 r"""
220 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.
222 Parameters
223 ----------
224 rhs : dtype_f
225 Right-hand side for the linear system.
226 factor : float
227 Abbrev. for the local stepsize (or any other factor required).
228 u0 : dtype_u
229 Initial guess for the iterative solver.
230 t : float
231 Current time (e.g. for time-dependent BCs).
233 Returns
234 -------
235 me : dtype_u
236 The solution as mesh.
237 """
238 me = self.dtype_u(self.init)
239 L = 1 - factor * self.lambdas_implicit
240 L[L == 0] = 1 # to avoid potential divisions by zeros
241 me[:] = rhs
242 me /= L
243 return me
245 def u_exact(self, t, u_init=None, t_init=None):
246 """
247 Routine to compute the exact solution at time t.
249 Parameters
250 ----------
251 t : float
252 Time of the exact solution.
253 u_init : pySDC.problem.testequation0d.dtype_u
254 Initial solution.
255 t_init : float
256 The initial time.
258 Returns
259 -------
260 me : dtype_u
261 The exact solution.
262 """
264 u_init = (self.u0 if u_init is None else u_init) * 1.0
265 t_init = 0.0 if t_init is None else t_init * 1.0
267 me = self.dtype_u(self.init)
268 me[:] = u_init * self.xp.exp((t - t_init) * (self.lambdas_implicit + self.lambdas_explicit))
269 return me