Support for SNESAdapt

[pbrubeck]

Implementation Plan: PETSc SNESAdapt with Multigrid Hierarchy in Firedrake

Disclaimer: This implementation plan was generated collaboratively by a Firedrake core developer interacting with Gemini. It reflects an architectural design tailored to Firedrake's immutable mesh paradigm, PETSc's adaptive solver components, and ngsPETSc boundary tracking hooks.

---

This document outlines the complete implementation plan for integrating PETSc's adaptive nonlinear solvers (SNESAdapt) within Firedrake, with specialized support for geometric multigrid hierarchies built on top of Netgen geometries via ngsPETSc.

Due to the immutable nature of Firedrake's mesh topologies, in-place $h$-adaptivity is replaced by an append-only, monotonic MeshHierarchy extension. Netgen is used exclusively as a background geometric oracle to snap newly introduced boundary vertices to the true CAD surface, guaranteeing strict topological nestedness and optimal performance for Geometric Multigrid (PCMG) transfer operators.

---

1. Architecture & Data Flow

[SNES Solver Loop] ──(Calls User Marking Callback)──> [Firedrake DG0 Indicator]
                                                               │
                                                    (Auto-Build DMLabel)
                                                               ▼
[MeshHierarchy Level N+1] <──(Snap Bnd)── [Topological Split via DMLabel (DMPlex)]
           │
           ▼ (Synchronize Preconditioner)
[Update PCMG Levels & Transfer Operators]

1. Error Estimation & Marking Callback: Prior to running a refinement loop step, the solver invokes a user-defined marking callback. This function computes a localized error metric and tags elements for refinement using standard Firedrake fields.
2. Automated DMLabel Production: Firedrake takes the flagged data array from the function space level and automatically populates a native PETSc DMLabel assigned to the mesh's underlying DMPlex.
3. Hierarchy Extension: Firedrake intercepts PETSc's DMRefine call via a Python callback, triggering a topological split of the mesh guided by the generated DMLabel, and appends a new discrete Mesh object to an append-only MeshHierarchy.
4. Geometric Alignment: Newly introduced boundary vertices are intercepted and projected onto the native Netgen CAD description to honor curved geometries.
5. Preconditioner Synchronization: PETSc's PCMG (Preconditioner Multigrid) is notified of the depth increase, and Firedrake's nested restriction and prolongation operators are cleanly mapped to the new solver level.

---

2. User-Facing API & Usage Sketch

This sketch demonstrates how a user defines a custom marking callback (e.g., executing a residual error estimation or threshold metric mapped to a $DG0$ space), connects it to the solver context, and solves using an automated DMLabel pipeline.

import netgen.gui
from netgen.geom2d import SplineGeometry
from firedrake import *

# 1. Define Netgen Geometry and construct initial coarse mesh
geo = SplineGeometry()
geo.AddCircle((0, 0), r=1.0)
ng_mesh = geo.GenerateMesh(maxh=0.5)
coarse_mesh = Mesh(ng_mesh)

# 2. Wrap into our new AdaptiveMeshHierarchy 
hierarchy = AdaptiveMeshHierarchy(coarse_mesh)

# 3. Setup Problem on the initial coarse level
V = FunctionSpace(hierarchy[0], "CG", 1)
u = Function(V)
v = TestFunction(V)
F = dot(grad(u), grad(v))*dx - f*v*dx

# 4. Define a custom Marking Callback using high-level Firedrake syntax
def my_marking_strategy(solver, current_u):
    """
    Computes an error indicator over the current solution and returns
    a DG0 Function where 1 means 'Refine' and 0 means 'Keep'.
    """
    mesh = current_u.function_space().mesh()
    DG0 = FunctionSpace(mesh, "DG", 0)
    indicator = Function(DG0)
    
    # Example: Simple gradient-magnitude based error indicator
    # (In production, this would evaluate a true residual or dual-weighted error estimator)
    assemble(sqrt(dot(grad(current_u), grad(current_u))) * dx, tensor=indicator)
    
    # Tag the top 10% highest error cells for refinement (Dörfler-like or direct threshold)
    threshold = 0.90 * indicator.dat.data.max()
    
    markers = Function(DG0)
    markers.dat.data[:] = np.where(indicator.dat.data >= threshold, 1, 0)
    return markers

# 5. Configure adaptive solver parameters targeting PCMG
params = {
    "snes_type": "newtonls",
    "snes_adapt_type": "basic",          # Instructs SNES to drive an adaptation step
    "snes_max_it": 50,
    "ksp_type": "cg",
    "pc_type": "mg",
    "pc_mg_levels": 1,
    "mg_levels_ksp_type": "chebyshev",
    "mg_levels_pc_type": "sor"
}

prob = NonlinearVariationalProblem(F, u)
solver = NonlinearVariationalSolver(prob, solver_parameters=params)

# Attach the marking callback directly to the Firedrake solver wrapper
solver.set_marking_callback(my_marking_strategy)

# 6. Solve. Internal loops execute your marking callback and produce PETSc DMLabels automatically
solver.solve()

# 7. Extract the solution on the dynamically refined finest mesh layer
u_fine = solver.get_adapted_solution()

print(f"Solver adapted grid successfully over {len(hierarchy)} levels.")
File("refined_solution.pvd").write(u_fine)

---

3. Code Implementation Blueprint

3.1 Extension of MeshHierarchy (firedrake/mesh.py)

Tracks the hierarchy depth and routes topological modifications guided by localized cell-marking indicators.

from firedrake.mesh import MeshHierarchy

class AdaptiveMeshHierarchy(MeshHierarchy):
    def __init__(self, base_mesh, **kwargs):
        super().__init__(base_mesh, refinements=0, **kwargs)
        
        self.netgen_geometry = None
        if hasattr(base_mesh, "_netgen_mesh") and hasattr(base_mesh._netgen_mesh, "geom"):
            self.netgen_geometry = base_mesh._netgen_mesh.geom
        elif hasattr(base_mesh, "_netgen_geom"):
            self.netgen_geometry = base_mesh._netgen_geom

    def append_marked_refined_level(self, dmlabel=None):
        """Topologically refines the current fine mesh using a PETSc DMLabel indicator."""
        current_fine = self[-1]
        
        # 1. Perform nested topological refinement guided by the DMLabel flags
        new_fine = current_fine.refine(marker=dmlabel)
        
        # 2. If a Netgen geometry context was extracted, snap boundary nodes outward
        if self.netgen_geometry is not None:
            from firedrake.netgen_projection import project_mesh_to_netgen_geometry
            project_mesh_to_netgen_geometry(new_fine, self.netgen_geometry)
            
            degree = current_fine.coordinates.function_space().ufl_element().degree()
            if degree > 1:
                new_fine = new_fine.curve_field(degree)
        
        self._meshes.append(new_fine)
        self._parent_child_relationships.append((current_fine, new_fine))
        return new_fine

3.2 Automated DMLabel Construction Hook (firedrake/dmhooks.py)

Converts high-level Firedrake element arrays into native low-level PETSc tracking structures mapped to DMPlex points.

import petsc4py.PETSc as PETSc

def build_petsc_dmlabel_from_markers(plex, marker_function):
    """
    Takes a Firedrake DG0 marking function (1 for refine, 0 for keep) 
    and returns a populated PETSc DMLabel configured for DMPlex point ranges.
    """
    # 1. Instantiate or clear an adaptation target label inside the mesh DMPlex
    label_name = "adapt"
    if plex.hasLabel(label_name):
        plex.removeLabel(label_name)
    
    dm_label = plex.createLabel(label_name)
    
    # 2. Extract cell stratum from the mesh plex
    cell_start, cell_end = plex.getHeightStratum(0)
    
    # 3. Read the underlying array data from the Firedrake field container
    marker_data = marker_function.dat.data_ro
    
    # 4. Map element indices sequentially to DMPlex cell points
    for i, flag in enumerate(marker_data):
        if flag == 1:
            plex_cell_point = cell_start + i
            dm_label.setValue(plex_cell_point, PETSc.DMAdaptFlag.REFINE)
            
    return dm_label

def configure_adaptive_dm_hooks(dm, hierarchy):
    """Binds the Firedrake hierarchy append action to PETSc's DM refinement step."""
    def adaptive_refine_callback(fine_dm, comm):
        # Read the generated adaptation tracking label from the active context
        petsc_label = fine_dm.getLabel("adapt") if fine_dm.hasLabel("adapt") else None
        new_mesh = hierarchy.append_marked_refined_level(dmlabel=petsc_label)
        
        new_dm = new_mesh.topology_dm
        new_dm.set_appctx(new_mesh)
        return new_dm

    dm.setRefine(adaptive_refine_callback)

3.3 Internal Solver Loop & Callback Bridge (firedrake/variational_solver.py)

Updates NonlinearVariationalSolver to execute the user's marking logic, construct the required DMLabel, and unwrap the final result field.

# Inside firedrake/variational_solver.py

class NonlinearVariationalSolver(object):
    def __init__(self, problem, solver_parameters=None, **kwargs):
        super().__init__(problem, solver_parameters=solver_parameters, **kwargs)
        self._marking_callback = None

    def set_marking_callback(self, callback):
        """Registers a user-defined callback that generates a DG0 marking function."""
        self._marking_callback = callback

    def _run_marking_pipeline(self, hierarchy):
        """Executes the user callback and converts the results to a PETSc DMLabel."""
        if self._marking_callback is None:
            return None
            
        current_fine_mesh = hierarchy[-1]
        
        # Pull the active solver solution wrapped as a function on the current finest layer
        # (For level 0, this is self._problem.u; for upper levels it maps to the active SNES vector)
        if len(hierarchy) == 1:
            current_u = self._problem.u
        else:
            # Generate a temporary wrapper mapping the current fine mesh
            V_current = self._problem.u.function_space().clone(mesh=current_fine_mesh)
            current_u = Function(V_current)
            with current_u.dat.vec as v_fd:
                self.snes.getSolution().copy(v_fd)
        
        # 1. Call user-defined logic to compute element flags
        markers_dg0 = self._marking_callback(self, current_u)
        
        # 2. Automate mapping onto the active DMPlex layout
        from firedrake.dmhooks import build_petsc_dmlabel_from_markers
        build_petsc_dmlabel_from_markers(current_fine_mesh.topology_dm, markers_dg0)

    def solve(self, bounds=None):
        mesh = self._problem.u.function_space().mesh()
        hierarchy = getattr(mesh, "hierarchy", None)
        
        if hierarchy and hasattr(hierarchy, "append_marked_refined_level"):
            from firedrake.dmhooks import configure_adaptive_dm_hooks
            configure_adaptive_dm_hooks(mesh.topology_dm, hierarchy)
            
            # Pre-adapt intercept: Generate DMLabel before PETSc starts a refinement block
            self._run_marking_pipeline(hierarchy)
            
            # Post-adaptation hook: Update Multigrid structures if depth increases
            self.snes.setUpdate(lambda snes, it: update_pcmg_hierarchy(snes, hierarchy))

        super().solve(bounds=bounds)

    def get_adapted_solution(self):
        """Reconstructs the function space layout on the finest level and populates it."""
        base_space = self._problem.u.function_space()
        mesh = base_space.mesh()
        hierarchy = getattr(mesh, "hierarchy", None)
        
        if not hierarchy or len(hierarchy) == 1:
            return self._problem.u
            
        finest_mesh = hierarchy[-1]
        fine_space = base_space.clone(mesh=finest_mesh)
        fine_u = Function(fine_space)
        
        with fine_u.dat.vec as v_fd:
            self.snes.getSolution().copy(v_fd)
            
        return fine_u

[pbrubeck]

Architectural Design: Goal-Oriented Dual Weighted Residual (DWR) Adaptivity via Bubble/Cone Localization

Disclaimer: This implementation plan was generated collaboratively by a Firedrake core developer interacting with Gemini. It builds upon and supersedes the framework from PR #4893 (Goal-Adaptive Nonlinear Variational Solver) by unifying the error estimation loops directly inside a streamlined, automated marking callback wrapper.

This framework integrates Goal-Oriented Dual Weighted Residual (DWR) error estimation into the SNESAdapt workflow following the a posteriori localization framework of Rognes and Logg (2012). The marking callback handles the higher-order primal problem along with the low- and high-order dual (adjoint) problems. To maximize efficiency, we expose front-end support for solving the adjoint system directly through PETSc via solver.snes.ksp.solveTranspose.

---

1. Architecture & Data Flow

[Primal Low-Order u_h] ───────────────────────> [solveTranspose -> z_h]
          │                                               │
          ▼                                               ▼
[High-Order Enriched Space V_H]                           │
          │                                               │
          ▼                                               ▼
[Solve Primal High-Order u_H] ────────────────> [solveTranspose -> z_H]
          │                                               │
          └───────────────────────┬───────────────────────┘
                                  ▼
              [Solve Localized DG0 Variational Problem]
                                  ▼
                    [Build and Return DG0 DMLabel]

1. Low-Order Adjoint Solver: The derivative of the Goal functional with respect to $u_h$ is evaluated and solved via solve_adjoint_transpose to yield $z_h$.
2. Space Enrichment: The low-order element configuration is expanded to allocate a structurally sound, higher-order function space $V_H$.
3. High-Order Enriched Primal/Adjoint Solver: The forward problem is evaluated on $V_H$ to resolve $u_H$, and its goal derivative is solved to yield $z_H$. The dual error weights are computed as $e^* = z_H - I_H z_h$.
4. Variational Localization (PR GoalAdaptiveNonlinearVariationalSolver #4893): The indicator field $\eta \in DG0$ is resolved by solving an invertibly decoupled mass matrix projection problem. The absolute value of the weak residual contributions evaluated against localized bubble and cone support expansions acts as the RHS of this variational system.

---

2. User-Facing API

The user API abstracts the underlying matrix transpositions, dual space mappings, and localization mechanics into a single setup call.

from firedrake import *

# 1. Mesh and Hierarchy Setup
geo = SplineGeometry()
geo.AddCircle((0, 0), r=1.0)
coarse_mesh = Mesh(geo.GenerateMesh(maxh=0.5))
hierarchy = AdaptiveMeshHierarchy(coarse_mesh)

# 2. Primal Problem Definition (Low-order Space V)
V = FunctionSpace(hierarchy[0], "CG", 1)
u = Function(V)
v = TestFunction(V)
F = dot(grad(u), grad(u))*dot(grad(u), grad(v))*dx - f*v*dx  # p-Laplacian example

# 3. Define the Goal Functional (Quantity of Interest)
J = u * dx

params = {
    "snes_type": "newtonls",
    "snes_adapt_type": "basic",
    "pc_type": "mg"
}

prob = NonlinearVariationalProblem(F, u)
solver = NonlinearVariationalSolver(prob, solver_parameters=params)

# 4. Attach the Rognes-Logg DWR marking callback via goal and enrichment configurations
solver.set_dwr_marking_strategy(goal=J, extra_degree=1, fraction=0.15)

# 5. Execute. Adjoint solves happen implicitly via solveTranspose behind the scenes
solver.solve()

u_fine = solver.get_adapted_solution()

---

3. Implementation Details

3.1 Front-End Transpose Solve Implementation (firedrake/variational_solver.py)

# Add to NonlinearVariationalSolver class inside firedrake/variational_solver.py

def solve_adjoint_transpose(self, cofunction_rhs, adjoint_output_function):
    """
    Solves the adjoint linear system utilizing the transposed Jacobian operator
    directly out of the active SNES linear solver context.
    """
    snes = self.snes
    ksp = snes.getKSP()
    
    with cofunction_rhs.dat.vec_ro as b_petsc, adjoint_output_function.dat.vec as x_petsc:
        ksp.solveTranspose(b_petsc, x_petsc)
        
    return adjoint_output_function

3.2 The Unified DWR Bubble/Cone Pipeline (firedrake/dwr_edge.py)

from firedrake import *
from firedrake.preconditioners.pmg import PMGBase
import numpy as np

class DWRMarkingStrategy(object):
    def __init__(self, goal_expr, extra_degree=1, fraction=0.10):
        self.goal_expr = goal_expr
        self.extra_degree = extra_degree
        self.fraction = fraction

    def __call__(self, solver, current_u):
        mesh = current_u.function_space().mesh()
        V_low = current_u.function_space()
        
        # --- Step 1: Solve Low-Order Adjoint z_h ---
        z_h = Function(V_low)
        active_goal_low = replace(self.goal_expr, {solver._problem.u: current_u})
        rhs_cofunction_low = assemble(derivative(active_goal_low, current_u))
        solver.solve_adjoint_transpose(rhs_cofunction_low, z_h)
        
        # --- Step 2: Reconstruct Enriched Space V_H via PMGBase ---
        low_element = V_low.ufl_element()
        base_degree = PMGBase.max_degree(low_element)
        target_degree = base_degree + self.extra_degree
        
        high_element = PMGBase.reconstruct_degree(low_element, target_degree)
        V_high = FunctionSpace(mesh, high_element)
        
        u_H = Function(V_high)
        v_H = TestFunction(V_high)
        prolong(current_u, u_H)
        
        F_high_expr = replace(solver._problem.F, {solver._problem.u: u_H, TestFunction(V_low): v_H})
        prob_high = NonlinearVariationalProblem(F_high_expr, u_H)
        solver_high = NonlinearVariationalSolver(prob_high, solver_parameters={"snes_type": "newtonls"})
        solver_high.solve()
        
        # --- Step 3: Solve High-Order Adjoint z_H ---
        z_H = Function(V_high)
        active_goal_high = replace(self.goal_expr, {solver._problem.u: u_H})
        rhs_cofunction_high = assemble(derivative(active_goal_high, u_H))
        solver_high.solve_adjoint_transpose(rhs_cofunction_high, z_H)
        
        # --- Step 4: Verbatim PR #4893 DWR Variational Localization Pipeline ---
        DG0 = FunctionSpace(mesh, "DG", 0)
        indicators = Function(DG0)
        eta = TrialFunction(DG0)
        w_dg0 = TestFunction(DG0)
        
        z_h_prolonged = Function(V_high)
        prolong(z_h, z_h_prolonged)
        error_weight = Function(V_high).assign(z_H - z_h_prolonged)
        
        v_bubble = TestFunction(V_high)
        
        # 1. Interior element residual localization using element bubbles
        F_interior = replace(solver._problem.F, {solver._problem.u: u_H, TestFunction(V_low): v_bubble * error_weight})
        
        # 2. Interface jump localization using facet cones
        F_facet_inner = replace(solver._problem.F, {solver._problem.u: u_H, TestFunction(V_low): v_bubble * error_weight})
        
        # Formulate the localized error linear system following the implementation in PR #4893
        a_dg0 = eta * w_dg0 * dx
        L_dg0 = abs(F_interior) * w_dg0 * dx + abs(F_facet_inner)('internal') * w_dg0('+') * dS
        
        # Solve the decoupled mass problem locally (since DG0 features a block-diagonal system)
        prob_indicator = LinearVariationalProblem(a_dg0, L_dg0, indicators)
        solver_indicator = LinearVariationalSolver(prob_indicator, solver_parameters={
            "ksp_type": "preonly",
            "pc_type": "jacobi"
        })
        solver_indicator.solve()
        
        # --- Step 5: Mark Cells via Thresholding Fraction ---
        markers = Function(DG0)
        cutoff = np.percentile(indicators.dat.data_ro, 100.0 * (1.0 - self.fraction))
        markers.dat.data[:] = np.where(indicators.dat.data_ro >= cutoff, 1, 0)
        
        return markers

3.3 Integration Hook in Solver Backend

# Add to class NonlinearVariationalSolver in firedrake/variational_solver.py

def set_dwr_marking_strategy(self, goal, extra_degree=1, fraction=0.10):
    """Enables and configures the automated Dual Weighted Residual goal-oriented pipeline."""
    from firedrake.dwr_edge import DWRMarkingStrategy
    strategy = DWRMarkingStrategy(goal, extra_degree, fraction)
    self.set_marking_callback(strategy)

---

4. Engineering Advantages

• Front-End solveTranspose Efficiency: Reuses the internal Jacobian storage structures directly from the primal solver line step, cutting assembly costs in half and avoiding parallel solver state duplication.
• Transient Memory Control: High-order field matrices are scoped strictly within the marking runtime callback, freeing GPU/CPU memory buffers before the DMLabel is committed to the main solver line.
• Robust Projection Mapping: Computing indicators via a local $DG0$ variational problem guarantees that the computed quantities correctly scale with the volume/area weight of elements, maintaining consistency without mesh-dependent metric drift.

[connorjward]

This is so much detail... I don't really want to read this. Can't you write a more concise summary (not AI) about what you want to do? We want to evaluate the high level ideas and this is a lot of effort to read and respond to.

The plan you have here is basically an implementation detail that we don't need to see.

[pbrubeck]

This is so much detail... I don't really want to read this. Can't you write a more concise summary (not AI) about what you want to do? We want to evaluate the high level ideas and this is a lot of effort to read and respond to.

The plan you have here is basically an implementation detail that we don't need to see.

I'm still iterating on this. I specifically asked it to show the user-facing API for both steps (separate PRs). That should give you a good summary

[pbrubeck]

Also, during the PETSc meeting we discussed this at the top level, and this approach is the one that we came up with.

[pefarrell]

In the DWR user-facing API: where does the tolerance on the goal functional error go? Or does it only do one refinement, and it's up to an outer loop to keep refining until the error tolerance is met? I confess I prefer the API in #4893 .