diff --git a/docs/source/appctx.rst b/docs/source/appctx.rst new file mode 100644 index 0000000..98c9625 --- /dev/null +++ b/docs/source/appctx.rst @@ -0,0 +1,119 @@ +The OptionsManager and the AppContext +------------------------------------- + +The PETSc options provide a simple but powerful DSL for configuring composable solvers. +However, their main limitation is that the values of each option is limited to primitive C types, e.g. ``str``, ``float``, ``int``, or ``complex``. +Sometimes more advanced data is useful or essential for building a particular solver. + +:class:`petsctools.AppContext <.appctx.AppContext>` fulfils this need by providing a means of passing arbitrary Python types through to Python PETSc types (e.g. Python type PCs). +In this demo we show how to use the :class:`~.appctx.AppContext` to pass data to a custom Python type PC using the variable coefficient diffusion equation as an example. + + +Diffusion equation with variable coefficients +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The diffusion equation with coefficient :math:`\sigma(x)` depending on the spatial coordinate is: + +.. math:: + + u - \nabla\cdot\left(\sigma(x)\nabla u\right) = b + +We will solve this matrix with finite differences with the standard 3 point central stencil. +The particular details of the discretisation are not essential for this demo so we will be brief in the description. + +If :math:`D` is the assembled matrix for the finite difference gradient stencil, and :math:`\Sigma` is a diagonal matrix with the value of the diffusion coefficient at each grid point, then the assembled matrix for the diffusion equation is: + +.. math:: + + Au = \left(I + D^{T}\Sigma D\right)u = b, + \quad + \Sigma_{ii} = \sigma(x_{i}) + +The following Python function takes a numpy array ``sigma`` with the value of :math:`\sigma` at each grid point and assembles a sparse (``aij``) PETSc Mat for the diffusion equation. +We will use it later to build the :class:`~petsc4py.PETSc.Mat` for a :class:`~petsc4py.PETSc.KSP` to solve the diffusion equation. + +.. literalinclude:: ../../tests/docs/test_appctx_docs.py + :language: python3 + :dedent: + :start-after: [appctx_docs create_mat-start] + :end-before: [appctx_docs create_mat-end] + +A PC needing Python data +~~~~~~~~~~~~~~~~~~~~~~~~ + +Imagine a scenario where we need to solve :math:`A` multiple times and :math:`\sigma` changes slightly each time, for example if we are solving the unsteady diffusion equation with time-varying coefficients. +Rather than recomputing a preconditioner every time :math:`A` changes, we might instead find a representative :math:`\sigma_{p}` and use that to compute a preconditioner which can be reused for all solves. + +For simplicity, in this demo the preconditioner :math:`P` will just be the diagonal of the assembled matrix :math:`A_{p}` for the diffusion equation with :math:`\sigma_{p}` with a simple scaling factor :math:`\omega`: + +.. math:: + + A_{p} = I + D^{T}\Sigma_{p} D, + \quad + P = \omega^{-1}\mathrm{diag}(A_{p}). + +The diagonal of a matrix is clearly not expensive to compute, but in practice we would use a factorisation of :math:`A_{p}` which would be more expensive to compute and so would be more worthwhile reusing. + +The preconditioner defined above is implemented with a Python type PC called ``DiffusionJacobiPC`` in the code below. +Constructing :math:`P` requires two values, :math:`\sigma_{p}` and :math:`\omega`, which must be provided by the user. + +1. The scaling factor :math:`\omega` is just a real number, and can therefore be passed as usual via the :class:`PETSc.Options ` using the ``"djacobi_scale"`` option. + +2. The diffusion coefficient at each grid point :math:`\sigma_{p}(x_{i})` is defined as a numpy array. + This is clearly not a primitive type and so cannot be passed via the :class:`PETSc.Options ` directly. + Instead, we access it via the :class:`~.appctx.AppContext` using the ``"djacobi_sigma"`` key. + +The :class:`~petsctools.appctx.AppContext` mimics the :class:`PETSc.Options ` very closely, but can contain arbitrary Python data. +We will see below how to add ``sigma`` into the :class:`~petsctools.appctx.AppContext` so that it is available to the ``DiffusionJacobiPC``. + +.. literalinclude:: ../../tests/docs/test_appctx_docs.py + :language: python3 + :dedent: + :start-after: [appctx_docs pc-start] + :end-before: [appctx_docs pc-end] + +Assembling the system +~~~~~~~~~~~~~~~~~~~~~ + +We specify the diffusion coefficient as some random variations :math:`\sigma'` around a constant value :math:`\overline{\sigma}`, i.e. :math:`\sigma(x) = \overline{\sigma} + \sigma'(x)`. +Assuming that :math:`\sigma'` is the component that may vary from solve to solve, we use :math:`\sigma_{p}=\overline{\sigma}`. + +.. literalinclude:: ../../tests/docs/test_appctx_docs.py + :language: python3 + :dedent: + :start-after: [appctx_docs create_ksp-start] + :end-before: [appctx_docs create_ksp-end] + +The Options and the AppContext +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Now we configure ``ksp`` by passing PETSc options as key-value pairs in the ``parameters`` dictionary to :func:`petsctools.set_from_options <.options.set_from_options>`. +This function will create a :class:`petsctools.OptionsManager <.options.OptionsManager>` and attach it to ``ksp``. + +We can see that common options, e.g. ``"ksp_type"`` are set as usual in the ``parameters`` dictionary. +However, when we come to the ``"djacobi_sigma"`` value we use the :class:`petsctools.AppContextManager <.appctx.AppContextManager>` class. +The :meth:`AppContextManager.add <.appctx.AppContextManager.add>` method returns a unique value which is used to associate a PETSc option to whatever data was passed to ``add``. +For example, adding the key-value pair ``"djacobi_sigma": appmngr.add(sigma_p)`` to the ``parameters`` dictionary means that, during the solve, we will be able to access ``sigma_p`` via ``AppContext()["djacobi_sigma"]``. + +We then pass the :class:`~.appctx.AppContextManager` to :func:`~.options.set_from_options` so that it can be attached to ``ksp`` and its data can be made available later on during the solve. + +.. literalinclude:: ../../tests/docs/test_appctx_docs.py + :language: python3 + :dedent: + :start-after: [appctx_docs set_from_options-start] + :end-before: [appctx_docs set_from_options-end] + +Solving the KSP +~~~~~~~~~~~~~~~ + +Now we come to actually solving the linear equation :math:`Au=b`. +To avoid memory leaks, :func:`~.options.set_from_options` does not permanently insert the contents of ``parameters`` and the ``appmngr`` into the global :class:`PETSc.Options ` and :class:`~.appctx.AppContext` databases respectively. +Instead, we use the :func:`petsctools.inserted_options <.options.inserted_options>` context manager. +On entry, this context manager inserts the contents of ``parameters`` and ``appmngr`` into the global databases, and on exit it removes them again. +This means that we need to use the :func:`~.options.inserted_options` context manager whenever these entries will be needed, for example during the solve when the KSP and PC are being set up. + +.. literalinclude:: ../../tests/docs/test_appctx_docs.py + :language: python3 + :dedent: + :start-after: [appctx_docs solve-start] + :end-before: [appctx_docs solve-end] diff --git a/docs/source/conf.py b/docs/source/conf.py index eadb2c7..257fe82 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -19,8 +19,14 @@ extensions = [ "sphinx.ext.apidoc", "sphinx.ext.napoleon", + "sphinx.ext.intersphinx" ] +# Document special methods +autodoc_default_options = { + 'special-members': '__getitem__', +} + templates_path = ["_templates"] exclude_patterns = [] @@ -32,6 +38,13 @@ # -- sphinx.ext.apidoc configuration ------------------------------------------ -apidoc_modules = [ - {"path": "../../petsctools", "destination": "generated"}, -] +apidoc_modules = [{ + "path": "../../petsctools", + "destination": "generated", +}] + +# -- sphinx.ext.intersphinx configuration ------------------------------------- + +intersphinx_mapping = { + 'petsc4py': ('https://petsc.org/release/petsc4py/', None), +} diff --git a/docs/source/index.rst b/docs/source/index.rst index 978da07..f4bb9f6 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -9,4 +9,5 @@ petsctools provides Pythonic extensions for petsc4py and slepc4py. examples cython + appctx generated/modules diff --git a/petsctools/__init__.py b/petsctools/__init__.py index 907379f..88df751 100644 --- a/petsctools/__init__.py +++ b/petsctools/__init__.py @@ -15,6 +15,11 @@ # is not available then attempting to access these attributes will raise an # informative error. if PETSC4PY_INSTALLED: + from .appctx import ( # noqa: F401 + AppContext, + AppContextManager, + PetscToolsAppctxException, + ) from .citation import ( # noqa: F401 add_citation, cite, @@ -65,6 +70,9 @@ def __getattr__(name): "set_default_parameter", "DefaultOptionSet", "PCBase", + "AppContext", + "AppContextManager", + "PetscToolsAppctxException", } if name in petsc4py_attrs: raise ImportError( diff --git a/petsctools/appctx.py b/petsctools/appctx.py new file mode 100644 index 0000000..661c3a1 --- /dev/null +++ b/petsctools/appctx.py @@ -0,0 +1,282 @@ +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. + + Warning + ------- + This type should not be instantiated directly by the user, it + will be generated as needed by the :class:`.AppContextManager`. + + See Also + -------- + .AppContext + .AppContextManager + """ + + _count = itertools.count() + + @classmethod + def _generate_key(cls): + return f"petsctools_appctx_key_{next(cls._count)}" + + +class AppContext: + """ + A dictionary-like object to pass Python data to solvers analogously to + passing primitive types with :class:`petsc4py.PETSc.Options`. + + The ``PETSc.Options`` dictionary can only contain primitive types (e.g. + str, int, float, bool) as values. The ``AppContext`` allows other Python + types to be passed into PETSc solvers while still making use of the + namespacing provided by options prefixing. + + This class *must* be used in conjunction with the + :class:`.AppContextManager`. The example below shows how to use these + classes together. + + A typical use case is a Python PC type, here called ``MyCustomPC``, which + requires some data which is a non-primitive Python type, here an instance + of ``MyCustomData``. Non-primitive types cannot be passed via the + ``PETSc.Options``, so instead it is accessed via the + ``petsctools.AppContext`` using a (fully prefixed) key. + + .. code-block:: python3 + + class MyCustomData: + ... + + class MyCustomPC: + def setUp(self, pc): + prefix = pc.getOptionsPrefix() + option_key = 'custompc_data' + self.data = petsctools.AppContext()[prefix+option_key] + # or: + # self.data = petsctools.AppContext(prefix)[option_key] + ... + + Data is added to the ``AppContext`` using an :class:`.AppContextManager`, + which stores Python data associated with a specific PETSc object (e.g. a + KSP), in conjunction with :func:`.set_from_options` (or with the + :class:`.OptionsManager` directly). + Note that data *cannot* be added to the ``AppContext`` directly. + + The ``AppContextManager`` is created before calling ``set_from_options``. + When specifying the option for the Python data (e.g. the + ``"custompc_data"`` option key), the value in the ``parameters`` dictionary + must be ``appmngr.add(data)`` rather than the ``data`` object itself. + The ``AppContextManager`` is then passed to ``set_from_options``. + + .. code-block:: python3 + + data = MyCustomData(...) + + appmngr = petsctools.AppContextManager() + + petsctools.set_from_options( + ksp, + parameters={ + 'pc_type': 'python', + 'pc_python_type': 'MyCustomPC', + 'custompc_data': appmngr.add(data)}, + options_prefix='solver', + appmngr=appmngr) + + When the :func:`.inserted_options` context manager is used for the PETSc + object, all entries added to the ``petsctools.AppContextManager`` will be + inserted into the global ``petsctools.AppContext`` database (with the + ``options_prefix`` prepended to the keys), and they can be accessed inside + the solver as shown above. + When the context manager exits, the entries are removed from the global + ``AppContext``. + + .. code-block:: python3 + + with petsctools.inserted_options(ksp): + ksp.solve(b, x) + + Parameters + ---------- + prefix : + If provided, all option keys passed to ``__getitem__`` or ``get`` will + be prepended with this prefix before searching in the global database. + i.e. ``AppContext("prefix_")["option"]`` is equivalent to + ``AppContext()["prefix_option"]``. + + See Also + -------- + .AppContextManager + .OptionsManager + petsc4py.PETSc.Options + """ + 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: + """The prefix prepended to all keys before + before searching the global database. + """ + return self._prefix + + @cached_property + def options_object(self): + """A :class:`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 corresponding to the key ``option`` in the + :class:`PETSc.Options ` dictionary. + + If this ``AppContext`` instance has a prefix then the value + corresponding to the key ``self.prefix + option`` will be returned. + + Parameters + ---------- + option : + The PETSc option or key. + + Returns + ------- + Any : + 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 get( + self, option: str | AppContextKey, default: Any | None = None + ) -> Any: + """ + Return the value corresponding to the key ``option`` in the + :class:`PETSc.Options ` dictionary, + or the ``default`` value if ``option`` is not found in the + global options database. + + If this ``AppContext`` instance has a prefix then the value + corresponding to the key ``self.prefix + option`` will be returned. + + Parameters + ---------- + option : + The PETSc option or key. + default : + The value to return if ``option`` is not in the ``AppContext`` + + Returns + ------- + Any : + The value for the key ``option``, or ``default``. + """ + try: + return self[option] + except PetscToolsAppctxException: + return default + + +class AppContextManager: + """ + Class for storing Python data associated with a particular PETSc object. + + This class must be used in conjunction with the :class:`.AppContext` + class and :func:`.set_from_options` (or :class:`.OptionsManager`). + See the documentation for the :class:`.AppContext` for a description + of how these classes are used together. + + See Also + -------- + .AppContext + .OptionsManager + petsc4py.PETSc.Options + """ + + def __init__(self): + self._data = {} + + def add(self, val: Any) -> AppContextKey: + """ + Add a value to be inserted into the global :class:`.AppContext` + database by the :func:`.inserted_options` context manager, or + by the ``AppContextManager`` directly with + :meth:`.AppContextManager.inserted_appctx`. + + The autogenerated key returned by this method must only be used + as the value for the corresponding entry in the parameters dictionary + passed to :func:`.set_from_options`. + + Parameters + ---------- + val + The value to be inserted into the ``AppContext``. + + Returns + ------- + AppContextKey + The key to put into the ``PETSc.Options`` dictionary. + """ + key = AppContextKey._generate_key() + self._data[key] = val + return key + + @contextmanager + def inserted_appctx(self): + """Context manager inside which the global :class:`.AppContext` + database contains the entries added to this ``AppContextManager``. + """ + # 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..be92fd5 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. + appmngr + 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, + appmngr: 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.appmngr = appmngr + # 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 ``appmngr`` + 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 ``appmngr`` 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.appmngr: + with self.appmngr.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, + appmngr: 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. + appmngr + 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, + appmngr=appmngr, ) 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, + appmngr: 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. + appmngr + 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, + appmngr=appmngr, ) 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 ``appmngr`` then all entries are + inserted into the :class:`AppContext`. Parameters ---------- diff --git a/pyproject.toml b/pyproject.toml index 1bf7d60..81efaec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,9 +32,9 @@ requires = ["setuptools>=77.0.3"] build-backend = "setuptools.build_meta" [dependency-groups] -docs = ["sphinx"] +docs = ["sphinx", "numpy"] lint = ["ruff"] -test = ["pytest"] +test = ["pytest", "numpy"] ci = [ {include-group = "docs"}, {include-group = "lint"}, diff --git a/tests/docs/test_appctx_docs.py b/tests/docs/test_appctx_docs.py new file mode 100644 index 0000000..3f5cdf1 --- /dev/null +++ b/tests/docs/test_appctx_docs.py @@ -0,0 +1,124 @@ +import pytest +# [appctx_docs create_mat-start] +import numpy as np +import petsctools + + +def diffusion_mat(sigma): + """ + AIJ Mat for the diffusion equation with a variable diffusion coefficient. + + (I + D.T@sigma@D)u = b + """ + from petsc4py import PETSc + n = sigma.shape[0] + dtype = sigma.dtype + + # index lists for CSR format + row_start = [0] + col_indices = [] + + # top row + idxs = [0, 1] + col_indices.extend(idxs) + row_start.append(row_start[-1] + len(idxs)) + + # interior rows + for j in range(1, n-1): + idxs = [j-1, j, j+1] + col_indices.extend(idxs) + row_start.append(row_start[-1] + len(idxs)) + + # bottom row + idxs = [n-2, n-1] + col_indices.extend(idxs) + row_start.append(row_start[-1] + len(idxs)) + + # values for leading and upper/lower diagonals + diagonal = 1 + sigma.copy() + diagonal[:-1] += sigma[1:] + offdiags = -sigma[1:] + + # interleave diagonal entries + Avals = np.zeros(3*n-2, dtype=dtype) + Avals[::3] = diagonal + Avals[1::3] = offdiags + Avals[2::3] = offdiags + + A = PETSc.Mat().createAIJWithArrays( + size=(n, n), + csr=(row_start, col_indices, Avals) + ) + return A +# [appctx_docs create_mat-end] + + +# [appctx_docs pc-start] +class DiffusionJacobiPC: + prefix = "djacobi_" + + def setFromOptions(self, pc): + from petsc4py import PETSc + prefix = (pc.getOptionsPrefix() or "") + self.prefix + + options = PETSc.Options() + scale = options.getReal(prefix + "scale", 1.0) + + appctx = petsctools.AppContext() + sigma = appctx[prefix + "sigma"] + + Ap = diffusion_mat(sigma) + P = Ap.getDiagonal() + P.scale(1/scale) + self.P = P + + def apply(self, pc, x, y): + y.pointwiseDivide(x, self.P) +# [appctx_docs pc-end] + + +@pytest.mark.skipnopetsc4py +def test_appctx_docs(): + # [appctx_docs create_ksp-start] + PETSc = petsctools.init() + np.random.seed(13) + n = 50 + + sigma_bar = 2*np.ones(n) + sigma_prime = -0.2 + 0.4*np.random.random_sample(n) + + sigma = sigma_bar + sigma_prime + sigma_p = sigma_bar + + A = diffusion_mat(sigma) + + ksp = PETSc.KSP().create() + ksp.setOperators(A) + # [appctx_docs create_ksp-end] + + # [appctx_docs set_from_options-start] + appmngr = petsctools.AppContextManager() + + petsctools.set_from_options( + ksp, + parameters={ + 'ksp_converged_reason': None, + 'ksp_type': 'richardson', + 'pc_type': 'python', + 'pc_python_type': f'{__name__}.DiffusionJacobiPC', + 'djacobi_scale': 0.9, + 'djacobi_sigma': appmngr.add(sigma_p), + }, + appmngr=appmngr, + options_prefix="", + ) + # [appctx_docs set_from_options-end] + + # [appctx_docs solve-start] + u, b = A.createVecs() + u.zeroEntries() + b.array[:] = np.random.random_sample(n) + + with petsctools.inserted_options(ksp): + ksp.solve(b, u) + # [appctx_docs solve-end] diff --git a/tests/test_appctx.py b/tests/test_appctx.py new file mode 100644 index 0000000..4e94eed --- /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 + + use_prefixed_appctx = PETSc.Options().getBool( + prefix + "use_prefixed_appctx") + + if use_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) + + appmngr = petsctools.AppContextManager() + + petsctools.set_from_options( + ksp, + parameters={ + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': f'{__name__}.JacobiTestPC', + 'jacobi_scale': appmngr.add(diag), + 'jacobi_use_prefixed_appctx': use_prefix == "with_prefix", + }, + options_prefix="myksp", + appmngr=appmngr, + ) + + 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 = (