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 d91a72d5e8..b40bfbd7ac 100644 --- a/firedrake/mg/ufl_utils.py +++ b/firedrake/mg/ufl_utils.py @@ -282,47 +282,45 @@ def coarsen_snescontext(context, self, coefficient_mapping=None): # Assume not something that needs coarsening (e.g. float) new_appctx[k] = v + # 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: + 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}_" + 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 + sub_pmat_type = None + for prefix in (levels_prefix, current_level_prefix): + 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 + + pmat_type = pmat_type or mat_type + sub_pmat_type = sub_pmat_type or sub_mat_type + coarse = context.reconstruct(problem=problem, + 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=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..db18865712 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,10 @@ 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)) + 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) @staticmethod diff --git a/tests/firedrake/multigrid/test_options_prefix.py b/tests/firedrake/multigrid/test_options_prefix.py new file mode 100644 index 0000000000..204b584c7f --- /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(W) + 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"