diff --git a/firedrake/__init__.py b/firedrake/__init__.py index cb57a882f8..6093f0c503 100644 --- a/firedrake/__init__.py +++ b/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 * diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 9ca4d17237..5fd2854bff 100644 --- a/firedrake/assemble.py +++ b/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 diff --git a/firedrake/preconditioners/assembled.py b/firedrake/preconditioners/assembled.py index b2f1ede79a..565d60c016 100644 --- a/firedrake/preconditioners/assembled.py +++ b/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) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 57580b6f86..e3149cbfc3 100644 --- a/firedrake/preconditioners/bddc.py +++ b/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 diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py index 26e3f984ca..b1c25f165a 100644 --- a/firedrake/preconditioners/facet_split.py +++ b/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") diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 4ba26a621f..db248d6de4 100644 --- a/firedrake/preconditioners/fdm.py +++ b/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) diff --git a/firedrake/preconditioners/gtmg.py b/firedrake/preconditioners/gtmg.py index 4d25ce7259..c3b3fb667c 100644 --- a/firedrake/preconditioners/gtmg.py +++ b/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. diff --git a/firedrake/preconditioners/hiptmair.py b/firedrake/preconditioners/hiptmair.py index f4071d7a10..2495c93130 100644 --- a/firedrake/preconditioners/hiptmair.py +++ b/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")) diff --git a/firedrake/preconditioners/hypre_ads.py b/firedrake/preconditioners/hypre_ads.py index b639c21f05..58bb1c3057 100644 --- a/firedrake/preconditioners/hypre_ads.py +++ b/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,12 +31,12 @@ 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: @@ -41,7 +44,7 @@ def initialize(self, obj): 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') diff --git a/firedrake/preconditioners/massinv.py b/firedrake/preconditioners/massinv.py index d29c704e8b..ea4f74e452 100644 --- a/firedrake/preconditioners/massinv.py +++ b/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 diff --git a/firedrake/preconditioners/pcd.py b/firedrake/preconditioners/pcd.py index dbc4671720..3f3b9367c2 100644 --- a/firedrake/preconditioners/pcd.py +++ b/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) diff --git a/firedrake/slate/static_condensation/hybridization.py b/firedrake/slate/static_condensation/hybridization.py index bd0fc761f0..29dd8add64 100644 --- a/firedrake/slate/static_condensation/hybridization.py +++ b/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) diff --git a/firedrake/slate/static_condensation/scpc.py b/firedrake/slate/static_condensation/scpc.py index a7536d7bda..0e831e6d35 100644 --- a/firedrake/slate/static_condensation/scpc.py +++ b/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()) diff --git a/firedrake/solving.py b/firedrake/solving.py index 480823d32d..2deafad008 100644 --- a/firedrake/solving.py +++ b/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, diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 393e5e2b8c..5a4bed115f 100644 --- a/firedrake/solving_utils.py +++ b/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 diff --git a/firedrake/variational_solver.py b/firedrake/variational_solver.py index 724c4257d3..37540e3760 100644 --- a/firedrake/variational_solver.py +++ b/firedrake/variational_solver.py @@ -179,8 +179,8 @@ def __init__(self, problem, *, solver_parameters=None, specify the near nullspace (for multigrid solvers). :kwarg solver_parameters: Solver parameters to pass to PETSc. This should be a dict mapping PETSc options to values. - :kwarg appctx: A dictionary containing application context that - is passed to the preconditioner if matrix-free. + :kwarg appctx: A petsctools.AppContext containing application + context that is passed to the preconditioner if matrix-free. :kwarg options_prefix: an optional prefix used to distinguish PETSc options. If not provided a unique prefix will be created. Use this option if you want to pass options @@ -422,8 +422,8 @@ class LinearVariationalSolver(NonlinearVariationalSolver): created. Use this option if you want to pass options to the solver from the command line in addition to through the ``solver_parameters`` dict. - :kwarg appctx: A dictionary containing application context that - is passed to the preconditioner if matrix-free. + :kwarg appctx: A petsctools.AppContext containing application + context that is passed to the preconditioner if matrix-free. :kwarg pre_jacobian_callback: A user-defined function that will be called immediately before Jacobian assembly. This can be used, for example, to update a coefficient function diff --git a/tests/firedrake/regression/test_fdm.py b/tests/firedrake/regression/test_fdm.py index 947ae43fbc..48beafedb6 100644 --- a/tests/firedrake/regression/test_fdm.py +++ b/tests/firedrake/regression/test_fdm.py @@ -312,6 +312,7 @@ def test_ipdg_direct_solver(fs): + inner(outer(u_exact, n), alpha(2*num_flux_b(v) - grad(v))) * ds_Dir) problem = LinearVariationalProblem(a, L, uh, bcs=bcs) + appctx = AppContext() solver = LinearVariationalSolver(problem, solver_parameters={ "mat_type": "matfree", "ksp_type": "cg", @@ -322,10 +323,11 @@ def test_ipdg_direct_solver(fs): "ksp_norm_type": "unpreconditioned", "pc_type": "python", "pc_python_type": "firedrake.PoissonFDMPC", + "fdm_eta": appctx.add(eta), "fdm_pc_type": "cholesky", "fdm_pc_factor_mat_solver_type": DEFAULT_DIRECT_SOLVER, "fdm_pc_factor_mat_ordering_type": "nd", - }, appctx={"eta": eta, }) + }, appctx=appctx) solver.solve() assert solver.snes.ksp.getIterationNumber() == 1 diff --git a/tests/firedrake/slate/test_hybrid_poisson_sphere.py b/tests/firedrake/slate/test_hybrid_poisson_sphere.py index 065f9cf125..11ba378fe6 100644 --- a/tests/firedrake/slate/test_hybrid_poisson_sphere.py +++ b/tests/firedrake/slate/test_hybrid_poisson_sphere.py @@ -26,24 +26,21 @@ def run_hybrid_poisson_sphere(MeshClass, refinement, hdiv_space): L = inner(f, v)*dx w = Function(W) + # Provide a callback to construct the trace nullspace + def nullspace_basis(T): + return VectorSpaceBasis(constant=True, comm=COMM_WORLD) + + appctx = AppContext() + params = { 'mat_type': 'matfree', 'ksp_type': 'preonly', 'pc_type': 'python', 'pc_python_type': 'firedrake.HybridizationPC', - 'hybridization': { - 'ksp_type': 'preonly', - 'pc_type': 'redundant', - 'redundant_pc_type': 'lu', - 'redundant_pc_factor': DEFAULT_DIRECT_SOLVER_PARAMETERS - } + 'hybridization': DEFAULT_DIRECT_SOLVER_PARAMETERS, + 'hybridization_trace_nullspace': appctx.add(nullspace_basis), } - # Provide a callback to construct the trace nullspace - def nullspace_basis(T): - return VectorSpaceBasis(constant=True) - - appctx = {'trace_nullspace': nullspace_basis} solve(a == L, w, solver_parameters=params, appctx=appctx) u_h, _ = w.subfunctions error = errornorm(u_exact, u_h)