Coverage for pySDC/implementations/problem_classes/TestEquation_0D.py: 67%
104 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-04-18 08:18 +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 re = self.xp.linspace(-30, 19, 50)
177 im = self.xp.linspace(-50, 49, 50)
178 lambdas_implicit = self.xp.array(
179 [[complex(re[i], im[j]) for i in range(len(re))] for j in range(len(im))]
180 ).reshape((len(re) * len(im)))
181 lambdas_implicit = self.xp.asarray(lambdas_implicit)
182 lambdas_explicit = self.xp.asarray(lambdas_explicit)
184 assert lambdas_implicit.ndim == 1, f'expect flat list here, got {lambdas_implicit}'
185 assert lambdas_explicit.shape == lambdas_implicit.shape
186 nvars = lambdas_implicit.size
187 assert nvars > 0, 'expect at least one lambda parameter here'
189 # invoke super init, passing number of dofs, dtype_u and dtype_f
190 super().__init__(init=(nvars, None, self.xp.dtype('complex128')))
192 self.A = self.xsp.diags(lambdas_implicit)
193 self._makeAttributeAndRegister(
194 'nvars', 'lambdas_implicit', 'lambdas_explicit', 'u0', localVars=locals(), readOnly=True
195 )
196 self.work_counters['rhs'] = WorkCounter()
198 def eval_f(self, u, t):
199 """
200 Routine to evaluate the right-hand side of the problem.
202 Parameters
203 ----------
204 u : dtype_u
205 Current values of the numerical solution.
206 t : float
207 Current time of the numerical solution is computed.
209 Returns
210 -------
211 f : dtype_f
212 The right-hand side of the problem.
213 """
215 f = self.dtype_f(self.init)
216 f.impl[:] = u * self.lambdas_implicit
217 f.expl[:] = u * self.lambdas_explicit
218 self.work_counters['rhs']()
219 return f
221 def solve_system(self, rhs, factor, u0, t):
222 r"""
223 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`.
225 Parameters
226 ----------
227 rhs : dtype_f
228 Right-hand side for the linear system.
229 factor : float
230 Abbrev. for the local stepsize (or any other factor required).
231 u0 : dtype_u
232 Initial guess for the iterative solver.
233 t : float
234 Current time (e.g. for time-dependent BCs).
236 Returns
237 -------
238 me : dtype_u
239 The solution as mesh.
240 """
241 me = self.dtype_u(self.init)
242 L = 1 - factor * self.lambdas_implicit
243 L[L == 0] = 1 # to avoid potential divisions by zeros
244 me[:] = rhs
245 me /= L
246 return me
248 def u_exact(self, t, u_init=None, t_init=None):
249 """
250 Routine to compute the exact solution at time t.
252 Parameters
253 ----------
254 t : float
255 Time of the exact solution.
256 u_init : pySDC.problem.testequation0d.dtype_u
257 Initial solution.
258 t_init : float
259 The initial time.
261 Returns
262 -------
263 me : dtype_u
264 The exact solution.
265 """
267 u_init = (self.u0 if u_init is None else u_init) * 1.0
268 t_init = 0.0 if t_init is None else t_init * 1.0
270 me = self.dtype_u(self.init)
271 me[:] = u_init * self.xp.exp((t - t_init) * (self.lambdas_implicit + self.lambdas_explicit))
272 return me