# Coverage for pySDC/implementations/problem_classes/TestEquation_0D.py: 100%

## 53 statements

, created at 2024-09-09 14:59 +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

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