From aaabace607c5c2c305f201e73d91c0750c944cf6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 4 Nov 2025 17:18:59 +0000 Subject: [PATCH 01/13] Fieldsplit: update MatNest Jacobian --- firedrake/dmhooks.py | 47 ++++++++----------- firedrake/formmanipulation.py | 20 ++++++-- firedrake/matrix.py | 10 ++-- firedrake/solving_utils.py | 37 +++++++++++---- .../test_nested_fieldsplit_solves.py | 34 ++++++++++++++ 5 files changed, 106 insertions(+), 42 deletions(-) diff --git a/firedrake/dmhooks.py b/firedrake/dmhooks.py index 046852b2e6..41b71661d2 100644 --- a/firedrake/dmhooks.py +++ b/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() diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index eb830493e1..bb8cd8bcc4 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -1,4 +1,3 @@ - import numpy import collections @@ -161,6 +160,7 @@ 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, EquationBC ises = [] args = [] for a in o.arguments(): @@ -180,8 +180,22 @@ def matrix(self, o): args.append(asplit) 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() + W = W.parent or W + + number = spaces.index(W) + V = args[number].function_space() + field = self.blocks[number] + if isinstance(bc, DirichletBC): + sub_bc = bc.reconstruct(field=field, V=V, g=bc.function_arg) + elif isinstance(bc, EquationBC): + raise NotImplementedError("Please get in touch if you need this") + if sub_bc is not None: + bcs.append(sub_bc) + return AssembledMatrix(tuple(args), tuple(bcs), submat) def zero_base_form(self, o): return ZeroBaseForm(tuple(map(self, o.arguments()))) diff --git a/firedrake/matrix.py b/firedrake/matrix.py index 2f33841289..d2c84451f7 100644 --- a/firedrake/matrix.py +++ b/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() diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 661f3ae218..999c193dfd 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -407,23 +407,33 @@ 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)) 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: + 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)) @@ -504,6 +514,15 @@ def form_jacobian(snes, X, J, P): ctx.set_nullspace(ctx._nullspace_T, ises, transpose=True, near=False) ctx.set_nullspace(ctx._near_nullspace, ises, transpose=False, near=True) + # Bump petsc matrix state of each split by assembling them. + # Ensures that if the matrix changed, the preconditioner is + # updated if necessary. + for field, splits in ctx._splits.items(): + for subctx in splits: + subctx._jac.assemble() + if subctx.Jp is not None: + subctx._pjac.assemble() + @staticmethod def compute_operators(ksp, J, P): r"""Form the Jacobian for this problem diff --git a/tests/firedrake/regression/test_nested_fieldsplit_solves.py b/tests/firedrake/regression/test_nested_fieldsplit_solves.py index 35de398216..d7784f7ea5 100644 --- a/tests/firedrake/regression/test_nested_fieldsplit_solves.py +++ b/tests/firedrake/regression/test_nested_fieldsplit_solves.py @@ -152,6 +152,40 @@ def test_nested_fieldsplit_solve_parallel(W, A, b, expect): assert norm(f) < 1e-11 +def test_nonlinear_fielsplit(): + 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": "nest", + "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", + }, + } + J = derivative(F, u) + solver = NonlinearVariationalSolver(NonlinearVariationalProblem(F, u), solver_parameters=sp) + solver.solve() + assert np.allclose(solver.snes.ksp.pc.getFieldSplitSubKSP()[1].computeEigenvalues(), 1) + + def test_matrix_types(W): a = inner(TrialFunction(W), TestFunction(W))*dx From 61b421d05f36af3de89e481b33eddde532e4fdc3 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 4 Nov 2025 17:39:25 +0000 Subject: [PATCH 02/13] Test split assembled matrix --- .../test_nested_fieldsplit_solves.py | 2 +- tests/firedrake/regression/test_split.py | 25 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/firedrake/regression/test_nested_fieldsplit_solves.py b/tests/firedrake/regression/test_nested_fieldsplit_solves.py index d7784f7ea5..8b6c80b299 100644 --- a/tests/firedrake/regression/test_nested_fieldsplit_solves.py +++ b/tests/firedrake/regression/test_nested_fieldsplit_solves.py @@ -152,7 +152,7 @@ def test_nested_fieldsplit_solve_parallel(W, A, b, expect): assert norm(f) < 1e-11 -def test_nonlinear_fielsplit(): +def test_nonlinear_fieldsplit(): mesh = UnitIntervalMesh(1) V = FunctionSpace(mesh, "DG", 0) Z = V * V * V diff --git a/tests/firedrake/regression/test_split.py b/tests/firedrake/regression/test_split.py index fcf8f221f8..3a49db9c07 100644 --- a/tests/firedrake/regression/test_split.py +++ b/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) From 469f583f9c81ee0b7b5b150ebbcb95386348eb27 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 4 Nov 2025 17:43:52 +0000 Subject: [PATCH 03/13] lint --- tests/firedrake/regression/test_nested_fieldsplit_solves.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/firedrake/regression/test_nested_fieldsplit_solves.py b/tests/firedrake/regression/test_nested_fieldsplit_solves.py index 8b6c80b299..5e4c726191 100644 --- a/tests/firedrake/regression/test_nested_fieldsplit_solves.py +++ b/tests/firedrake/regression/test_nested_fieldsplit_solves.py @@ -180,8 +180,8 @@ def test_nonlinear_fieldsplit(): "pc_type": "jacobi", }, } - J = derivative(F, u) - solver = NonlinearVariationalSolver(NonlinearVariationalProblem(F, u), solver_parameters=sp) + problem = NonlinearVariationalProblem(F, u) + solver = NonlinearVariationalSolver(problem, solver_parameters=sp) solver.solve() assert np.allclose(solver.snes.ksp.pc.getFieldSplitSubKSP()[1].computeEigenvalues(), 1) From 6d9f609496d58f71feb93096d6dd294334f129a2 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 5 Nov 2025 10:41:15 +0000 Subject: [PATCH 04/13] fixes --- firedrake/solving_utils.py | 19 +++++++++++++------ .../test_nested_fieldsplit_solves.py | 15 ++++++++++++--- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 999c193dfd..5c154f8e02 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -353,9 +353,15 @@ 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._jac.petscmat.type == "nest" else problem.J + Jpbig = self._pjac if self._pjac.petscmat.type == "nest" else problem.Jp + 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 +403,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. @@ -408,13 +413,15 @@ def split(self, fields): 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 if isinstance(J, MatrixBase) and J.has_bcs: + # The BCs of the problem are already encoded in the Jacobian bcs = None else: bcs = [] @@ -517,7 +524,7 @@ def form_jacobian(snes, X, J, P): # Bump petsc matrix state of each split by assembling them. # Ensures that if the matrix changed, the preconditioner is # updated if necessary. - for field, splits in ctx._splits.items(): + for fields, splits in ctx._splits.items(): for subctx in splits: subctx._jac.assemble() if subctx.Jp is not None: diff --git a/tests/firedrake/regression/test_nested_fieldsplit_solves.py b/tests/firedrake/regression/test_nested_fieldsplit_solves.py index 5e4c726191..ec7dd9f57b 100644 --- a/tests/firedrake/regression/test_nested_fieldsplit_solves.py +++ b/tests/firedrake/regression/test_nested_fieldsplit_solves.py @@ -152,7 +152,8 @@ def test_nested_fieldsplit_solve_parallel(W, A, b, expect): assert norm(f) < 1e-11 -def test_nonlinear_fieldsplit(): +@pytest.mark.parametrize("mat_type,pmat_type", [("nest", "nest"), ("matfree", "nest")]) +def test_nonlinear_fieldsplit(mat_type, pmat_type): mesh = UnitIntervalMesh(1) V = FunctionSpace(mesh, "DG", 0) Z = V * V * V @@ -167,7 +168,8 @@ def test_nonlinear_fieldsplit(): u.subfunctions[1].assign(Constant(1)) sp = { - "mat_type": "nest", + "mat_type": mat_type, + "pmat_type": pmat_type, "snes_max_it": 10, "ksp_type": "fgmres", "pc_type": "fieldsplit", @@ -182,8 +184,15 @@ def test_nonlinear_fieldsplit(): } 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() - assert np.allclose(solver.snes.ksp.pc.getFieldSplitSubKSP()[1].computeEigenvalues(), 1) def test_matrix_types(W): From c7c10cf1faecc5199d28c60e870bd8091f095a2a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 5 Nov 2025 11:03:54 +0000 Subject: [PATCH 05/13] fixes --- firedrake/formmanipulation.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index bb8cd8bcc4..22642edb65 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -160,7 +160,7 @@ 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, EquationBC + from firedrake.bcs import DirichletBC ises = [] args = [] for a in o.arguments(): @@ -184,15 +184,17 @@ def matrix(self, o): spaces = [a.function_space() for a in o.arguments()] for bc in o.bcs: W = bc.function_space() - W = W.parent or W + 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): sub_bc = bc.reconstruct(field=field, V=V, g=bc.function_arg) - elif isinstance(bc, EquationBC): - raise NotImplementedError("Please get in touch if you need this") + else: + raise NotImplementedError(f"Extracting matrix subblocks not supported with {type(bc).__name__}. " + "Please get in touch if you need this.") if sub_bc is not None: bcs.append(sub_bc) return AssembledMatrix(tuple(args), tuple(bcs), submat) From 537c6cd07bb602bdfbb5f233f24648b38f1c42a6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 5 Nov 2025 11:10:08 +0000 Subject: [PATCH 06/13] fixes --- firedrake/solving_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 5c154f8e02..9ad0ef1e76 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -358,6 +358,8 @@ def split(self, fields): # Reuse the submatrices if we are splitting a MatNest Jbig = self._jac if self._jac.petscmat.type == "nest" else problem.J Jpbig = self._pjac if self._pjac.petscmat.type == "nest" else problem.Jp + if Jpbig is Jbig: + Jpbig = None for field in fields: F = splitter.split(Fbig, argument_indices=(field, )) From 1d753e356b0340669ee6765c34be77ec0bd55e50 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 5 Nov 2025 11:18:21 +0000 Subject: [PATCH 07/13] lint --- firedrake/formmanipulation.py | 3 +-- tests/firedrake/regression/test_nested_fieldsplit_solves.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index 22642edb65..7fc438fd43 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -193,8 +193,7 @@ def matrix(self, o): if isinstance(bc, DirichletBC): sub_bc = bc.reconstruct(field=field, V=V, g=bc.function_arg) else: - raise NotImplementedError(f"Extracting matrix subblocks not supported with {type(bc).__name__}. " - "Please get in touch if you need this.") + raise NotImplementedError(f"Extracting matrix subblocks not supported with {type(bc).__name__}. Please get in touch if you need this.") if sub_bc is not None: bcs.append(sub_bc) return AssembledMatrix(tuple(args), tuple(bcs), submat) diff --git a/tests/firedrake/regression/test_nested_fieldsplit_solves.py b/tests/firedrake/regression/test_nested_fieldsplit_solves.py index ec7dd9f57b..e061a810ad 100644 --- a/tests/firedrake/regression/test_nested_fieldsplit_solves.py +++ b/tests/firedrake/regression/test_nested_fieldsplit_solves.py @@ -152,7 +152,7 @@ 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")]) +@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) From 4ab9490edbf7356965801a932b5fbe1c808fdf99 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 21 Nov 2025 09:45:24 +0000 Subject: [PATCH 08/13] Split EquationBCSplit --- firedrake/formmanipulation.py | 16 ++++++++++------ firedrake/solving_utils.py | 9 --------- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index 7fc438fd43..8926c395c5 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -160,7 +160,7 @@ 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 + from firedrake.bcs import DirichletBC, EquationBCSplit ises = [] args = [] for a in o.arguments(): @@ -180,6 +180,7 @@ def matrix(self, o): args.append(asplit) submat = o.petscmat.createSubMatrix(*ises) + bcs = [] spaces = [a.function_space() for a in o.arguments()] for bc in o.bcs: @@ -191,11 +192,14 @@ def matrix(self, o): V = args[number].function_space() field = self.blocks[number] if isinstance(bc, DirichletBC): - sub_bc = bc.reconstruct(field=field, V=V, g=bc.function_arg) - else: - raise NotImplementedError(f"Extracting matrix subblocks not supported with {type(bc).__name__}. Please get in touch if you need this.") - if sub_bc is not None: - bcs.append(sub_bc) + bc_temp = bc.reconstruct(field=field, V=V, g=bc.function_arg, use_split=True) + elif isinstance(bc, EquationBCSplit): + row_field = self.blocks[0] + col_field = self.blocks[1] + bc_temp = bc.reconstruct(field=field, V=V, row_field=row_field, col_field=col_field, use_split=True) + if bc_temp is not None: + bcs.append(bc_temp) + return AssembledMatrix(tuple(args), tuple(bcs), submat) def zero_base_form(self, o): diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 9ad0ef1e76..c2ffb4238a 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -523,15 +523,6 @@ def form_jacobian(snes, X, J, P): ctx.set_nullspace(ctx._nullspace_T, ises, transpose=True, near=False) ctx.set_nullspace(ctx._near_nullspace, ises, transpose=False, near=True) - # Bump petsc matrix state of each split by assembling them. - # Ensures that if the matrix changed, the preconditioner is - # updated if necessary. - for fields, splits in ctx._splits.items(): - for subctx in splits: - subctx._jac.assemble() - if subctx.Jp is not None: - subctx._pjac.assemble() - @staticmethod def compute_operators(ksp, J, P): r"""Form the Jacobian for this problem From b7ea78c02dce350bea1e8a9b215c348db98e12f7 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 25 Nov 2025 09:23:08 +0000 Subject: [PATCH 09/13] cleanup --- firedrake/solving_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index c2ffb4238a..523faf38d1 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -356,8 +356,8 @@ def split(self, fields): Fbig = problem.F # Reuse the submatrices if we are splitting a MatNest - Jbig = self._jac if self._jac.petscmat.type == "nest" else problem.J - Jpbig = self._pjac if self._pjac.petscmat.type == "nest" else problem.Jp + 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 From 866fda211f6f199ed3ab9fdbb74c69a1359c9b6f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 25 Nov 2025 14:38:52 +0000 Subject: [PATCH 10/13] WIP --- firedrake/formmanipulation.py | 13 +++++++++---- tests/firedrake/equation_bcs/test_equation_bcs.py | 4 ++-- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index 8926c395c5..2e30cd7858 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -2,7 +2,7 @@ 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 @@ -183,6 +183,8 @@ def matrix(self, o): bcs = [] spaces = [a.function_space() for a in o.arguments()] + row_field = self.blocks[0] + col_field = self.blocks[1] for bc in o.bcs: W = bc.function_space() while W.parent is not None: @@ -194,13 +196,16 @@ def matrix(self, o): 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 = self.blocks[0] - col_field = self.blocks[1] 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(tuple(args), tuple(bcs), submat) + if isinstance(o.a, Form): + a = self.split(o.a, argument_indices=(row_field, col_field)) + else: + a = args + return AssembledMatrix(a, tuple(bcs), submat) def zero_base_form(self, o): return ZeroBaseForm(tuple(map(self, o.arguments()))) diff --git a/tests/firedrake/equation_bcs/test_equation_bcs.py b/tests/firedrake/equation_bcs/test_equation_bcs.py index 53daf97a54..c3151e46e5 100644 --- a/tests/firedrake/equation_bcs/test_equation_bcs.py +++ b/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 From 561480f174f5a459726ca9830dae1ba64b2a54bc Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 26 Nov 2025 10:35:47 +0000 Subject: [PATCH 11/13] Split the form --- firedrake/formmanipulation.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index 2e30cd7858..558593f86c 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -161,6 +161,15 @@ def cofunction(self, o): def matrix(self, o): from firedrake.bcs import DirichletBC, EquationBCSplit + row_field = self.blocks[0] + col_field = self.blocks[1] + if isinstance(o.a, Form): + form = self.split(o.a, argument_indices=(row_field, col_field)) + if isinstance(form, ZeroBaseForm): + return form + else: + form = None + ises = [] args = [] for a in o.arguments(): @@ -183,8 +192,6 @@ def matrix(self, o): bcs = [] spaces = [a.function_space() for a in o.arguments()] - row_field = self.blocks[0] - col_field = self.blocks[1] for bc in o.bcs: W = bc.function_space() while W.parent is not None: @@ -201,11 +208,7 @@ def matrix(self, o): if bc_temp is not None: bcs.append(bc_temp) - if isinstance(o.a, Form): - a = self.split(o.a, argument_indices=(row_field, col_field)) - else: - a = args - return AssembledMatrix(a, tuple(bcs), submat) + return AssembledMatrix(form or tuple(args), tuple(bcs), submat) def zero_base_form(self, o): return ZeroBaseForm(tuple(map(self, o.arguments()))) From 4688ff1c5be985e05b82c0c79ce74b167ed83372 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 26 Nov 2025 11:47:34 +0000 Subject: [PATCH 12/13] fixup --- firedrake/formmanipulation.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index 558593f86c..0ee187b8bc 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -161,32 +161,34 @@ def cofunction(self, o): def matrix(self, o): from firedrake.bcs import DirichletBC, EquationBCSplit - row_field = self.blocks[0] - col_field = self.blocks[1] - if isinstance(o.a, Form): - form = self.split(o.a, argument_indices=(row_field, col_field)) - if isinstance(form, ZeroBaseForm): - return form - else: - form = None - 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) @@ -203,6 +205,7 @@ def matrix(self, o): 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: From 0a3fa466f76fef5f1811683d79cafdb13ce759a4 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 15 Dec 2025 14:58:11 +0000 Subject: [PATCH 13/13] Bump PETSc version --- firedrake/__init__.py | 2 +- pyproject.toml | 4 ++-- scripts/firedrake-configure | 2 +- setup.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/firedrake/__init__.py b/firedrake/__init__.py index 70dec1d6ca..07ff341863 100644 --- a/firedrake/__init__.py +++ b/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(): diff --git a/pyproject.toml b/pyproject.toml index 6376531907..06536d0a3c 100644 --- a/pyproject.toml +++ b/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" diff --git a/scripts/firedrake-configure b/scripts/firedrake-configure index 1b2ffcb42b..c3d52639af 100755 --- a/scripts/firedrake-configure +++ b/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(): diff --git a/setup.py b/setup.py index 27015d23cc..76ea9bf8fd 100644 --- a/setup.py +++ b/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