Coverage for pySDC/implementations/problem_classes/HeatEquation_ND_FD_CuPy.py: 0%

1import numpy as np 

2import cupy as cp 

3import cupyx.scipy.sparse as csp 

4from cupyx.scipy.sparse.linalg import spsolve, cg # , gmres, minres 

5 

6from pySDC.core.errors import ProblemError 

7from pySDC.core.problem import Problem 

8from pySDC.helpers import problem_helper 

9from pySDC.implementations.datatype_classes.cupy_mesh import cupy_mesh, imex_cupy_mesh 

10 

11 

12class heatNd_forced(Problem): # pragma: no cover 

13 r""" 

14 This class implements the unforced :math:`N`-dimensional heat equation with periodic boundary conditions 

15 

16 .. math:: 

17 \frac{\partial u}{\partial t} = \nu 

18 \left( 

19 \frac{\partial^2 u}{\partial x^2_1} + .. + \frac{\partial^2 u}{\partial x^2_N} 

20 \right) 

21 

22 for :math:`(x_1,..,x_N) \in [0, 1]^{N}` with :math:`N \leq 3`. The initial solution is of the form 

23 

24 .. math:: 

25 u({\bf x},0) = \prod_{i=1}^N \sin(\pi k_i x_i). 

26 

27 The spatial term is discretized using finite differences. 

28 

29 This class uses the ``CuPy`` package in order to make ``pySDC`` available for GPUs. 

30 

31 Parameters 

32 ---------- 

33 nvars : int, optional 

34 Spatial resolution (same in all dimensions). Using a tuple allows to 

35 consider several dimensions, e.g ``nvars=(16,16)`` for a 2D problem. 

36 nu : float, optional 

37 Diffusion coefficient :math:`\nu`. 

38 freq : int, optional 

39 Spatial frequency :math:`k_i` of the initial conditions, can be tuple. 

40 stencil_type : str, optional 

41 Type of the finite difference stencil. 

42 order : int, optional 

43 Order of the finite difference discretization. 

44 lintol : float, optional 

45 Tolerance for spatial solver. 

46 liniter : int, optional 

47 Max. iterations number for spatial solver. 

48 solver_type : str, optional 

49 Solve the linear system directly or using CG. 

50 bc : str, optional 

51 Boundary conditions, either ``'periodic'`` or ``'dirichlet'``. 

52 sigma : float, optional 

53 If ``freq=-1`` and ``ndim=1``, uses a Gaussian initial solution of the form 

54 

55 .. math:: 

56 u(x,0) = e^{ 

57 \frac{\displaystyle 1}{\displaystyle 2} 

58 \left( 

59 \frac{\displaystyle x-1/2}{\displaystyle \sigma} 

60 \right)^2 

61 } 

62 

63 Attributes 

64 ---------- 

65 A : sparse matrix (CSC) 

66 FD discretization matrix of the ND grad operator. 

67 Id : sparse matrix (CSC) 

68 Identity matrix of the same dimension as A 

69 """ 

70 

71 dtype_u = cupy_mesh 

72 dtype_f = imex_cupy_mesh 

73 

74 def __init__( 

75 self, 

76 nvars=512, 

77 nu=0.1, 

78 freq=2, 

79 bc='periodic', 

80 order=2, 

81 stencil_type='center', 

82 lintol=1e-12, 

83 liniter=10000, 

84 solver_type='direct', 

85 ): 

86 """Initialization routine""" 

87 

88 # make sure parameters have the correct types 

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

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

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

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

93 

94 # transforms nvars into a tuple 

95 if type(nvars) is int: 

96 nvars = (nvars,) 

97 

98 # automatically determine ndim from nvars 

99 ndim = len(nvars) 

100 if ndim > 3: 

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

102 

103 # eventually extend freq to other dimension 

104 if type(freq) is int: 

105 freq = (freq,) * ndim 

106 if len(freq) != ndim: 

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

108 

109 # check values for freq and nvars 

110 for f in freq: 

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

112 # use Gaussian initial solution in 1D 

113 bc = 'periodic' 

114 break 

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

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

117 for nvar in nvars: 

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

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

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

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

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

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

124 

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

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

127 self._makeAttributeAndRegister( 

128 'nvars', 

129 'nu', 

130 'freq', 

131 'bc', 

132 'order', 

133 'stencil_type', 

134 'lintol', 

135 'liniter', 

136 'solver_type', 

137 localVars=locals(), 

138 readOnly=True, 

139 ) 

140 

141 # compute dx (equal in both dimensions) and get discretization matrix A 

142 if self.bc == 'periodic': 

143 self.dx = 1.0 / self.nvars[0] 

144 xvalues = cp.array([i * self.dx for i in range(self.nvars[0])]) 

145 elif self.bc == 'dirichlet-zero': 

146 self.dx = 1.0 / (self.nvars[0] + 1) 

147 xvalues = cp.array([(i + 1) * self.dx for i in range(self.nvars[0])]) 

148 else: 

149 raise ProblemError(f'Boundary conditions {self.bc} not implemented.') 

150 

151 self.A, _ = problem_helper.get_finite_difference_matrix( 

152 derivative=2, 

153 order=self.order, 

154 stencil_type=self.stencil_type, 

155 dx=self.dx, 

156 size=self.nvars[0], 

157 dim=self.ndim, 

158 bc=self.bc, 

159 cupy=True, 

160 ) 

161 self.A *= self.nu 

162 

163 self.xv = cp.meshgrid(*[xvalues for _ in range(self.ndim)]) 

164 self.Id = csp.eye(np.prod(self.nvars), format='csc') 

165 

166 @property 

167 def ndim(self): 

168 """Number of dimensions of the spatial problem""" 

169 return len(self.nvars) 

170 

171 def eval_f(self, u, t): 

172 """ 

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

174 

175 Parameters 

176 ---------- 

177 u : dtype_u 

178 Current values of the numerical solution. 

179 t : float 

180 Current time of the numerical solution is computed. 

181 

182 Returns 

183 ------- 

184 f : dtype_f 

185 The right-hand side of the problem. 

186 """ 

187 

188 f = self.dtype_f(self.init) 

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

190 if self.ndim == 1: 

191 f.expl[:] = cp.sin(np.pi * self.freq[0] * self.xv[0]) * ( 

192 self.nu * np.pi**2 * sum([freq**2 for freq in self.freq]) * cp.cos(t) - cp.sin(t) 

193 ) 

194 elif self.ndim == 2: 

195 f.expl[:] = ( 

196 cp.sin(np.pi * self.freq[0] * self.xv[0]) 

197 * cp.sin(np.pi * self.freq[1] * self.xv[1]) 

198 * (self.nu * np.pi**2 * sum([freq**2 for freq in self.freq]) * cp.cos(t) - cp.sin(t)) 

199 ) 

200 elif self.ndim == 3: 

201 f.expl[:] = ( 

202 cp.sin(np.pi * self.freq[0] * self.xv[0]) 

203 * cp.sin(np.pi * self.freq[1] * self.xv[1]) 

204 * cp.sin(np.pi * self.freq[2] * self.xv[2]) 

205 * (self.nu * np.pi**2 * sum([freq**2 for freq in self.freq]) * cp.cos(t) - cp.sin(t)) 

206 ) 

207 

208 return f 

209 

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

211 r""" 

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

213 

214 Parameters 

215 ---------- 

216 rhs : dtype_f 

217 Right-hand side for the linear system. 

218 factor : float 

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

220 u0 : dtype_u 

221 Initial guess for the iterative solver. 

222 t : float 

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

224 

225 Returns 

226 ------- 

227 me : dtype_u 

228 Solution. 

229 """ 

230 

231 me = self.dtype_u(self.init) 

232 

233 if self.solver_type == 'direct': 

234 me[:] = spsolve(self.Id - factor * self.A, rhs.flatten()).reshape(self.nvars) 

235 elif self.solver_type == 'CG': 

236 me[:] = cg( 

237 self.Id - factor * self.A, 

238 rhs.flatten(), 

239 x0=u0.flatten(), 

240 tol=self.lintol, 

241 maxiter=self.liniter, 

242 atol=0, 

243 )[0].reshape(self.nvars) 

244 else: 

245 raise NotImplementedError(f'Solver {self.solver_type} not implemented in GPU heat equation!') 

246 return me 

247 

248 def u_exact(self, t): 

249 r""" 

250 Routine to compute the exact solution at time :math:`t`. 

251 

252 Parameters 

253 ---------- 

254 t : float 

255 Time of the exact solution. 

256 

257 Returns 

258 ------- 

259 me : dtype_u 

260 Exact solution. 

261 """ 

262 

263 me = self.dtype_u(self.init) 

264 

265 if self.ndim == 1: 

266 me[:] = cp.sin(np.pi * self.freq[0] * self.xv[0]) * cp.cos(t) 

267 elif self.ndim == 2: 

268 me[:] = cp.sin(np.pi * self.freq[0] * self.xv[0]) * cp.sin(np.pi * self.freq[1] * self.xv[1]) * cp.cos(t) 

269 elif self.ndim == 3: 

270 me[:] = ( 

271 cp.sin(np.pi * self.freq[0] * self.xv[0]) 

272 * cp.sin(np.pi * self.freq[1] * self.xv[1]) 

273 * cp.sin(np.pi * self.freq[2] * self.xv[2]) 

274 * cp.cos(t) 

275 ) 

276 return me 

277 

278 

279class heatNd_unforced(heatNd_forced): 

280 r""" 

281 This class implements the forced :math:`N`-dimensional heat equation with periodic boundary conditions 

282 

283 .. math:: 

284 \frac{\partial u}{\partial t} = \nu 

285 \left( 

286 \frac{\partial^2 u}{\partial x^2_1} + .. + \frac{\partial^2 u}{\partial x^2_N} 

287 \right) + f({\bf x}, t) 

288 

289 for :math:`(x_1,..,x_N) \in [0, 1]^{N}` with :math:`N \leq 3`, and forcing term 

290 

291 .. math:: 

292 f({\bf x}, t) = \prod_{i=1}^N \sin(\pi k_i x_i) \left( 

293 \nu \pi^2 \sum_{i=1}^N k_i^2 \cos(t) - \sin(t) 

294 \right), 

295 

296 where :math:`k_i` denotes the frequency in the :math:`i^{th}` dimension. The exact solution is 

297 

298 .. math:: 

299 u({\bf x}, t) = \prod_{i=1}^N \sin(\pi k_i x_i) \cos(t). 

300 

301 The spatial term is discretized using finite differences. 

302 

303 The implementation is this class uses the ``CuPy`` package in order to make ``pySDC`` available for GPUs. 

304 """ 

305 

306 dtype_f = cupy_mesh 

307 

308 def eval_f(self, u, t): 

309 """ 

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

311 

312 Parameters 

313 ---------- 

314 u : dtype_u 

315 Current values of the numerical solution. 

316 t : float 

317 Current time of the numerical solution is computed. 

318 

319 Returns 

320 ------- 

321 f : dtype_f 

322 The right-hand side of the problem. 

323 """ 

324 

325 f = self.dtype_f(self.init) 

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

327 

328 return f 

329 

330 def u_exact(self, t): 

331 r""" 

332 Routine to compute the exact solution at time :math:`t`. 

333 

334 Parameters 

335 ---------- 

336 t : float 

337 Time of the exact solution. 

338 

339 Returns 

340 ------- 

341 me : dtype_u 

342 The exact solution. 

343 """ 

344 

345 me = self.dtype_u(self.init) 

346 

347 if self.ndim == 1: 

348 rho = (2.0 - 2.0 * cp.cos(np.pi * self.freq[0] * self.dx)) / self.dx**2 

349 me[:] = cp.sin(np.pi * self.freq[0] * self.xv[0]) * cp.exp(-t * self.nu * rho) 

350 elif self.ndim == 2: 

351 rho = (2.0 - 2.0 * cp.cos(np.pi * self.freq[0] * self.dx)) / self.dx**2 + ( 

352 2.0 - 2.0 * cp.cos(np.pi * self.freq[1] * self.dx) 

353 ) / self.dx**2 

354 

355 me[:] = ( 

356 cp.sin(np.pi * self.freq[0] * self.xv[0]) 

357 * cp.sin(np.pi * self.freq[1] * self.xv[1]) 

358 * cp.exp(-t * self.nu * rho) 

359 ) 

360 elif self.ndim == 3: 

361 rho = ( 

362 (2.0 - 2.0 * cp.cos(np.pi * self.freq[0] * self.dx)) / self.dx**2 

363 + (2.0 - 2.0 * cp.cos(np.pi * self.freq[1] * self.dx)) 

364 + (2.0 - 2.0 * cp.cos(np.pi * self.freq[2] * self.dx)) / self.dx**2 

365 ) 

366 me[:] = ( 

367 cp.sin(np.pi * self.freq[0] * self.xv[0]) 

368 * cp.sin(np.pi * self.freq[1] * self.xv[1]) 

369 * cp.sin(np.pi * self.freq[2] * self.xv[2]) 

370 * cp.exp(-t * self.nu * rho) 

371 ) 

372 

373 return me