RestrictedFunctionSpace: support p-multigrid

Author: pbrubeck

firedrake/mg/embedded.py

@@ -90,14 +90,17 @@ def cache(self, V):
         except KeyError:
             return self.caches.setdefault(key, TransferManager.Cache(*key))
 
+    def cache_key(self, V):
+        return (V.dim(), tuple(V.boundary_set))
+
     def V_dof_weights(self, V):
         """Dof weights for averaging projection.
 
         :arg V: function space to compute weights for.
         :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:
@@ -122,7 +125,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:
@@ -153,7 +156,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:
@@ -171,7 +174,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:
@@ -193,7 +196,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:
@@ -210,7 +213,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

@@ -127,7 +127,7 @@ def restrict(fine_dual, coarse_dual):
             next = coarse_dual
         else:
             Vc = Vf.reconstruct(mesh=meshes[next_level])
-            next = firedrake.Cofunction(Vc.dual())
+            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.

firedrake/mg/utils.py

@@ -25,10 +25,7 @@ def fine_node_to_coarse_node_map(Vf, Vc):
     if levelc + increment != levelf:
         raise ValueError("Can't map between level %s and level %s" % (levelc, levelf))
 
-    key = (entity_dofs_key(Vc.finat_element.entity_dofs())
-           + entity_dofs_key(Vf.finat_element.entity_dofs())
-           + (levelc, levelf))
-
+    key = _cache_key(Vc, Vf)
     cache = mesh._shared_data_cache["hierarchy_fine_node_to_coarse_node_map"]
     try:
         return cache[key]
@@ -63,10 +60,7 @@ def coarse_node_to_fine_node_map(Vc, Vf):
     if levelc + increment != levelf:
         raise ValueError("Can't map between level %s and level %s" % (levelc, levelf))
 
-    key = (entity_dofs_key(Vc.finat_element.entity_dofs())
-           + entity_dofs_key(Vf.finat_element.entity_dofs())
-           + (levelc, levelf))
-
+    key = _cache_key(Vc, Vf)
     cache = mesh._shared_data_cache["hierarchy_coarse_node_to_fine_node_map"]
     try:
         return cache[key]
@@ -101,7 +95,7 @@ def coarse_cell_to_fine_node_map(Vc, Vf):
     if levelc + increment != levelf:
         raise ValueError("Can't map between level %s and level %s" % (levelc, levelf))
 
-    key = (entity_dofs_key(Vf.finat_element.entity_dofs()) + (levelc, levelf))
+    key = _cache_key(Vc, Vf)
     cache = mesh._shared_data_cache["hierarchy_coarse_cell_to_fine_node_map"]
     try:
         return cache[key]
@@ -142,7 +136,7 @@ def physical_node_locations(V):
     # This is a defaultdict, so the first time we access the key we
     # get a fresh dict for the cache.
     cache = mesh._geometric_shared_data_cache["hierarchy_physical_node_locations"]
-    key = element
+    key = (element, tuple(V.boundary_set))
     try:
         return cache[key]
     except KeyError:
@@ -172,3 +166,10 @@ def get_level(obj):
 def has_level(obj):
     """Does the provided object have level info?"""
     return hasattr(obj.topological, "__level_info__")
+
+
+def _cache_key(Vc, Vf):
+    return (entity_dofs_key(Vc.finat_element.entity_dofs())
+            + entity_dofs_key(Vf.finat_element.entity_dofs())
+            + (Vc.dim(), Vf.dim())
+            + (tuple(Vc.boundary_set), tuple(Vf.boundary_set)))

firedrake/preconditioners/pmg.py

@@ -181,21 +181,21 @@ def coarsen(self, fdm, comm):
         assert parent is not None
 
         test, trial = fctx.J.arguments()
-        fV = test.function_space()
+        fV = trial.function_space()
         cele = self.coarsen_element(fV.ufl_element())
 
         # Have we already done this?
         cctx = fctx._coarse
         if cctx is not None:
-            cV = cctx.J.arguments()[0].function_space()
-            if (cV.ufl_element() == cele) and (cV.mesh() == fV.mesh()):
+            cV = cctx.J.arguments()[1].function_space()
+            if (cV.ufl_element() == cele) and (cV.mesh() == fV.mesh()) and all(cV_.boundary_set == fV_.boundary_set for cV_, fV_ in zip(cV, fV)):
                 return cV.dm
 
-        cV = firedrake.FunctionSpace(fV.mesh(), cele)
+        cV = fV.reconstruct(element=cele)
         cdm = cV.dm
 
         fproblem = fctx._problem
-        fu = fproblem.u
+        fu = fproblem.u_restrict
         cu = firedrake.Function(cV)
 
         fdeg = PMGBase.max_degree(fV.ufl_element())
@@ -370,8 +370,8 @@ def create_transfer(self, mat_type, cctx, fctx, cbcs, fbcs):
                 construct_mat = prolongation_matrix_aij
             else:
                 raise ValueError("Unknown matrix type")
-            cV = cctx.J.arguments()[0].function_space()
-            fV = fctx.J.arguments()[0].function_space()
+            cV = cctx._problem.u_restrict.function_space()
+            fV = fctx._problem.u_restrict.function_space()
             cbcs = tuple(cctx._problem.bcs) if cbcs else tuple()
             fbcs = tuple(fctx._problem.bcs) if fbcs else tuple()
             return cache.setdefault(key, construct_mat(cV, fV, cbcs, fbcs))
@@ -1179,7 +1179,7 @@ def make_permutation_code(V, vshape, pshape, t_in, t_out, array_name):
 
 def reference_value_space(V):
     element = finat.ufl.WithMapping(V.ufl_element(), mapping="identity")
-    return firedrake.FunctionSpace(V.mesh(), element)
+    return V.collapse().reconstruct(element=element)
 
 
 class StandaloneInterpolationMatrix(object):
@@ -1206,13 +1206,13 @@ def __init__(self, Vc, Vf, Vc_bcs, Vf_bcs):
             self.Vf = reference_value_space(self.Vf)
             self.uc = firedrake.Function(self.Vc, val=self.uc.dat)
             self.uf = firedrake.Function(self.Vf, val=self.uf.dat)
-            self.Vc_bcs = [bc.reconstruct(V=self.Vc) for bc in self.Vc_bcs]
-            self.Vf_bcs = [bc.reconstruct(V=self.Vf) for bc in self.Vf_bcs]
+            self.Vc_bcs = [bc.reconstruct(V=self.Vc, g=0) for bc in self.Vc_bcs]
+            self.Vf_bcs = [bc.reconstruct(V=self.Vf, g=0) for bc in self.Vf_bcs]
 
     def work_function(self, V):
         if isinstance(V, firedrake.Function):
             return V
-        key = (V.ufl_element(), V.mesh())
+        key = (V.ufl_element(), V.mesh(), tuple(V.boundary_set))
         try:
             return self._cache_work[key]
         except KeyError:
@@ -1337,17 +1337,14 @@ def make_blas_kernels(self, Vf, Vc):
             restrict = [""]*5
             # get embedding element for Vf with identity mapping and collocated vector component DOFs
             try:
-                qelem = felem
-                if qelem.mapping() != "identity":
-                    qelem = qelem.reconstruct(mapping="identity")
-                Qf = Vf if qelem == felem else firedrake.FunctionSpace(Vf.mesh(), qelem)
+                Qf = Vf if felem.mapping() == "identity" else Vf.reconstruct(mapping="identity")
                 mapping_output = make_mapping_code(Qf, cmapping, fmapping, "t0", "t1")
                 in_place_mapping = True
             except Exception:
                 qelem = finat.ufl.FiniteElement("DQ", cell=felem.cell, degree=PMGBase.max_degree(felem))
                 if Vf.value_shape:
                     qelem = finat.ufl.TensorElement(qelem, shape=Vf.value_shape, symmetry=felem.symmetry())
-                Qf = firedrake.FunctionSpace(Vf.mesh(), qelem)
+                Qf = Vf.reconstruct(element=qelem)
                 mapping_output = make_mapping_code(Qf, cmapping, fmapping, "t0", "t1")
 
             qshape = (Qf.block_size, Qf.finat_element.space_dimension())

tests/firedrake/multigrid/test_p_multigrid.py

@@ -137,7 +137,8 @@ def mat_type(request):
     return request.param
 
 
-def test_p_multigrid_scalar(mesh, mat_type):
+@pytest.mark.parametrize("restrict", [False, True], ids=("unrestrict", "restrict"))
+def test_p_multigrid_scalar(mesh, mat_type, restrict):
     V = FunctionSpace(mesh, "CG", 4)
 
     u = Function(V)
@@ -175,7 +176,7 @@ def test_p_multigrid_scalar(mesh, mat_type):
           "pmg_mg_coarse_mg_coarse_ksp_monitor": None,
           "pmg_mg_coarse_mg_coarse_pc_type": "gamg",
           "pmg_mg_coarse_mg_coarse_pc_gamg_threshold": 0}
-    problem = NonlinearVariationalProblem(F, u, bcs)
+    problem = NonlinearVariationalProblem(F, u, bcs, restrict=restrict)
     solver = NonlinearVariationalSolver(problem, solver_parameters=sp)
     solver.solve()