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