AV formulation with fieldsplit Hypre AMS

Hi,

I am trying to solve the low frequency magnetoquasistatic equation with AV formulation with Coulomb gauge. I use Nedelec1 for A, CG1 for the scalar potential V and the lagrange multipliers used to force the A gauge. In steady state, with monolitic mumps LU, I get a reasonable solution, but I would like to use an iterative solver to limit memory. I am trying a fieldsplit with Schur, using Hypre AMS for the Nedelec elements providing the discrete gradient and edge vectors, while I use mumps lu for the remaining part, which should handle the singular lagrange multiplier block. I always get nan on the AMS iteration 0. Am I missing something or is there something wrong in the solver setup? Here is the minimal example, with the mesh generation script :

import gmsh

gmsh.initialize()
gmsh.model.add("cube_with_two_tori")

occ = gmsh.model.occ

# -------------------------
# Parameters
# -------------------------
L = 1.0
cube = occ.addBox(-L/2, -L/2, 0, L, L, L)

cx = 0.0
cy = 0.0

# Torus centers and radii
z1 = 0.35
z2 = 0.65

R1 = 0.18   # major radius
R2 = 0.30

r1 = 0.05   # tube radius
r2 = 0.05   # tube radius

torus1 = occ.addTorus(cx, cy, z1, R1, r1)
torus2 = occ.addTorus(cx, cy, z2, R2, r2)
bill = occ.addCylinder(0,0,z1,0,0,0.5*(z2-z1),R1*0.5)


# Make all volumes conformal
occ.fragment([(3, cube)], [(3, torus1), (3, torus2), (3, bill)])
occ.synchronize()
vols = gmsh.model.getEntities(3)

air_vols = []
loop1_vols = []
loop2_vols = []

for dim, tag in vols:
    x, y, z = occ.getCenterOfMass(dim, tag)
    rho = ((x - cx) ** 2 + (y - cy) ** 2) ** 0.5

    d1 = ((rho - R1) ** 2 + (z - z1) ** 2) ** 0.5
    d2 = ((rho - R2) ** 2 + (z - z2) ** 2) ** 0.5

    if d1 < r1 * 0.7:
        loop1_vols.append(tag)
    elif d2 < r2 * 0.7:
        loop2_vols.append(tag)
    else:
        air_vols.append(tag)

# Physical groups
ct = 1
for dim, tag in gmsh.model.getEntities(dim=3):
    gmsh.model.addPhysicalGroup(3, [tag], ct)
    ct += 1

ct = 100
for dim, tag in gmsh.model.getEntities(dim=2):
    gmsh.model.addPhysicalGroup(2, [tag], ct)
    ct += 1

ct = 200
for dim, tag in gmsh.model.getEntities(dim=1):
    gmsh.model.addPhysicalGroup(1, [tag], ct)
    ct += 1

ct = 300
for dim, tag in gmsh.model.getEntities(dim=0):
    gmsh.model.addPhysicalGroup(0, [tag], ct)
    ct += 1
occ.synchronize()
# Mesh options
gmsh.option.setNumber("Mesh.CharacteristicLengthMin", 0.03)
gmsh.option.setNumber("Mesh.CharacteristicLengthMax", 0.07)

gmsh.model.mesh.generate(3)
gmsh.model.mesh.optimize()
gmsh.write("cube_with_two_tori.msh")
gmsh.finalize()
from __future__ import annotations

import math
import time

import basix.ufl
import numpy as np
import ufl
from dolfinx import fem
from dolfinx.fem.petsc import (
    apply_lifting,
    assemble_matrix,
    assemble_vector,
    discrete_gradient,
    set_bc,
)
from dolfinx.io.gmsh import read_from_msh
from mpi4py import MPI
from petsc4py import PETSc


def build_piecewise_sigma(mesh, cell_tags):
    """Create DG0 conductivity field matching compare_solvers materials.

    - air tags (1, 3, 4): sigma = 10 S/m
    - copper tag (2): sigma = 5e7 S/m
    """
    q0 = fem.functionspace(mesh, ("DG", 0))
    sigma = fem.Function(q0)
    sigma.x.array[:] = 0.0

    air_value = PETSc.ScalarType(10.0)
    copper_value = PETSc.ScalarType(5.0e7)

    for tag in (1, 3, 4):
        cells = cell_tags.find(tag)
        sigma.x.array[cells] = air_value

    copper_cells = cell_tags.find(2)
    sigma.x.array[copper_cells] = copper_value

    sigma.x.scatter_forward()
    return sigma


def build_current_density_function(v_hcurl):
    """Interpolate source J = 1e5 * (-sin(atan2(y,x)), cos(atan2(y,x)), 0)."""
    j_fun = fem.Function(v_hcurl)

    def j_expr(x):
        theta = np.arctan2(x[1], x[0])
        values = np.zeros((3, x.shape[1]), dtype=PETSc.ScalarType)
        values[0, :] = -1.0e5 * np.sin(theta)
        values[1, :] = 1.0e5 * np.cos(theta)
        return values

    j_fun.interpolate(j_expr)
    j_fun.x.scatter_forward()
    return j_fun


def assemble_system(mesh, cell_tags, point_tags):
    """Assemble mixed AV-lambda linear system and return matrix, rhs, metadata.

    Problem: Magnetostatic AV-lambda formulation with Coulomb gauge constraint.

    This is a saddle-point problem arising from:
    - curl(1/mu * curl(A)) = J (current source) + grad(V) (electric coupling)
    - div(sigma * grad(V)) = -div(sigma * grad(lambda))  (Coulomb gauge constraint)
    - div(A) = lambda  (Lagrange multiplier constraint)

    Mixed FE spaces:
    - A in H(curl) via Nedelec elements (order 1)
    - V, lambda in H^1 via Lagrange elements (order 1)

    Materials: air (tags 1,3,4), copper (tag 2) with different conductivities.
    Source: current density J on volume tag 1 only.
    BCs: V = lambda = 0 on point tag 300.
    """
    # Mixed function space: (A, V, lambda)
    a_el = basix.ufl.element("N1curl", mesh.basix_cell(), 1)
    v_el = basix.ufl.element("Lagrange", mesh.basix_cell(), 1)
    l_el = basix.ufl.element("Lagrange", mesh.basix_cell(), 1)
    w_space = fem.functionspace(mesh, basix.ufl.mixed_element([a_el, v_el, l_el]))

    # Unknown and test functions
    (a_trial, v_trial, l_trial) = ufl.TrialFunctions(w_space)
    (w_test, q_test, r_test) = ufl.TestFunctions(w_space)

    # Collapse maps needed for fieldsplit IS definitions
    w0, map_a = w_space.sub(0).collapse()
    _, map_v = w_space.sub(1).collapse()
    _, map_l = w_space.sub(2).collapse()

    # Material coefficients
    mu0 = 4.0e-7 * math.pi
    nu = PETSc.ScalarType(1.0 / mu0)
    sigma = build_piecewise_sigma(mesh, cell_tags)

    # Domain measures with subdomain tags
    dx = ufl.Measure("dx", domain=mesh, subdomain_data=cell_tags)

    # AV-lambda mixed formulation (magnetostatic with Coulomb gauge constraint):
    # - A: magnetic vector potential (N1curl / H(curl), test fn w)
    # - V: electric scalar potential (CG1 / H^1, test fn q)
    # - lambda: Lagrange multiplier enforcing gauge (CG1 / H^1, test fn r)
    #
    # Weak form breakdown:
    # 1. nu * <curl(A), curl(w)>           : curl-curl bilinear form (stiffness from reluctivity)
    # 2. <sigma·grad(V), w + grad(q)>      : conductivity coupling (electric & constraint mixing)
    # 3. <grad(lambda), w>                 : Lagrange multiplier to A (enforces constraint on A)
    # 4. <A, grad(r)>                      : A to Lagrange multiplier (dual constraint coupling)
    #
    # This gives the saddle-point system:
    #   [ K_A      0        G^T   ] [ A      ]   [ J ]
    #   [ 0        K_V      D     ] [ V      ] = [ 0 ]
    #   [ G        D^T      0     ] [ lambda ]   [ 0 ]
    # where K_A = curl-curl, K_V = sigma·grad, G = grad, D = mixed conductivity term
    a_form = fem.form(
        (
            # Term 1: Magnetic energy (curl-curl)
            nu * ufl.inner(ufl.curl(a_trial), ufl.curl(w_test))
            # Term 2: Electric conductivity (appears on both A and V rows via Coulomb constraint)
            + ufl.inner(sigma * ufl.grad(v_trial), w_test + ufl.grad(q_test))
            # Term 3: Lagrange multiplier constraint on A (gradient of lambda couples to w)
            + ufl.inner(ufl.grad(l_trial), w_test)
            # Term 4: Constraint dual form (A couples to gradient of r via Lagrange)
            + ufl.inner(a_trial, ufl.grad(r_test))
        )
        * dx,
    )

    # Current source only on volume tag 1, as in compare_solvers
    j = build_current_density_function(w0)
    l_form = fem.form(ufl.inner(j, w_test) * dx(1))

    # Point constraints on V and lambda at point tag 300
    point_entities = point_tags.find(300)
    dofs_v = fem.locate_dofs_topological(w_space.sub(1), 0, point_entities)
    dofs_l = fem.locate_dofs_topological(w_space.sub(2), 0, point_entities)
    bc_v = fem.dirichletbc(PETSc.ScalarType(0.0), dofs_v, w_space.sub(1))
    bc_l = fem.dirichletbc(PETSc.ScalarType(0.0), dofs_l, w_space.sub(2))
    bcs = [bc_v, bc_l]

    # Assemble linear system
    a_mat = assemble_matrix(a_form, bcs=bcs)
    a_mat.assemble()

    b_vec = assemble_vector(l_form)
    apply_lifting(b_vec, [a_form], [bcs])
    b_vec.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
    set_bc(b_vec, bcs)

    metadata = {
        "W": w_space,
        "W0": w0,
        "map_a": map_a.astype(np.int32),
        "map_v": map_v.astype(np.int32),
        "map_l": map_l.astype(np.int32),
    }
    return a_mat, b_vec, metadata


def solve_monolithic_lu(a_mat, b_vec):
    """Reference solve with preonly+LU(MUMPS)."""
    x = a_mat.createVecRight()

    ksp = PETSc.KSP().create(a_mat.comm)
    ksp.setOperators(a_mat)
    ksp.setType("preonly")
    pc = ksp.getPC()
    pc.setType("lu")
    pc.setFactorSolverType("mumps")
    # ICNTL(24)=1: detect null pivots (handles gauge singularity in curl-curl block)
    pc.setFactorSetUpSolverType()
    pc.getFactorMatrix().setMumpsIcntl(24, 1)

    t0 = time.perf_counter()
    ksp.solve(b_vec, x)
    elapsed = time.perf_counter() - t0

    return {
        "x": x,
        "ksp": ksp,
        "time": elapsed,
        "reason": ksp.getConvergedReason(),
        "its": ksp.getIterationNumber(),
    }


def solve_fieldsplit_ams(a_mat, b_vec, meta):
    """FGMRES + FieldSplit Schur with AMS on A and LU on (V,lambda).

    This solver uses a block preconditioner:
    - Outer: FGMRES (flexible restarted GMRES, allows varying preconditioner)
    - Outer PC: Fieldsplit with Schur complement factorization (UPPER form)
      * Block (A): HYPRE AMS (Auxiliary Space Maxwell Solver for H(curl) problems)
        - Uses discrete gradient (curl operator) and edge constant vectors
        - Designed for curl-curl systems in H(curl) spaces
      * Block (V,lambda): MUMPS LU (direct solver on Schur complement)


    """
    x = a_mat.createVecRight()

    ksp = PETSc.KSP().create(a_mat.comm)
    ksp.setOperators(a_mat)
    ksp.setType("fgmres")
    ksp.setTolerances(rtol=1e-8, atol=1e-50, divtol=1e3, max_it=10000)

    pc = ksp.getPC()
    pc.setType("fieldsplit")
    pc.setFieldSplitType(PETSc.PC.CompositeType.SCHUR)
    pc.setFieldSplitSchurFactType(PETSc.PC.SchurFactType.UPPER)
    pc.setFieldSplitSchurPreType(PETSc.PC.SchurPreType.A11)

    # A vs (V, lambda) split
    # Index set for A-block: all DOFs from Nedelec space (field 0)
    # Index set for (V,lambda)-block: all DOFs from both CG1 spaces (fields 1 & 2)
    is_a = PETSc.IS().createGeneral(meta["map_a"], comm=a_mat.comm)
    map_vl = np.concatenate((meta["map_v"], meta["map_l"]))
    is_vl = PETSc.IS().createGeneral(map_vl.astype(np.int32), comm=a_mat.comm)
    pc.setFieldSplitIS(("field_0", is_a), ("field_1", is_vl))

    # Materialize sub-KSPs after split setup
    ksp.setUp()
    sub_ksps = pc.getFieldSplitSubKSP()
    ksp_a = sub_ksps[0]
    ksp_vl = sub_ksps[1]

    # A-block: preonly + HYPRE AMS
    # AMS (Auxiliary Space Maxwell Solver) is designed for H(curl) discretizations.
    ksp_a.setType("preonly")
    pc_a = ksp_a.getPC()
    pc_a.setType("hypre")
    pc_a.setHYPREType("ams")
    opts = PETSc.Options()
    v_h1 = fem.functionspace(meta["W"].mesh, ("Lagrange", 1))
    g_mat = discrete_gradient(v_h1, meta["W0"])
    g_mat.assemble()

    # Edge constant vectors: coordinate basis vectors in H(curl)
    cvec0 = fem.Function(meta["W0"])
    cvec1 = fem.Function(meta["W0"])
    cvec2 = fem.Function(meta["W0"])
    cvec0.interpolate(
        lambda x: np.vstack(
            (np.ones_like(x[0]), np.zeros_like(x[1]), np.zeros_like(x[2])),
        ),
    )
    cvec1.interpolate(
        lambda x: np.vstack(
            (np.zeros_like(x[0]), np.ones_like(x[1]), np.zeros_like(x[2])),
        ),
    )
    cvec2.interpolate(
        lambda x: np.vstack(
            (np.zeros_like(x[0]), np.zeros_like(x[1]), np.ones_like(x[2])),
        ),
    )
    cvec0.x.scatter_forward()
    cvec1.x.scatter_forward()
    cvec2.x.scatter_forward()

    # Pass auxiliary data to AMS preconditioner
    pc_a.setHYPREDiscreteGradient(g_mat)
    pc_a.setHYPRESetEdgeConstantVectors(
        cvec0.x.petsc_vec,
        cvec1.x.petsc_vec,
        cvec2.x.petsc_vec,
    )
    # Apply options and trigger HYPRE's internal AMS setup
    ksp_a.setFromOptions()
    # Explicit PC.setUp() triggers HYPRE's internal AMS setup NOW
    # (without this, AMS setup may be deferred and skipped in fieldsplit context)
    pc_a.setUp()

    # (V, lambda)-block: preonly + LU(MUMPS)
    ksp_vl.setType("preonly")
    pc_vl = ksp_vl.getPC()
    pc_vl.setType("lu")
    pc_vl.setFactorSolverType("mumps")
    # ICNTL(24)=1: detect null pivots (Schur complement may inherit singularity)
    pc_vl.setFactorSetUpSolverType()
    pc_vl.getFactorMatrix().setMumpsIcntl(24, 1)

    t0 = time.perf_counter()
    ksp.solve(b_vec, x)
    elapsed = time.perf_counter() - t0

    return {
        "x": x,
        "ksp": ksp,
        "time": elapsed,
        "reason": ksp.getConvergedReason(),
        "its": ksp.getIterationNumber(),
    }


def main():
    comm = MPI.COMM_WORLD
    rank = comm.rank

    mesh_file = "examples/data/cube_with_two_tori.msh"
    if rank == 0:
        print("=" * 80)
        print("Minimal DOLFINx AV-lambda test (small mesh)")
        print("=" * 80)
        print(f"Loading mesh: {mesh_file}")

    # read_from_msh returns (mesh, cell_tags, facet_tags, edge_tags, point_tags, ...)
    loaded = read_from_msh(mesh_file, comm=comm, rank=0, gdim=3)
    mesh = loaded[0]
    cell_tags = loaded[1]
    point_tags = loaded[4]

    a_mat, b_vec, meta = assemble_system(mesh, cell_tags, point_tags)

    if rank == 0:
        print("System assembled")
        print(f"  Global size: {a_mat.getSize()[0]}")
        print(f"  A dofs:      {meta['W0'].dofmap.index_map.size_global}")
        print(f"  V dofs:      {len(meta['map_v'])}")
        print(f"  Lambda dofs: {len(meta['map_l'])}")

    # Baseline
    if rank == 0:
        print("\n[1/2] Monolithic LU (MUMPS)")
    lu_res = solve_monolithic_lu(a_mat, b_vec)
    if rank == 0:
        print(
            f"  reason={lu_res['reason']}, its={lu_res['its']}, "
            f"time={lu_res['time']:.3f}s",
        )

    # AMS fieldsplit
    if rank == 0:
        print("\n[2/2] FieldSplit Schur + AMS(A) + LU(V,lambda)")
    ams_res = solve_fieldsplit_ams(a_mat, b_vec, meta)
    if rank == 0:
        print(
            f"  reason={ams_res['reason']}, its={ams_res['its']}, "
            f"time={ams_res['time']:.3f}s",
        )

    if rank == 0:
        print("\nDone.")


if __name__ == "__main__":
    main()

Thank you,
Mattia

Following the thread https://fenicsproject.discourse.group/t/solving-a-rank-deficient-curl-curl-matrix/18423, the nan i get is due to the singular curl-curl block, adding a small regularization on the A part solved the issue.

I am told that the hypre ams preconditioner (with a right discrete gradient) can handle the null space of the curl curl operator on the Nedelec element of the first kind. But when I tried, it didn’t work out as it is supposed to.

Could you elaborate?
There is an implementation at GitHub - jpdean/maxwell · GitHub that works well for Nedelec order k elements.

Just putting this resource here for future reference on this topic. AMS and ADS preconditioning implementations for both Raviart-Thomas and BDM elements, for 2D and 3D, and for k=1 and higher order can be found here:

https://gitlab.com/fenicsx4flow/projects/preconditioning/-/blob/main/vector_laplace/advection_diffusion_reaction_Hdiv.py?ref_type=heads

When I worked on it last, there appeared to be a bug in the petsc4py handle that didn’t permit the correct passing of arguments for the 3D higher order case (discussed from lines 292 onward), so that works a little suboptimally. Maybe that has been patched since.