Coverage for pySDC/implementations/problem_classes/TestEquation_0D.py: 85%

1import numpy as np 

2import scipy.sparse as nsp 

3 

4from pySDC.core.Problem import ptype, WorkCounter 

5from pySDC.implementations.datatype_classes.mesh import mesh 

6 

7 

8class testequation0d(ptype): 

9 r""" 

10 This class implements the simple test equation of the form 

11 

12 .. math:: 

13 \frac{d u(t)}{dt} = A u(t) 

14 

15 for :math:`A = diag(\lambda_1, .. ,\lambda_n)`. 

16 

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. 

23 

24 Attributes 

25 ---------- 

26 A : scipy.sparse.csc_matrix 

27 Diagonal matrix containing :math:`\lambda_1,..,\lambda_n`. 

28 """ 

29 

30 xp = np 

31 xsp = nsp 

32 dtype_u = mesh 

33 dtype_f = mesh 

34 

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 

43 

44 cls.xp = cp 

45 cls.xsp = csp 

46 cls.dtype_u = cupy_mesh 

47 cls.dtype_f = cupy_mesh 

48 

49 def __init__(self, lambdas=None, u0=0.0, useGPU=False): 

50 """Initialization routine""" 

51 if useGPU: 

52 self.setup_GPU() 

53 

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' 

64 

65 # invoke super init, passing number of dofs, dtype_u and dtype_f 

66 super().__init__(init=(nvars, None, self.xp.dtype('complex128'))) 

67 

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() 

72 

73 def eval_f(self, u, t): 

74 """ 

75 Routine to evaluate the right-hand side of the problem. 

76 

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. 

83 

84 Returns 

85 ------- 

86 f : dtype_f 

87 The right-hand side of the problem. 

88 """ 

89 

90 f = self.dtype_f(self.init) 

91 f[:] = u 

92 f *= self.lambdas 

93 self.work_counters['rhs']() 

94 return f 

95 

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}`. 

99 

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). 

110 

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 

122 

123 def u_exact(self, t, u_init=None, t_init=None): 

124 """ 

125 Routine to compute the exact solution at time t. 

126 

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. 

135 

136 Returns 

137 ------- 

138 me : dtype_u 

139 The exact solution. 

140 """ 

141 

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 

144 

145 me = self.dtype_u(self.init) 

146 me[:] = u_init * self.xp.exp((t - t_init) * self.lambdas) 

147 return me