Coverage for pySDC/implementations/problem_classes/HeatEquation_2D_PETSc_forced.py: 100%
90 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-29 09:02 +0000
1import numpy as np
2from petsc4py import PETSc
4from pySDC.core.Errors import ParameterError, ProblemError
5from pySDC.core.Problem import ptype
6from pySDC.implementations.datatype_classes.petsc_vec import petsc_vec, petsc_vec_imex
9# noinspection PyUnusedLocal
10class heat2d_petsc_forced(ptype):
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(),
95 readOnly=True,
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