From 386538f778e03a8657fda1ca9a32241ff8901a31 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 9 May 2026 18:25:36 +0100 Subject: [PATCH 1/7] Fix _SNESContext.options_prefix --- firedrake/mg/ufl_utils.py | 68 +++++++++++++++++--------------------- firedrake/solving_utils.py | 8 ++--- 2 files changed, 33 insertions(+), 43 deletions(-) diff --git a/firedrake/mg/ufl_utils.py b/firedrake/mg/ufl_utils.py index d5a26a36e6..0671cc9765 100644 --- a/firedrake/mg/ufl_utils.py +++ b/firedrake/mg/ufl_utils.py @@ -278,47 +278,39 @@ def coarsen_snescontext(context, self, coefficient_mapping=None): # Assume not something that needs coarsening (e.g. float) new_appctx[k] = v + opts = PETSc.Options(context.options_prefix) + if opts.getString("snes_type", "") == "fas": + solver_prefix = "fas_" + else: + solver_prefix = "mg_" _, level = utils.get_level(problem.u_restrict.function_space().mesh()) if level == 0: - # Use different mat_type on coarsest level - opts = PETSc.Options(context.options_prefix) - if opts.getString("snes_type", "") == "fas": - solver_prefix = "fas_" - else: - solver_prefix = "mg_" - - coarse_mat_type = opts.getString(f"{solver_prefix}coarse_mat_type", "") - if coarse_mat_type == "": - coarse_mat_type = context.mat_type - default_pmat_type = context.pmat_type - else: - default_pmat_type = coarse_mat_type - coarse_pmat_type = opts.getString(f"{solver_prefix}coarse_pmat_type", - default_pmat_type) - - coarse_sub_mat_type = opts.getString(f"{solver_prefix}coarse_sub_mat_type", "") - if coarse_sub_mat_type == "": - coarse_sub_mat_type = context.sub_mat_type - default_sub_pmat_type = context.sub_pmat_type - else: - default_sub_pmat_type = coarse_sub_mat_type - coarse_sub_pmat_type = opts.getString(f"{solver_prefix}coarse_sub_pmat_type", - default_sub_pmat_type or "") or None + levels_prefix = f"{solver_prefix}coarse_" else: - coarse_mat_type = context.mat_type - coarse_pmat_type = context.pmat_type - coarse_sub_mat_type = context.sub_mat_type - coarse_sub_pmat_type = context.sub_pmat_type - - coarse = _SNESContext(problem, - mat_type=coarse_mat_type, - pmat_type=coarse_pmat_type, - sub_mat_type=coarse_sub_mat_type, - sub_pmat_type=coarse_sub_pmat_type, - appctx=new_appctx, - options_prefix=context.options_prefix, - transfer_manager=context.transfer_manager, - pre_apply_bcs=context.pre_apply_bcs) + levels_prefix = f"{solver_prefix}levels_" + current_level_prefix = f"{solver_prefix}levels_{level}_" + new_options_prefix = f"{context.options_prefix}{current_level_prefix}" + + new_mat_type = None + new_pmat_type = None + new_sub_mat_type = None + new_sub_pmat_type = None + for prefix in (levels_prefix, current_level_prefix): + new_mat_type = opts.getString(f"{prefix}mat_type", "") or new_mat_type + new_pmat_type = opts.getString(f"{prefix}pmat_type", "") or new_pmat_type + new_sub_mat_type = opts.getString(f"{prefix}sub_mat_type", "") or new_sub_mat_type + new_sub_pmat_type = opts.getString(f"{prefix}sub_pmat_type", "") or new_sub_pmat_type + + new_pmat_type = new_pmat_type or new_mat_type + new_sub_pmat_type = new_sub_pmat_type or new_sub_mat_type + coarse = context.reconstruct(problem=problem, + mat_type=new_mat_type, + pmat_type=new_pmat_type, + sub_mat_type=new_sub_mat_type, + sub_pmat_type=new_sub_pmat_type, + appctx=new_appctx, + options_prefix=new_options_prefix, + ) coarse._coefficient_mapping = coefficient_mapping coarse._fine = context context._coarse = coarse diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index b1090d1111..3440d2af19 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -373,7 +373,7 @@ def split(self, fields): splits = [] problem = self._problem splitter = ExtractSubBlock() - for field in fields: + for field_num, field in enumerate(fields): F = splitter.split(problem.F, argument_indices=(field, )) J = splitter.split(problem.J, argument_indices=(field, field)) us = problem.u_restrict.subfunctions @@ -443,10 +443,8 @@ def split(self, fields): 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, - appctx=self.appctx, - transfer_manager=self.transfer_manager, - pre_apply_bcs=self.pre_apply_bcs)) + options_prefix = f"{self.options_prefix}fieldsplit_{field_num}_" + splits.append(self.reconstruct(new_problem, options_prefix=options_prefix)) return self._splits.setdefault(tuple(fields), splits) @staticmethod From 26047865971bd50504a24e7f6ba2c38bbab6acc5 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 11 May 2026 11:14:39 +0100 Subject: [PATCH 2/7] review suggestions --- firedrake/dmhooks.py | 2 +- firedrake/functionspaceimpl.py | 2 +- firedrake/mg/ufl_utils.py | 31 +++++++++++++++---------------- firedrake/solving_utils.py | 3 ++- 4 files changed, 19 insertions(+), 19 deletions(-) diff --git a/firedrake/dmhooks.py b/firedrake/dmhooks.py index 5b1562e84c..447d76d205 100644 --- a/firedrake/dmhooks.py +++ b/firedrake/dmhooks.py @@ -338,7 +338,7 @@ def create_field_decomposition(dm, *args, **kwargs): the appropriate subdms as well. """ W = get_function_space(dm) - # Don't pass split number if name is None (this way the + # Don't pass split number if name is not None (this way the # recursively created splits have the names you want) names = [s.name for s in W] dms = [V.dm for V in W] diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index b5b2e90a0a..36c2652615 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -926,7 +926,7 @@ def local_to_global_map(self, bcs, lgmap=None, mat_type=None): return PETSc.LGMap().create(indices, bsize=bsize, comm=lgmap.comm) def collapse(self): - return type(self)(self.mesh(), self.ufl_element()) + return type(self)(self.mesh(), self.ufl_element(), name=self.name) class RestrictedFunctionSpace(FunctionSpace): diff --git a/firedrake/mg/ufl_utils.py b/firedrake/mg/ufl_utils.py index 0671cc9765..4436503d4c 100644 --- a/firedrake/mg/ufl_utils.py +++ b/firedrake/mg/ufl_utils.py @@ -289,27 +289,26 @@ def coarsen_snescontext(context, self, coefficient_mapping=None): else: levels_prefix = f"{solver_prefix}levels_" current_level_prefix = f"{solver_prefix}levels_{level}_" - new_options_prefix = f"{context.options_prefix}{current_level_prefix}" - new_mat_type = None - new_pmat_type = None - new_sub_mat_type = None - new_sub_pmat_type = None + mat_type = None + pmat_type = None + sub_mat_type = None + sub_pmat_type = None for prefix in (levels_prefix, current_level_prefix): - new_mat_type = opts.getString(f"{prefix}mat_type", "") or new_mat_type - new_pmat_type = opts.getString(f"{prefix}pmat_type", "") or new_pmat_type - new_sub_mat_type = opts.getString(f"{prefix}sub_mat_type", "") or new_sub_mat_type - new_sub_pmat_type = opts.getString(f"{prefix}sub_pmat_type", "") or new_sub_pmat_type + mat_type = opts.getString(f"{prefix}mat_type", "") or mat_type + pmat_type = opts.getString(f"{prefix}pmat_type", "") or pmat_type + sub_mat_type = opts.getString(f"{prefix}sub_mat_type", "") or sub_mat_type + sub_pmat_type = opts.getString(f"{prefix}sub_pmat_type", "") or sub_pmat_type - new_pmat_type = new_pmat_type or new_mat_type - new_sub_pmat_type = new_sub_pmat_type or new_sub_mat_type + pmat_type = pmat_type or mat_type + sub_pmat_type = sub_pmat_type or sub_mat_type coarse = context.reconstruct(problem=problem, - mat_type=new_mat_type, - pmat_type=new_pmat_type, - sub_mat_type=new_sub_mat_type, - sub_pmat_type=new_sub_pmat_type, + mat_type=mat_type, + pmat_type=pmat_type, + sub_mat_type=sub_mat_type, + sub_pmat_type=sub_pmat_type, appctx=new_appctx, - options_prefix=new_options_prefix, + options_prefix=context.options_prefix, ) coarse._coefficient_mapping = coefficient_mapping coarse._fine = context diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 3440d2af19..335e12e111 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -443,7 +443,8 @@ def split(self, fields): 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 - options_prefix = f"{self.options_prefix}fieldsplit_{field_num}_" + field_prefix = f"fieldsplit_{V.name or field_num}_" + options_prefix = f"{self.options_prefix}{field_prefix}" splits.append(self.reconstruct(new_problem, options_prefix=options_prefix)) return self._splits.setdefault(tuple(fields), splits) From 285db1ff88e3985fe752aed8d7ea2938e401d02b Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 11 May 2026 15:52:36 +0100 Subject: [PATCH 3/7] different options_prefix for mg_levels --- firedrake/dmhooks.py | 2 +- firedrake/mg/ufl_utils.py | 11 +++++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/firedrake/dmhooks.py b/firedrake/dmhooks.py index 447d76d205..5b1562e84c 100644 --- a/firedrake/dmhooks.py +++ b/firedrake/dmhooks.py @@ -338,7 +338,7 @@ def create_field_decomposition(dm, *args, **kwargs): the appropriate subdms as well. """ W = get_function_space(dm) - # Don't pass split number if name is not None (this way the + # Don't pass split number if name is None (this way the # recursively created splits have the names you want) names = [s.name for s in W] dms = [V.dm for V in W] diff --git a/firedrake/mg/ufl_utils.py b/firedrake/mg/ufl_utils.py index 4436503d4c..a5b80b8069 100644 --- a/firedrake/mg/ufl_utils.py +++ b/firedrake/mg/ufl_utils.py @@ -278,7 +278,12 @@ def coarsen_snescontext(context, self, coefficient_mapping=None): # Assume not something that needs coarsening (e.g. float) new_appctx[k] = v - opts = PETSc.Options(context.options_prefix) + # Get options prefix for current level + parent_context = context + while parent_context._fine: + parent_context = parent_context._fine + parent_prefix = parent_context.options_prefix + opts = PETSc.Options(parent_prefix) if opts.getString("snes_type", "") == "fas": solver_prefix = "fas_" else: @@ -289,7 +294,9 @@ def coarsen_snescontext(context, self, coefficient_mapping=None): else: levels_prefix = f"{solver_prefix}levels_" current_level_prefix = f"{solver_prefix}levels_{level}_" + options_prefix = f"{parent_prefix}{current_level_prefix}" + # Use different mat_type on each level mat_type = None pmat_type = None sub_mat_type = None @@ -308,7 +315,7 @@ def coarsen_snescontext(context, self, coefficient_mapping=None): sub_mat_type=sub_mat_type, sub_pmat_type=sub_pmat_type, appctx=new_appctx, - options_prefix=context.options_prefix, + options_prefix=options_prefix, ) coarse._coefficient_mapping = coefficient_mapping coarse._fine = context From 8d22ff2eb1e4ec08e8929d9ddc0e03c0749bbc39 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 13 May 2026 10:30:23 +0100 Subject: [PATCH 4/7] add a test --- .../multigrid/test_options_prefix.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 tests/firedrake/multigrid/test_options_prefix.py diff --git a/tests/firedrake/multigrid/test_options_prefix.py b/tests/firedrake/multigrid/test_options_prefix.py new file mode 100644 index 0000000000..81a8a0dbba --- /dev/null +++ b/tests/firedrake/multigrid/test_options_prefix.py @@ -0,0 +1,57 @@ +from firedrake import * +import pytest + + +@pytest.mark.parametrize("named", [False, True], ids=["unnamed", "named"]) +def test_fieldsplit_mg_options_prefix(named): + if named: + names = ("V1", "V2") + else: + names = (None, None) + refine = 2 + base = UnitIntervalMesh(10) + mh = MeshHierarchy(base, refine) + mesh = mh[-1] + V1 = FunctionSpace(mesh, 'CG', 1, name=names[0]) + V2 = FunctionSpace(mesh, 'CG', 2, name=names[1]) + W = MixedFunctionSpace([V1, V2]) + params0 = { + "pc_type": "mg", + "mg_levels_ksp_type": "chebyshev", + "mg_levels_pc_type": "jacobi", + "mg_coarse_mat_type": "aij", + "mg_coarse_ksp_type": "preonly", + "mg_coarse_pc_type": "lu", + } + params1 = { + "pc_type": "mg", + "mg_levels_ksp_type": "chebyshev", + "mg_levels_pc_type": "jacobi", + "mg_levels_0_mat_type": "aij", + "mg_levels_0_ksp_type": "preonly", + "mg_levels_0_pc_type": "lu", + } + sp = { + "mat_type": "matfree", + "ksp_rtol": 1E-12, + "ksp_max_it": 12, + "ksp_monitor": None, + "ksp_type": "cg", + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "additive", + "fieldsplit_ksp_type": "preonly", + f"fieldsplit_{names[0] or 0}": params0, + f"fieldsplit_{names[1] or 1}": params1, + } + + x = SpatialCoordinate(mesh) + z_exact = as_vector([x[0], x[0]**2]) + z = Function(W) + w = TrialFunction(W) + v = TestFunction(W) + a = inner(grad(w), grad(v)) * dx + L = a(v, z_exact) + bcs = [DirichletBC(W.sub(i), z_exact[i], "on_boundary") for i in range(len(W))] + solve(a == L, z, bcs=bcs, solver_parameters=sp) + + assert errornorm(z_exact, z)/norm(z_exact) < 1E-12 From e7bd3243991bfa3525309bc37c3edfc789e69bf1 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 13 May 2026 10:30:23 +0100 Subject: [PATCH 5/7] review suggestions --- .../multigrid/test_options_prefix.py | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 tests/firedrake/multigrid/test_options_prefix.py diff --git a/tests/firedrake/multigrid/test_options_prefix.py b/tests/firedrake/multigrid/test_options_prefix.py new file mode 100644 index 0000000000..9d75e41d1a --- /dev/null +++ b/tests/firedrake/multigrid/test_options_prefix.py @@ -0,0 +1,69 @@ +from firedrake import * +import pytest + + +@pytest.mark.parametrize("named", [False, True], ids=["unnamed", "named"]) +def test_fieldsplit_mg_options_prefix(named): + if named: + names = ("V1", "V2") + else: + names = (None, None) + refine = 2 + base = UnitIntervalMesh(10) + mh = MeshHierarchy(base, refine) + mesh = mh[-1] + V1 = FunctionSpace(mesh, 'CG', 1, name=names[0]) + V2 = FunctionSpace(mesh, 'CG', 2, name=names[1]) + W = MixedFunctionSpace([V1, V2]) + params0 = { + "pc_type": "mg", + "mg_levels_ksp_type": "chebyshev", + "mg_levels_pc_type": "jacobi", + "mg_coarse_mat_type": "aij", + "mg_coarse_ksp_type": "preonly", + "mg_coarse_pc_type": "lu", + } + params1 = { + "pc_type": "mg", + "mg_levels_ksp_type": "chebyshev", + "mg_levels_pc_type": "jacobi", + "mg_levels_0_mat_type": "aij", + "mg_levels_0_ksp_type": "preonly", + "mg_levels_0_pc_type": "lu", + } + sp = { + "mat_type": "matfree", + "ksp_rtol": 1E-12, + "ksp_max_it": 12, + "ksp_type": "cg", + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "additive", + "fieldsplit_ksp_type": "preonly", + f"fieldsplit_{names[0] or 0}": params0, + f"fieldsplit_{names[1] or 1}": params1, + } + + x = SpatialCoordinate(mesh) + z_exact = as_vector([x[0], x[0]**2]) + z = Function(W) + w = TrialFunction(W) + v = TestFunction(W) + a = inner(grad(w), grad(v)) * dx + L = a(v, z_exact) + bcs = [DirichletBC(W.sub(i), z_exact[i], "on_boundary") for i in range(len(W))] + problem = LinearVariationalProblem(a, L, z, bcs=bcs) + solver = LinearVariationalSolver(problem, solver_parameters=sp) + solver.solve() + assert errornorm(z_exact, z) / norm(z_exact) < 1E-12 + + assert solver.snes.ksp.pc.getOperators()[0].getType() == "python" + fsplit = solver.snes.ksp.pc.getFieldSplitSubKSP() + assert len(fsplit) == len(V) + for ksp in fsplit: + coarse = ksp.pc.getMGSmoother(0) + assert coarse.pc.getType() == "lu" + assert coarse.getOperators()[0].getType().endswith("aij") + for l in range(1, len(mh)): + level = ksp.pc.getMGSmoother(l) + assert level.pc.getType() == "jacobi" + assert level.getOperators()[0].getType() == "python" From 7d4287a4d508a4d3b0aabc892939c708614bb202 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 13 May 2026 21:41:03 +0100 Subject: [PATCH 6/7] Apply suggestion from @pbrubeck --- tests/firedrake/multigrid/test_options_prefix.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/firedrake/multigrid/test_options_prefix.py b/tests/firedrake/multigrid/test_options_prefix.py index 9d75e41d1a..204b584c7f 100644 --- a/tests/firedrake/multigrid/test_options_prefix.py +++ b/tests/firedrake/multigrid/test_options_prefix.py @@ -58,7 +58,7 @@ def test_fieldsplit_mg_options_prefix(named): assert solver.snes.ksp.pc.getOperators()[0].getType() == "python" fsplit = solver.snes.ksp.pc.getFieldSplitSubKSP() - assert len(fsplit) == len(V) + assert len(fsplit) == len(W) for ksp in fsplit: coarse = ksp.pc.getMGSmoother(0) assert coarse.pc.getType() == "lu" From 76a6e55f4c1d46fd9a0637f075df40fb08dc179c Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 14 May 2026 13:25:14 +0100 Subject: [PATCH 7/7] Update firedrake/solving_utils.py --- firedrake/solving_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 335e12e111..db18865712 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -443,7 +443,8 @@ def split(self, fields): 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 - field_prefix = f"fieldsplit_{V.name or field_num}_" + name = V.name if len(V) == 1 else None + field_prefix = f"fieldsplit_{name or field_num}_" options_prefix = f"{self.options_prefix}{field_prefix}" splits.append(self.reconstruct(new_problem, options_prefix=options_prefix)) return self._splits.setdefault(tuple(fields), splits)