Restrict Vanka patches to the adaptively refined region

firedrake/cython/patchimpl.pyx

@@ -30,9 +30,9 @@ def create_star_points(PETSc.DM dm, seeds, offsets):
     Returns
     -------
     tuple of numpy.ndarray
-        The star of every owned seed by decreasing topological dimension,
-        the offsets saying where each star starts in them, and the offsets
-        saying which stars make up each patch.
+        The star of every owned seed by decreasing topological dimension, the
+        offsets saying where each star starts in them, the offsets saying which
+        stars make up each patch, and the owned seeds themselves, one per star.
 
     """
     cdef:
@@ -46,6 +46,7 @@ def create_star_points(PETSc.DM dm, seeds, offsets):
         PetscInt[::1] cseeds = np.asarray(seeds, dtype=IntType)
         PetscInt[::1] coffsets = np.asarray(offsets, dtype=IntType)
         PetscInt[::1] cstar_offsets = np.zeros(cseeds.shape[0] + 1, dtype=IntType)
+        PetscInt[::1] cowned_seeds = np.zeros(cseeds.shape[0], dtype=IntType)
         PetscInt[::1] cpatch_offsets = np.zeros(npatch + 1, dtype=IntType)
         PetscInt[::1] view
         DMLabel ghost = NULL
@@ -84,6 +85,7 @@ def create_star_points(PETSc.DM dm, seeds, offsets):
                     points[npoints] = star[2*s]
                     npoints += 1
                 CHKERR(DMPlexRestoreTransitiveClosure(dm.dm, seed, PETSC_FALSE, &starSize, &star))
+                cowned_seeds[nstars] = seed
                 nstars += 1
                 cstar_offsets[nstars] = npoints
         cpatch_offsets[npatch] = nstars
@@ -99,21 +101,109 @@ def create_star_points(PETSc.DM dm, seeds, offsets):
             CHKERR(DMLabelDestroyIndex(ghost))
     return (star_points,
             np.asarray(cstar_offsets)[:nstars + 1],
-            np.asarray(cpatch_offsets))
+            np.asarray(cpatch_offsets),
+            np.asarray(cowned_seeds)[:nstars])
 
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def create_patch_ises(points, offsets, sections, bsizes, indices):
-    """Gather the degrees of freedom carried by the mesh points of each patch.
+def create_closure_points(PETSc.DM dm, points, offsets):
+    """Collect the mesh points in the closure of each group of points.
+
+    Applied to the stars of `create_star_points`, these are the points PCPatch
+    solves for with ``-pc_patch_construct_type vanka``.
 
     Parameters
     ----------
+    dm : PETSc.DM
+        The mesh topology.
     points : numpy.ndarray
-        The mesh points of the patches, patch by patch.
+        The mesh points to close over, group by group.
     offsets : numpy.ndarray
-        Where each patch starts in ``points``, of length one more than the
-        number of patches.
+        Where each group starts in ``points``, of length one more than the
+        number of groups.
+
+    Returns
+    -------
+    tuple of numpy.ndarray
+        The closure of every group by decreasing topological dimension, and the
+        offsets saying where each closure starts in them.
+
+    """
+    cdef:
+        PetscInt ngroups = len(offsets) - 1
+        PetscInt npoints = 0, maxpoints = 0, start = 0
+        PetscInt g, k, q, p, closureSize = 0
+        PetscInt *closure = NULL
+        PetscInt *points_ = NULL
+        PetscInt *newpoints = NULL
+        PetscInt *seen = NULL
+        PetscInt pStart = 0, pEnd = 0
+        PetscInt[::1] cpoints = np.asarray(points, dtype=IntType)
+        PetscInt[::1] coffsets = np.asarray(offsets, dtype=IntType)
+        PetscInt[::1] cclosure_offsets = np.zeros(ngroups + 1, dtype=IntType)
+        PetscInt[::1] view
+
+    CHKERR(DMPlexGetChart(dm.dm, &pStart, &pEnd))
+    # A point already taken carries the group that took it, one-based so that
+    # the zeroed initial state marks a point no group has taken
+    CHKERR(PetscCalloc1(pEnd - pStart, &seen))
+    maxpoints = 16 * (cpoints.shape[0] + 1)
+    CHKERR(PetscMalloc1(maxpoints, &points_))
+    try:
+        for g in range(ngroups):
+            start = npoints
+            # Against the order of the group, so that the closure of its lowest
+            # dimensional point comes first and survives the deduplication
+            for k in range(coffsets[g+1] - 1, coffsets[g] - 1, -1):
+                CHKERR(DMPlexGetTransitiveClosure(dm.dm, cpoints[k], PETSC_TRUE, &closureSize, &closure))
+                while npoints + closureSize > maxpoints:
+                    maxpoints *= 2
+                    CHKERR(PetscMalloc1(maxpoints, &newpoints))
+                    for q in range(npoints):
+                        newpoints[q] = points_[q]
+                    CHKERR(PetscFree(points_))
+                    points_ = newpoints
+                    newpoints = NULL
+                for q in range(closureSize):
+                    p = closure[2*q]
+                    if seen[p - pStart] == g + 1:
+                        continue
+                    seen[p - pStart] = g + 1
+                    points_[npoints] = p
+                    npoints += 1
+                CHKERR(DMPlexRestoreTransitiveClosure(dm.dm, cpoints[k], PETSC_TRUE, &closureSize, &closure))
+            # Undo the reversed walk, leaving the closure by decreasing topological dimension
+            for k in range((npoints - start) // 2):
+                p = points_[start + k]
+                points_[start + k] = points_[npoints - 1 - k]
+                points_[npoints - 1 - k] = p
+            cclosure_offsets[g+1] = npoints
+
+        closure_points = np.empty(npoints, dtype=IntType)
+        if npoints > 0:
+            view = closure_points
+            for k in range(npoints):
+                view[k] = points_[k]
+    finally:
+        CHKERR(PetscFree(points_))
+        CHKERR(PetscFree(seen))
+    return closure_points, np.asarray(cclosure_offsets)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def create_patch_ises(points, offsets, sections, bsizes, indices):
+    """Gather the degrees of freedom carried by the mesh points of each patch.
+
+    Parameters
+    ----------
+    points : list of numpy.ndarray
+        The mesh points each subspace reads, patch by patch. Subspaces reading
+        the same points may share one array.
+    offsets : list of numpy.ndarray
+        Where each patch starts in the points of each subspace, each of length
+        one more than the number of patches.
     sections : list of PETSc.Section
         The local section of each subspace.
     bsizes : list of int
@@ -130,33 +220,47 @@ def create_patch_ises(points, offsets, sections, bsizes, indices):
     Raises
     ------
     ValueError
-        If a mesh point appears more than once in a patch.
+        If the subspaces disagree on the number of patches, or if a mesh point
+        appears more than once in a patch.
 
     """
     cdef:
         PetscInt nsub = len(sections)
-        PetscInt npatch = len(offsets) - 1
+        PetscInt npatch = 0
         PetscInt ndofs = 0, maxdofs = 0
         PetscInt i, k, p, v, dof, off, bs, index
         PetscInt *dofs = NULL
         PetscInt *cbs = NULL
         PetscInt **cindices = NULL
-        PetscInt[::1] cpoints = np.asarray(points, dtype=IntType)
-        PetscInt[::1] coffsets = np.asarray(offsets, dtype=IntType)
+        PetscInt **cpoints = NULL
+        PetscInt **coffsets = NULL
         PETSc.PetscSection *csections = NULL
         PetscInt[::1] view
 
+    if nsub > 0:
+        npatch = len(offsets[0]) - 1
+        if any(len(offsets[i]) - 1 != npatch for i in range(nsub)):
+            raise ValueError("The subspaces disagree on the number of patches")
+
     CHKERR(PetscMalloc1(nsub, &csections))
     CHKERR(PetscMalloc1(nsub, &cindices))
+    CHKERR(PetscMalloc1(nsub, &cpoints))
+    CHKERR(PetscMalloc1(nsub, &coffsets))
     CHKERR(PetscMalloc1(nsub, &cbs))
     # The memoryviews must outlive the pointers taken out of them
     views = [np.asarray(indices[i], dtype=IntType) for i in range(nsub)]
+    point_views = [np.asarray(points[i], dtype=IntType) for i in range(nsub)]
+    offset_views = [np.asarray(offsets[i], dtype=IntType) for i in range(nsub)]
     for i in range(nsub):
         csections[i] = (<PETSc.Section?>sections[i]).sec
         cbs[i] = bsizes[i]
         view = views[i]
         cindices[i] = &view[0] if view.shape[0] > 0 else NULL
         maxdofs += view.shape[0]
+        view = point_views[i]
+        cpoints[i] = &view[0] if view.shape[0] > 0 else NULL
+        view = offset_views[i]
+        coffsets[i] = &view[0]
 
     # No mesh point appears twice in a patch, so a patch cannot hold more degrees
     # of freedom than the process has
@@ -168,11 +272,11 @@ def create_patch_ises(points, offsets, sections, bsizes, indices):
             ndofs = 0
             for i in range(nsub):
                 bs = cbs[i]
-                for p in range(coffsets[v], coffsets[v+1]):
-                    CHKERR(PetscSectionGetDof(csections[i], cpoints[p], &dof))
+                for p in range(coffsets[i][v], coffsets[i][v+1]):
+                    CHKERR(PetscSectionGetDof(csections[i], cpoints[i][p], &dof))
                     if dof <= 0:
                         continue
-                    CHKERR(PetscSectionGetOffset(csections[i], cpoints[p], &off))
+                    CHKERR(PetscSectionGetOffset(csections[i], cpoints[i][p], &off))
                     if ndofs + dof*bs > maxdofs:
                         raise ValueError("A mesh point appears more than once in a patch")
                     for k in range(off*bs, (off + dof)*bs):
@@ -191,6 +295,8 @@ def create_patch_ises(points, offsets, sections, bsizes, indices):
         CHKERR(PetscFree(dofs))
         CHKERR(PetscFree(csections))
         CHKERR(PetscFree(cindices))
+        CHKERR(PetscFree(cpoints))
+        CHKERR(PetscFree(coffsets))
         CHKERR(PetscFree(cbs))
     return ises
 

firedrake/cython/petschdr.pxi

@@ -29,6 +29,7 @@ cdef extern from "petsc.h":
 cdef extern from "petscsys.h" nogil:
     PetscErrorCode PetscMalloc1(PetscInt,void*)
     PetscErrorCode PetscMalloc2(PetscInt,void*,PetscInt,void*)
+    PetscErrorCode PetscCalloc1(PetscInt,void*)
     PetscErrorCode PetscFree(void*)
     PetscErrorCode PetscFree2(void*,void*)
     PetscErrorCode PetscSortIntWithArray(PetscInt,PetscInt[],PetscInt[])