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

## 90 statements

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

1import numpy as np

2from petsc4py import PETSc

4from pySDC.core.errors import ParameterError, ProblemError

5from pySDC.core.problem import Problem

6from pySDC.implementations.datatype_classes.petsc_vec import petsc_vec, petsc_vec_imex

9# noinspection PyUnusedLocal

10class heat2d_petsc_forced(Problem):

11 r"""

12 Example implementing the forced two-dimensional heat equation with Dirichlet boundary conditions

13 :math:(x, y) \in [0,1]^2

15 .. math::

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

17 \left(

18 \frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2}

19 \right) + f(x, y, t)

21 and forcing term :math:f(x, y, t) such that the exact solution is

23 .. math::

24 u(x, y, t) = \sin(2 \pi x) \sin(2 \pi y) \cos(t).

26 The spatial discretization uses central finite differences and is realized with PETSc [1]_, [2]_.

28 Parameters

29 ----------

30 cnvars : tuple, optional

31 Spatial resolution for the 2D problem, e.g. cnvars=(16, 16).

32 nu : float, optional

33 Diffusion coefficient :math:\nu.

34 freq : int, optional

35 Spatial frequency of the initial conditions (equal for both dimensions).

36 refine : int, optional

37 Defines the refinement of the mesh, e.g. refine=2 means the mesh is refined with factor 2.

38 comm : COMM_WORLD

39 Communicator for PETSc.

40 sol_tol : float, optional

41 Tolerance that the solver needs to satisfy for termination.

42 sol_maxiter : int, optional

43 Maximum number of iterations for the solver to terminate.

45 Attributes

46 ----------

47 A : PETSc matrix object

48 Second-order FD discretization of the 2D Laplace operator.

49 Id : PETSc matrix object

50 Identity matrix.

51 dx : float

52 Distance between two spatial nodes in x direction.

53 dy : float

54 Distance between two spatial nodes in y direction.

55 ksp : object

56 PETSc linear solver object.

57 ksp_ncalls : int

58 Calls of PETSc's linear solver object.

59 ksp_itercount : int

60 Iterations done by PETSc's linear solver object.

62 References

63 ----------

64 .. [1] PETSc Web page. Satish Balay et al. https://petsc.org/ (2023).

65 .. [2] Parallel distributed computing using Python. Lisandro D. Dalcin, Rodrigo R. Paz, Pablo A. Kler,

66 Alejandro Cosimo. Advances in Water Resources (2011).

67 """

69 dtype_u = petsc_vec

70 dtype_f = petsc_vec_imex

72 def __init__(self, cnvars, nu, freq, refine, comm=PETSc.COMM_WORLD, sol_tol=1e-10, sol_maxiter=None):

73 """Initialization routine"""

74 # make sure parameters have the correct form

75 if len(cnvars) != 2:

76 raise ProblemError('this is a 2d example, got %s' % cnvars)

78 # create DMDA object which will be used for all grid operations

79 da = PETSc.DMDA().create([cnvars[0], cnvars[1]], stencil_width=1, comm=comm)

80 for _ in range(refine):

81 da = da.refine()

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

84 super().__init__(init=da)

85 self._makeAttributeAndRegister(

86 'cnvars',

87 'nu',

88 'freq',

89 'comm',

90 'refine',

91 'comm',

92 'sol_tol',

93 'sol_maxiter',

94 localVars=locals(),

96 )

98 # compute dx, dy and get local ranges

99 self.dx = 1.0 / (self.init.getSizes()[0] - 1)

100 self.dy = 1.0 / (self.init.getSizes()[1] - 1)

101 (self.xs, self.xe), (self.ys, self.ye) = self.init.getRanges()

103 # compute discretization matrix A and identity

104 self.A = self.__get_A()

105 self.Id = self.__get_Id()

107 # setup solver

108 self.ksp = PETSc.KSP()

109 self.ksp.create(comm=self.comm)

110 self.ksp.setType('gmres')

111 pc = self.ksp.getPC()

112 pc.setType('none')

113 # pc.setType('hypre')

114 # self.ksp.setInitialGuessNonzero(True)

115 self.ksp.setFromOptions()

116 self.ksp.setTolerances(rtol=self.sol_tol, atol=self.sol_tol, max_it=self.sol_maxiter)

118 self.ksp_ncalls = 0

119 self.ksp_itercount = 0

121 def __get_A(self):

122 r"""

123 Helper function to assemble PETSc matrix A.

125 Returns

126 -------

127 A : PETSc matrix object

128 Matrix A defining the 2D Laplace operator.

129 """

130 # create matrix and set basic options

131 A = self.init.createMatrix()

132 A.setType('aij') # sparse

133 A.setFromOptions()

134 A.setPreallocationNNZ((5, 5))

135 A.setUp()

137 # fill matrix

138 A.zeroEntries()

139 row = PETSc.Mat.Stencil()

140 col = PETSc.Mat.Stencil()

141 mx, my = self.init.getSizes()

142 (xs, xe), (ys, ye) = self.init.getRanges()

143 for j in range(ys, ye):

144 for i in range(xs, xe):

145 row.index = (i, j)

146 row.field = 0

147 if i == 0 or j == 0 or i == mx - 1 or j == my - 1:

148 A.setValueStencil(row, row, 1.0)

149 else:

150 diag = self.nu * (-2.0 / self.dx**2 - 2.0 / self.dy**2)

151 for index, value in [

152 ((i, j - 1), self.nu / self.dy**2),

153 ((i - 1, j), self.nu / self.dx**2),

154 ((i, j), diag),

155 ((i + 1, j), self.nu / self.dx**2),

156 ((i, j + 1), self.nu / self.dy**2),

157 ]:

158 col.index = index

159 col.field = 0

160 A.setValueStencil(row, col, value)

161 A.assemble()

163 return A

165 def __get_Id(self):

166 r"""

167 Helper function to assemble PETSc identity matrix.

169 Returns

170 -------

171 Id : PETSc matrix object

172 Identity matrix.

173 """

175 # create matrix and set basic options

176 Id = self.init.createMatrix()

177 Id.setType('aij') # sparse

178 Id.setFromOptions()

179 Id.setPreallocationNNZ((1, 1))

180 Id.setUp()

182 # fill matrix

183 Id.zeroEntries()

184 row = PETSc.Mat.Stencil()

185 (xs, xe), (ys, ye) = self.init.getRanges()

186 for j in range(ys, ye):

187 for i in range(xs, xe):

188 row.index = (i, j)

189 row.field = 0

190 Id.setValueStencil(row, row, 1.0)

192 Id.assemble()

194 return Id

196 def eval_f(self, u, t):

197 """

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

200 Parameters

201 ----------

202 u : dtype_u

203 Current values of the numerical solution.

204 t : float

205 Current time at which the numerical solution is computed at.

207 Returns

208 -------

209 f : dtype_f

210 Right-hand side of the problem.

211 """

213 f = self.dtype_f(self.init)

214 # evaluate Au for implicit part

215 self.A.mult(u, f.impl)

217 # evaluate forcing term for explicit part

218 fa = self.init.getVecArray(f.expl)

219 xv, yv = np.meshgrid(range(self.xs, self.xe), range(self.ys, self.ye), indexing='ij')

220 fa[self.xs : self.xe, self.ys : self.ye] = (

221 -np.sin(np.pi * self.freq * xv * self.dx)

222 * np.sin(np.pi * self.freq * yv * self.dy)

223 * (np.sin(t) - self.nu * 2.0 * (np.pi * self.freq) ** 2 * np.cos(t))

224 )

226 return f

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

229 r"""

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

232 Parameters

233 ----------

234 rhs : dtype_f

235 Right-hand side for the linear system.

236 factor : float

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

238 u0 : dtype_u

239 Initial guess for the iterative solver.

240 t : float

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

243 Returns

244 -------

245 me : dtype_u

246 Solution.

247 """

249 me = self.dtype_u(u0)

250 self.ksp.setOperators(self.Id - factor * self.A)

251 self.ksp.solve(rhs, me)

252 self.ksp_ncalls += 1

253 self.ksp_itercount += int(self.ksp.getIterationNumber())

254 return me

256 def u_exact(self, t):

257 r"""

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

260 Parameters

261 ----------

262 t : float

263 Time of the exact solution.

265 Returns

266 -------

267 me : dtype_u

268 Exact solution.

269 """

271 me = self.dtype_u(self.init)

272 xa = self.init.getVecArray(me)

273 for i in range(self.xs, self.xe):

274 for j in range(self.ys, self.ye):

275 xa[i, j] = np.sin(np.pi * self.freq * i * self.dx) * np.sin(np.pi * self.freq * j * self.dy) * np.cos(t)

277 return me