WIP: Namespaced AppContext

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,

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()

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."""

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):

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']