diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index bd8b74c816..48d50e1526 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -1,6 +1,7 @@ from itertools import chain import numpy +import petsctools import ufl from pyop2 import op2 @@ -131,7 +132,7 @@ def check_snes_convergence(snes): %s""" % (snes.getIterationNumber(), msg)) -class _SNESContext(object): +class _SNESContext: """Context holding information for SNES callbacks. Parameters @@ -152,9 +153,6 @@ class _SNESContext(object): sub_pmat_type Indicates the matrix type for the sparse blocks in the preconditioner if pmat_type='nest', ignored otherwise. - appctx - Any extra information used in the assembler. For the - matrix-free case this will contain the Newton state in ``"state"``. pre_jacobian_callback User-defined function called immediately before Jacobian assembly. post_jacobian_callback @@ -171,6 +169,8 @@ class _SNESContext(object): pre_apply_bcs If `False`, the problem is linearised around the initial guess before imposing the boundary conditions. + legacy_appctx_prefix + Options prefix used to find appctx entries passed in the legacy fashion. The idea here is that the SNES holds a shell DM which contains this object as "user context". When the SNES calls back to the @@ -184,12 +184,12 @@ def __init__(self, problem, mat_type: str, pmat_type: str, sub_mat_type: str | None = None, sub_pmat_type: str | None = None, - appctx: dict | None = None, pre_jacobian_callback=None, pre_function_callback=None, post_jacobian_callback=None, post_function_callback=None, options_prefix: str | None = None, transfer_manager=None, - pre_apply_bcs: bool = True): + pre_apply_bcs: bool = True, + legacy_appctx_prefix: str | None = None): from firedrake.assemble import get_assembler if pmat_type is None: @@ -216,17 +216,7 @@ def __init__(self, problem, # Function to hold current guess self._x = problem.u_restrict - if appctx is None: - appctx = {} - # 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.legacy_appctx_prefix = legacy_appctx_prefix self.matfree = matfree self.pmatfree = pmatfree self.F = problem.F @@ -286,10 +276,10 @@ def reconstruct(self, problem=None, mat_type=None, pmat_type=None, **kwargs): default_options = { "sub_mat_type": self.sub_mat_type, "sub_pmat_type": self.sub_pmat_type, - "appctx": self.appctx, "options_prefix": self.options_prefix, "transfer_manager": self.transfer_manager, "pre_apply_bcs": self.pre_apply_bcs, + "legacy_appctx_prefix": self.legacy_appctx_prefix, } for k, v in default_options.items(): if kwargs.get(k) is None: @@ -342,6 +332,23 @@ def transfer_manager(self, manager): raise ValueError("Must set transfer manager before first use.") self._transfer_manager = manager + # TODO: remove this method eventually + @cached_property + def appctx(self) -> dict: + ctx = petsctools.AppContext(self.legacy_appctx_prefix).getAll() + + # 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. + if "state" not in ctx: + ctx["state"] = self._x + if "form_compiler_parameters" not in ctx: + ctx["form_compiler_parameters"] = self.fcp + + return ctx + def set_function(self, snes): r"""Set the residual evaluation function""" with self._F.dat.vec_wo as v: diff --git a/firedrake/variational_solver.py b/firedrake/variational_solver.py index 708b314e99..fc6ce96dc0 100644 --- a/firedrake/variational_solver.py +++ b/firedrake/variational_solver.py @@ -170,6 +170,17 @@ def compute_bc_lifting(J: ufl.BaseForm | slate.TensorBase, return F +_solver_id = -1 +"""A unique number for each variational solver. + +This is used to uniquely key appctx entries (now deprecated practice). + +The fact that this value is not parallel safe is not a concern because PETSc +options are only read rank-locally. + +""" + + class NonlinearVariationalSolver(NonlinearVariationalSolverMixin): r"""Solves a :class:`NonlinearVariationalProblem`.""" @@ -259,6 +270,9 @@ def update_diffusivity(current_solution): pre_jacobian_callback=update_diffusivity) """ + global _solver_id + _solver_id += 1 + assert isinstance(problem, NonlinearVariationalProblem) solver_parameters = flatten_parameters(solver_parameters or {}) @@ -278,10 +292,17 @@ def update_diffusivity(current_solution): ksp_defaults=self.DEFAULT_KSP_PARAMETERS, snes_defaults=self.DEFAULT_SNES_PARAMETERS) + # Move entries from the deprecated appctx into the parameter dictionary + if appctx is not None: + legacy_appctx_prefix = f"firedrake_legacy_appctx_{_solver_id}_" + for key, value in appctx: + solver_parameters[f"{legacy_appctx_prefix}_{key}"] = value + else: + legacy_appctx_prefix = None + self.options_manager = OptionsManager(solver_parameters, options_prefix, default_prefix="firedrake") - # Now the correct parameters live in self.parameters (via the - # OptionsManager mixin) + # Now the correct parameters live in self.parameters mat_type = self.parameters.get("mat_type") pmat_type = self.parameters.get("pmat_type") sub_mat_type = self.parameters.get("sub_mat_type") @@ -291,13 +312,13 @@ def update_diffusivity(current_solution): pmat_type=pmat_type, sub_mat_type=sub_mat_type, sub_pmat_type=sub_pmat_type, - appctx=appctx, pre_jacobian_callback=pre_jacobian_callback, pre_function_callback=pre_function_callback, post_jacobian_callback=post_jacobian_callback, post_function_callback=post_function_callback, options_prefix=self.options_prefix, - pre_apply_bcs=pre_apply_bcs) + pre_apply_bcs=pre_apply_bcs, + legacy_appctx_prefix=legacy_appctx_prefix) self.snes = PETSc.SNES().create(comm=problem.dm.comm) diff --git a/pyproject.toml b/pyproject.toml index 74d38ccf8e..33f1934b1a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "packaging", # TODO RELEASE # "petsc4py==3.25.0", - "petsctools>=2026.0", + "petsctools @ git+https://github.com/firedrakeproject/petsctools@connorjward/appctx-getall", "pkgconfig", "progress", "pyadjoint-ad>=2026.4.0", diff --git a/requirements-build.txt b/requirements-build.txt index 7916a00db7..3105a344a8 100644 --- a/requirements-build.txt +++ b/requirements-build.txt @@ -6,7 +6,7 @@ mpi4py>3; python_version >= '3.13' mpi4py; python_version < '3.13' numpy pkgconfig -petsctools +petsctools @ git+https://github.com/firedrakeproject/petsctools@connorjward/appctx-getall pybind11 setuptools>=77.0.3