Move the operator registry from individual backends to the registry class to allow flexible addition of operators (without loading the backend).

Author: david-zwicker

pde/backends/__init__.py

@@ -11,7 +11,7 @@
 .. codeauthor:: David Zwicker <[email protected]>
 """
 
-# load the registry, which manages the backends
+# load the registry, which manages all backends
 from .registry import backends  # noqa: I001
 
 # load and register the default backend

pde/backends/base.py

@@ -7,8 +7,7 @@
 
 import inspect
 import logging
-from collections import defaultdict
-from typing import Any, Callable, Literal, NamedTuple
+from typing import Any, Callable, Literal
 
 from ..fields import DataFieldBase
 from ..grids import BoundariesBase, GridBase
@@ -22,31 +21,21 @@
     NumberOrArray,
     NumericArray,
     OperatorFactory,
+    OperatorInfo,
     TField,
 )
 
 _base_logger = logging.getLogger(__name__.rsplit(".", 1)[0])
 """:class:`logging.Logger`: Base logger for backends."""
 
 
-class OperatorInfo(NamedTuple):
-    """Stores information about an operator."""
-
-    factory: OperatorFactory
-    rank_in: int
-    rank_out: int
-    name: str = ""  # attach a unique name to help caching
-
-
 class BackendBase:
     """Basic backend from which all other backends inherit."""
 
-    _operators: dict[type[GridBase], dict[str, OperatorInfo]]
     _logger: logging.Logger  # logger instance to output information
 
     def __init__(self, name: str = "numpy"):
         self.name = name
-        self._operators = defaultdict(dict)
 
     def __init_subclass__(cls, **kwargs) -> None:
         """Initialize class-level attributes of subclasses."""
@@ -59,8 +48,7 @@ def register_operator(
         grid_cls: type[GridBase],
         name: str,
         factory_func: OperatorFactory | None = None,
-        rank_in: int = 0,
-        rank_out: int = 0,
+        **kwargs,
     ):
         """Register an operator for a particular grid.
 
@@ -95,21 +83,9 @@ def make_operator(grid: GridBase): ...
             rank_out (int):
                 The rank of the field that is returned by the operator
         """
+        from .registry import backends
 
-        def register_operator(factor_func_arg: OperatorFactory):
-            """Helper function to register the operator."""
-            self._operators[grid_cls][name] = OperatorInfo(
-                factory=factor_func_arg, rank_in=rank_in, rank_out=rank_out, name=name
-            )
-            return factor_func_arg
-
-        if factory_func is None:
-            # method is used as a decorator, so return the helper function
-            return register_operator
-        else:
-            # method is used directly
-            register_operator(factory_func)
-            return None
+        backends.register_operator(self.name, grid_cls, name, factory_func, **kwargs)
 
     def get_registered_operators(self, grid_id: GridBase | type[GridBase]) -> set[str]:
         """Returns all operators defined for a grid.
@@ -118,14 +94,16 @@ def get_registered_operators(self, grid_id: GridBase | type[GridBase]) -> set[st
             grid (:class:`~pde.grid.base.GridBase` or its type):
                 Grid for which the operator need to be returned
         """
+        from . import backends
+
         grid_cls = grid_id if inspect.isclass(grid_id) else grid_id.__class__
 
         # get all operators registered on the class
         operators = set()
         # add all custom defined operators
         classes = inspect.getmro(grid_cls)[:-1]  # type: ignore
         for cls in classes:
-            operators |= set(self._operators[cls].keys())
+            operators |= set(backends._operators[self.name][cls].keys())
 
         return operators
 
@@ -149,11 +127,13 @@ def get_operator_info(
             return operator
         assert isinstance(operator, str)
 
+        from . import backends
+
         # look for defined operators on all parent grid classes (except `object`)
         classes = inspect.getmro(grid.__class__)[:-1]
         for cls in classes:
-            if operator in self._operators[cls]:
-                return self._operators[cls][operator]
+            if operator in backends._operators[self.name][cls]:
+                return backends._operators[self.name][cls][operator]
 
         # throw an error since operator was not found
         raise NotImplementedError(

pde/backends/numba/__init__.py

@@ -12,8 +12,7 @@
 from .backend import NumbaBackend
 
 # add the loaded numba backend to the registry
-numba_backend = NumbaBackend("numba")
-backends.add(numba_backend)
+backends.add(NumbaBackend("numba"))
 
 # register all the standard operators
 from . import operators

pde/backends/numba/backend.py

@@ -728,7 +728,7 @@ def make_inner_stepper(
             time `t_end`. The function call signature is `(state: numpy.ndarray,
             t_start: float, t_end: float)`
         """
-        assert solver.backend == "numba"
+        assert solver.backend == self.name
 
         from .solvers import make_adaptive_stepper, make_fixed_stepper
 

pde/backends/numba/operators/cartesian.py

@@ -26,7 +26,7 @@
 from ....tools.misc import module_available
 from ....tools.numba import jit
 from ....tools.typing import NumericArray, OperatorType
-from .. import numba_backend
+from ...registry import backends
 
 
 def make_corner_point_setter_2d(grid: CartesianGrid) -> Callable[[NumericArray], None]:
@@ -132,13 +132,13 @@ def laplace(arr: NumericArray, out: NumericArray) -> None:
     else:
         # use 9-point stencil with interpolated boundary conditions
         w = corner_weight
-        numba_backend._logger.info(
+        backends["numba"]._logger.info(
             "Create 2D Cartesian Laplacian with 9-point stencil (w=%.3g)", w
         )
 
         if not np.isclose(*grid.discretization):
             # we have not yet found a good expression for the 9-point stencil for dx!=dy
-            numba_backend._logger.warning(
+            backends["numba"]._logger.warning(
                 "9-point stencils with anisotropic grids are not tested and might "
                 "produce wrong results."
             )
@@ -292,7 +292,7 @@ def laplace(arr: NumericArray, out: NumericArray) -> None:
     return laplace  # type: ignore
 
 
-@numba_backend.register_operator(CartesianGrid, "laplace", rank_in=0, rank_out=0)
+@backends.register_operator("numba", CartesianGrid, "laplace", rank_in=0, rank_out=0)
 def make_laplace(
     grid: CartesianGrid, *, spectral: bool | None = None, **kwargs
 ) -> OperatorType:
@@ -484,7 +484,7 @@ def gradient(arr: NumericArray, out: NumericArray) -> None:
     return gradient  # type: ignore
 
 
-@numba_backend.register_operator(CartesianGrid, "gradient", rank_in=0, rank_out=1)
+@backends.register_operator("numba", CartesianGrid, "gradient", rank_in=0, rank_out=1)
 def make_gradient(
     grid: CartesianGrid,
     *,
@@ -681,8 +681,8 @@ def gradient_squared(arr: NumericArray, out: NumericArray) -> None:
     return gradient_squared  # type: ignore
 
 
-@numba_backend.register_operator(
-    CartesianGrid, "gradient_squared", rank_in=0, rank_out=0
+@backends.register_operator(
+    "numba", CartesianGrid, "gradient_squared", rank_in=0, rank_out=0
 )
 def make_gradient_squared(grid: CartesianGrid, *, central: bool = True) -> OperatorType:
     """Make a gradient operator on a Cartesian grid.
@@ -841,7 +841,7 @@ def divergence(arr: NumericArray, out: NumericArray) -> None:
     return divergence  # type: ignore
 
 
-@numba_backend.register_operator(CartesianGrid, "divergence", rank_in=1, rank_out=0)
+@backends.register_operator("numba", CartesianGrid, "divergence", rank_in=1, rank_out=0)
 def make_divergence(
     grid: CartesianGrid,
     *,
@@ -904,8 +904,8 @@ def vectorized_operator(arr: NumericArray, out: NumericArray) -> None:
     return register_jitable(vectorized_operator)  # type: ignore
 
 
-@numba_backend.register_operator(
-    CartesianGrid, "vector_gradient", rank_in=1, rank_out=2
+@backends.register_operator(
+    "numba", CartesianGrid, "vector_gradient", rank_in=1, rank_out=2
 )
 def make_vector_gradient(
     grid: CartesianGrid,
@@ -927,7 +927,9 @@ def make_vector_gradient(
     return _vectorize_operator(make_gradient, grid, method=method)
 
 
-@numba_backend.register_operator(CartesianGrid, "vector_laplace", rank_in=1, rank_out=1)
+@backends.register_operator(
+    "numba", CartesianGrid, "vector_laplace", rank_in=1, rank_out=1
+)
 def make_vector_laplace(grid: CartesianGrid) -> OperatorType:
     """Make a vector Laplacian on a Cartesian grid.
 
@@ -941,8 +943,8 @@ def make_vector_laplace(grid: CartesianGrid) -> OperatorType:
     return _vectorize_operator(make_laplace, grid)
 
 
-@numba_backend.register_operator(
-    CartesianGrid, "tensor_divergence", rank_in=2, rank_out=1
+@backends.register_operator(
+    "numba", CartesianGrid, "tensor_divergence", rank_in=2, rank_out=1
 )
 def make_tensor_divergence(
     grid: CartesianGrid,

pde/backends/numba/operators/cylindrical_sym.py

@@ -22,10 +22,12 @@
 from ....tools.docstrings import fill_in_docstring
 from ....tools.numba import jit
 from ....tools.typing import NumericArray, OperatorType
-from .. import numba_backend
+from ...registry import backends
 
 
-@numba_backend.register_operator(CylindricalSymGrid, "laplace", rank_in=0, rank_out=0)
+@backends.register_operator(
+    "numba", CylindricalSymGrid, "laplace", rank_in=0, rank_out=0
+)
 @fill_in_docstring
 def make_laplace(grid: CylindricalSymGrid) -> OperatorType:
     """Make a discretized laplace operator for a cylindrical grid.
@@ -64,7 +66,9 @@ def laplace(arr: NumericArray, out: NumericArray) -> None:
     return laplace  # type: ignore
 
 
-@numba_backend.register_operator(CylindricalSymGrid, "gradient", rank_in=0, rank_out=1)
+@backends.register_operator(
+    "numba", CylindricalSymGrid, "gradient", rank_in=0, rank_out=1
+)
 @fill_in_docstring
 def make_gradient(grid: CylindricalSymGrid) -> OperatorType:
     """Make a discretized gradient operator for a cylindrical grid.
@@ -97,8 +101,8 @@ def gradient(arr: NumericArray, out: NumericArray) -> None:
     return gradient  # type: ignore
 
 
-@numba_backend.register_operator(
-    CylindricalSymGrid, "gradient_squared", rank_in=0, rank_out=0
+@backends.register_operator(
+    "numba", CylindricalSymGrid, "gradient_squared", rank_in=0, rank_out=0
 )
 @fill_in_docstring
 def make_gradient_squared(
@@ -154,8 +158,8 @@ def gradient_squared(arr: NumericArray, out: NumericArray) -> None:
     return gradient_squared  # type: ignore
 
 
-@numba_backend.register_operator(
-    CylindricalSymGrid, "divergence", rank_in=1, rank_out=0
+@backends.register_operator(
+    "numba", CylindricalSymGrid, "divergence", rank_in=1, rank_out=0
 )
 @fill_in_docstring
 def make_divergence(grid: CylindricalSymGrid) -> OperatorType:
@@ -194,8 +198,8 @@ def divergence(arr: NumericArray, out: NumericArray) -> None:
     return divergence  # type: ignore
 
 
-@numba_backend.register_operator(
-    CylindricalSymGrid, "vector_gradient", rank_in=1, rank_out=2
+@backends.register_operator(
+    "numba", CylindricalSymGrid, "vector_gradient", rank_in=1, rank_out=2
 )
 @fill_in_docstring
 def make_vector_gradient(grid: CylindricalSymGrid) -> OperatorType:
@@ -244,8 +248,8 @@ def vector_gradient(arr: NumericArray, out: NumericArray) -> None:
     return vector_gradient  # type: ignore
 
 
-@numba_backend.register_operator(
-    CylindricalSymGrid, "vector_laplace", rank_in=1, rank_out=1
+@backends.register_operator(
+    "numba", CylindricalSymGrid, "vector_laplace", rank_in=1, rank_out=1
 )
 @fill_in_docstring
 def make_vector_laplace(grid: CylindricalSymGrid) -> OperatorType:
@@ -306,8 +310,8 @@ def vector_laplace(arr: NumericArray, out: NumericArray) -> None:
     return vector_laplace  # type: ignore
 
 
-@numba_backend.register_operator(
-    CylindricalSymGrid, "tensor_divergence", rank_in=2, rank_out=1
+@backends.register_operator(
+    "numba", CylindricalSymGrid, "tensor_divergence", rank_in=2, rank_out=1
 )
 @fill_in_docstring
 def make_tensor_divergence(grid: CylindricalSymGrid) -> OperatorType:

pde/backends/numba/operators/polar_sym.py

@@ -21,10 +21,10 @@
 from ....tools.docstrings import fill_in_docstring
 from ....tools.numba import jit
 from ....tools.typing import NumericArray, OperatorType
-from .. import numba_backend
+from ...registry import backends
 
 
-@numba_backend.register_operator(PolarSymGrid, "laplace", rank_in=0, rank_out=0)
+@backends.register_operator("numba", PolarSymGrid, "laplace", rank_in=0, rank_out=0)
 @fill_in_docstring
 def make_laplace(grid: PolarSymGrid) -> OperatorType:
     """Make a discretized laplace operator for a polar grid.
@@ -56,7 +56,7 @@ def laplace(arr: NumericArray, out: NumericArray) -> None:
     return laplace  # type: ignore
 
 
-@numba_backend.register_operator(PolarSymGrid, "gradient", rank_in=0, rank_out=1)
+@backends.register_operator("numba", PolarSymGrid, "gradient", rank_in=0, rank_out=1)
 @fill_in_docstring
 def make_gradient(
     grid: PolarSymGrid, *, method: Literal["central", "forward", "backward"] = "central"
@@ -101,8 +101,8 @@ def gradient(arr: NumericArray, out: NumericArray) -> None:
     return gradient  # type: ignore
 
 
-@numba_backend.register_operator(
-    PolarSymGrid, "gradient_squared", rank_in=0, rank_out=0
+@backends.register_operator(
+    "numba", PolarSymGrid, "gradient_squared", rank_in=0, rank_out=0
 )
 @fill_in_docstring
 def make_gradient_squared(grid: PolarSymGrid, *, central: bool = True) -> OperatorType:
@@ -152,7 +152,7 @@ def gradient_squared(arr: NumericArray, out: NumericArray) -> None:
     return gradient_squared  # type: ignore
 
 
-@numba_backend.register_operator(PolarSymGrid, "divergence", rank_in=1, rank_out=0)
+@backends.register_operator("numba", PolarSymGrid, "divergence", rank_in=1, rank_out=0)
 @fill_in_docstring
 def make_divergence(grid: PolarSymGrid) -> OperatorType:
     """Make a discretized divergence operator for a polar grid.
@@ -185,7 +185,9 @@ def divergence(arr: NumericArray, out: NumericArray) -> None:
     return divergence  # type: ignore
 
 
-@numba_backend.register_operator(PolarSymGrid, "vector_gradient", rank_in=1, rank_out=2)
+@backends.register_operator(
+    "numba", PolarSymGrid, "vector_gradient", rank_in=1, rank_out=2
+)
 @fill_in_docstring
 def make_vector_gradient(grid: PolarSymGrid) -> OperatorType:
     """Make a discretized vector gradient operator for a polar grid.
@@ -224,8 +226,8 @@ def vector_gradient(arr: NumericArray, out: NumericArray) -> None:
     return vector_gradient  # type: ignore
 
 
-@numba_backend.register_operator(
-    PolarSymGrid, "tensor_divergence", rank_in=2, rank_out=1
+@backends.register_operator(
+    "numba", PolarSymGrid, "tensor_divergence", rank_in=2, rank_out=1
 )
 @fill_in_docstring
 def make_tensor_divergence(grid: PolarSymGrid) -> OperatorType: