RestrictedFunctionSpace: support Fieldsplit, multigrid, and python PC (#4169)

* dmhooks: support RestrictedFunctionSpace

* RestrictedFunctionSpace: support geometric multigrid

* RestrictedFunctionSpace: support p-multigrid

* PC: replace FunctionSpace() constructor with V.reconstruct() to preserve boundary restriction ids.

Author: pbrubeck

firedrake/bcs.py

@@ -319,6 +319,8 @@ def reconstruct(self, field=None, V=None, g=None, sub_domain=None, use_split=Fal
             V = V.sub(index)
         if g is None:
             g = self._original_arg
+            if isinstance(g, firedrake.Function) and g.function_space() != V:
+                g = firedrake.Function(V).interpolate(g)
         if sub_domain is None:
             sub_domain = self.sub_domain
         if field is not None:
@@ -744,11 +746,11 @@ def restricted_function_space(V, ids):
         return V
 
     assert len(ids) == len(V)
-    spaces = [Vsub if len(boundary_set) == 0 else
-              firedrake.RestrictedFunctionSpace(Vsub, boundary_set=boundary_set)
-              for Vsub, boundary_set in zip(V, ids)]
+    spaces = [V_ if len(boundary_set) == 0 else
+              firedrake.RestrictedFunctionSpace(V_, boundary_set=boundary_set, name=V_.name)
+              for V_, boundary_set in zip(V, ids)]
 
     if len(spaces) == 1:
         return spaces[0]
     else:
-        return firedrake.MixedFunctionSpace(spaces)
+        return firedrake.MixedFunctionSpace(spaces, name=V.name)

firedrake/dmhooks.py

@@ -53,11 +53,13 @@ def get_function_space(dm):
     :raises RuntimeError: if no function space was found.
     """
     info = dm.getAttr("__fs_info__")
-    meshref, element, indices, (name, names) = info
+    meshref, element, indices, (name, names), boundary_sets = info
     mesh = meshref()
     if mesh is None:
         raise RuntimeError("Somehow your mesh was collected, this should never happen")
     V = firedrake.FunctionSpace(mesh, element, name=name)
+    if any(boundary_sets):
+        V = firedrake.bcs.restricted_function_space(V, boundary_sets)
     if len(V) > 1:
         for V_, name in zip(V, names):
             V_.topological.name = name
@@ -93,8 +95,8 @@ def set_function_space(dm, V):
     if len(V) > 1:
         names = tuple(V_.name for V_ in V)
     element = V.ufl_element()
-
-    info = (weakref.ref(mesh), element, tuple(reversed(indices)), (V.name, names))
+    boundary_sets = tuple(V_.boundary_set for V_ in V)
+    info = (weakref.ref(mesh), element, tuple(reversed(indices)), (V.name, names), boundary_sets)
     dm.setAttr("__fs_info__", info)
 
 
@@ -457,7 +459,7 @@ def refine(dm, comm):
     if hasattr(V, "_fine"):
         fdm = V._fine.dm
     else:
-        V._fine = firedrake.FunctionSpace(hierarchy[level + 1], V.ufl_element())
+        V._fine = V.reconstruct(mesh=hierarchy[level + 1])
         fdm = V._fine.dm
     V._fine._coarse = V
     return fdm

firedrake/functionspaceimpl.py

@@ -377,17 +377,20 @@ def make_function_space(cls, mesh, element, name=None):
             new = cls.create(new, mesh)
         return new
 
-    def reconstruct(self, mesh=None, name=None, **kwargs):
+    def reconstruct(self, mesh=None, element=None, name=None, **kwargs):
         r"""Reconstruct this :class:`.WithGeometryBase` .
 
         :kwarg mesh: the new :func:`~.Mesh` (defaults to same mesh)
+        :kwarg element: the new :class:`finat.ufl.FiniteElement` (defaults to same element)
         :kwarg name: the new name (defaults to None)
         :returns: the new function space of the same class as ``self``.
 
         Any extra kwargs are used to reconstruct the finite element.
         For details see :meth:`finat.ufl.finiteelement.FiniteElement.reconstruct`.
         """
+        from firedrake.bcs import restricted_function_space
         V_parent = self
+
         # Deal with ProxyFunctionSpace
         indices = []
         while True:
@@ -402,13 +405,21 @@ def reconstruct(self, mesh=None, name=None, **kwargs):
 
         if mesh is None:
             mesh = V_parent.mesh()
+        if element is None:
+            element = V_parent.ufl_element()
 
-        element = V_parent.ufl_element()
         cell = mesh.topology.ufl_cell()
         if len(kwargs) > 0 or element.cell != cell:
             element = element.reconstruct(cell=cell, **kwargs)
 
+        # Reconstruct the parent space
         V = type(self).make_function_space(mesh, element, name=name)
+
+        # Deal with RestrictedFunctionSpace
+        boundary_sets = [V_.boundary_set for V_ in V_parent]
+        if any(boundary_sets):
+            V = restricted_function_space(V, boundary_sets)
+
         for i in reversed(indices):
             V = V.sub(i)
         return V
@@ -899,8 +910,7 @@ def __init__(self, function_space, boundary_set=frozenset(), name=None):
                                                      function_space.ufl_element(),
                                                      label=self._label)
         self.function_space = function_space
-        self.name = name or (function_space.name or "Restricted" + "_"
-                             + "_".join(sorted(map(str, self.boundary_set))))
+        self.name = name or function_space.name
 
     def set_shared_data(self):
         sdata = get_shared_data(self._mesh, self.ufl_element(), self.boundary_set)

firedrake/mg/embedded.py

@@ -87,11 +87,14 @@ def _native_transfer(self, element, op):
         return None
 
     def cache(self, V):
-        key = (V.ufl_element(), V.value_shape)
+        key = (V.ufl_element(), V.value_shape, V.boundary_set)
         try:
             return self.caches[key]
         except KeyError:
-            return self.caches.setdefault(key, TransferManager.Cache(*key))
+            return self.caches.setdefault(key, TransferManager.Cache(*key[:2]))
+
+    def cache_key(self, V):
+        return (V.dim(),)
 
     def V_dof_weights(self, V):
         """Dof weights for averaging projection.
@@ -100,7 +103,7 @@ def V_dof_weights(self, V):
         :returns: A PETSc Vec.
         """
         cache = self.cache(V)
-        key = V.dim()
+        key = self.cache_key(V)
         try:
             return cache._V_dof_weights[key]
         except KeyError:
@@ -125,7 +128,7 @@ def V_DG_mass(self, V, DG):
         :returns: A PETSc Mat mapping from V -> DG
         """
         cache = self.cache(V)
-        key = V.dim()
+        key = self.cache_key(V)
         try:
             return cache._V_DG_mass[key]
         except KeyError:
@@ -156,7 +159,7 @@ def V_approx_inv_mass(self, V, DG):
         :returns: A PETSc Mat mapping from V -> DG.
         """
         cache = self.cache(V)
-        key = V.dim()
+        key = self.cache_key(V)
         try:
             return cache._V_approx_inv_mass[key]
         except KeyError:
@@ -174,7 +177,7 @@ def V_inv_mass_ksp(self, V):
         :returns: A PETSc KSP for inverting (V, V).
         """
         cache = self.cache(V)
-        key = V.dim()
+        key = self.cache_key(V)
         try:
             return cache._V_inv_mass_ksp[key]
         except KeyError:
@@ -196,7 +199,7 @@ def DG_work(self, V):
         """
         needs_dual = ufl.duals.is_dual(V)
         cache = self.cache(V)
-        key = (V.dim(), needs_dual)
+        key = self.cache_key(V) + (needs_dual,)
         try:
             return cache._DG_work[key]
         except KeyError:
@@ -213,7 +216,7 @@ def work_vec(self, V):
         :returns: A PETSc Vec for V.
         """
         cache = self.cache(V)
-        key = V.dim()
+        key = self.cache_key(V)
         try:
             return cache._work_vec[key]
         except KeyError:

firedrake/mg/interface.py

@@ -58,18 +58,17 @@ def prolong(coarse, fine):
     repeat = (fine_level - coarse_level)*refinements_per_level
     next_level = coarse_level * refinements_per_level
 
-    element = Vc.ufl_element()
     meshes = hierarchy._meshes
     for j in range(repeat):
         next_level += 1
         if j == repeat - 1:
             next = fine
             Vf = fine.function_space()
         else:
-            Vf = firedrake.FunctionSpace(meshes[next_level], element)
+            Vf = Vc.reconstruct(mesh=meshes[next_level])
             next = firedrake.Function(Vf)
 
-        coarse_coords = Vc.mesh().coordinates
+        coarse_coords = get_coordinates(Vc)
         fine_to_coarse = utils.fine_node_to_coarse_node_map(Vf, Vc)
         fine_to_coarse_coords = utils.fine_node_to_coarse_node_map(Vf, coarse_coords.function_space())
         kernel = kernels.prolong_kernel(coarse)
@@ -119,7 +118,6 @@ def restrict(fine_dual, coarse_dual):
     repeat = (fine_level - coarse_level)*refinements_per_level
     next_level = fine_level * refinements_per_level
 
-    element = Vc.ufl_element()
     meshes = hierarchy._meshes
 
     for j in range(repeat):
@@ -128,15 +126,15 @@ def restrict(fine_dual, coarse_dual):
             coarse_dual.dat.zero()
             next = coarse_dual
         else:
-            Vc = firedrake.FunctionSpace(meshes[next_level], element)
-            next = firedrake.Cofunction(Vc.dual())
+            Vc = Vf.reconstruct(mesh=meshes[next_level])
+            next = firedrake.Cofunction(Vc)
         Vc = next.function_space()
         # XXX: Should be able to figure out locations by pushing forward
         # reference cell node locations to physical space.
         # x = \sum_i c_i \phi_i(x_hat)
-        node_locations = utils.physical_node_locations(Vf)
+        node_locations = utils.physical_node_locations(Vf.dual())
 
-        coarse_coords = Vc.mesh().coordinates
+        coarse_coords = get_coordinates(Vc.dual())
         fine_to_coarse = utils.fine_node_to_coarse_node_map(Vf, Vc)
         fine_to_coarse_coords = utils.fine_node_to_coarse_node_map(Vf, coarse_coords.function_space())
         # Have to do this, because the node set core size is not right for
@@ -195,7 +193,6 @@ def inject(fine, coarse):
     repeat = (fine_level - coarse_level)*refinements_per_level
     next_level = fine_level * refinements_per_level
 
-    element = Vc.ufl_element()
     meshes = hierarchy._meshes
 
     for j in range(repeat):
@@ -205,12 +202,12 @@ def inject(fine, coarse):
             next = coarse
             Vc = next.function_space()
         else:
-            Vc = firedrake.FunctionSpace(meshes[next_level], element)
+            Vc = Vf.reconstruct(mesh=meshes[next_level])
             next = firedrake.Function(Vc)
         if not dg:
             node_locations = utils.physical_node_locations(Vc)
 
-            fine_coords = Vf.mesh().coordinates
+            fine_coords = get_coordinates(Vf)
             coarse_node_to_fine_nodes = utils.coarse_node_to_fine_node_map(Vc, Vf)
             coarse_node_to_fine_coords = utils.coarse_node_to_fine_node_map(Vc, fine_coords.function_space())
 
@@ -242,3 +239,11 @@ def inject(fine, coarse):
         fine = next
         Vf = Vc
     return coarse
+
+
+def get_coordinates(V):
+    coords = V.mesh().coordinates
+    if V.boundary_set:
+        W = V.reconstruct(element=coords.function_space().ufl_element())
+        coords = firedrake.Function(W).interpolate(coords)
+    return coords

firedrake/mg/ufl_utils.py

@@ -195,7 +195,7 @@ def inject_on_restrict(fine, restriction, rscale, injection, coarse):
             for bc in cctx._problem.dirichlet_bcs():
                 bc.apply(cctx._x)
 
-    dm = problem.u.function_space().dm
+    dm = problem.u_restrict.function_space().dm
     if not dm.getAttr("_coarsen_hook"):
         # The hook is persistent and cumulative, but also problem-independent.
         # Therefore, we are only adding it once.
@@ -209,7 +209,7 @@ def inject_on_restrict(fine, restriction, rscale, injection, coarse):
     F = self(problem.F, self, coefficient_mapping=coefficient_mapping)
     J = self(problem.J, self, coefficient_mapping=coefficient_mapping)
     Jp = self(problem.Jp, self, coefficient_mapping=coefficient_mapping)
-    u = coefficient_mapping[problem.u]
+    u = coefficient_mapping[problem.u_restrict]
 
     fine = problem
     problem = firedrake.NonlinearVariationalProblem(F, u, bcs=bcs, J=J, Jp=Jp, is_linear=problem.is_linear,
@@ -283,7 +283,7 @@ def coarsen_snescontext(context, self, coefficient_mapping=None):
         if isinstance(val, (firedrake.Function, firedrake.Cofunction)):
             V = val.function_space()
             coarseneddm = V.dm
-            parentdm = get_parent(context._problem.u.function_space().dm)
+            parentdm = get_parent(context._problem.u_restrict.function_space().dm)
 
             # Now attach the hook to the parent DM
             if get_appctx(coarseneddm) is None:
@@ -303,11 +303,11 @@ def coarsen_snescontext(context, self, coefficient_mapping=None):
 
 
 class Interpolation(object):
-    def __init__(self, coarse, fine, manager, cbcs=None, fbcs=None):
-        self.cprimal = coarse
-        self.fprimal = fine
-        self.cdual = coarse.riesz_representation(riesz_map="l2")
-        self.fdual = fine.riesz_representation(riesz_map="l2")
+    def __init__(self, Vcoarse, Vfine, manager, cbcs=None, fbcs=None):
+        self.cprimal = firedrake.Function(Vcoarse)
+        self.fprimal = firedrake.Function(Vfine)
+        self.cdual = firedrake.Cofunction(Vcoarse.dual())
+        self.fdual = firedrake.Cofunction(Vfine.dual())
         self.cbcs = cbcs or []
         self.fbcs = fbcs or []
         self.manager = manager
@@ -352,9 +352,9 @@ def multTransposeAdd(self, mat, x, y, w):
 
 
 class Injection(object):
-    def __init__(self, cfn, ffn, manager, cbcs=None):
-        self.cfn = cfn
-        self.ffn = ffn
+    def __init__(self, Vcoarse, Vfine, manager, cbcs=None):
+        self.cfn = firedrake.Function(Vcoarse)
+        self.ffn = firedrake.Function(Vfine)
         self.cbcs = cbcs or []
         self.manager = manager
 
@@ -374,18 +374,15 @@ def create_interpolation(dmc, dmf):
 
     manager = get_transfer_manager(dmf)
 
-    V_c = cctx._problem.u.function_space()
-    V_f = fctx._problem.u.function_space()
+    V_c = cctx._problem.u_restrict.function_space()
+    V_f = fctx._problem.u_restrict.function_space()
 
     row_size = V_f.dof_dset.layout_vec.getSizes()
     col_size = V_c.dof_dset.layout_vec.getSizes()
-
-    cfn = firedrake.Function(V_c)
-    ffn = firedrake.Function(V_f)
     cbcs = tuple(cctx._problem.dirichlet_bcs())
     fbcs = tuple(fctx._problem.dirichlet_bcs())
 
-    ctx = Interpolation(cfn, ffn, manager, cbcs, fbcs)
+    ctx = Interpolation(V_c, V_f, manager, cbcs, fbcs)
     mat = PETSc.Mat().create(comm=dmc.comm)
     mat.setSizes((row_size, col_size))
     mat.setType(mat.Type.PYTHON)
@@ -400,16 +397,13 @@ def create_injection(dmc, dmf):
 
     manager = get_transfer_manager(dmf)
 
-    V_c = cctx._problem.u.function_space()
-    V_f = fctx._problem.u.function_space()
+    V_c = cctx._problem.u_restrict.function_space()
+    V_f = fctx._problem.u_restrict.function_space()
 
     row_size = V_c.dof_dset.layout_vec.getSizes()
     col_size = V_f.dof_dset.layout_vec.getSizes()
 
-    cfn = firedrake.Function(V_c)
-    ffn = firedrake.Function(V_f)
-
-    ctx = Injection(cfn, ffn, manager)
+    ctx = Injection(V_c, V_f, manager)
     mat = PETSc.Mat().create(comm=dmc.comm)
     mat.setSizes((row_size, col_size))
     mat.setType(mat.Type.PYTHON)