Fieldsplit: update MatNest Jacobian and support bcs on AssembledMatrix

firedrake/__init__.py

@@ -3,7 +3,7 @@
 # the specific version, here we are more permissive. This is to catch the
 # case where users don't update their PETSc for a really long time or
 # accidentally install a too-new release that isn't yet supported.
-PETSC_SUPPORTED_VERSIONS = ">=3.23.0"
+PETSC_SUPPORTED_VERSIONS = ">=3.24.2"
 
 
 def init_petsc():

firedrake/dmhooks.py

@@ -363,43 +363,36 @@ def create_subdm(dm, fields, *args, **kwargs):
     :arg fields: The fields in the new sub-DM.
     """
     W = get_function_space(dm)
-    ctx = get_appctx(dm)
-    coarsen = get_ctx_coarsener(dm)
-    parent = get_parent(dm)
     if len(fields) == 1:
         # Subspace is just a single FunctionSpace.
         idx, = fields
-        subdm = W[idx].dm
+        subspace = W[idx]
         iset = W._ises[idx]
-        add_hook(parent, setup=partial(push_parent, subdm, parent), teardown=partial(pop_parent, subdm, parent),
-                 call_setup=True)
-
-        if ctx is not None:
-            ctx, = ctx.split([(idx, )])
-            add_hook(parent, setup=partial(push_appctx, subdm, ctx), teardown=partial(pop_appctx, subdm, ctx),
-                     call_setup=True)
-            add_hook(parent, setup=partial(push_ctx_coarsener, subdm, coarsen), teardown=partial(pop_ctx_coarsener, subdm, coarsen),
-                     call_setup=True)
-        return iset, subdm
     else:
         # Need to build an MFS for the subspace
         subspace = firedrake.MixedFunctionSpace([W[f] for f in fields])
 
-        add_hook(parent, setup=partial(push_parent, subspace.dm, parent), teardown=partial(pop_parent, subspace.dm, parent),
-                 call_setup=True)
         # Index set mapping from W into subspace.
-        iset = PETSc.IS().createGeneral(numpy.concatenate([W._ises[f].indices
-                                                           for f in fields]),
+        iset = PETSc.IS().createGeneral(numpy.concatenate([W.dof_dset.field_ises[f].indices for f in fields]),
                                         comm=W._comm)
-        if ctx is not None:
-            ctx, = ctx.split([fields])
-            add_hook(parent, setup=partial(push_appctx, subspace.dm, ctx),
-                     teardown=partial(pop_appctx, subspace.dm, ctx),
-                     call_setup=True)
-            add_hook(parent, setup=partial(push_ctx_coarsener, subspace.dm, coarsen),
-                     teardown=partial(pop_ctx_coarsener, subspace.dm, coarsen),
-                     call_setup=True)
-        return iset, subspace.dm
+
+    subdm = subspace.dm
+    parent = get_parent(dm)
+    add_hook(parent, setup=partial(push_parent, subdm, parent),
+             teardown=partial(pop_parent, subdm, parent),
+             call_setup=True)
+
+    ctx = get_appctx(dm)
+    coarsen = get_ctx_coarsener(dm)
+    if ctx is not None:
+        ctx, = ctx.split([fields])
+        add_hook(parent, setup=partial(push_appctx, subdm, ctx),
+                 teardown=partial(pop_appctx, subdm, ctx),
+                 call_setup=True)
+        add_hook(parent, setup=partial(push_ctx_coarsener, subdm, coarsen),
+                 teardown=partial(pop_ctx_coarsener, subdm, coarsen),
+                 call_setup=True)
+    return iset, subdm
 
 
 @PETSc.Log.EventDecorator()

firedrake/formmanipulation.py

@@ -1,9 +1,8 @@
-
 import numpy
 import collections
 
 from ufl import as_tensor, as_vector, split
-from ufl.classes import Zero, FixedIndex, ListTensor, ZeroBaseForm
+from ufl.classes import Form, Zero, FixedIndex, ListTensor, ZeroBaseForm
 from ufl.algorithms.map_integrands import map_integrand_dags
 from ufl.algorithms import expand_derivatives
 from ufl.corealg.map_dag import MultiFunction, map_expr_dags
@@ -161,27 +160,58 @@ def cofunction(self, o):
             return Cofunction(W, val=MixedDat(o.dat[i] for i in indices))
 
     def matrix(self, o):
+        from firedrake.bcs import DirichletBC, EquationBCSplit
         ises = []
         args = []
+        argument_indices = []
         for a in o.arguments():
             V = a.function_space()
             iset = PETSc.IS()
             if a.number() in self.blocks:
+                fields = self.blocks[a.number()]
                 asplit = self._subspace_argument(a)
-                for f in self.blocks[a.number()]:
+                for f in fields:
                     fset = V.dof_dset.field_ises[f]
                     iset = iset.expand(fset)
             else:
+                fields = tuple(range(len(V)))
                 asplit = a
                 for fset in V.dof_dset.field_ises:
                     iset = iset.expand(fset)
 
             ises.append(iset)
             args.append(asplit)
+            argument_indices.append(fields)
+
+        if isinstance(o.a, Form):
+            form = self.split(o.a, argument_indices=argument_indices)
+            if isinstance(form, ZeroBaseForm):
+                return form
+        else:
+            form = None
 
         submat = o.petscmat.createSubMatrix(*ises)
-        bcs = ()
-        return AssembledMatrix(tuple(args), bcs, submat)
+
+        bcs = []
+        spaces = [a.function_space() for a in o.arguments()]
+        for bc in o.bcs:
+            W = bc.function_space()
+            while W.parent is not None:
+                W = W.parent
+
+            number = spaces.index(W)
+            V = args[number].function_space()
+            field = self.blocks[number]
+            if isinstance(bc, DirichletBC):
+                bc_temp = bc.reconstruct(field=field, V=V, g=bc.function_arg, use_split=True)
+            elif isinstance(bc, EquationBCSplit):
+                row_field, col_field = argument_indices
+                bc_temp = bc.reconstruct(field=field, V=V, row_field=row_field, col_field=col_field, use_split=True)
+                bc_temp = None
+            if bc_temp is not None:
+                bcs.append(bc_temp)
+
+        return AssembledMatrix(form or tuple(args), tuple(bcs), submat)
 
     def zero_base_form(self, o):
         return ZeroBaseForm(tuple(map(self, o.arguments())))

firedrake/matrix.py

@@ -170,9 +170,7 @@ def __init__(self, a, bcs, mat_type, *args, **kwargs):
         self.mat_type = mat_type
 
     def assemble(self):
-        raise NotImplementedError("API compatibility to apply bcs after 'assemble(a)'\
-                                  has been removed.  Use 'assemble(a, bcs=bcs)', which\
-                                  now returns an assembled matrix.")
+        self.M.assemble()
 
 
 class ImplicitMatrix(MatrixBase):
@@ -250,3 +248,9 @@ def __init__(self, a, bcs, petscmat, *args, **kwargs):
 
     def mat(self):
         return self.petscmat
+
+    def assemble(self):
+        # Bump petsc matrix state by assembling it.
+        # Ensures that if the matrix changed, the preconditioner is
+        # updated if necessary.
+        self.petscmat.assemble()

firedrake/solving_utils.py

@@ -353,9 +353,17 @@ def split(self, fields):
         splits = []
         problem = self._problem
         splitter = ExtractSubBlock()
+
+        Fbig = problem.F
+        # Reuse the submatrices if we are splitting a MatNest
+        Jbig = self._jac if self.mat_type == "nest" else problem.J
+        Jpbig = self._pjac if self.pmat_type == "nest" else problem.Jp
+        if Jpbig is Jbig:
+            Jpbig = None
+
         for field in fields:
-            F = splitter.split(problem.F, argument_indices=(field, ))
-            J = splitter.split(problem.J, argument_indices=(field, field))
+            F = splitter.split(Fbig, argument_indices=(field, ))
+            J = splitter.split(Jbig, argument_indices=(field, field))
             us = problem.u_restrict.subfunctions
             V = F.arguments()[0].function_space()
             # Exposition:
@@ -397,7 +405,6 @@ def split(self, fields):
             # solving for, and some spaces that have just become
             # coefficients in the new form.
             u = as_vector(vec)
-            J = replace(J, {problem.u_restrict: u})
             if problem.is_linear and isinstance(J, MatrixBase):
                 # The BC lifting term is action(MatrixBase, u).
                 # We cannot replace u with the split solution, as action expects a Function.
@@ -407,23 +414,35 @@ def split(self, fields):
                 F += problem.compute_bc_lifting(J, subu)
             else:
                 F = replace(F, {problem.u_restrict: u})
-            if problem.Jp is not None:
-                Jp = splitter.split(problem.Jp, argument_indices=(field, field))
+
+            J = replace(J, {problem.u_restrict: u})
+            if Jpbig is not None:
+                Jp = splitter.split(Jpbig, argument_indices=(field, field))
                 Jp = replace(Jp, {problem.u_restrict: u})
             else:
                 Jp = None
-            bcs = []
-            for bc in problem.bcs:
-                if isinstance(bc, DirichletBC):
-                    bc_temp = bc.reconstruct(field=field, V=V, g=bc.function_arg, sub_domain=bc.sub_domain)
-                elif isinstance(bc, EquationBC):
-                    bc_temp = bc.reconstruct(V, subu, u, field, problem.is_linear)
-                if bc_temp is not None:
-                    bcs.append(bc_temp)
+
+            if isinstance(J, MatrixBase) and J.has_bcs:
+                # The BCs of the problem are already encoded in the Jacobian
+                bcs = None
+            else:
+                bcs = []
+                for bc in problem.bcs:
+                    if isinstance(bc, DirichletBC):
+                        bc_temp = bc.reconstruct(field=field, V=V, g=bc.function_arg)
+                    elif isinstance(bc, EquationBC):
+                        bc_temp = bc.reconstruct(V, subu, u, field, problem.is_linear)
+                    if bc_temp is not None:
+                        bcs.append(bc_temp)
+
             new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=Jp, is_linear=problem.is_linear,
                                form_compiler_parameters=problem.form_compiler_parameters)
             new_problem._constant_jacobian = problem._constant_jacobian
-            splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type,
+            splits.append(type(self)(new_problem,
+                                     mat_type=self.mat_type,
+                                     pmat_type=self.pmat_type,
+                                     sub_mat_type=self.sub_mat_type,
+                                     sub_pmat_type=self.sub_pmat_type,
                                      appctx=self.appctx,
                                      transfer_manager=self.transfer_manager,
                                      pre_apply_bcs=self.pre_apply_bcs))

pyproject.toml

@@ -30,7 +30,7 @@ dependencies = [
   "loopy>2024.1",
   "numpy",
   "packaging",
-  "petsc4py==3.24.0",
+  "petsc4py==3.24.2",
   "petsctools",
   "pkgconfig",
   "progress",
@@ -150,7 +150,7 @@ requires = [
   "pkgconfig",
   "pybind11",
   "setuptools>=77.0.3",
-  "petsc4py==3.24.0",
+  "petsc4py==3.24.2",
   "rtree>=1.2",
 ]
 build-backend = "setuptools.build_meta"

scripts/firedrake-configure

@@ -39,7 +39,7 @@ ARCH_DEFAULT = FiredrakeArch.DEFAULT
 ARCH_COMPLEX = FiredrakeArch.COMPLEX
 
 
-SUPPORTED_PETSC_VERSION = "v3.24.0"
+SUPPORTED_PETSC_VERSION = "v3.24.2"
 
 
 def main():

setup.py

@@ -20,7 +20,7 @@
 
 
 # Ensure that the PETSc getting linked against is compatible
-petsctools.init(version_spec=">=3.23.0")
+petsctools.init(version_spec=">=3.24.2")
 import petsc4py
 
 

tests/firedrake/equation_bcs/test_equation_bcs.py

@@ -296,8 +296,8 @@ def test_EquationBC_mixedpoisson_matrix(eq_type):
     assert abs(math.log2(err[0][0]) - math.log2(err[1][0]) - (porder+1)) < 0.05
 
 
-def test_EquationBC_mixedpoisson_matrix_fieldsplit():
-    mat_type = "aij"
+@pytest.mark.parametrize("mat_type", ("aij", "nest"))
+def test_EquationBC_mixedpoisson_matrix_fieldsplit(mat_type):
     eq_type = "linear"
     porder = 0
     # Mixed poisson with EquationBCs

tests/firedrake/regression/test_nested_fieldsplit_solves.py

@@ -152,6 +152,49 @@ def test_nested_fieldsplit_solve_parallel(W, A, b, expect):
     assert norm(f) < 1e-11
 
 
+@pytest.mark.parametrize("mat_type,pmat_type", [("nest", "nest"), ("matfree", "nest"), ("matfree", "aij")])
+def test_nonlinear_fieldsplit(mat_type, pmat_type):
+    mesh = UnitIntervalMesh(1)
+    V = FunctionSpace(mesh, "DG", 0)
+    Z = V * V * V
+
+    u = Function(Z)
+    u0, u1, u2 = split(u)
+    v0, v1, v2 = TestFunctions(Z)
+
+    F = inner(u0, v0) * dx
+    F += inner(0.5*u1**2 + u1, v1) * dx
+    F += inner(u2, v2) * dx
+    u.subfunctions[1].assign(Constant(1))
+
+    sp = {
+        "mat_type": mat_type,
+        "pmat_type": pmat_type,
+        "snes_max_it": 10,
+        "ksp_type": "fgmres",
+        "pc_type": "fieldsplit",
+        "pc_fieldsplit_type": "additive",
+        "pc_fieldsplit_0_fields": "0",
+        "pc_fieldsplit_1_fields": "1,2",
+        "fieldsplit_1_ksp_view_eigenvalues": None,
+        "fieldsplit": {
+            "ksp_type": "gmres",
+            "pc_type": "jacobi",
+        },
+    }
+    problem = NonlinearVariationalProblem(F, u)
+    solver = NonlinearVariationalSolver(problem, solver_parameters=sp)
+
+    def mymonitor(snes, it, fnorm):
+        if it == 0:
+            # This call happens before the first linear solve
+            return
+        assert np.allclose(snes.ksp.pc.getFieldSplitSubKSP()[1].computeEigenvalues(), 1)
+
+    solver.snes.setMonitor(mymonitor)
+    solver.solve()
+
+
 def test_matrix_types(W):
     a = inner(TrialFunction(W), TestFunction(W))*dx
 

tests/firedrake/regression/test_split.py

@@ -91,6 +91,31 @@ def test_assemble_split_mixed_derivative():
     assert np.allclose(actual.M.values, expect.M.values)
 
 
+@pytest.mark.parametrize("mat_type", ("aij", "nest"))
+@pytest.mark.parametrize("bcs", (True, False))
+def test_split_assembled_matrix(mat_type, bcs):
+    mesh = UnitSquareMesh(2, 2)
+    V = FunctionSpace(mesh, "CG", 1)
+    Q = FunctionSpace(mesh, "DG", 0)
+    Z = V * Q
+    bcs = [DirichletBC(Z.sub(0), 0, "on_boundary")] if bcs else []
+
+    test = TestFunction(Z)
+    trial = TrialFunction(Z)
+
+    a = inner(test, trial)*dx
+    A = assemble(a, bcs=bcs, mat_type=mat_type)
+
+    splitter = ExtractSubBlock()
+    actual = splitter.split(A, (0, 0))
+
+    bcs = [bc.reconstruct(V=V) for bc in bcs]
+    expect = assemble(splitter.split(a, (0, 0)), bcs=bcs)
+
+    expect.petscmat.axpy(-1, actual.petscmat)
+    assert np.allclose(expect.petscmat[:, :], 0)
+
+
 def test_split_coordinate_derivative():
     mesh = UnitSquareMesh(1, 1)
     V = FunctionSpace(mesh, "P", 1)