From 6c8264962f5e21dc00b15625f0ad70d39937dc51 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 6 Aug 2026 15:57:02 +0100 Subject: [PATCH 1/2] Build ASMVankaPC patches in Cython, with coloring and adaptive restriction ASMVankaPC now shares the same Cython IS-construction mechanism as ASMStarPC: create_star_points() takes the star of each owned seed, and the new create_closure_points() takes the closure of a group of points which, applied to those stars, are the points -pc_patch_construct_type vanka solves for. This replaces the pure-Python build_vanka_indices()/ get_entity_dofs() path that only handled the general, non-colored case. create_patch_ises() now takes points and offsets per subspace instead of a single array shared by all of them, since a Vanka patch's excluded subspaces (pc_vanka_exclude_subspaces) read the star, or with pc_vanka_include_type entity just the seed, while the included ones read its closure. The seeds may be restricted to the entities marked by a DMLabel, named by pc_vanka_construct_label, and the coloring then colors only those. With pc_vanka_adaptive the label comes from adapt.mark_refined_entities(), the same restriction ASMStarPC applies, so that a smoother on an adaptively refined level only relaxes the entities whose patch meets the refined region. PatchPC already builds the same restricted patches for -patch_pc_patch_construct_type vanka, since its adaptive option is construct-type agnostic. get_colors() no longer forces the mat_coloring_type to "power" for the distance-3 separation Vanka patches need: DMPlexCreateColoringLabel() now folds the requested distance into the graph it builds, so the default greedy coloring, applied to that graph at distance one, is correct for any requested distance. Co-Authored-By: Claude Sonnet 5 --- firedrake/cython/patchimpl.pyx | 138 ++++++++++++++++++++++++++---- firedrake/cython/petschdr.pxi | 1 + firedrake/preconditioners/asm.py | 142 ++++++++++--------------------- 3 files changed, 170 insertions(+), 111 deletions(-) diff --git a/firedrake/cython/patchimpl.pyx b/firedrake/cython/patchimpl.pyx index c16e1c38ed..93d4b7fc7e 100644 --- a/firedrake/cython/patchimpl.pyx +++ b/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] = (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 diff --git a/firedrake/cython/petschdr.pxi b/firedrake/cython/petschdr.pxi index 42ac97e24d..2c6e31f66e 100644 --- a/firedrake/cython/petschdr.pxi +++ b/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[]) diff --git a/firedrake/preconditioners/asm.py b/firedrake/preconditioners/asm.py index 4575326c67..a26a881352 100644 --- a/firedrake/preconditioners/asm.py +++ b/firedrake/preconditioners/asm.py @@ -1,5 +1,4 @@ import abc -import petsctools from pyop2.datatypes import IntType from firedrake.cython import patchimpl @@ -192,7 +191,7 @@ def get_patches(self, V): # Take the star of every seed point, one patch at a time seeds, seed_offsets = get_seeds(mesh_dm, use_coloring, depth, 1, label, value) - points, star_offsets, patch_offsets = patchimpl.create_star_points(mesh_dm, seeds, seed_offsets) + points, star_offsets, patch_offsets, _ = patchimpl.create_star_points(mesh_dm, seeds, seed_offsets) if ordering != "natural": # Each star is reordered on its own, since it is a patch of its own # unless a coloring grouped it with others @@ -204,8 +203,9 @@ def get_patches(self, V): # so we need to cache these for efficiency V_local_ises_indices = get_local_ises_indices(V) - # Build index sets for the patches - return patchimpl.create_patch_ises(points, star_offsets[patch_offsets], + # Build index sets for the patches, every subspace reading the same points + return patchimpl.create_patch_ises([points] * len(V), + [star_offsets[patch_offsets]] * len(V), [W.dm.getLocalSection() for W in V], [W.block_size for W in V], V_local_ises_indices) @@ -223,6 +223,19 @@ class ASMVankaPC(ASMPatchPC): coloring of the mesh entities. This is specified via the option `pc_vanka_use_coloring`. + The mesh entities that get a patch may be restricted to those marked + by a `PETSc.DMLabel` on the mesh's DMPlex, named by the option + `pc_vanka_construct_label` and holding them in the stratum + `pc_vanka_construct_label_value`. The option `pc_vanka_adaptive` builds + that label with `~firedrake.adapt.mark_refined_entities`, so that a + smoother on an adaptively refined level only relaxes the entities whose + patch meets the refined region. The coloring, if requested, then colors + only the marked entities. + + The subspaces listed in `pc_vanka_exclude_subspaces` contribute only the + DoFs on the star of the mesh entity, or on the entity alone when + `pc_vanka_include_type` is `entity` rather than the default `star`. + The mesh entities in the patches may be reordered by applying a matrix reordering to the connectivity graph with the option `pc_vanka_mat_ordering_type`. @@ -250,29 +263,41 @@ def get_patches(self, V): validate_overlap(mesh, depth, "vanka") exclude_subspaces = opts.getIntArray("exclude_subspaces", default=[]) - include_subspaces = [i for i in range(len(V)) if i not in exclude_subspaces] include_type = opts.getString("include_type", default="star").lower() if include_type not in ["star", "entity"]: raise ValueError(f"{self.prefix}include_type must be either 'star' or 'entity', not {include_type}") - include_star = include_type == "star" use_coloring = opts.getBool("use_coloring", default=False) ordering = opts.getString("mat_ordering_type", default="natural") + label, value = get_construct_label(mesh, self.prefix) - def splitting(V): - return (tuple(V[i] for i in include_subspaces), tuple(V[i] for i in exclude_subspaces)) + # Take the closure of the star of every seed point, one patch at a time + seeds, seed_offsets = get_seeds(mesh_dm, use_coloring, depth, 3, label, value) + star, star_offsets, patch_offsets, seeds = patchimpl.create_star_points(mesh_dm, seeds, seed_offsets) + if ordering != "natural": + # Each star is reordered on its own, since it is a patch of its own + # unless a coloring grouped it with others, and the closure follows it + for s in range(len(star_offsets) - 1): + points = slice(star_offsets[s], star_offsets[s+1]) + star[points] = order_points(mesh_dm, star[points], ordering, self.prefix) + closure, closure_offsets = patchimpl.create_closure_points(mesh_dm, star, star_offsets) + + # The included subspaces span the closure of the star, the excluded ones + # only the star, or the entity it was built around + included = (closure, closure_offsets[patch_offsets]) + excluded = (star, star_offsets[patch_offsets]) if include_type == "star" else (seeds, patch_offsets) + groups = [excluded if i in exclude_subspaces else included for i in range(len(V))] - Z = splitting(V) # Accessing .indices causes the allocation of a global array, # so we need to cache these for efficiency V_local_ises_indices = get_local_ises_indices(V) - Z_local_ises_indices = splitting(V_local_ises_indices) # Build index sets for the patches - colors = get_colors(mesh_dm, use_coloring, depth, distance=3) - ises = [build_vanka_indices(Z, Z_local_ises_indices, mesh_dm, ordering, self.prefix, - include_star, color) for color in colors] - return ises + return patchimpl.create_patch_ises([points for points, _ in groups], + [offsets for _, offsets in groups], + [W.dm.getLocalSection() for W in V], + [W.block_size for W in V], + V_local_ises_indices) class ASMLinesmoothPC(ASMPatchPC): @@ -623,7 +648,8 @@ def get_colors(mesh_dm: PETSc.DMPlex, use_coloring: bool, depth: int, distance: depth The topological dimension of the entities. distance - The coloring distance. + How far through the mesh an entity reaches, 1 to separate star patches + and 3 to separate Vanka patches. label A label selecting the entities, or `None` for all of them. value @@ -636,14 +662,10 @@ def get_colors(mesh_dm: PETSc.DMPlex, use_coloring: bool, depth: int, distance: """ if use_coloring: - # Greedy, the default, only supports distances 1 and 2 - parameters = {"mat_coloring_type": "power" if distance > 2 else "greedy"} - with petsctools.inserted_options(parameters=parameters, - options_prefix="dm_plex_coloring_"): - # Colors the subgraph the selected entities induce, so restricting - # them can only merge colors, never split one - colors = mesh_dm.createColoringLabel(depth=depth, distance=distance, - label=label, value=value) + # Colors the subgraph the selected entities induce, so restricting + # them can only merge colors, never split one + colors = mesh_dm.createColoringLabel(depth=depth, distance=distance, + label=label, value=value) elif label is None: colors = range(*mesh_dm.getDepthStratum(depth)) else: @@ -671,7 +693,8 @@ def get_seeds(mesh_dm: PETSc.DMPlex, use_coloring: bool, depth: int, distance: i depth The topological dimension of the entities. distance - The coloring distance. + How far through the mesh an entity reaches, 1 to separate star patches + and 3 to separate Vanka patches. label A label selecting the entities, or `None` for all of them. value @@ -698,29 +721,6 @@ def get_seeds(mesh_dm: PETSc.DMPlex, use_coloring: bool, depth: int, distance: i return seeds, numpy.arange(len(seeds) + 1, dtype=IntType) -def get_entity_dofs(V, V_local_ises_indices, points): - """Return degrees of freedom associated with mesh entities (points of the DMPlex). - - :arg V: the FunctionSpace to extract DOFs from - :arg V_local_ises_indices: V.local_ises.indices - :points: an iterable of mesh entities - - :returns: a list with the DOFs of V associated with the mesh entities - """ - indices = [] - for (i, W) in enumerate(V): - section = W.dm.getLocalSection() - for p in points: - dof = section.getDof(p) - if dof <= 0: - continue - off = section.getOffset(p) - # Local indices within W - W_slice = slice(off*W.block_size, W.block_size * (off + dof)) - indices.extend(V_local_ises_indices[i][W_slice]) - return indices - - def get_star_points(mesh_dm, ordering, prefix, seed_points): """Get DMPlex points in the star of each point in seed_points. @@ -746,51 +746,3 @@ def get_star_points(mesh_dm, ordering, prefix, seed_points): star = order_points(mesh_dm, star[::-1], ordering, prefix) points.extend(star) return points - - -def build_vanka_indices(Z, Z_local_ises_indices, mesh_dm, ordering, prefix, include_star, seed_points): - """Return DOFs in the Vanka patches constructed at each point in seed_points. - - :arg Z: a tuple of the included/excluded FunctionSpaces to extract DOFs from - :arg Z_local_ises_indices: (Z[0].local_ises.indices, Z[1].local_ises.indices) - :arg mesh_dm: the DMPlex - :arg ordering: a Mat.OrderingType indicating the ordering type - :arg prefix: the PETSc.Options prefix to further specify the ordering - :arg include_star: whether to include DOFs of Z[1] in the star or just the entity - :seed_points: an iterable of point indices to construct the Vanka patches - - :returns: A PETSc.IS with the degrees of freedom in the Vanka patches - """ - if isinstance(seed_points, PETSc.IS): - seed_points = seed_points.indices - elif numpy.isscalar(seed_points): - seed_points = (seed_points,) - indices = [] - for seed in seed_points: - V_points = [] - Q_points = [] - # Only build patches over owned DoFs - if mesh_dm.getLabelValue("pyop2_ghost", seed) != -1: - continue - # Create point list from mesh DM, by decreasing topological dimension - # (interiors, faces, edges, vertices) - star, _ = mesh_dm.getTransitiveClosure(seed, useCone=False) - star = order_points(mesh_dm, star[::-1], ordering, prefix) - if include_star: - Q_points.extend(star) - else: - Q_points.append(seed) - closure = [] - for s in reversed(star): - cs, _ = mesh_dm.getTransitiveClosure(s, useCone=True) - closure.extend(cs) - # Grab unique points with stable ordering - closure = reversed(dict.fromkeys(closure)) - V_points.extend(closure) - indices.extend(get_entity_dofs(Z[0], Z_local_ises_indices[0], V_points)) - indices.extend(get_entity_dofs(Z[1], Z_local_ises_indices[1], Q_points)) - - indices = numpy.array(indices, dtype=PETSc.IntType) - indices = indices[indices >= 0] - iset = PETSc.IS().createGeneral(indices, comm=PETSc.COMM_SELF) - return iset From bb4fad03b719abbb64e368e306745785d5a4083f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 6 Aug 2026 19:07:52 +0100 Subject: [PATCH 2/2] Test adaptive restriction and coloring for ASMVankaPC Adds test_adaptive_vanka_equivalence, test_adaptive_vanka_coloring, and test_adaptive_vanka_is_restricted, mirroring the existing test_adaptive_star_* tests: they drive a manufactured Taylor-Hood Stokes problem through ASMVankaPC and PatchPC(construct_type=vanka), each with pc_vanka_adaptive/pc_vanka_use_coloring, and check that the two implementations take the same number of iterations and hold the expected number of patches. Unlike a star patch, a Vanka patch reaches every cell that shares a point with the seed's closure, so validate_overlap() requires overlap distance two rather than one. The existing adaptive_mesh fixture only requests distance one, which is enough for ASMStarPC but leaves ASMVankaPC's patches inconsistent across ranks; adaptive_vanka_mesh requests the extra overlap, sharing the corner-marking logic with adaptive_mesh through the new _refined_corner_mesh() helper. Co-Authored-By: Claude Sonnet 5 --- tests/firedrake/regression/test_star_pc.py | 141 ++++++++++++++++++++- 1 file changed, 137 insertions(+), 4 deletions(-) diff --git a/tests/firedrake/regression/test_star_pc.py b/tests/firedrake/regression/test_star_pc.py index 55eab8fe83..f3b1816926 100644 --- a/tests/firedrake/regression/test_star_pc.py +++ b/tests/firedrake/regression/test_star_pc.py @@ -4,7 +4,11 @@ import petsctools from firedrake import * from firedrake.adapt import mark_refined_entities, REFINED_LABEL +from firedrake.functionspaceimpl import WithGeometry +from firedrake.mesh import MeshGeometry from firedrake.petsc import DEFAULT_DIRECT_SOLVER +from firedrake.variational_solver import LinearVariationalSolver +from ufl.core.expr import Expr @pytest.fixture(params=["scalar", @@ -542,6 +546,15 @@ def test_vanka_coloring(): assert its[True] == its[False] +def _refined_corner_mesh(base: MeshGeometry) -> MeshGeometry: + """Refine the corner of ``base`` where the coordinates sum to less than + one half, so that most of the mesh is untouched.""" + x = SpatialCoordinate(base) + marker = Function(FunctionSpace(base, "DG", 0)) + marker.interpolate(conditional(sum(x) < 0.5, 1, 0)) + return base.refine_marked_elements(marker) + + @pytest.fixture(params=["square", "cube"]) def adaptive_mesh(request): """A mesh with a single refined corner, so that most of it is untouched.""" @@ -550,11 +563,18 @@ def adaptive_mesh(request): base = UnitSquareMesh(6, 6, distribution_parameters=dparams) else: base = UnitCubeMesh(3, 3, 3, distribution_parameters=dparams) + return _refined_corner_mesh(base) - x = SpatialCoordinate(base) - marker = Function(FunctionSpace(base, "DG", 0)) - marker.interpolate(conditional(sum(x) < 0.5, 1, 0)) - return base.refine_marked_elements(marker) + +@pytest.fixture(params=["square", "cube"]) +def adaptive_vanka_mesh(request): + """Like ``adaptive_mesh``, but with the extra overlap Vanka patches need.""" + dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 2)} + if request.param == "square": + base = UnitSquareMesh(6, 6, distribution_parameters=dparams) + else: + base = UnitCubeMesh(3, 3, 3, distribution_parameters=dparams) + return _refined_corner_mesh(base) def adaptive_solver(V, parameters): @@ -614,6 +634,83 @@ def adaptive_patch_parameters(**kwargs): **{f"patch_pc_patch_{k}": v for k, v in kwargs.items()}} +def adaptive_vanka_parameters(**kwargs: object) -> dict: + """Parameters driving ASMVankaPC, with ``pc_vanka_`` prefixing every keyword.""" + return {"mat_type": "aij", + "ksp_type": "gmres", + "ksp_rtol": 1E-8, + "pc_type": "python", + "pc_python_type": "firedrake.ASMVankaPC", + "pc_vanka_construct_dim": 0, + "pc_vanka_exclude_subspaces": "1", + **{f"pc_vanka_{k}": v for k, v in kwargs.items()}} + + +def adaptive_vanka_patch_parameters(**kwargs: object) -> dict: + """Parameters driving PatchPC with vanka patches, with ``patch_pc_patch_`` prefixing every keyword.""" + return {"mat_type": "matfree", + "ksp_type": "gmres", + "ksp_rtol": 1E-8, + "pc_type": "python", + "pc_python_type": "firedrake.PatchPC", + "patch_pc_patch_save_operators": True, + "patch_pc_patch_construct_type": "vanka", + "patch_pc_patch_construct_dim": 0, + "patch_pc_patch_exclude_subspaces": "1", + "patch_pc_patch_sub_mat_type": "seqdense", + "patch_sub_ksp_type": "preonly", + "patch_sub_pc_type": "lu", + **{f"patch_pc_patch_{k}": v for k, v in kwargs.items()}} + + +def taylor_hood_space(mesh: MeshGeometry) -> WithGeometry: + """A Taylor-Hood (CG2, CG1) mixed velocity-pressure space on ``mesh``.""" + V = VectorFunctionSpace(mesh, "CG", 2) + Q = FunctionSpace(mesh, "CG", 1) + return V * Q + + +def stokes_manufactured_solution(mesh: MeshGeometry) -> tuple[Expr, Expr]: + """A divergence-free velocity, built as the curl of a potential that + vanishes on the boundary, paired with a linear pressure.""" + gdim = mesh.geometric_dimension + x = SpatialCoordinate(mesh) + if gdim == 2: + psi = x[0]*(1-x[0]) * x[1]*(1-x[1]) + else: + bubble = x[0]*(1-x[0]) * x[1]*(1-x[1]) * x[2]*(1-x[2]) + psi = as_vector([bubble, bubble, bubble]) + uexact = curl(psi) + pexact = x[0] + return uexact, pexact + + +def adaptive_stokes_solver(Z: WithGeometry, parameters: dict) -> LinearVariationalSolver: + """Solve a manufactured Stokes problem on the Taylor-Hood space ``Z``, + returning the solver so that its iteration count and its patches can be + inspected.""" + mesh = Z.mesh() + gdim = mesh.geometric_dimension + uexact, pexact = stokes_manufactured_solution(mesh) + + u, p = TrialFunctions(Z) + v, q = TestFunctions(Z) + a = (inner(grad(u), grad(v)) * dx + - inner(p, div(v)) * dx + - inner(div(u), q) * dx) + test, trial = a.arguments() + L = a(test, as_vector([uexact[i] for i in range(gdim)] + [pexact])) + bcs = DirichletBC(Z.sub(0), uexact, "on_boundary") + + zh = Function(Z) + problem = LinearVariationalProblem(a, L, zh, bcs=bcs) + nsp = MixedVectorSpaceBasis(Z, [Z.sub(0), VectorSpaceBasis(constant=True, comm=mesh.comm)]) + solver = LinearVariationalSolver(problem, solver_parameters=parameters, + nullspace=nsp, transpose_nullspace=nsp) + solver.solve() + return solver + + def num_asm_patches(solver): """The number of patches the ASM preconditioner of a solver holds.""" return len(solver.snes.ksp.pc.getPythonContext().asmpc.getASMSubKSP()) @@ -684,6 +781,42 @@ def test_adaptive_star_is_restricted(adaptive_mesh): assert 0 < num_asm_patches(restricted) < num_asm_patches(every) +@pytest.mark.parallel([1, 3]) +def test_adaptive_vanka_equivalence(adaptive_vanka_mesh): + Z = taylor_hood_space(adaptive_vanka_mesh) + vanka = adaptive_stokes_solver(Z, adaptive_vanka_parameters(adaptive=True)) + patch = adaptive_stokes_solver(Z, adaptive_vanka_patch_parameters(adaptive=True)) + assert vanka.snes.getLinearSolveIterations() == patch.snes.getLinearSolveIterations() + + +@pytest.mark.parallel([1, 3]) +def test_adaptive_vanka_coloring(adaptive_vanka_mesh): + Z = taylor_hood_space(adaptive_vanka_mesh) + plain = adaptive_stokes_solver(Z, adaptive_vanka_parameters(adaptive=True)) + colored = adaptive_stokes_solver(Z, adaptive_vanka_parameters(adaptive=True, use_coloring=True)) + assert plain.snes.getLinearSolveIterations() == colored.snes.getLinearSolveIterations() + + dm = adaptive_vanka_mesh.topology_dm + label = dm.getLabel(REFINED_LABEL) + vstart, vend = dm.getDepthStratum(0) + marked = label.getStratumIS(1).indices + vertices = marked[(marked >= vstart) & (marked < vend)] + assert num_asm_patches(plain) == len(vertices) + + # Coloring groups the same seeds into fewer, larger patches + colors = dm.createColoringLabel(depth=0, distance=3, label=label, value=1) + assert num_asm_patches(colored) == len(colors) + assert sum(len(color.indices) for color in colors) <= len(vertices) + + +@pytest.mark.parallel([1, 3]) +def test_adaptive_vanka_is_restricted(adaptive_vanka_mesh): + Z = taylor_hood_space(adaptive_vanka_mesh) + every = adaptive_stokes_solver(Z, adaptive_vanka_parameters()) + restricted = adaptive_stokes_solver(Z, adaptive_vanka_parameters(adaptive=True)) + assert 0 < num_asm_patches(restricted) < num_asm_patches(every) + + def test_uniform_mesh_has_no_refined_region(): # A mesh with no adaptive parent, such as the coarsest mesh of a hierarchy or a # uniformly refined level, has no refined region to single out