Fix _SNESContext.options_prefix

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):

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

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

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"