Use petsctools.AppContext

firedrake/__init__.py

@@ -55,6 +55,8 @@ def init_petsc():
 from pyop2 import op2                        # noqa: F401
 from pyop2.mpi import COMM_WORLD, COMM_SELF  # noqa: F401
 
+from petsctools import AppContext  # noqa: F401
+
 from firedrake.assemble import *
 from firedrake.bcs import *
 from firedrake.checkpointing import *

firedrake/assemble.py

@@ -78,7 +78,7 @@ def assemble(expr, *args, **kwargs):
         not supplied, defaults to ``parameters["default_sub_matrix_type"]``.
     options_prefix : str
         PETSc options prefix to apply to matrices.
-    appctx : dict
+    appctx : petsctools.AppContext
         Additional information to hang on the assembled
         matrix if an implicit matrix is requested (mat_type ``"matfree"``).
     zero_bc_nodes : bool

firedrake/preconditioners/assembled.py

@@ -18,13 +18,13 @@ class AssembledPC(PCBase):
 
     def initialize(self, pc):
         from firedrake.assemble import get_assembler
+        from dmhooks import get_appctx
         A, P = pc.getOperators()
 
         if pc.type != "python":
             raise ValueError("Expecting PC type python")
         opc = pc
-        appctx = self.get_appctx(pc)
-        fcp = appctx.get("form_compiler_parameters")
+        fcp = get_appctx(pc).fcp
 
         V = get_function_space(pc.getDM()).collapse()
         test = TestFunction(V)

firedrake/preconditioners/bddc.py

@@ -53,7 +53,8 @@ def initialize(self, pc):
         bddcpc.setOperators(*pc.getOperators())
         bddcpc.setType(PETSc.PC.Type.BDDC)
 
-        opts = PETSc.Options(bddcpc.getOptionsPrefix())
+        prefix = bddcpc.getOptionsPrefix()
+        opts = PETSc.Options(prefix)
         if V.ufl_element().variant() == "fdm" and "pc_bddc_use_local_mat_graph" not in opts:
             # Disable computation of disconected components of subdomain interfaces
             opts["pc_bddc_use_local_mat_graph"] = False
@@ -84,7 +85,7 @@ def initialize(self, pc):
         tdim = V.mesh().topological_dimension()
         degree = max(as_tuple(V.ufl_element().degree()))
         if tdim >= 2 and V.finat_element.formdegree == tdim-1:
-            B = appctx.get("divergence_mat", None)
+            B = appctx.get(prefix+"divergence_mat", None)
             if B is None:
                 from firedrake.assemble import assemble
                 d = {HCurl: curl, HDiv: div}[sobolev_space]
@@ -99,7 +100,7 @@ def initialize(self, pc):
                 B = assemble(b, mat_type="matfree")
             bddcpc.setBDDCDivergenceMat(B.petscmat)
         elif sobolev_space == HCurl:
-            gradient = appctx.get("discrete_gradient", None)
+            gradient = appctx.get(prefix+"discrete_gradient", None)
             if gradient is None:
                 from firedrake.preconditioners.fdm import tabulate_exterior_derivative
                 from firedrake.preconditioners.hiptmair import curl_to_grad

firedrake/preconditioners/facet_split.py

@@ -49,8 +49,7 @@ def initialize(self, pc):
         domains = domains.split(",")
 
         a, bcs = self.form(pc)
-        appctx = self.get_appctx(pc)
-        fcp = appctx.get("form_compiler_parameters")
+        fcp = dmhooks.get_appctx(pc).fcp
         V = a.arguments()[-1].function_space()
         if len(V) != 1:
             raise ValueError("Decomposition of mixed elements is not supported")

firedrake/preconditioners/fdm.py

@@ -2,6 +2,7 @@
 from functools import partial
 from itertools import chain, product
 from firedrake.petsc import PETSc
+from firedrake.dmhooks import get_appctx
 from firedrake.preconditioners.base import PCBase
 from firedrake.preconditioners.patch import bcdofs
 from firedrake.preconditioners.pmg import (prolongation_matrix_matfree,
@@ -100,6 +101,8 @@ def initialize(self, pc):
         prefix = pc.getOptionsPrefix() or ""
         options_prefix = prefix + self._prefix
         options = PETSc.Options(options_prefix)
+        self.options = options
+        self.options_prefix = options_prefix
 
         use_amat = options.getBool("pc_use_amat", True)
         use_static_condensation = options.getBool("static_condensation", False)
@@ -111,9 +114,8 @@ def initialize(self, pc):
             allow_repeated = options.getBool("mat_is_allow_repeated", True)
         self.allow_repeated = allow_repeated
 
-        appctx = self.get_appctx(pc)
-        fcp = appctx.get("form_compiler_parameters") or {}
-        self.appctx = appctx
+        fcp = get_appctx(pc).fcp
+        self.appctx = self.get_appctx(pc)
 
         # Get original Jacobian form and bcs
         J, bcs = self.form(pc)
@@ -1908,7 +1910,9 @@ def assemble_reference_tensor(self, V):
         axes_shifts, = shifts
 
         degree = max(e.degree() for e in line_elements)
-        eta = float(self.appctx.get("eta", degree*(degree+1)))
+        eta = float(self.appctx.get(self.options_prefix+"eta",
+                                    default=degree*(degree+1)))
+
         is_dg = V.finat_element.is_dg()
         Afdm = []  # sparse interval mass and stiffness matrices for each direction
         Dfdm = []  # tabulation of normal derivatives at the boundary for each direction
@@ -2075,7 +2079,7 @@ def cell_to_global(lgmap, cell_to_local, cell_index, result=None):
                 raise NotImplementedError("Static condensation for SIPG not implemented")
             if tdim < V.mesh().geometric_dimension():
                 raise NotImplementedError("SIPG on immersed meshes is not implemented")
-            eta = float(self.appctx.get("eta"))
+            eta = float(self.appctx.get(self.options_prefix+"eta"))
 
             lgmap = self.lgmaps[V]
             index_facet, local_facet_data, nfacets = extrude_interior_facet_maps(V)

firedrake/preconditioners/gtmg.py

@@ -66,11 +66,14 @@ def initialize(self, pc):
         :arg pc: PETSc preconditioner instance
         """
         from firedrake.assemble import assemble, get_assembler
+        from dmhooks import get_appctx
 
         Citations().register("Gopalakrishnan2009")
         _, P = pc.getOperators()
-        appctx = self.get_appctx(pc)
-        fcp = appctx.get("form_compiler_parameters")
+        ctx = get_appctx(pc)
+        fcp = ctx.fcp
+
+        appctx = ctx.appctx
 
         if pc.getType() != "python":
             raise ValueError("Expecting PC type python")
@@ -119,21 +122,21 @@ def initialize(self, pc):
         coarse_mat_type = opts.getString(coarse_options_prefix + "mat_type",
                                          parameters["default_matrix_type"])
 
-        get_coarse_space = appctx.get("get_coarse_space", None)
+        get_coarse_space = appctx.get(options_prefix+"get_coarse_space", None)
         if not get_coarse_space:
             raise ValueError("Need to provide a callback which provides the coarse space.")
         coarse_space = get_coarse_space()
 
-        get_coarse_operator = appctx.get("get_coarse_operator", None)
+        get_coarse_operator = appctx.get(options_prefix+"get_coarse_operator", None)
         if not get_coarse_operator:
             raise ValueError("Need to provide a callback which provides the coarse operator.")
         coarse_operator = get_coarse_operator()
 
-        coarse_space_bcs = appctx.get("coarse_space_bcs", None)
+        coarse_space_bcs = appctx.get(options_prefix+"coarse_space_bcs", None)
 
         # These should be callbacks which return the relevant nullspaces
-        get_coarse_nullspace = appctx.get("get_coarse_op_nullspace", None)
-        get_coarse_transpose_nullspace = appctx.get("get_coarse_op_transpose_nullspace", None)
+        get_coarse_nullspace = appctx.get(options_prefix+"get_coarse_op_nullspace", None)
+        get_coarse_transpose_nullspace = appctx.get(options_prefix+"get_coarse_op_transpose_nullspace", None)
 
         coarse_form_assembler = get_assembler(coarse_operator, bcs=coarse_space_bcs, form_compiler_parameters=fcp, mat_type=coarse_mat_type, options_prefix=coarse_options_prefix)
         self.coarse_op = coarse_form_assembler.allocate()
@@ -150,14 +153,14 @@ def initialize(self, pc):
             tnsp = get_coarse_transpose_nullspace()
             coarse_opmat.setTransposeNullSpace(tnsp.nullspace())
 
-        interp_petscmat = appctx.get("interpolation_matrix", None)
+        interp_petscmat = appctx.get(options_prefix+"interpolation_matrix", None)
         if interp_petscmat is None:
             # Create interpolation matrix from coarse space to fine space
             fine_space = ctx.J.arguments()[0].function_space()
             coarse_test, coarse_trial = coarse_operator.arguments()
             interp = assemble(Interpolate(coarse_trial, fine_space))
             interp_petscmat = interp.petscmat
-        restr_petscmat = appctx.get("restriction_matrix", None)
+        restr_petscmat = appctx.get(options_prefix+"restriction_matrix", None)
 
         # We set up a PCMG object that uses the constructed interpolation
         # matrix to generate the restriction/prolongation operators.

firedrake/preconditioners/hiptmair.py

@@ -37,9 +37,9 @@ def coarsen(self, pc):
 
     def initialize(self, pc):
         from firedrake.assemble import get_assembler
+        from dmhooks import get_appctx
         A, P = pc.getOperators()
-        appctx = self.get_appctx(pc)
-        fcp = appctx.get("form_compiler_parameters")
+        fcp = get_appctx(pc).fcp
 
         prefix = pc.getOptionsPrefix() or ""
         options_prefix = prefix + self._prefix
@@ -168,13 +168,13 @@ def coarsen(self, pc):
         coarse_space_bcs = [bc.reconstruct(V=coarse_space, g=0) for bc in bcs]
 
         if element.sobolev_space == ufl.HDiv:
-            G_callback = appctx.get("get_curl", None)
+            G_callback = appctx.get(options_prefix+"get_curl", None)
             dminus = ufl.curl
             if V.shape:
                 dminus = lambda u: ufl.as_vector([ufl.curl(u[k, ...])
                                                   for k in range(u.ufl_shape[0])])
         else:
-            G_callback = appctx.get("get_gradient", None)
+            G_callback = appctx.get(options_prefix+"get_gradient", None)
             dminus = ufl.grad
 
         # Get only the zero-th order term of the form
@@ -197,7 +197,7 @@ def coarsen(self, pc):
 
         cdegree = max(as_tuple(celement.degree()))
         if formdegree > 1 and cdegree > 1:
-            shift = appctx.get("hiptmair_shift", None)
+            shift = appctx.get(options_prefix+"shift", None)
             if shift is not None:
                 b = beta(test, shift * trial)
                 coarse_operator += ufl.Form(b.integrals_by_type("cell"))

firedrake/preconditioners/hypre_ads.py

@@ -13,10 +13,13 @@
 
 
 class HypreADS(PCBase):
+    _prefix = "hypre_ads_"
+
     def initialize(self, obj):
         A, P = obj.getOperators()
         appctx = self.get_appctx(obj)
         prefix = obj.getOptionsPrefix() or ""
+        options_prefix = prefix + self.prefix
         V = get_function_space(obj.getDM())
         mesh = V.mesh()
 
@@ -28,20 +31,20 @@ def initialize(self, obj):
 
         P1 = FunctionSpace(mesh, "Lagrange", 1)
         NC1 = FunctionSpace(mesh, "N1curl" if mesh.ufl_cell().is_simplex() else "NCE", 1)
-        G_callback = appctx.get("get_gradient", None)
+        G_callback = appctx.get(options_prefix+"get_gradient", None)
         if G_callback is None:
             G = chop(assemble.assemble(interpolate(grad(TestFunction(P1)), NC1)).petscmat)
         else:
             G = G_callback(P1, NC1)
-        C_callback = appctx.get("get_curl", None)
+        C_callback = appctx.get(options_prefix+"get_curl", None)
         if C_callback is None:
             C = chop(assemble.assemble(interpolate(curl(TestFunction(NC1)), V)).petscmat)
         else:
             C = C_callback(NC1, V)
 
         pc = PETSc.PC().create(comm=obj.comm)
         pc.incrementTabLevel(1, parent=obj)
-        pc.setOptionsPrefix(prefix + "hypre_ads_")
+        pc.setOptionsPrefix(options_prefix)
         pc.setOperators(A, P)
 
         pc.setType('hypre')

firedrake/preconditioners/massinv.py

@@ -1,5 +1,5 @@
 from firedrake.preconditioners.assembled import AssembledPC
-from firedrake import inner, dx
+from firedrake import inner, dx, Constant
 
 __all__ = ("MassInvPC", )
 
@@ -23,7 +23,8 @@ def form(self, pc, test, trial):
         _, bcs = super(MassInvPC, self).form(pc)
 
         appctx = self.get_appctx(pc)
-        mu = appctx.get("mu", 1.0)
+        options_prefix = (pc.getOptionsPrefix() or "") + self._prefix
+        mu = appctx.get(options_prefix+"mu", Constant(1.0))
         a = inner((1/mu) * trial, test) * dx
         return a, bcs
 

firedrake/preconditioners/pcd.py

@@ -101,17 +101,19 @@ def initialize(self, pc):
         Kksp.setFromOptions()
         self.Kksp = Kksp
 
-        state = context.appctx["state"]
+        state = context._problem.u_restrict
 
-        Re = context.appctx.get("Re", 1.0)
+        Re = context.appctx.get(prefix+"re", Constant(1.0))
 
-        velid = context.appctx["velocity_space"]
+        vel_idx = opts.getInt(prefix+"velocity_space")
 
-        u0 = split(state)[velid]
+        u0 = split(state)[vel_idx]
         fp = 1.0/Re * inner(grad(p), grad(q))*dx + inner(u0, grad(p))*q*dx
 
         self.Re = Re
-        form_assembler = get_assembler(fp, bcs=None, form_compiler_parameters=context.fc_params, mat_type=self.Fp_mat_type, options_prefix=prefix + "Fp_")
+        form_assembler = get_assembler(
+            fp, bcs=None, form_compiler_parameters=context.fc_params,
+            mat_type=self.Fp_mat_type, options_prefix=prefix + "Fp_")
         self.Fp = form_assembler.allocate()
         self._assemble_Fp = form_assembler.assemble
         self._assemble_Fp(tensor=self.Fp)

firedrake/slate/static_condensation/hybridization.py

@@ -211,7 +211,7 @@ def initialize(self, pc):
 
         Smat = self.S.petscmat
 
-        nullspace = self.ctx.appctx.get("trace_nullspace", None)
+        nullspace = self.ctx.appctx.get(prefix+"trace_nullspace", None)
         if nullspace is not None:
             nsp = nullspace(TraceSpace)
             Smat.setNullSpace(nsp.nullspace())
@@ -238,7 +238,7 @@ def initialize(self, pc):
         trace_ksp.setOperators(Smat, Smat)
 
         # Option to add custom monitor
-        monitor = self.ctx.appctx.get('custom_monitor', None)
+        monitor = self.ctx.appctx.get(prefix+'custom_monitor', None)
         if monitor:
             monitor.add_reconstructor(self.backward_substitution)
             trace_ksp.setMonitor(monitor)

firedrake/slate/static_condensation/scpc.py

@@ -138,7 +138,7 @@ def initialize(self, pc):
         # Get nullspace for the condensed operator (if any).
         # This is provided as a user-specified callback which
         # returns the basis for the nullspace.
-        nullspace = self.cxt.appctx.get("condensed_field_nullspace", None)
+        nullspace = self.cxt.appctx.get(prefix+"nullspace", None)
         if nullspace is not None:
             nsp = nullspace(Vc)
             Smat.setNullSpace(nsp.nullspace())

firedrake/solving.py

@@ -20,6 +20,7 @@
 __all__ = ["solve"]
 
 import ufl
+from petsctools import AppContext
 
 import firedrake.linear_solver as ls
 from firedrake.matrix import MatrixBase
@@ -163,7 +164,7 @@ def _solve_varproblem(*args, **kwargs):
         form_compiler_parameters = {}
     form_compiler_parameters['scalar_type'] = ScalarType
 
-    appctx = kwargs.get("appctx", {})
+    appctx = kwargs.get("appctx", AppContext())
     if isinstance(eq.lhs, (ufl.Form, MatrixBase)) and isinstance(eq.rhs, ufl.BaseForm):
         # Create linear variational problem
         problem = vs.LinearVariationalProblem(eq.lhs, eq.rhs, u, bcs, Jp,
@@ -238,7 +239,7 @@ def _la_solve(A, x, b, **kwargs):
     if bcs is not None:
         raise RuntimeError("It is no longer possible to apply or change boundary conditions after assembling the matrix `A`; pass any necessary boundary conditions to `assemble` when assembling `A`.")
 
-    appctx = solver_parameters.get("appctx", {})
+    appctx = solver_parameters.get("appctx", AppContext())
     solver = ls.LinearSolver(A=A, P=P, solver_parameters=solver_parameters,
                              nullspace=nullspace,
                              transpose_nullspace=nullspace_T,

firedrake/solving_utils.py

@@ -141,9 +141,8 @@ class _SNESContext(object):
     :arg pmat_type: Indicates whether the preconditioner (if present) is assembled
         monolithically ('aij'), as a block sparse matrix ('nest') or
         matrix-free (as :class:`~.ImplicitMatrix`, 'matfree').
-    :arg appctx: Any extra information used in the assembler.  For the
-        matrix-free case this will contain the Newton state in
-        ``"state"``.
+    :arg appctx: A petsctools.AppContext containing application
+        context that is passed to the preconditioner if matrix-free.
     :arg pre_jacobian_callback: User-defined function called immediately
         before Jacobian assembly
     :arg post_jacobian_callback: User-defined function called immediately
@@ -194,14 +193,13 @@ def __init__(self, problem, mat_type, pmat_type, appctx=None,
         self._x = problem.u_restrict
 
         if appctx is None:
-            appctx = {}
+            from petsctools import AppContext
+            appctx = AppContext()
         # A split context will already get the full state.
         # TODO, a better way of doing this.
         # Now we don't have a temporary state inside the snes
         # context we could just require the user to pass in the
         # full state on the outside.
-        appctx.setdefault("state", self._x)
-        appctx.setdefault("form_compiler_parameters", self.fcp)
 
         self.appctx = appctx
         self.matfree = matfree