Coverage for pySDC/implementations/problem_classes/HeatFiredrake.py: 98%

1from pySDC.core.problem import Problem, WorkCounter 

2from pySDC.implementations.datatype_classes.firedrake_mesh import firedrake_mesh, IMEX_firedrake_mesh 

3import firedrake as fd 

4import numpy as np 

5from mpi4py import MPI 

6 

7 

8class Heat1DForcedFiredrake(Problem): 

9 r""" 

10 Forced 1D heat equation with Dirichlet BCs and Firedrake finite elements, IMEX with diffusion implicit. 

11 

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

13 

14 .. math:: 

15 \frac{d u}{d t} = \nu \frac{d^2 u}{d x^2} + f 

16 

17 for :math:`x \in \Omega:=[0,1]`, where the forcing term :math:`f` is defined by 

18 

19 .. math:: 

20 f(x, t) = -\sin(\pi x) (\sin(t) - \nu \pi^2 \cos(t)). 

21 

22 For initial conditions with constant c 

23 

24 .. math:: 

25 u(x, 0) = \sin(\pi x) + c, 

26 

27 the exact solution is given by 

28 

29 .. math:: 

30 u(x, t) = \sin(\pi x)\cos(t) + c. 

31 

32 Here, the problem is discretized with finite elements using firedrake. Hence, the problem 

33 is reformulated to the *weak formulation* 

34 

35 .. math: 

36 \int_\Omega u_t v\,dx = - \nu \int_\Omega \nabla u \nabla v\,dx + \int_\Omega f v\,dx. 

37 

38 We invert the Laplacian implicitly and treat the forcing term explicitly. 

39 The solvers for the arising variational problems are cached for multiple collocation nodes and step sizes. 

40 

41 Parameters 

42 ---------- 

43 nvars : int, optional 

44 Spatial resolution, i.e., numbers of degrees of freedom in space. 

45 nu : float, optional 

46 Diffusion coefficient :math:`\nu`. 

47 c: float, optional 

48 Constant for the Dirichlet boundary condition :math: `c` 

49 LHS_cache_size : int, optional 

50 Cache size for variational problem solvers 

51 comm : MPI communicator, optional 

52 Supply an MPI communicator for spatial parallelism 

53 """ 

54 

55 dtype_u = firedrake_mesh 

56 dtype_f = IMEX_firedrake_mesh 

57 

58 def __init__(self, n=30, nu=0.1, c=0.0, order=4, LHS_cache_size=12, mesh=None, comm=None): 

59 """ 

60 Initialization 

61 

62 Args: 

63 n (int): Number of degrees of freedom 

64 nu (float): Diffusion parameter 

65 c (float): Boundary condition constant 

66 order (int): Order of finite elements 

67 LHS_cache_size (int): Size of the cache for solvers 

68 mesh (Firedrake mesh, optional): Give a custom mesh, for instance from a hierarchy 

69 comm (mpi4pi.Intracomm, optional): MPI communicator for spatial parallelism 

70 """ 

71 comm = MPI.COMM_WORLD if comm is None else comm 

72 

73 # prepare Firedrake mesh and function space 

74 self.mesh = fd.UnitIntervalMesh(n, comm=comm) if mesh is None else mesh 

75 self.V = fd.FunctionSpace(self.mesh, "CG", order) 

76 

77 # prepare pySDC problem class infrastructure by passing the function space to super init 

78 super().__init__(self.V) 

79 self._makeAttributeAndRegister( 

80 'n', 'nu', 'c', 'order', 'LHS_cache_size', 'comm', localVars=locals(), readOnly=True 

81 ) 

82 

83 # prepare caches and IO variables for solvers 

84 self.solvers = {} 

85 self.tmp_in = fd.Function(self.V) 

86 self.tmp_out = fd.Function(self.V) 

87 

88 # prepare work counters 

89 self.work_counters['solver_setup'] = WorkCounter() 

90 self.work_counters['solves'] = WorkCounter() 

91 self.work_counters['rhs'] = WorkCounter() 

92 

93 def eval_f(self, u, t): 

94 """ 

95 Evaluate the right hand side. 

96 The forcing term is simply interpolated to the grid. 

97 The Laplacian is evaluated via a variational problem, where the mass matrix is inverted and homogeneous boundary conditions are applied. 

98 Note that we cache the solver to obtain much better performance. 

99 

100 Parameters 

101 ---------- 

102 u : dtype_u 

103 Solution at which to evaluate 

104 t : float 

105 Time at which to evaluate 

106 

107 Returns 

108 ------- 

109 f : dtype_f 

110 The evaluated right hand side 

111 """ 

112 # construct and cache a solver for the implicit part of the right hand side evaluation 

113 if not hasattr(self, '_solv_eval_f_implicit'): 

114 v = fd.TestFunction(self.V) 

115 u_trial = fd.TrialFunction(self.V) 

116 

117 a = u_trial * v * fd.dx 

118 L_impl = -fd.inner(self.nu * fd.nabla_grad(self.tmp_in), fd.nabla_grad(v)) * fd.dx 

119 

120 bcs = [fd.bcs.DirichletBC(self.V, fd.Constant(0), area) for area in [1, 2]] 

121 

122 prob = fd.LinearVariationalProblem(a, L_impl, self.tmp_out, bcs=bcs) 

123 self._solv_eval_f_implicit = fd.LinearVariationalSolver(prob) 

124 

125 # copy the solution we want to evaluate at into the input buffer 

126 self.tmp_in.assign(u.functionspace) 

127 

128 # perform the solve using the cached solver 

129 self._solv_eval_f_implicit.solve() 

130 

131 me = self.dtype_f(self.init) 

132 

133 # copy the result of the solver from the output buffer to the variable this function returns 

134 me.impl.assign(self.tmp_out) 

135 

136 # evaluate explicit part. 

137 # Because it does not depend on the current solution, we can simply interpolate the expression 

138 x = fd.SpatialCoordinate(self.mesh) 

139 me.expl.interpolate(-(np.sin(t) - self.nu * np.pi**2 * np.cos(t)) * fd.sin(np.pi * x[0])) 

140 

141 self.work_counters['rhs']() 

142 

143 return me 

144 

145 def solve_system(self, rhs, factor, *args, **kwargs): 

146 r""" 

147 Linear solver for :math:`(M - factor nu * Lap) u = rhs`. 

148 

149 Parameters 

150 ---------- 

151 rhs : dtype_f 

152 Right-hand side for the nonlinear system. 

153 factor : float 

154 Abbrev. for the node-to-node stepsize (or any other factor required). 

155 *args, **kwargs 

156 Take the initial guess ``u0`` and the current time ``t`` of the generic interface, neither is used here. 

157 

158 Returns 

159 ------- 

160 u : dtype_u 

161 Solution. 

162 """ 

163 

164 # construct and cache a solver for the current factor (preconditioner entry times step size) 

165 if factor not in self.solvers.keys(): 

166 

167 # check if we need to evict something from the cache 

168 if len(self.solvers) >= self.LHS_cache_size: 

169 self.solvers.pop(list(self.solvers.keys())[0]) 

170 

171 u = fd.TrialFunction(self.V) 

172 v = fd.TestFunction(self.V) 

173 

174 a = u * v * fd.dx + fd.Constant(factor) * fd.inner(self.nu * fd.nabla_grad(u), fd.nabla_grad(v)) * fd.dx 

175 L = fd.inner(self.tmp_in, v) * fd.dx 

176 

177 bcs = [fd.bcs.DirichletBC(self.V, fd.Constant(self.c), area) for area in [1, 2]] 

178 

179 prob = fd.LinearVariationalProblem(a, L, self.tmp_out, bcs=bcs) 

180 self.solvers[factor] = fd.LinearVariationalSolver(prob) 

181 

182 self.work_counters['solver_setup']() 

183 

184 # copy solver rhs to the input buffer. Copying also to the output buffer uses it as initial guess 

185 self.tmp_in.assign(rhs.functionspace) 

186 self.tmp_out.assign(rhs.functionspace) 

187 

188 # call the cached solver 

189 self.solvers[factor].solve() 

190 

191 # copy from output buffer to return variable 

192 me = self.dtype_u(self.init) 

193 me.assign(self.tmp_out) 

194 

195 self.work_counters['solves']() 

196 return me 

197 

198 @fd.utils.cached_property 

199 def x(self): 

200 """ 

201 Spatial coordinate of the mesh, cached after the first access. 

202 

203 Returns 

204 ------- 

205 x : ufl.SpatialCoordinate 

206 Symbolic spatial coordinate, used to interpolate expressions on the mesh. 

207 """ 

208 return fd.SpatialCoordinate(self.mesh) 

209 

210 def u_exact(self, t): 

211 r""" 

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

213 

214 Parameters 

215 ---------- 

216 t : float 

217 Time of the exact solution. 

218 

219 Returns 

220 ------- 

221 me : dtype_u 

222 Exact solution. 

223 """ 

224 me = self.u_init 

225 me.interpolate(np.cos(t) * fd.sin(np.pi * self.x[0]) + self.c) 

226 return me