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

## 26 statements

, created at 2024-09-09 14:59 +0000

1import numpy as np

2import cupy as cp

3import cupyx.scipy.sparse as csp

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

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

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

13 r"""

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

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)

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

24 .. math::

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

27 The spatial term is discretized using finite differences.

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

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

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 }

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 """

71 dtype_u = cupy_mesh

72 dtype_f = imex_cupy_mesh

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"""

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

94 # transforms nvars into a tuple

95 if type(nvars) is int:

96 nvars = (nvars,)

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}')

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}')

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)

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

139 )

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

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

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

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

166 @property

167 def ndim(self):

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

169 return len(self.nvars)

171 def eval_f(self, u, t):

172 """

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

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.

182 Returns

183 -------

184 f : dtype_f

185 The right-hand side of the problem.

186 """

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 )

208 return f

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

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

225 Returns

226 -------

227 me : dtype_u

228 Solution.

229 """

231 me = self.dtype_u(self.init)

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

248 def u_exact(self, t):

249 r"""

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

252 Parameters

253 ----------

254 t : float

255 Time of the exact solution.

257 Returns

258 -------

259 me : dtype_u

260 Exact solution.

261 """

263 me = self.dtype_u(self.init)

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

279class heatNd_unforced(heatNd_forced):

280 r"""

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

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)

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

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

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

298 .. math::

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

301 The spatial term is discretized using finite differences.

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

304 """

306 dtype_f = cupy_mesh

308 def eval_f(self, u, t):

309 """

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

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.

319 Returns

320 -------

321 f : dtype_f

322 The right-hand side of the problem.

323 """

325 f = self.dtype_f(self.init)

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

328 return f

330 def u_exact(self, t):

331 r"""

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

334 Parameters

335 ----------

336 t : float

337 Time of the exact solution.

339 Returns

340 -------

341 me : dtype_u

342 The exact solution.

343 """

345 me = self.dtype_u(self.init)

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

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 )

373 return me