Coverage for pySDC/implementations/problem_classes/generic_ND_FD.py: 94%

71 statements  

« prev     ^ index     » next       coverage.py v7.6.7, created at 2024-11-16 14:51 +0000

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4Created on Sat Feb 11 22:39:30 2023 

5""" 

6import numpy as np 

7import scipy.sparse as sp 

8from scipy.sparse.linalg import gmres, spsolve, cg 

9 

10from pySDC.core.errors import ProblemError 

11from pySDC.core.problem import Problem, WorkCounter 

12from pySDC.helpers import problem_helper 

13from pySDC.implementations.datatype_classes.mesh import mesh 

14 

15 

16class GenericNDimFinDiff(Problem): 

17 r""" 

18 Base class for finite difference spatial discretisation in :math:`N` dimensions 

19 

20 .. math:: 

21 \frac{d u}{dt} = A u, 

22 

23 where :math:`A \in \mathbb{R}^{nN \times nN}` is a matrix arising from finite difference discretisation of spatial 

24 derivatives with :math:`n` degrees of freedom per dimension and :math:`N` dimensions. This generic class follows the MOL 

25 (method-of-lines) approach and can be used to discretize partial differential equations such as the advection 

26 equation and the heat equation. 

27 

28 Parameters 

29 ---------- 

30 nvars : int, optional 

31 Spatial resolution for the ND problem. For :math:`N = 2`, 

32 set ``nvars=(16, 16)``. 

33 coeff : float, optional 

34 Factor for finite difference matrix :math:`A`. 

35 derivative : int, optional 

36 Order of the spatial derivative. 

37 freq : tuple of int, optional 

38 Spatial frequency, can be a tuple. 

39 stencil_type : str, optional 

40 Stencil type for finite differences. 

41 order : int, optional 

42 Order of accuracy of the finite difference discretization. 

43 lintol : float, optional 

44 Tolerance for spatial solver. 

45 liniter : int, optional 

46 Maximum number of iterations for linear solver. 

47 solver_type : str, optional 

48 Type of solver. Can be ``'direct'``, ``'GMRES'`` or ``'CG'``. 

49 bc : str or tuple of 2 string, optional 

50 Type of boundary conditions. Default is ``'periodic'``. 

51 To define two different types of boundary condition for each side, 

52 you can use a tuple, for instance ``bc=("dirichlet", "neumann")`` 

53 uses Dirichlet BC on the left side, and Neumann BC on the right side. 

54 bcParams : dict, optional 

55 Parameters for boundary conditions, that can contains those keys : 

56 

57 - **val** : value for the boundary value (Dirichlet) or derivative 

58 (Neumann), default to 0 

59 - **reduce** : if true, reduce the order of the A matrix close to the 

60 boundary. If false (default), use shifted stencils close to the 

61 boundary. 

62 - **neumann_bc_order** : finite difference order that should be used 

63 for the neumann BC derivative. If None (default), uses the same 

64 order as the discretization for A. 

65 

66 Default is None, which takes the default values for each parameters. 

67 You can also define a tuple to set different parameters for each 

68 side. 

69 

70 Attributes 

71 ---------- 

72 A : sparse matrix (CSC) 

73 FD discretization matrix of the ND operator. 

74 Id : sparse matrix (CSC) 

75 Identity matrix of the same dimension as A. 

76 xvalues : np.1darray 

77 Values of spatial grid. 

78 """ 

79 

80 dtype_u = mesh 

81 dtype_f = mesh 

82 

83 def __init__( 

84 self, 

85 nvars=512, 

86 coeff=1.0, 

87 derivative=1, 

88 freq=2, 

89 stencil_type='center', 

90 order=2, 

91 lintol=1e-12, 

92 liniter=10000, 

93 solver_type='direct', 

94 bc='periodic', 

95 bcParams=None, 

96 ): 

97 # make sure parameters have the correct types 

98 if type(nvars) not in [int, tuple]: 

99 raise ProblemError('nvars should be either tuple or int') 

100 if type(freq) not in [int, tuple]: 

101 raise ProblemError('freq should be either tuple or int') 

102 

103 # transforms nvars into a tuple 

104 if type(nvars) is int: 

105 nvars = (nvars,) 

106 

107 # automatically determine ndim from nvars 

108 ndim = len(nvars) 

109 if ndim > 3: 

110 raise ProblemError(f'can work with up to three dimensions, got {ndim}') 

111 

112 # eventually extend freq to other dimension 

113 if type(freq) is int: 

114 freq = (freq,) * ndim 

115 if len(freq) != ndim: 

116 raise ProblemError(f'len(freq)={len(freq)}, different to ndim={ndim}') 

117 

118 # check values for freq and nvars 

119 for f in freq: 

120 if ndim == 1 and f == -1: 

121 # use Gaussian initial solution in 1D 

122 bc = 'periodic' 

123 break 

124 if f % 2 != 0 and bc == 'periodic': 

125 raise ProblemError('need even number of frequencies due to periodic BCs') 

126 for nvar in nvars: 

127 if nvar % 2 != 0 and bc == 'periodic': 

128 raise ProblemError('the setup requires nvars = 2^p per dimension') 

129 if (nvar + 1) % 2 != 0 and bc == 'dirichlet-zero': 

130 raise ProblemError('setup requires nvars = 2^p - 1') 

131 if ndim > 1 and nvars[1:] != nvars[:-1]: 

132 raise ProblemError('need a square domain, got %s' % nvars) 

133 

134 # invoke super init, passing number of dofs 

135 super().__init__(init=(nvars[0] if ndim == 1 else nvars, None, np.dtype('float64'))) 

136 

137 dx, xvalues = problem_helper.get_1d_grid(size=nvars[0], bc=bc, left_boundary=0.0, right_boundary=1.0) 

138 

139 self.A, _ = problem_helper.get_finite_difference_matrix( 

140 derivative=derivative, 

141 order=order, 

142 stencil_type=stencil_type, 

143 dx=dx, 

144 size=nvars[0], 

145 dim=ndim, 

146 bc=bc, 

147 ) 

148 self.A *= coeff 

149 

150 self.xvalues = xvalues 

151 self.Id = sp.eye(np.prod(nvars), format='csc') 

152 

153 # store attribute and register them as parameters 

154 self._makeAttributeAndRegister('nvars', 'stencil_type', 'order', 'bc', localVars=locals(), readOnly=True) 

155 self._makeAttributeAndRegister('freq', 'lintol', 'liniter', 'solver_type', localVars=locals()) 

156 

157 if self.solver_type != 'direct': 

158 self.work_counters[self.solver_type] = WorkCounter() 

159 

160 @property 

161 def ndim(self): 

162 """Number of dimensions of the spatial problem""" 

163 return len(self.nvars) 

164 

165 @property 

166 def dx(self): 

167 """Size of the mesh (in all dimensions)""" 

168 return self.xvalues[1] - self.xvalues[0] 

169 

170 @property 

171 def grids(self): 

172 """ND grids associated to the problem""" 

173 x = self.xvalues 

174 if self.ndim == 1: 

175 return x 

176 if self.ndim == 2: 

177 return x[None, :], x[:, None] 

178 if self.ndim == 3: 

179 return x[None, :, None], x[:, None, None], x[None, None, :] 

180 

181 @classmethod 

182 def get_default_sweeper_class(cls): 

183 from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit 

184 

185 return generic_implicit 

186 

187 def eval_f(self, u, t): 

188 """ 

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

190 

191 Parameters 

192 ---------- 

193 u : dtype_u 

194 Current values. 

195 t : float 

196 Current time. 

197 

198 Returns 

199 ------- 

200 f : dtype_f 

201 Values of the right-hand side of the problem. 

202 """ 

203 f = self.f_init 

204 f[:] = self.A.dot(u.flatten()).reshape(self.nvars) 

205 return f 

206 

207 def solve_system(self, rhs, factor, u0, t): 

208 r""" 

209 Simple linear solver for :math:`(I-factor\cdot A)\vec{u}=\vec{rhs}`. 

210 

211 Parameters 

212 ---------- 

213 rhs : dtype_f 

214 Right-hand side for the linear system. 

215 factor : float 

216 Abbrev. for the local stepsize (or any other factor required). 

217 u0 : dtype_u 

218 Initial guess for the iterative solver. 

219 t : float 

220 Current time (e.g. for time-dependent BCs). 

221 

222 Returns 

223 ------- 

224 sol : dtype_u 

225 The solution of the linear solver. 

226 """ 

227 solver_type, Id, A, nvars, lintol, liniter, sol = ( 

228 self.solver_type, 

229 self.Id, 

230 self.A, 

231 self.nvars, 

232 self.lintol, 

233 self.liniter, 

234 self.u_init, 

235 ) 

236 

237 if solver_type == 'direct': 

238 sol[:] = spsolve(Id - factor * A, rhs.flatten()).reshape(nvars) 

239 elif solver_type == 'GMRES': 

240 sol[:] = gmres( 

241 Id - factor * A, 

242 rhs.flatten(), 

243 x0=u0.flatten(), 

244 rtol=lintol, 

245 maxiter=liniter, 

246 atol=0, 

247 callback=self.work_counters[solver_type], 

248 callback_type='legacy', 

249 )[0].reshape(nvars) 

250 elif solver_type == 'CG': 

251 sol[:] = cg( 

252 Id - factor * A, 

253 rhs.flatten(), 

254 x0=u0.flatten(), 

255 rtol=lintol, 

256 maxiter=liniter, 

257 atol=0, 

258 callback=self.work_counters[solver_type], 

259 )[0].reshape(nvars) 

260 else: 

261 raise ValueError(f'solver type "{solver_type}" not known in generic advection-diffusion implementation!') 

262 

263 return sol