Merge pull request #1328 from firedrakeproject/lm/bary-multigrid

* lm/bary-multigrid:
  mg: Add test of "non-nested" hierarchies
  mg: Fix level indexing for skipping hierarchies
  mg: Save nestedness of hierarchies, error for DG injection if not nested
  Fix a stack smash in the restriction code
  And the same for prolongation, just in case
  mg: Fix transfer kernels for non-nested
  mg: Make fine_to_coarse_cells a multimap

firedrake/mg/impl.pyx

@@ -105,7 +105,7 @@ def coarse_to_fine_nodes(Vc, Vf, np.ndarray[PetscInt, ndim=2, mode="c"] coarse_t
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def fine_to_coarse_nodes(Vf, Vc, np.ndarray[PetscInt, ndim=1, mode="c"] fine_to_coarse_cells):
+def fine_to_coarse_nodes(Vf, Vc, np.ndarray[PetscInt, ndim=2, mode="c"] fine_to_coarse_cells):
     cdef:
         np.ndarray[PetscInt, ndim=2, mode="c"] fine_map, coarse_map, fine_to_coarse_map
         np.ndarray[PetscInt, ndim=1, mode="c"] coarse_offset, fine_offset
@@ -122,25 +122,28 @@ def fine_to_coarse_nodes(Vf, Vc, np.ndarray[PetscInt, ndim=1, mode="c"] fine_to_
         coarse_offset = Vc.offset
         fine_offset = Vf.offset
         layers = Vc.mesh().layers - 1
-    fine_cells = fine_map.shape[0]
+
+    fine_cells = fine_to_coarse_cells.shape[0]
+    coarse_per_fine = fine_to_coarse_cells.shape[1]
     coarse_per_cell = coarse_map.shape[1]
     fine_per_cell = fine_map.shape[1]
     fine_to_coarse_map = np.full((Vf.dof_dset.total_size,
-                                  coarse_per_cell),
+                                  coarse_per_fine*coarse_per_cell),
                                  -1,
                                  dtype=IntType)
+
     for i in range(fine_cells):
-        coarse_cell = fine_to_coarse_cells[i]
-        for j in range(fine_per_cell):
-            node = fine_map[i, j]
-            if extruded:
-                for layer in range(layers):
+        for l, coarse_cell in enumerate(fine_to_coarse_cells[i, :]):
+            for j in range(fine_per_cell):
+                node = fine_map[i, j]
+                if extruded:
+                    for layer in range(layers):
+                        for k in range(coarse_per_cell):
+                            fine_to_coarse_map[node + fine_offset[j]*layer, k] = (coarse_map[coarse_cell, k] +
+                                                                                  coarse_offset[k]*layer)
+                else:
                     for k in range(coarse_per_cell):
-                        fine_to_coarse_map[node + fine_offset[j]*layer, k] = (coarse_map[coarse_cell, k] +
-                                                                              coarse_offset[k]*layer)
-            else:
-                for k in range(coarse_per_cell):
-                    fine_to_coarse_map[node, k] = coarse_map[coarse_cell, k]
+                        fine_to_coarse_map[node, coarse_per_cell*l + k] = coarse_map[coarse_cell, k]
 
     return fine_to_coarse_map
 
@@ -230,7 +233,7 @@ def coarse_to_fine_cells(mc, mf, clgmaps, flgmaps):
         PetscInt cStart, cEnd, c, val, dim, nref, ncoarse
         PetscInt i, ccell, fcell, nfine
         np.ndarray[PetscInt, ndim=2, mode="c"] coarse_to_fine
-        np.ndarray[PetscInt, ndim=1, mode="c"] fine_to_coarse
+        np.ndarray[PetscInt, ndim=2, mode="c"] fine_to_coarse
         np.ndarray[PetscInt, ndim=1, mode="c"] co2n, fn2o, idx
 
     cdm = mc._plex
@@ -242,7 +245,7 @@ def coarse_to_fine_cells(mc, mf, clgmaps, flgmaps):
     co2n, _ = get_entity_renumbering(cdm, mc._cell_numbering, "cell")
     _, fn2o = get_entity_renumbering(fdm, mf._cell_numbering, "cell")
     coarse_to_fine = np.full((ncoarse, nref), -1, dtype=PETSc.IntType)
-    fine_to_coarse = np.full(nfine, -1, dtype=PETSc.IntType)
+    fine_to_coarse = np.full((nfine, 1), -1, dtype=PETSc.IntType)
     # Walk owned fine cells:
     cStart, cEnd = 0, nfine
 
@@ -276,7 +279,7 @@ def coarse_to_fine_cells(mc, mf, clgmaps, flgmaps):
         # cells should map into owned coarse cells)
         ccell = co2n[fcell // nref]
         assert 0 <= ccell < ncoarse
-        fine_to_coarse[c] = ccell
+        fine_to_coarse[c, 0] = ccell
         for i in range(nref):
             if coarse_to_fine[ccell, i] == -1:
                 coarse_to_fine[ccell, i] = c

firedrake/mg/interface.py

@@ -178,6 +178,8 @@ def inject(fine, coarse):
 
     kernel, dg = kernels.inject_kernel(Vf, Vc)
     hierarchy, coarse_level = utils.get_level(coarse.ufl_domain())
+    if dg and not hierarchy.nested:
+        raise NotImplementedError("Sorry, we can't do supermesh projections yet!")
     _, fine_level = utils.get_level(fine.ufl_domain())
     refinements_per_level = hierarchy.refinements_per_level
     repeat = (fine_level - coarse_level)*refinements_per_level

firedrake/mg/kernels.py

@@ -1,5 +1,6 @@
 import numpy
 import string
+from fractions import Fraction
 from pyop2 import op2
 from pyop2.datatypes import IntType, as_cstr
 from firedrake.functionspacedata import entity_dofs_key
@@ -197,7 +198,8 @@ def predicate(index):
 
 
 def prolong_kernel(expression):
-    hierarchy, _ = utils.get_level(expression.ufl_domain())
+    hierarchy, level = utils.get_level(expression.ufl_domain())
+    levelf = level + Fraction(1 / hierarchy.refinements_per_level)
     cache = hierarchy._shared_data_cache["transfer_kernels"]
     coordinates = expression.ufl_domain().coordinates
     key = (("prolong", )
@@ -212,34 +214,50 @@ def prolong_kernel(expression):
         to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
         element = create_element(expression.ufl_element())
         eval_args = evaluate_kernel.args[:-1]
+        coords_element = create_element(coordinates.ufl_element())
 
-        args = ", ".join(a.gencode(not_scope=True) for a in eval_args)
-        arg_names = ", ".join(a.sym.symbol for a in eval_args)
+        args = eval_args[-1].gencode(not_scope=True)
+        R, coarse = (a.sym.symbol for a in eval_args)
         my_kernel = """
         %(to_reference)s
         %(evaluate)s
-        void prolong_kernel(%(args)s, const double *X, const double *Xc)
+        void prolong_kernel(double *R, %(args)s, const double *X, const double *Xc)
         {
             double Xref[%(tdim)d];
-            to_reference_coords_kernel(Xref, X, Xc);
+            int cell = -1;
+            for (int i = 0; i < %(ncandidate)d; i++) {
+                const double *Xci = Xc + i*%(Xc_cell_inc)d;
+                to_reference_coords_kernel(Xref, X, Xci);
+                if (%(inside_cell)s) {
+                    cell = i;
+                    break;
+                }
+            }
+            if (cell == -1) abort();
+            const double *coarsei = %(coarse)s + cell*%(coarse_cell_inc)d;
             for ( int i = 0; i < %(Rdim)d; i++ ) {
                 %(R)s[i] = 0;
             }
-            evaluate_kernel(%(arg_names)s, Xref);
+            evaluate_kernel(%(R)s, coarsei, Xref);
         }
         """ % {"to_reference": str(to_reference_kernel),
                "evaluate": str(evaluate_kernel),
                "args": args,
-               "R": arg_names[0],
+               "R": R,
+               "coarse": coarse,
+               "ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1],
                "Rdim": numpy.prod(element.value_shape),
-               "arg_names": arg_names,
+               "inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"),
+               "Xc_cell_inc": coords_element.space_dimension(),
+               "coarse_cell_inc": element.space_dimension(),
                "tdim": mesh.topological_dimension()}
 
         return cache.setdefault(key, op2.Kernel(my_kernel, name="prolong_kernel"))
 
 
 def restrict_kernel(Vf, Vc):
-    hierarchy, _ = utils.get_level(Vc.ufl_domain())
+    hierarchy, level = utils.get_level(Vc.ufl_domain())
+    levelf = level + Fraction(1 / hierarchy.refinements_per_level)
     cache = hierarchy._shared_data_cache["transfer_kernels"]
     coordinates = Vc.ufl_domain().coordinates
     key = (("restrict", )
@@ -253,24 +271,38 @@ def restrict_kernel(Vf, Vc):
         mesh = coordinates.ufl_domain()
         evaluate_kernel = compile_element(firedrake.TestFunction(Vc), Vf)
         to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
-
+        coords_element = create_element(coordinates.ufl_element())
+        element = create_element(Vc.ufl_element())
         eval_args = evaluate_kernel.args[:-1]
-
-        args = ", ".join(a.gencode(not_scope=True) for a in eval_args)
-        arg_names = ", ".join(a.sym.symbol for a in eval_args)
+        args = eval_args[-1].gencode(not_scope=True)
+        R, fine = (a.sym.symbol for a in eval_args)
         my_kernel = """
         %(to_reference)s
         %(evaluate)s
-        void restrict_kernel(%(args)s, const double *X, const double *Xc)
+        void restrict_kernel(double *R, %(args)s, const double *X, const double *Xc)
         {
             double Xref[%(tdim)d];
-            to_reference_coords_kernel(Xref, X, Xc);
-            evaluate_kernel(%(arg_names)s, Xref);
+            int cell = -1;
+            for (int i = 0; i < %(ncandidate)d; i++) {
+                const double *Xci = Xc + i*%(Xc_cell_inc)d;
+                to_reference_coords_kernel(Xref, X, Xci);
+                if (%(inside_cell)s) {
+                    cell = i;
+                    const double *Ri = %(R)s + cell*%(coarse_cell_inc)d;
+                    evaluate_kernel(Ri, %(fine)s, Xref);
+                    break;
+                }
+            }
         }
         """ % {"to_reference": str(to_reference_kernel),
                "evaluate": str(evaluate_kernel),
+               "ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1],
+               "inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"),
+               "Xc_cell_inc": coords_element.space_dimension(),
+               "coarse_cell_inc": element.space_dimension(),
                "args": args,
-               "arg_names": arg_names,
+               "R": R,
+               "fine": fine,
                "tdim": mesh.topological_dimension()}
 
         return cache.setdefault(key, op2.Kernel(my_kernel, name="restrict_kernel"))

firedrake/mg/mesh.py

@@ -21,6 +21,7 @@ class HierarchyBase(object):
        pair, mapping each fine cell into coarse cells it intersects.
     :arg refinements_per_level: number of mesh refinements each
        multigrid level should "see".
+    :arg nested: Is this mesh hierarchy nested?
 
     .. note::
 
@@ -29,14 +30,15 @@ class HierarchyBase(object):
        :func:`ExtrudedMeshHierarchy`, or :func:`NonNestedHierarchy`.
     """
     def __init__(self, meshes, coarse_to_fine_cells, fine_to_coarse_cells,
-                 refinements_per_level=1):
+                 refinements_per_level=1, nested=False):
         from firedrake_citations import Citations
         Citations().register("Mitchell2016")
         self._meshes = tuple(meshes)
         self.meshes = tuple(meshes[::refinements_per_level])
         self.coarse_to_fine_cells = coarse_to_fine_cells
         self.fine_to_coarse_cells = fine_to_coarse_cells
         self.refinements_per_level = refinements_per_level
+        self.nested = nested
         for level, m in enumerate(meshes):
             set_level(m, self, Fraction(level, refinements_per_level))
         for level, m in enumerate(self):
@@ -166,7 +168,7 @@ def MeshHierarchy(mesh, refinement_levels,
     fine_to_coarse_cells = dict((Fraction(i, refinements_per_level), f2c)
                                 for i, f2c in enumerate(fine_to_coarse_cells))
     return HierarchyBase(meshes, coarse_to_fine_cells, fine_to_coarse_cells,
-                         refinements_per_level)
+                         refinements_per_level, nested=True)
 
 
 def ExtrudedMeshHierarchy(base_hierarchy, layers, kernel=None, layer_height=None,
@@ -191,8 +193,10 @@ def ExtrudedMeshHierarchy(base_hierarchy, layers, kernel=None, layer_height=None
     return HierarchyBase(meshes,
                          base_hierarchy.coarse_to_fine_cells,
                          base_hierarchy.fine_to_coarse_cells,
-                         refinements_per_level=base_hierarchy.refinements_per_level)
+                         refinements_per_level=base_hierarchy.refinements_per_level,
+                         nested=base_hierarchy.nested)
 
 
 def NonNestedHierarchy(*meshes):
-    return HierarchyBase(meshes, [None for _ in meshes], [None for _ in meshes])
+    return HierarchyBase(meshes, [None for _ in meshes], [None for _ in meshes],
+                         nested=False)

tests/multigrid/test_grid_transfer.py

@@ -3,7 +3,9 @@
 from firedrake import *
 
 
-@pytest.fixture(params=["interval", "triangle", "quadrilateral", "tetrahedron",
+@pytest.fixture(params=["interval", "triangle",
+                        "triangle-nonnested",  # no parameterized fixtures, UGH!
+                        "quadrilateral", "tetrahedron",
                         "prism", "hexahedron"], scope="module")
 def cell(request):
     return request.param
@@ -27,7 +29,7 @@ def hierarchy(cell, refinements_per_level):
     if cell == "interval":
         mesh = UnitIntervalMesh(3)
         return MeshHierarchy(mesh, 2)
-    elif cell in {"triangle", "prism"}:
+    elif cell in {"triangle", "triangle-nonnested", "prism"}:
         mesh = UnitSquareMesh(3, 3, quadrilateral=False)
     elif cell in {"quadrilateral", "hexahedron"}:
         mesh = UnitSquareMesh(3, 3, quadrilateral=True)
@@ -39,7 +41,20 @@ def hierarchy(cell, refinements_per_level):
 
     if cell in {"prism", "hexahedron"}:
         hierarchy = ExtrudedMeshHierarchy(hierarchy, layers=3)
-
+    if cell == "triangle-nonnested":
+        c2f = {}
+        for k, v in hierarchy.coarse_to_fine_cells.items():
+            v = numpy.hstack([v, numpy.roll(v, 1, axis=0)])
+            c2f[k] = v
+        f2c = {}
+
+        for k, v in hierarchy.fine_to_coarse_cells.items():
+            if v is not None:
+                v = numpy.hstack([v, numpy.roll(v, 4, axis=0)])
+            f2c[k] = v
+        hierarchy = HierarchyBase(tuple(hierarchy), c2f, f2c,
+                                  refinements_per_level=refinements_per_level,
+                                  nested=False)
     return hierarchy
 
 
@@ -50,8 +65,11 @@ def vector(request):
 
 
 @pytest.fixture(params=["injection", "restriction", "prolongation"])
-def transfer_type(request):
-    return request.param
+def transfer_type(request, hierarchy):
+    if not hierarchy.nested and request.param == "injection":
+        return pytest.mark.xfail(reason="Supermesh projections not implemented yet")(request.param)
+    else:
+        return request.param
 
 
 @pytest.fixture
@@ -114,29 +132,38 @@ def run_prolongation(hierarchy, vector, space, degrees):
 
 
 def run_restriction(hierarchy, vector, space, degrees):
-    def exact_dual(V):
-        if V.shape:
-            c = Constant([1] * V.value_size)
-        else:
-            c = Constant(1)
-        return assemble(inner(c, TestFunction(V))*dx)
+    def victim(V):
+        return Function(V).assign(1)
+
+    def dual(V):
+        f = Function(V).assign(1)
+        return assemble(inner(f, TestFunction(V))*dx)
+
+    def functional(victim, dual):
+        with victim.dat.vec_ro as v, dual.dat.vec_ro as dv:
+            return dv.dot(v)
 
     for degree in degrees:
         Ve = element(space, hierarchy[0].ufl_cell(), degree, vector)
-        mesh = hierarchy[-1]
-        V = FunctionSpace(mesh, Ve)
+        for cmesh, fmesh in zip(hierarchy[:-1], hierarchy[1:]):
+            Vc = FunctionSpace(cmesh, Ve)
+            Vf = FunctionSpace(fmesh, Ve)
+            fine_dual = dual(Vf)
+            coarse_primal = victim(Vc)
 
-        actual = exact_dual(V)
-        for mesh in reversed(hierarchy[:-1]):
-            V = FunctionSpace(mesh, Ve)
-            expect = exact_dual(V)
-            tmp = Function(V)
-            restrict(actual, tmp)
-            actual = tmp
-            assert numpy.allclose(expect.dat.data_ro, actual.dat.data_ro)
+            coarse_dual = Function(Vc)
+            fine_primal = Function(Vf)
+            restrict(fine_dual, coarse_dual)
+            prolong(coarse_primal, fine_primal)
+            coarse_functional = functional(coarse_primal, coarse_dual)
+            fine_functional = functional(fine_primal, fine_dual)
+
+            assert numpy.allclose(fine_functional, coarse_functional)
 
 
 def test_grid_transfer(hierarchy, vector, space, degrees, transfer_type):
+    if not hierarchy.nested and transfer_type == "injection":
+        pytest.skip("Not implemented")
     if transfer_type == "injection":
         run_injection(hierarchy, vector, space, degrees)
     elif transfer_type == "restriction":
@@ -150,6 +177,8 @@ def test_grid_transfer_parallel(hierarchy, transfer_type):
     space = "CG"
     degs = degrees(space)
     vector = False
+    if not hierarchy.nested and hierarchy.refinements_per_level > 1:
+        pytest.skip("Not implemented")
     if transfer_type == "injection":
         run_injection(hierarchy, vector, space, degs)
     elif transfer_type == "restriction":