Adaptive Multigrid

firedrake/__init__.py

@@ -100,7 +100,8 @@ def init_petsc():
     HierarchyBase, MeshHierarchy, ExtrudedMeshHierarchy,
     NonNestedHierarchy, SemiCoarsenedExtrudedHierarchy,
     prolong, restrict, inject, TransferManager,
-    OpenCascadeMeshHierarchy
+    OpenCascadeMeshHierarchy, AdaptiveMeshHierarchy,
+    AdaptiveTransferManager
 )
 from firedrake.norms import errornorm, norm  # noqa: F401
 from firedrake.nullspace import VectorSpaceBasis, MixedVectorSpaceBasis  # noqa: F401

firedrake/mg/__init__.py

@@ -7,3 +7,5 @@
 )
 from firedrake.mg.embedded import TransferManager  # noqa F401
 from firedrake.mg.opencascade_mh import OpenCascadeMeshHierarchy  # noqa F401
+from firedrake.mg.adaptive_hierarchy import AdaptiveMeshHierarchy  # noqa F401
+from firedrake.mg.adaptive_transfer_manager import AdaptiveTransferManager  # noqa: F401

firedrake/mg/adaptive_hierarchy.py

@@ -0,0 +1,79 @@
+"""
+This module contains the class for the AdaptiveMeshHierarchy and
+related helper functions
+"""
+
+from collections import defaultdict
+
+from firedrake.cofunction import Cofunction
+from firedrake.function import Function
+from firedrake.mg import HierarchyBase
+from firedrake.mg.utils import set_level
+
+__all__ = ["AdaptiveMeshHierarchy"]
+
+
+class AdaptiveMeshHierarchy(HierarchyBase):
+    """
+    HierarchyBase for hierarchies of adaptively refined meshes
+    """
+    def __init__(self, base_mesh, refinements_per_level=1, nested=True):
+        self.meshes = []
+        self._meshes = []
+        self.refinements_per_level = refinements_per_level
+        self.nested = nested
+        self.add_mesh(base_mesh)
+
+    def add_mesh(self, mesh):
+        """
+        Adds newly refined mesh into hierarchy.
+        """
+        self._meshes.append(mesh)
+        self.meshes.append(mesh)
+        level = len(self.meshes)
+        set_level(self.meshes[-1], self, level - 1)
+        self._shared_data_cache = defaultdict(dict)
+
+    def adapt(self, eta: Function | Cofunction, theta: float):
+        """
+        Add a refinement level to the hierarchy by local refinement
+        with a simplified variant of Dorfler marking.
+
+        Parameters
+        ----------
+        eta
+            A DG0 :class:`~firedrake.function.Function` with the local error estimator.
+        theta
+            The threshold for marking as a fraction of the maximum error.
+
+        Note
+        ----
+        Dorfler marking involves sorting all of the elements by decreasing
+        error estimator and taking the minimal set that exceeds some fixed
+        fraction of the total error. What this code implements is the simpler
+        variant that doesn't have a proof of convergence (as far as I know)
+        but works as well in practice.
+
+        """
+        if not isinstance(eta, (Function, Cofunction)):
+            raise TypeError(f"eta must be a Function or Cofunction, not a {type(eta).__name__}")
+        M = eta.function_space()
+        if M.finat_element.space_dimension() != 1:
+            raise ValueError("eta must be a Function or Cofunction in DG0")
+        mesh = self.meshes[-1]
+        if M.mesh() is not mesh:
+            raise ValueError("eta must be defined on the finest mesh of the hierarchy")
+
+        # Take the maximum over all processes
+        with eta.dat.vec_ro as eta_:
+            eta_max = eta_.max()[1]
+
+        threshold = theta * eta_max
+        should_refine = eta.dat.data_ro > threshold
+
+        markers = Function(M)
+        markers.dat.data_wo[should_refine] = 1
+
+        refined_mesh = mesh.refine_marked_elements(markers)
+        self.add_mesh(refined_mesh)
+        return refined_mesh

firedrake/mg/adaptive_transfer_manager.py

@@ -0,0 +1,94 @@
+"""
+This module contains the AdaptiveTransferManager used to perform
+transfer operations on AdaptiveMeshHierarchies
+"""
+from firedrake.mg.embedded import TransferManager
+from firedrake.ufl_expr import action, TrialFunction
+from firedrake.functionspace import FunctionSpace, TensorFunctionSpace
+from firedrake.interpolation import interpolate
+from firedrake.preconditioners.bddc import is_lagrange
+from finat.quadrature import QuadratureRule
+from functools import partial
+
+import numpy
+
+
+__all__ = ("AdaptiveTransferManager",)
+
+
+class AdaptiveTransferManager(TransferManager):
+    """
+    TransferManager for adaptively refined mesh hierarchies
+    """
+    def __init__(self, *, native_transfers=None, use_averaging=True):
+        super().__init__(native_transfers=native_transfers, use_averaging=use_averaging)
+        self.cache = {}
+
+    def get_operators(self, Vc, Vf):
+        key = (Vc, Vf)
+        try:
+            return self.cache[key]
+        except KeyError:
+            ops = get_mg_interpolator(Vc, Vf)
+            return self.cache.setdefault(key, ops)
+
+    def forward(self, uc, uf):
+        from firedrake.assemble import assemble
+        Vc = uc.function_space()
+        Vf = uf.function_space()
+        ops = self.get_operators(Vc, Vf)
+
+        expr = uc
+        for op in ops:
+            expr = action(op, expr)
+        return assemble(expr, tensor=uf)
+
+    def adjoint(self, uf, uc):
+        from firedrake.assemble import assemble
+        Vc = uc.function_space().dual()
+        Vf = uf.function_space().dual()
+        ops = self.get_operators(Vc, Vf)
+
+        expr = uf
+        for op in reversed(ops):
+            expr = action(expr, op)
+        return assemble(expr, tensor=uc)
+
+    def prolong(self, uf, uc):
+        return self.forward(uf, uc)
+
+    def inject(self, uc, uf):
+        return self.forward(uc, uf)
+
+    def restrict(self, uc, uf):
+        return self.adjoint(uc, uf)
+
+
+def make_quadrature_space(V):
+    fe = V.finat_element
+    _, ps = fe.dual_basis
+    wts = numpy.full(len(ps.points), numpy.nan)
+    scheme = QuadratureRule(ps, wts, ref_el=fe.cell)
+    if V.value_shape == ():
+        make_space = FunctionSpace
+    else:
+        make_space = partial(TensorFunctionSpace, shape=V.value_shape)
+    return make_space(V.mesh(), "Quadrature", degree=fe.degree, quad_scheme=scheme)
+
+
+def get_mg_interpolator(V1, V2):
+    from firedrake.assemble import assemble
+    if is_lagrange(V2.finat_element):
+        spaces = (V1, V2)
+    else:
+        Q2 = make_quadrature_space(V2)
+        spaces = (V1, Q2, V2)
+
+    ops = []
+    for i in range(len(spaces)-1):
+        Vsrc = spaces[i]
+        Vdest = spaces[i+1]
+        Iexpr = interpolate(TrialFunction(Vsrc), Vdest)
+        op = assemble(Iexpr, mat_type="aij")
+        ops.append(op)
+    return ops