RestrictedFunctionSpace: support geometric multigrid

Author: pbrubeck

firedrake/dmhooks.py

@@ -459,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

@@ -357,7 +357,7 @@ def collapse(self):
         return type(self).create(self.topological.collapse(), self.mesh())
 
     @classmethod
-    def make_function_space(cls, mesh, element, name=None):
+    def make_function_space(cls, mesh, element, name=None, boundary_set=None):
         r"""Factory method for :class:`.WithGeometryBase`."""
         mesh.init()
         topology = mesh.topology
@@ -376,12 +376,18 @@ def make_function_space(cls, mesh, element, name=None):
         if mesh is not topology:
             # Create a concrete WithGeometry or FiredrakeDualSpace on this mesh
             new = cls.create(new, mesh)
+
+        if boundary_set:
+            new = RestrictedFunctionSpace(new, boundary_set=boundary_set)
+            if mesh is not topology:
+                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``.
 
@@ -404,12 +410,14 @@ def reconstruct(self, mesh=None, name=None, **kwargs):
         if mesh is None:
             mesh = V_parent.mesh()
 
-        element = V_parent.ufl_element()
+        if element is None:
+            element = V_parent.ufl_element()
         cell = mesh.topology.ufl_cell()
         if len(kwargs) > 0 or element.cell != cell:
             element = element.reconstruct(cell=cell, **kwargs)
 
-        V = type(self).make_function_space(mesh, element, name=name)
+        V = type(self).make_function_space(mesh, element, name=name,
+                                           boundary_set=V_parent.boundary_set)
         for i in reversed(indices):
             V = V.sub(i)
         return V

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)
+            Vc = Vf.reconstruct(mesh=meshes[next_level])
             next = firedrake.Cofunction(Vc.dual())
         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

@@ -186,7 +186,7 @@ def inject_on_restrict(fine, restriction, rscale, injection, coarse):
                 if isinstance(g, firedrake.Function) and hasattr(g, "_child"):
                     manager.inject(g, g._child)
 
-    V = problem.u.function_space()
+    V = problem.u_restrict.function_space()
     if not hasattr(V, "_coarse"):
         # The hook is persistent and cumulative, but also problem-independent.
         # Therefore, we are only adding it once.
@@ -201,7 +201,7 @@ def inject_on_restrict(fine, restriction, rscale, injection, coarse):
     for c in coefficients:
         coefficient_mapping[c] = self(c, self, coefficient_mapping=coefficient_mapping)
 
-    u = coefficient_mapping[problem.u]
+    u = coefficient_mapping[problem.u_restrict]
 
     bcs = [self(bc, self) for bc in problem.bcs]
     J = self(problem.J, self, coefficient_mapping=coefficient_mapping)
@@ -277,7 +277,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:
@@ -369,8 +369,8 @@ 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()
@@ -395,8 +395,8 @@ 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_f.dof_dset.layout_vec.getSizes()
     col_size = V_c.dof_dset.layout_vec.getSizes()

firedrake/mg/utils.py

@@ -11,7 +11,7 @@
 def fine_node_to_coarse_node_map(Vf, Vc):
     if len(Vf) > 1:
         assert len(Vf) == len(Vc)
-        return op2.MixedMap(fine_node_to_coarse_node_map(f, c) for f, c in zip(Vf, Vc))
+        return op2.MixedMap(map(fine_node_to_coarse_node_map, Vf, Vc))
     mesh = Vf.mesh()
     assert hasattr(mesh, "_shared_data_cache")
     hierarchyf, levelf = get_level(Vf.mesh())
@@ -49,7 +49,7 @@ def fine_node_to_coarse_node_map(Vf, Vc):
 def coarse_node_to_fine_node_map(Vc, Vf):
     if len(Vf) > 1:
         assert len(Vf) == len(Vc)
-        return op2.MixedMap(coarse_node_to_fine_node_map(f, c) for f, c in zip(Vf, Vc))
+        return op2.MixedMap(map(coarse_node_to_fine_node_map, Vf, Vc))
     mesh = Vc.mesh()
     assert hasattr(mesh, "_shared_data_cache")
     hierarchyf, levelf = get_level(Vf.mesh())
@@ -146,7 +146,8 @@ def physical_node_locations(V):
     try:
         return cache[key]
     except KeyError:
-        Vc = firedrake.VectorFunctionSpace(mesh, element)
+        Vc = V.reconstruct(element=finat.ufl.VectorElement(element, dim=mesh.geometric_dimension()))
+
         # FIXME: This is unsafe for DG coordinates and CG target spaces.
         locations = firedrake.assemble(firedrake.Interpolate(firedrake.SpatialCoordinate(mesh), Vc))
         return cache.setdefault(key, locations)

pyop2/types/dat.py

@@ -105,7 +105,6 @@ def _wrapper_cache_key_(self):
     @utils.validate_in(('access', _modes, ex.ModeValueError))
     def __call__(self, access, path=None):
         from pyop2.parloop import DatLegacyArg
-
         if conf.configuration["type_check"] and path and path.toset != self.dataset.set:
             raise ex.MapValueError("To Set of Map does not match Set of Dat.")
         return DatLegacyArg(self, path, access)

tests/firedrake/regression/test_restricted_function_space.py

@@ -401,11 +401,11 @@ def test_restrict_fieldsplit(names):
 
 def test_restrict_python_pc():
     mesh = UnitSquareMesh(2, 2)
-    x, y = SpatialCoordinate(mesh)
     V = FunctionSpace(mesh, "CG", 1)
-
     u = Function(V)
     test = TestFunction(V)
+
+    x, y = SpatialCoordinate(mesh)
     u_exact = x + y
     g = Function(V).interpolate(u_exact)
 
@@ -419,8 +419,40 @@ def test_restrict_python_pc():
         "ksp_type": "preonly",
         "pc_type": "python",
         "pc_python_type": "firedrake.AssembledPC",
-        "assembled_pc_type": "lu"},
-        options_prefix="")
+        "assembled_pc_type": "lu"})
+    solver.solve()
+
+    assert errornorm(u_exact, u) < 1E-10
+
+
+def test_restrict_multigrid():
+    base = UnitSquareMesh(2, 2)
+    refine = 2
+    mh = MeshHierarchy(base, refine)
+    mesh = mh[-1]
+
+    V = FunctionSpace(mesh, "CG", 1)
+    u = Function(V)
+    test = TestFunction(V)
+
+    x, y = SpatialCoordinate(mesh)
+    u_exact = x + y
+    g = Function(V).interpolate(u_exact)
+
+    F = inner(grad(u - u_exact), grad(test)) * dx
+    bcs = [DirichletBC(V, g, 1), DirichletBC(V, u_exact, 2)]
+
+    problem = NonlinearVariationalProblem(F, u, bcs=bcs, restrict=True)
+    solver = NonlinearVariationalSolver(problem, solver_parameters={
+        "snes_type": "ksponly",
+        "ksp_type": "cg",
+        "ksp_rtol": 1E-10,
+        "ksp_max_it": 10,
+        "ksp_monitor": None,
+        "pc_type": "mg",
+        "mg_levels_ksp_type": "chebyshev",
+        "mg_levels_pc_type": "jacobi",
+        "mg_coarse_pc_type": "lu"})
     solver.solve()
 
     assert errornorm(u_exact, u) < 1E-10