Source code for firedrake.preconditioners.hiptmair

import abc

from firedrake.petsc import PETSc
from firedrake.preconditioners.base import PCBase
from firedrake.functionspace import FunctionSpace
from firedrake.ufl_expr import TestFunction, TrialFunction
from firedrake.preconditioners.hypre_ams import chop
from firedrake.parameters import parameters
from firedrake_citations import Citations
from firedrake.interpolation import Interpolator
from import expand_derivatives
import firedrake.dmhooks as dmhooks
import ufl
import finat.ufl

__all__ = ("TwoLevelPC", "HiptmairPC")

[docs] class TwoLevelPC(PCBase): """ PC for two-level methods should implement: - :meth:`coarsen` """ needs_python_pmat = False
[docs] @abc.abstractmethod def coarsen(self, pc): """Return a tuple with coarse bilinear form, coarse boundary conditions, and coarse-to-fine interpolation matrix """ raise NotImplementedError
[docs] def initialize(self, pc): from firedrake.assemble import allocate_matrix, TwoFormAssembler A, P = pc.getOperators() appctx = self.get_appctx(pc) fcp = appctx.get("form_compiler_parameters") prefix = pc.getOptionsPrefix() options_prefix = prefix + self._prefix opts = PETSc.Options() coarse_operator, coarse_space_bcs, interp_petscmat = self.coarsen(pc) # Handle the coarse operator coarse_options_prefix = options_prefix + "mg_coarse_" coarse_mat_type = opts.getString(coarse_options_prefix + "mat_type", parameters["default_matrix_type"]) self.coarse_op = allocate_matrix(coarse_operator, bcs=coarse_space_bcs, form_compiler_parameters=fcp, mat_type=coarse_mat_type, options_prefix=coarse_options_prefix) self._assemble_coarse_op = TwoFormAssembler(coarse_operator, tensor=self.coarse_op, form_compiler_parameters=fcp, bcs=coarse_space_bcs).assemble self._assemble_coarse_op() coarse_opmat = self.coarse_op.petscmat # We set up a PCMG object that uses the constructed interpolation # matrix to generate the restriction/prolongation operators. # This is a two-level multigrid preconditioner. pcmg = PETSc.PC().create(comm=pc.comm) pcmg.incrementTabLevel(1, parent=pc) pcmg.setType(pc.Type.MG) pcmg.setOptionsPrefix(options_prefix) pcmg.setMGLevels(2) pcmg.setMGType(pc.MGType.ADDITIVE) pcmg.setMGCycleType(pc.MGCycleType.V) pcmg.setMGInterpolation(1, interp_petscmat) # FIXME the default for MGRScale is created with the wrong shape when dim(coarse) > dim(fine) # FIXME there is no need for injection in a KSP context, probably this comes from the snes_ctx below # as workaround define injection as the restriction of the solution times a zero vector pcmg.setMGRScale(1, interp_petscmat.createVecRight()) pcmg.setOperators(A=A, P=P) coarse_solver = pcmg.getMGCoarseSolve() coarse_solver.setOperators(A=coarse_opmat, P=coarse_opmat) # coarse space dm coarse_space = coarse_operator.arguments()[-1].function_space() coarse_dm = coarse_solver.setDM(coarse_dm) coarse_solver.setDMActive(False) pcmg.setDM(pc.getDM()) pcmg.setFromOptions() self.pc = pcmg self._dm = coarse_dm prefix = coarse_solver.getOptionsPrefix() # Create new appctx self._ctx_ref = self.new_snes_ctx(pc, coarse_operator, coarse_space_bcs, coarse_mat_type, fcp, options_prefix=prefix) with dmhooks.add_hooks(coarse_dm, self, appctx=self._ctx_ref, save=False): coarse_solver.setFromOptions()
[docs] def update(self, pc): self._assemble_coarse_op() self.pc.setUp()
[docs] def apply(self, pc, X, Y): dm = self._dm with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref): self.pc.apply(X, Y)
[docs] def applyTranspose(self, pc, X, Y): dm = self._dm with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref): self.pc.applyTranspose(X, Y)
[docs] def view(self, pc, viewer=None): super(TwoLevelPC, self).view(pc, viewer) if hasattr(self, "pc"): viewer.printfASCII("Two level PC\n") self.pc.view(viewer)
[docs] class HiptmairPC(TwoLevelPC): """A two-level method for H(curl) or H(div) problems with an auxiliary potential space in H^1 or H(curl), respectively. Internally this creates a PETSc PCMG object that can be controlled by options using the extra options prefix ``hiptmair_mg_``. This allows for effective multigrid relaxation methods with patch solves centered around vertices for H^1, edges for H(curl), or faces for H(div). For the lowest-order spaces this corresponds to point-Jacobi. The H(div) auxiliary vector potential problem in H(curl) is singular for high-order. This can be overcome by pertubing the problem by a multiple of the mass matrix. The scaling factor can be provided (defaulting to 0) by providing a scalar in the application context, keyed on ``"hiptmair_shift"``. """ _prefix = "hiptmair_"
[docs] def coarsen(self, pc): Citations().register("Hiptmair1998") appctx = self.get_appctx(pc) _, P = pc.getOperators() if P.getType() == "python": ctx = P.getPythonContext() a = ctx.a bcs = tuple(ctx.bcs) else: ctx = dmhooks.get_appctx(pc.getDM()) a = ctx.Jp or ctx.J bcs = tuple(ctx._problem.bcs) V = a.arguments()[-1].function_space() mesh = V.mesh() element = V.ufl_element() degree = try: degree = max(degree) except TypeError: pass formdegree = V.finat_element.formdegree if formdegree == 1: celement = curl_to_grad(element) dminus = ufl.grad G_callback = appctx.get("get_gradient", None) elif formdegree == 2: celement = div_to_curl(element) dminus = ufl.curl if V.shape: dminus = lambda u: ufl.as_vector([ufl.curl(u[k, ...]) for k in range(u.ufl_shape[0])]) G_callback = appctx.get("get_curl", None) else: raise ValueError("Hiptmair decomposition not available for", element) coarse_space = FunctionSpace(mesh, celement) assert coarse_space.finat_element.formdegree + 1 == formdegree coarse_space_bcs = tuple(bc.reconstruct(V=coarse_space, g=0) for bc in bcs) # Get only the zero-th order term of the form replace_dict = {ufl.grad(t): for t in a.arguments()} beta = ufl.replace(expand_derivatives(a), replace_dict) test = TestFunction(coarse_space) trial = TrialFunction(coarse_space) coarse_operator = beta(dminus(test), dminus(trial), coefficients={}) if formdegree > 1 and degree > 1: shift = appctx.get("hiptmair_shift", None) if shift is not None: coarse_operator += beta(test, shift*trial, coefficients={}) if G_callback is None: interp_petscmat = chop(Interpolator(dminus(test), V, bcs=bcs + coarse_space_bcs).callable().handle) else: interp_petscmat = G_callback(coarse_space, V, coarse_space_bcs, bcs) return coarse_operator, coarse_space_bcs, interp_petscmat
def curl_to_grad(ele): if isinstance(ele, finat.ufl.VectorElement): return type(ele)(curl_to_grad(ele._sub_element), dim=ele.num_sub_elements) elif isinstance(ele, finat.ufl.TensorElement): return type(ele)(curl_to_grad(ele._sub_element), shape=ele.value_shape, symmetry=ele.symmetry()) elif isinstance(ele, finat.ufl.MixedElement): return type(ele)(*(curl_to_grad(e) for e in ele.sub_elements)) elif isinstance(ele, finat.ufl.RestrictedElement): return finat.ufl.RestrictedElement(curl_to_grad(ele._element), ele.restriction_domain()) else: cell = ele.cell family = variant = ele.variant() degree = if family.startswith("Sminus"): family = "S" else: family = "CG" if isinstance(degree, tuple) and isinstance(cell, ufl.TensorProductCell): cells = ele.cell.sub_cells() elems = [finat.ufl.FiniteElement(family, cell=c, degree=d, variant=variant) for c, d in zip(cells, degree)] return finat.ufl.TensorProductElement(*elems, cell=cell) return finat.ufl.FiniteElement(family, cell=cell, degree=degree, variant=variant) def div_to_curl(ele): if isinstance(ele, finat.ufl.VectorElement): return type(ele)(div_to_curl(ele._sub_element), dim=ele.num_sub_elements) elif isinstance(ele, finat.ufl.TensorElement): return type(ele)(div_to_curl(ele._sub_element), shape=ele.value_shape, symmetry=ele.symmetry()) elif isinstance(ele, finat.ufl.MixedElement): return type(ele)(*(div_to_curl(e) for e in ele.sub_elements)) elif isinstance(ele, finat.ufl.RestrictedElement): return finat.ufl.RestrictedElement(div_to_curl(ele._element), ele.restriction_domain()) elif isinstance(ele, finat.ufl.EnrichedElement): return type(ele)(*(div_to_curl(e) for e in reversed(ele._elements))) elif isinstance(ele, finat.ufl.TensorProductElement): return type(ele)(*(div_to_curl(e) for e in ele.sub_elements), cell=ele.cell) elif isinstance(ele, finat.ufl.WithMapping): return type(ele)(div_to_curl(ele.wrapee), ele.mapping()) elif isinstance(ele, finat.ufl.BrokenElement): return type(ele)(div_to_curl(ele._element)) elif isinstance(ele, finat.ufl.HDivElement): return finat.ufl.HCurlElement(div_to_curl(ele._element)) elif isinstance(ele, finat.ufl.HCurlElement): raise ValueError("Expecting an H(div) element") else: degree = family = if family in ["Lagrange", "CG", "Q"]: family = "DG" if ele.cell.is_simplex() else "DQ" degree = degree-1 elif family in ["Discontinuous Lagrange", "DG", "DQ"]: family = "CG" degree = degree+1 else: replace_dict = { "Raviart-Thomas": "N1curl", "Brezzi-Douglas-Marini": "N2curl", "RTCF": "RTCE", "NCF": "NCE", "SminusF": "SminusE", "SminusDiv": "SminusCurl", } family = replace_dict.get(family, None) if family is None: raise ValueError("Unexpected family %s" % family) return ele.reconstruct(degree=degree, family=family)