diff --git a/petsctools/__init__.py b/petsctools/__init__.py index 03b3682..caf4f23 100644 --- a/petsctools/__init__.py +++ b/petsctools/__init__.py @@ -1,3 +1,8 @@ +from .appctx import ( # noqa: F401 + AppContext, + push_appctx, + get_appctx, +) from .config import ( # noqa: F401 MissingPetscException, get_config, diff --git a/petsctools/appctx.py b/petsctools/appctx.py new file mode 100644 index 0000000..e3f3ddc --- /dev/null +++ b/petsctools/appctx.py @@ -0,0 +1,166 @@ +import itertools +from functools import cached_property +from contextlib import contextmanager +from petsctools.exceptions import PetscToolsAppctxException + +_global_appctx_stack = [] + +@contextmanager +def push_appctx(appctx): + _global_appctx_stack.append(appctx) + yield + _global_appctx_stack.pop() + + +def get_appctx(): + return _global_appctx_stack[-1] + + +class AppContextKey(int): + """A custom key type for AppContext.""" + + +class AppContext: + """ + Class for passing non-primitive types to PETSc python contexts. + + The PETSc.Options dictionary can only contain primitive types (str, + int, float, bool) as values. The AppContext allows other types to be + passed into PETSc solvers while still making use of the namespacing + provided by options prefixing. + + A typical usage is shown below. In this example we have a python PC + type `MyCustomPC` which requires additional data in the form of a + `MyCustomData` instance. + We can add the data to the AppContext with the `appctx.add` method, + but we need to tell `MyCustomPC` how to retrieve that data. The + `add` method returns a key which is a valid PETSc.Options entry, + i.e. a primitive type instance. This key is passed via PETSc.Options + with the 'custompc_somedata' prefix. + + NB: The user should never handle this key directly, it should only + ever be placed directly into the options dictionary. + + The data can be retrieved by giving the AppContext the (fully + prefixed) option for the key, in which case the AppContext will + internally fetch the key from the PETSc.Options and return the data. + + .. code-block:: python3 + + appctx = AppContext() + some_data = MyCustomData(5) + + opts = OptionsManager( + parameters={ + 'pc_type': 'python', + 'pc_python_type': 'MyCustomPC', + 'custompc_somedata': appctx.add(some_data)}, + options_prefix='solver') + + with opts.inserted_options(): + default = MyCustomData(10) + data = appctx.get('solver_custompc_somedata', default) + """ + + def __init__(self): + self._count = itertools.count(start=0) + self._data = {} + + def _keygen(self): + """ + Generate a new unique internal key. + + This should not called directly by the user. + """ + return AppContextKey(next(self._count)) + + def _key_from_option(self, option): + """ + Return the internal key for the PETSc option `option`. + + Parameters + ---------- + option : str + The PETSc option. + + Returns + ------- + key : AppContextKey + An internal key corresponding to `option`. + """ + return AppContextKey(self.options_object.getInt(option)) + + def add(self, val): + """ + Add a value to the application context and + return the autogenerated key for that value. + + The autogenerated key should be used as the value for the + corresponding entry in the solver_parameters dictionary. + + Parameters + ---------- + val : Any + The value to add to the AppContext. + + Returns + ------- + key : AppContextKey + The key to put into the PETSc Options dictionary. + """ + key = self._keygen() + self._data[key] = val + return key + + def __getitem__(self, option): + """ + Return the value with the key saved in `PETSc.Options()[option]`. + + Parameters + ---------- + option : Union[str, AppContextKey] + The PETSc option or key. + + Returns + ------- + val : Any + The value for the key `option`. + + Raises + ------ + PetscToolsAppctxException + If the AppContext does contain a value for `option`. + """ + try: + return self._data[self._key_from_option(option)] + except KeyError: + raise PetscToolsAppctxException( + f"AppContext does not have an entry for {option}") + + def get(self, option, default=None): + """ + Return the value with the key saved in PETSc.Options()[option], + or if it does not exist return default. + + Parameters + ---------- + option : Union[str, AppContextKey] + The PETSc option or key. + default : Any + The value to return if `option` is not in the AppContext + + Returns + ------- + val : Any + The value for the key `option`, or `default`. + """ + try: + return self[option] + except PetscToolsAppctxException: + return default + + @cached_property + def options_object(self): + """A PETSc.Options instance.""" + from petsc4py import PETSc + return PETSc.Options() diff --git a/petsctools/exceptions.py b/petsctools/exceptions.py index 70b0926..5d35e4c 100644 --- a/petsctools/exceptions.py +++ b/petsctools/exceptions.py @@ -6,5 +6,9 @@ class PetscToolsNotInitialisedException(PetscToolsException): """Exception raised when petsctools should have been initialised.""" +class PetscToolsAppctxException(PetscToolsException): + """Exception raised when the Appctx is missing an entry.""" + + class PetscToolsWarning(UserWarning): """Generic base class for petsctools warnings.""" diff --git a/petsctools/options.py b/petsctools/options.py index eaf972e..f955fd5 100644 --- a/petsctools/options.py +++ b/petsctools/options.py @@ -2,6 +2,7 @@ import functools import itertools import warnings +from petsctools.appctx import push_appctx from petsctools.exceptions import ( PetscToolsException, PetscToolsWarning, PetscToolsNotInitialisedException) @@ -408,7 +409,8 @@ def set_default_parameter(obj, key, val): def set_from_options(obj, parameters=None, - options_prefix=None): + options_prefix=None, + appctx=None): """Set up a PETSc object from the options in its OptionsManager. Calls ``obj.setOptionsPrefix`` and ``obj.setFromOptions`` whilst @@ -467,7 +469,11 @@ def set_from_options(obj, parameters=None, f" called for {petscobj2str(obj)}", PetscToolsWarning) - get_options(obj).set_from_options(obj) + if appctx is None: + get_options(obj).set_from_options(obj) + else: + with push_appctx(appctx): + get_options(obj).set_from_options(obj) def is_set_from_options(obj): diff --git a/tests/test_appctx.py b/tests/test_appctx.py new file mode 100644 index 0000000..544d43b --- /dev/null +++ b/tests/test_appctx.py @@ -0,0 +1,82 @@ +import pytest +import petsctools +from petsctools.exceptions import PetscToolsAppctxException + + +class JacobiTestPC: + prefix = "jacobi_" + def setFromOptions(self, pc): + appctx = petsctools.get_appctx() + prefix = (pc.getOptionsPrefix() or "") + self.prefix + self.scale = appctx[prefix + "scale"] + + def apply(self, pc, x, y): + y.pointwiseMult(x, self.scale) + + +@pytest.mark.skipnopetsc4py +def test_get_appctx(): + from numpy import allclose + PETSc = petsctools.init() + n = 4 + sizes = (n, n) + + appctx = petsctools.AppContext() + + diag = PETSc.Vec().createSeq(sizes) + diag.setSizes((n, n)) + diag.array[:] = [1, 2, 3, 4] + + mat = PETSc.Mat().createConstantDiagonal((sizes, sizes), 1.0) + + ksp = PETSc.KSP().create() + ksp.setOperators(mat, mat) + petsctools.set_from_options( + ksp, + parameters={ + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': f'{__name__}.JacobiTestPC', + 'jacobi_scale': appctx.add(diag) + }, + options_prefix="myksp", + appctx=appctx, + ) + + x, b = mat.createVecs() + b.setRandom() + + xcheck = x.duplicate() + xcheck.pointwiseMult(b, diag) + + with petsctools.inserted_options(ksp), petsctools.push_appctx(appctx): + ksp.solve(b, x) + + assert allclose(x.array_r, xcheck.array_r) + + +@pytest.mark.skipnopetsc4py +def test_appctx_key(): + PETSc = petsctools.init() + + appctx = petsctools.AppContext() + + param = 10 + options = PETSc.Options() + options['solver_param'] = appctx.add(param) + + # Can we access param via the prefixed option? + prm = appctx.get('solver_param') + assert prm is param + + prm = appctx['solver_param'] + assert prm is param + + # Can we set a default value? + default = 20 + prm = appctx.get('param', default) + assert prm is default + + # Will an invalid key raise an error + with pytest.raises(PetscToolsAppctxException): + appctx['param']