Fix coarse-to-fine cell maps for hierarchies built on an adaptively refined mesh

firedrake/cython/mgimpl.pyx

@@ -381,64 +381,100 @@ def adaptive_parent_child_cell_maps(PETSc.DM coarse_dm,
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def coarse_to_fine_cells(mc, mf, clgmaps, flgmaps):
-    """Return a map from (renumbered) cells in a coarse mesh to those
-    in a refined fine mesh.
-
-    :arg mc: the coarse mesh to create the map from.
-    :arg mf: the fine mesh to map to.
-    :arg clgmaps: coarse lgmaps (non-overlapped and overlapped)
-    :arg flgmaps: fine lgmaps (non-overlapped and overlapped)
-    :returns: Two arrays, one mapping coarse to fine cells, the second fine to coarse cells.
+    """Map the cells of a coarse mesh to those of its uniform refinement.
+
+    Parameters
+    ----------
+    mc : MeshGeometry
+        The coarse mesh.
+    mf : MeshGeometry
+        The fine mesh, obtained by uniformly refining the non-overlapped
+        plex of ``mc``.
+    clgmaps : tuple
+        The coarse ``(non-overlapped, overlapped)`` point local-to-global maps.
+    flgmaps : tuple
+        The fine ``(non-overlapped, overlapped)`` point local-to-global maps.
+
+    Returns
+    -------
+    numpy.ndarray
+        Map from each owned coarse cell to the fine cells it was split into.
+    numpy.ndarray
+        Map from each owned fine cell to the coarse cell it came from.
+
+    Notes
+    -----
+    Three numberings of the same cells meet here:
+
+    1. Firedrake numbering, which lists owned cells before halo cells. The
+       returned maps are indexed by, and contain, these numbers.
+    2. Overlapped plex numbering, that of ``mesh.topology_dm``.
+       `get_entity_renumbering` translates between 1 and 2.
+    3. Non-overlapped plex numbering, that of the halo-free plex that was
+       refined. Only here does the parent relation hold: uniform refinement
+       splits cell ``p`` into cells ``p*nref`` to ``p*nref + nref - 1``.
+
+    Applying an overlapped local-to-global map and then a non-overlapped
+    global-to-local one translates between 2 and 3. Those two numberings are
+    genuinely different orders, not a common prefix plus a halo: a plex built
+    by ``DMPlexTransform`` (an adaptively refined one) numbers its cells by
+    refinement case, so its owned cells are interleaved with its halo cells.
     """
     cdef:
         PETSc.DM cdm, fdm
         PetscInt cStart, cEnd, c, val, dim, nref, ncoarse
         PetscInt i, ccell, fcell, nfine
         np.ndarray coarse_to_fine
         np.ndarray fine_to_coarse
-        np.ndarray co2n, fn2o, idx
+        np.ndarray co2n, fn2o, idx, found, permuted
 
     cdm = mc.topology_dm
     fdm = mf.topology_dm
     dim = cdm.getDimension()
     nref = <PetscInt> 2 ** dim
     ncoarse = mc.cell_set.size
     nfine = mf.cell_set.size
+    # co2n: coarse overlapped plex cell -> coarse Firedrake cell
+    # fn2o: fine Firedrake cell -> fine overlapped plex cell
     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), -1, dtype=PETSc.IntType)
     # Walk owned fine cells:
     cStart, cEnd = 0, nfine
 
+    # In serial the overlapped and non-overlapped plexes are the same plex,
+    # so both maps already speak the numbering the parent relation holds in.
     if mc.comm.size > 1:
+        # Cells are the leading points of a plex chart, so these point maps
+        # can be applied to cell numbers directly.
         cno, co = clgmaps
         fno, fo = flgmaps
-        # Compute global numbers of original cell numbers
+        # Rebase fn2o onto the fine non-overlapped plex, one map per arrow:
+        # fine Firedrake cell -> overlapped -> global -> non-overlapped.
         fo.apply(fn2o, result=fn2o)
-        # Compute local numbers of original cells on non-overlapped mesh
         fn2o = fno.applyInverse(fn2o, PETSc.LGMap.MapMode.MASK)
-        # Need to permute order of co2n so it maps from non-overlapped
-        # cells to new cells (these may have changed order).  Need to
-        # map all known cells through.
+        # Rebase co2n the same way, but here it is the *index* that changes
+        # numbering, not the value, so send every local coarse cell through
+        # the translation. MASK gives -1 for cells the non-overlapped plex
+        # does not have.
         idx = np.arange(mc.cell_set.total_size, dtype=PETSc.IntType)
-        # LocalToGlobal
         co.apply(idx, result=idx)
-        # GlobalToLocal
-        # Drop values that did not exist on non-overlapped mesh
-        idx = cno.applyInverse(idx, PETSc.LGMap.MapMode.DROP)
-        co2n = co2n[idx]
+        idx = cno.applyInverse(idx, PETSc.LGMap.MapMode.MASK)
+        # idx[i] is where overlapped cell i lands, so scatter rather than
+        # slice: the surviving cells need not be the leading ones.
+        found = idx >= 0
+        permuted = np.empty(found.sum(), dtype=PETSc.IntType)
+        permuted[idx[found]] = co2n[found]
+        co2n = permuted
 
     for c in range(cStart, cEnd):
-        # get original (overlapped) cell number
+        # Every owned fine cell exists on the non-overlapped plex.
         fcell = fn2o[c]
-        # The owned cells should map into non-overlapped cell numbers
-        # (due to parallel growth strategy)
         assert 0 <= fcell < cEnd
 
-        # Find original coarse cell (fcell / nref) and then map
-        # forward to renumbered coarse cell (again non-overlapped
-        # cells should map into owned coarse cells)
+        # Uniform refinement numbers the nref children of a cell
+        # consecutively, so integer division recovers the parent.
         ccell = co2n[fcell // nref]
         assert 0 <= ccell < ncoarse
         fine_to_coarse[c, 0] = ccell

tests/firedrake/multigrid/test_adaptive_multigrid.py

@@ -297,6 +297,38 @@ def test_adapt_after_uniform_refinement(coarse_mesh, refine):
     _assert_adapt_after_uniform_refinement(mh)
 
 
+@pytest.mark.parallel([1, 2, 4])
+@pytest.mark.parametrize("refine", [1, 2])
+def test_adapt_before_uniform_refinement(coarse_mesh, refine):
+    """An adaptively refined mesh can be uniformly refined into a hierarchy.
+    Its plex numbers cells by refinement case, so its owned cells are
+    interleaved with its halo cells, which the cell maps must not assume away.
+    """
+    netgen_flags = {} if hasattr(coarse_mesh, "netgen_mesh") else None
+
+    M = FunctionSpace(coarse_mesh, "DG", 0)
+    markers = Function(M)
+    markers.dat.data_wo[:1] = 1
+    mesh = coarse_mesh.refine_marked_elements(markers)
+
+    mh = MeshHierarchy(mesh, refine, netgen_flags=netgen_flags)
+    assert len(mh) == refine + 1
+    assert np.allclose(assemble(1*dx(mh[-1])), assemble(1*dx(coarse_mesh)))
+
+    nref = 2 ** mesh.topological_dimension
+    for level in range(refine):
+        coarse_to_fine = mh.coarse_to_fine_cells[level]
+        fine_to_coarse = mh.fine_to_coarse_cells[level + 1]
+        assert coarse_to_fine.shape == (mh[level].cell_set.size, nref)
+        assert fine_to_coarse.shape == (mh[level + 1].cell_set.size, 1)
+        # Uniform refinement splits every owned coarse cell into nref owned
+        # fine cells, each of which points back at the cell it came from.
+        assert (coarse_to_fine >= 0).all()
+        assert (fine_to_coarse >= 0).all()
+        parents = np.arange(coarse_to_fine.shape[0]).reshape(-1, 1)
+        assert (fine_to_coarse[coarse_to_fine, 0] == parents).all()
+
+
 @pytest.mark.parallel([1, 2, 4])
 @pytest.mark.parametrize("operator", ["prolong", "inject"])
 def test_DG0(mh, operator):