diff --git a/petsctools/__init__.py b/petsctools/__init__.py index 907379f..9e3cc67 100644 --- a/petsctools/__init__.py +++ b/petsctools/__init__.py @@ -15,6 +15,10 @@ # is not available then attempting to access these attributes will raise an # informative error. if PETSC4PY_INSTALLED: + from .appctx import ( # noqa: F401 + AppContext, + AppContextManager, + ) from .citation import ( # noqa: F401 add_citation, cite, @@ -65,6 +69,8 @@ def __getattr__(name): "set_default_parameter", "DefaultOptionSet", "PCBase", + "AppContext", + "AppContextManager", } if name in petsc4py_attrs: raise ImportError( diff --git a/petsctools/appctx.py b/petsctools/appctx.py new file mode 100644 index 0000000..1b3e9bc --- /dev/null +++ b/petsctools/appctx.py @@ -0,0 +1,197 @@ +from typing import Any +import itertools +from functools import cached_property +from contextlib import contextmanager +from petsctools.exceptions import PetscToolsAppctxException + +_global_appctx_data = {} +"""The global storage for user data with arbitrary python types.""" + + +class AppContextKey(str): + """A custom key type for AppContext.""" + + _count = itertools.count() + + @classmethod + def _generate_key(cls): + return f"petsctools_appctx_key_{next(cls._count)}" + + +class AppContext: + def __init__(self, prefix: str | None = None): + from petsctools.options import _validate_prefix + + # possibly append underscore or cast to str + self._prefix = _validate_prefix(prefix or "") + + @property + def prefix(self) -> str: + return self._prefix + + @cached_property + def options_object(self): + """A PETSc.Options instance.""" + from petsc4py import PETSc + + return PETSc.Options() + + def _key_from_option(self, option: str) -> AppContextKey: + """ + Return the internal key for the PETSc option `option`. + + Parameters + ---------- + option + The PETSc option. + + Returns + ------- + key + An internal key corresponding to ``option``. + """ + return AppContextKey( + self.options_object.getString(self.prefix + option) + ) + + def __getitem__(self, option: str | AppContextKey, /) -> Any: + """ + Return the value with the key saved in ``PETSc.Options()[option]``. + + Parameters + ---------- + option : + The PETSc option or key. + + Returns + ------- + val : + The value for the key `option`. + + Raises + ------ + PetscToolsAppctxException + If the AppContext does contain a value for `option`. + """ + try: + return _global_appctx_data[self._key_from_option(option)] + except KeyError: + raise PetscToolsAppctxException( + f"AppContext does not have an entry for {option}" + ) + + def __setitem__(self, option: str, value: Any, /): + key = AppContextKey._generate_key() + self.options_object[self.prefix + option] = key + _global_appctx_data[key] = value + + def get( + self, option: str | AppContextKey, default: Any | None = None + ) -> Any: + """ + Return the value with the key saved in ``PETSc.Options()[option]``, + or if it does not exist return default. + + Parameters + ---------- + option : + The PETSc option or key. + default : + The value to return if ``option`` is not in the ``AppContext`` + + Returns + ------- + val : + The value for the key ``option``, or ``default``. + """ + try: + return self[option] + except PetscToolsAppctxException: + return default + + +class AppContextManager: + """ + 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._data = {} + + def add(self, val: Any) -> AppContextKey: + """ + 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 + The value to add to the AppContext. + + Returns + ------- + key + The key to put into the PETSc Options dictionary. + """ + key = AppContextKey._generate_key() + self._data[key] = val + return key + + @contextmanager + def inserted_appctx(self): + # We don't overwrite existing entries in the global data, + # so we need to keep track of what we do actually put in + # so we don't accidentally remove something we shouldn't. + to_delete = set() + try: + for k, v in self._data.items(): + if k not in _global_appctx_data: + _global_appctx_data[k] = v + to_delete.add(k) + yield + finally: + for k in self._data: + if k in to_delete: + del _global_appctx_data[k] + to_delete.remove(k) + assert len(to_delete) == 0 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 1c6b078..a06f1ad 100644 --- a/petsctools/options.py +++ b/petsctools/options.py @@ -15,6 +15,7 @@ PetscToolsWarning, PetscToolsNotInitialisedException, ) +from petsctools.appctx import AppContextManager _commandline_options = None @@ -388,6 +389,8 @@ class OptionsManager: default_options_set The prefix set for any default shared with other solvers. See :class:`DefaultOptionSet` for more information. + appctx + The :class:`AppContextManager` containing user python data. See Also -------- @@ -398,6 +401,8 @@ class OptionsManager: is_set_from_options inserted_options DefaultOptionSet + AppContext + AppContextManager """ count = itertools.count() @@ -405,7 +410,8 @@ class OptionsManager: def __init__(self, parameters: dict, options_prefix: str | None = None, default_prefix: str | None = None, - default_options_set: DefaultOptionSet | None = None): + default_options_set: DefaultOptionSet | None = None, + appctx: AppContextManager | None = None): super().__init__() if parameters is None: parameters = {} @@ -471,6 +477,10 @@ def __init__(self, parameters: dict, self.parameters[k[len(self.options_prefix):]] = v self._setfromoptions = False + + # user data + self.appctx = appctx + # Keep track of options used between invocations of inserted_options(). self._used_options = set() @@ -503,18 +513,29 @@ def set_default_parameter(self, key: str, val: Any) -> None: def set_from_options(self, petsc_obj): """Set up petsc_obj from the options database. - :arg petsc_obj: The PETSc object to call setFromOptions on. + Before calling ``petsc_obj.setFromOptions``, the options from + this OptionsManager's ``parameters`` are inserted into the global + :class:`PETSc.Options`, and if this OptionsManager has an ``appctx`` + then all entries are inserted into the :class:`AppContext`. - Raises PetscToolsWarning if this method has already been called. + Parameters + ---------- + petsc_obj + The PETSc object to call setFromOptions on. + + Raises + ------ + PetscToolsWarning + If this method has already been called. - Matt says: "Only ever call setFromOptions once". This - function ensures we do so. """ + # Matt says: "Only ever call setFromOptions once". This + # function ensures we do so. if not self._setfromoptions: + # Call setfromoptions inserting appropriate options + # and user data into the global databases. with self.inserted_options(): petsc_obj.setOptionsPrefix(self.options_prefix) - # Call setfromoptions inserting appropriate options into - # the options database. petsc_obj.setFromOptions() self._setfromoptions = True else: @@ -525,11 +546,18 @@ def set_from_options(self, petsc_obj): @contextlib.contextmanager def inserted_options(self): """Context manager inside which the petsc options database - contains the parameters from this object.""" + contains the parameters from this object. + If this OptionsManager has an ``appctx`` then all entries + are inserted into the :class:`AppContext`. + """ try: for k, v in self.parameters.items(): self.options_object[self.options_prefix + k] = v - yield + if self.appctx: + with self.appctx.inserted_appctx(): + yield + else: + yield finally: for k in self.to_delete: if self.options_object.used(self.options_prefix + k): @@ -563,7 +591,8 @@ def attach_options( parameters: dict | None = None, options_prefix: str | None = None, default_prefix: str | None = None, - default_options_set: DefaultOptionSet | None = None + default_options_set: DefaultOptionSet | None = None, + appctx: AppContextManager | None = None, ) -> None: """Set up an :class:`OptionsManager` and attach it to a PETSc Object. @@ -579,6 +608,8 @@ def attach_options( Base string for autogenerated default prefixes. default_options_set The prefix set for any default shared with other solvers. + appctx + The :class:`AppContextManager` containing user python data. See Also -------- @@ -596,7 +627,8 @@ def attach_options( parameters=parameters, options_prefix=options_prefix, default_prefix=default_prefix, - default_options_set=default_options_set + default_options_set=default_options_set, + appctx=appctx, ) obj.setAttr("options", options) @@ -690,7 +722,8 @@ def set_from_options( parameters: dict | None = None, options_prefix: str | None = None, default_prefix: str | None = None, - default_options_set: DefaultOptionSet | None = None + default_options_set: DefaultOptionSet | None = None, + appctx: AppContextManager | None = None, ) -> None: """Set up a PETSc object from the options in its :class:`OptionsManager`. @@ -716,6 +749,8 @@ def set_from_options( Base string for autogenerated default prefixes. default_options_set The prefix set for any default shared with other solvers. + appctx + An application context for passing non-native python types. Raises ------ @@ -734,6 +769,7 @@ def set_from_options( OptionsManager.set_from_options attach_options DefaultOptionSet + AppContextManager """ if has_options(obj): if parameters is not None or options_prefix is not None: @@ -753,7 +789,8 @@ def set_from_options( obj, parameters=parameters, options_prefix=options_prefix, default_prefix=default_prefix, - default_options_set=default_options_set + default_options_set=default_options_set, + appctx=appctx, ) if is_set_from_options(obj): @@ -797,6 +834,8 @@ def is_set_from_options(obj: petsc4py.PETSc.Object) -> bool: def inserted_options(obj): """Context manager inside which the PETSc options database contains the parameters from this object's :class:`OptionsManager`. + If the OptionsManager has an ``appctx`` then all entries are + inserted into the :class:`AppContext`. Parameters ---------- diff --git a/tests/test_appctx.py b/tests/test_appctx.py new file mode 100644 index 0000000..b30f01d --- /dev/null +++ b/tests/test_appctx.py @@ -0,0 +1,125 @@ +import pytest +import petsctools +from petsctools.exceptions import PetscToolsAppctxException + + +class JacobiTestPC: + prefix = "jacobi_" + + def setFromOptions(self, pc): + from petsc4py import PETSc + prefix = (pc.getOptionsPrefix() or "") + self.prefix + + prefixed_appctx = PETSc.Options().getBool( + prefix + "prefixed_appctx") + + if prefixed_appctx: + appctx = petsctools.AppContext(prefix) + self.scale = appctx["scale"] + else: + appctx = petsctools.AppContext() + self.scale = appctx[prefix + "scale"] + + def apply(self, pc, x, y): + y.pointwiseMult(x, self.scale) + + +@pytest.mark.skipnopetsc4py +@pytest.mark.parametrize("use_prefix", ["with_prefix", "without_prefix"]) +def test_appctx_context_manager(use_prefix): + PETSc = petsctools.init() + n = 4 + sizes = (n, n) + + 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) + + appctx = petsctools.AppContextManager() + + petsctools.set_from_options( + ksp, + parameters={ + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': f'{__name__}.JacobiTestPC', + 'jacobi_scale': appctx.add(diag), + 'jacobi_prefixed_appctx': use_prefix == "with_prefix", + }, + options_prefix="myksp", + appctx=appctx, + ) + + x, b = mat.createVecs() + b.setRandom() + + xcheck = x.duplicate() + xcheck.pointwiseMult(b, diag) + + with petsctools.inserted_options(ksp): + ksp.solve(b, x) + + assert (x - xcheck).norm() < 1e-14 + + +@pytest.mark.skipnopetsc4py +def test_appctx_key(): + PETSc = petsctools.init() + + manager = petsctools.AppContextManager() + + prefix0_param = 10 + options = PETSc.Options() + options['prefix0_param'] = manager.add(prefix0_param) + + appctx = petsctools.AppContext() + + # The param shouldn't be in the global dictionary yet + with pytest.raises(PetscToolsAppctxException): + appctx['param'] + + # Can we access param via the prefixed option? + with manager.inserted_appctx(): + prm = appctx.get('prefix0_param') + assert prm is prefix0_param + + prm = appctx['prefix0_param'] + assert prm is prefix0_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'] + + # Now try with a prefixed AppContext + + # First add a param option with a different prefix + prefix1_param = 20 + options['prefix1_param'] = manager.add(prefix1_param) + + appctx0 = petsctools.AppContext('prefix0') + appctx1 = petsctools.AppContext('prefix1') + + with manager.inserted_appctx(): + # This should only see prefix0 entries + prm = appctx0.get('param') + assert prm is prefix0_param + + prm = appctx0['param'] + assert prm is prefix0_param + + # This should only see prefix1 entries + prm = appctx1.get('param') + assert prm is prefix1_param + + prm = appctx1['param'] + assert prm is prefix1_param diff --git a/tests/test_config.py b/tests/test_config.py index 708b38a..82aef55 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -24,7 +24,7 @@ def test_get_petsc_dirs(): petsc_dir = petsctools.get_petsc_dir() petsc_arch = petsctools.get_petsc_arch() - expected = (petsc_dir, f"{petsc_dir}/{petsc_arch}") + expected = (petsc_dir, f"{petsc_dir}/{petsc_arch}") assert petsctools.get_petsc_dirs() == expected expected = (