Add PartialDerivativeCoefficients representation to lindiffops

src/linpde_gp/linfuncops/__init__.py

@@ -7,4 +7,4 @@
 from ._identity import Identity
 from ._linfuncop import LinearFunctionOperator
 from ._select_output import SelectOutput
-from .diffops import LambdaLinearDifferentialOperator, LinearDifferentialOperator
+from .diffops import LinearDifferentialOperator

src/linpde_gp/linfuncops/diffops/__init__.py

@@ -1,4 +1,5 @@
 from ._arithmetic import ScaledLinearDifferentialOperator
+from ._coefficients import MultiIndex, PartialDerivativeCoefficients
 from ._derivative import Derivative
 from ._directional_derivative import (
     DirectionalDerivative,
@@ -7,7 +8,7 @@
 )
 from ._heat import HeatOperator
 from ._laplacian import Laplacian, SpatialLaplacian, WeightedLaplacian
-from ._lindiffop import LambdaLinearDifferentialOperator, LinearDifferentialOperator
+from ._lindiffop import LinearDifferentialOperator
 
 # isort: off
 from . import _functions

src/linpde_gp/linfuncops/diffops/_arithmetic.py

@@ -13,16 +13,16 @@ def __init__(
     ) -> None:
         self._lindiffop = lindiffop
 
-        super().__init__(
-            input_shapes=self._lindiffop.input_shapes,
-            output_codomain_shape=self._lindiffop.output_codomain_shape,
-        )
-
         if not np.ndim(scalar) == 0:
             raise ValueError()
 
         self._scalar = np.asarray(scalar, dtype=np.double)
 
+        super().__init__(
+            coefficients=float(self._scalar) * self._lindiffop.coefficients,
+            input_shapes=self._lindiffop.input_shapes,
+        )
+
     @property
     def lindiffop(self) -> LinearDifferentialOperator:
         return self._lindiffop

src/linpde_gp/linfuncops/diffops/_coefficients.py

@@ -0,0 +1,183 @@
+from collections.abc import Iterator, Mapping
+from copy import deepcopy
+import functools
+
+import numpy as np
+from probnum.typing import ArrayLike, ShapeType
+
+
+class MultiIndex:
+    r"""Multi-index representation of a partial derivative.
+
+    A multi-index is an array of non-negative integers which is used to represent
+    a partial derivative of a function.
+
+    For example, the multi-index :math:`(1, 2, 0)` represents the partial derivative
+    :math:`\frac{\partial^3}{\partial x_1 \partial x_2^2}`
+    of a function :math:`f: \mathbb{R}^3 \to \mathbb{R}`.
+    """
+
+    def __init__(self, multi_index: ArrayLike) -> None:
+        self._multi_index = np.asarray(multi_index, dtype=int)
+        if np.any(self._multi_index < 0):
+            raise ValueError(f"Multi-index {multi_index} contains negative entries.")
+        self._multi_index.setflags(write=False)
+
+    @classmethod
+    def from_index(
+        cls, index: tuple[int, ...], shape: ShapeType, order: int
+    ) -> "MultiIndex":
+        multi_index = np.zeros(shape, dtype=int)
+        multi_index[index] = order
+        return cls(multi_index)
+
+    @functools.cached_property
+    def order(self) -> int:
+        return np.sum(self._multi_index)
+
+    @property
+    def array(self) -> np.ndarray:
+        return self._multi_index
+
+    @property
+    def shape(self) -> ShapeType:
+        return self._multi_index.shape
+
+    def __getitem__(self, index: tuple[int, ...]) -> int:
+        return self._multi_index[index]
+
+    def __hash__(self) -> int:
+        return hash(self._multi_index.data.tobytes())
+
+    def __eq__(self, __o: object) -> bool:
+        if not isinstance(__o, MultiIndex):
+            return NotImplemented
+        return np.all(self.array == __o.array)
+
+
+class PartialDerivativeCoefficients(Mapping[ShapeType, Mapping[MultiIndex, float]]):
+    r"""Partial derivative coefficients of a linear differential operator.
+
+    Any linear differential operator can be written as a sum of partial derivatives.
+
+    The coefficients are stored in a dictionary of the form
+    {input_codomain_index: {multi_index: coefficient}}.
+
+    For example, the Laplacian operator in 2D can be written as
+
+    .. math::
+        \Delta = \frac{\partial^2}{\partial x^2} + \frac{\partial^2}{\partial y^2}
+
+    and is stored as
+
+    .. code-block:: python
+
+        {
+            (): {
+                MultiIndex((2, 0)): 1.0,
+                MultiIndex((0, 2)): 1.0,
+            }
+        }
+
+    Parameters
+    ----------
+    coefficient_dict :
+        Dictionary of coefficients.
+    """
+
+    def __init__(
+        self,
+        coefficient_dict: Mapping[ShapeType, Mapping[MultiIndex, float]],
+        input_domain_shape: ShapeType,
+        input_codomain_shape: ShapeType,
+    ) -> None:
+        self._num_entries = 0
+        for codomain_idx in coefficient_dict.keys():
+            if len(codomain_idx) != len(input_codomain_shape) or not all(
+                x < y for x, y in zip(codomain_idx, input_codomain_shape)
+            ):
+                raise ValueError(
+                    f"Codomain index {codomain_idx} does not match shape"
+                    f"{input_codomain_shape}."
+                )
+            for multi_index in coefficient_dict[codomain_idx].keys():
+                if multi_index.shape != input_domain_shape:
+                    raise ValueError(
+                        f"Multi-index shape {multi_index.shape} does not match "
+                        f"input domain shape {input_domain_shape}."
+                    )
+                self._num_entries += 1
+
+        self._coefficient_dict = coefficient_dict
+        self._input_domain_shape = input_domain_shape
+        self._input_codomain_shape = input_codomain_shape
+
+    @property
+    def num_entries(self) -> int:
+        return self._num_entries
+
+    @property
+    def input_domain_shape(self) -> ShapeType:
+        return self._input_domain_shape
+
+    @property
+    def input_codomain_shape(self) -> ShapeType:
+        return self._input_codomain_shape
+
+    def __getitem__(self, codomain_idx: ShapeType) -> Mapping[MultiIndex, float]:
+        return self._coefficient_dict[codomain_idx]
+
+    def __len__(self) -> int:
+        return len(self._coefficient_dict)
+
+    def __iter__(self) -> Iterator[ShapeType]:
+        return iter(self._coefficient_dict)
+
+    def __neg__(self) -> "PartialDerivativeCoefficients":
+        return -1.0 * self
+
+    def __add__(self, other) -> "PartialDerivativeCoefficients":
+        if isinstance(other, PartialDerivativeCoefficients):
+            if self.input_domain_shape != other.input_domain_shape:
+                raise ValueError(
+                    "Cannot add coefficients with input domain shapes"
+                    f"{self.input_domain_shape} != {other.input_domain_shape}"
+                )
+            if self.input_codomain_shape != other.input_codomain_shape:
+                raise ValueError(
+                    "Cannot add coefficients with input codomain shapes"
+                    f"{self.input_codomain_shape} != {other.input_codomain_shape}"
+                )
+
+            new_dict = deepcopy(self._coefficient_dict)
+            for codomain_idx in other.keys():
+                if codomain_idx in new_dict.keys():
+                    for multi_index in other[codomain_idx].keys():
+                        if multi_index in new_dict[codomain_idx].keys():
+                            new_dict[codomain_idx][multi_index] += other[codomain_idx][
+                                multi_index
+                            ]
+                        else:
+                            new_dict[codomain_idx][multi_index] = other[codomain_idx][
+                                multi_index
+                            ]
+                else:
+                    new_dict[codomain_idx] = deepcopy(other[codomain_idx])
+            return PartialDerivativeCoefficients(
+                new_dict, self.input_domain_shape, self.input_codomain_shape
+            )
+        return NotImplemented
+
+    def __sub__(self, other) -> "PartialDerivativeCoefficients":
+        return self + (-other)
+
+    def __rmul__(self, other) -> "PartialDerivativeCoefficients":
+        if np.ndim(other) == 0:
+            scaled_dict = deepcopy(self._coefficient_dict)
+            for codomain_idx in scaled_dict.keys():
+                for multi_index in scaled_dict[codomain_idx].keys():
+                    scaled_dict[codomain_idx][multi_index] *= other
+            return PartialDerivativeCoefficients(
+                scaled_dict, self.input_domain_shape, self.input_codomain_shape
+            )
+        return NotImplemented

src/linpde_gp/linfuncops/diffops/_derivative.py

@@ -6,6 +6,7 @@
 
 import linpde_gp  # pylint: disable=unused-import # for type hints
 
+from ._coefficients import MultiIndex, PartialDerivativeCoefficients
 from ._lindiffop import LinearDifferentialOperator
 
 
@@ -17,7 +18,12 @@ def __init__(
         if order < 0:
             raise ValueError(f"Order must be >= 0, but got {order}.")
 
-        super().__init__(input_shapes=((), ()))
+        super().__init__(
+            coefficients=PartialDerivativeCoefficients(
+                {(): {MultiIndex(order): 1.0}}, (), ()
+            ),
+            input_shapes=((), ()),
+        )
 
         self._order = order
 

src/linpde_gp/linfuncops/diffops/_directional_derivative.py

@@ -6,14 +6,33 @@
 import probnum as pn
 from probnum.typing import ArrayLike, ShapeLike
 
+import linpde_gp  # pylint: disable=unused-import # for type hints
+
+from ._coefficients import MultiIndex, PartialDerivativeCoefficients
 from ._lindiffop import LinearDifferentialOperator
 
 
 class DirectionalDerivative(LinearDifferentialOperator):
     def __init__(self, direction: ArrayLike):
-        self._direction = np.asarray(direction)
+        direction = np.asarray(direction)
+
+        coefficients = PartialDerivativeCoefficients(
+            {
+                (): {
+                    MultiIndex.from_index(domain_index, direction.shape, 1): coefficient
+                    for domain_index, coefficient in np.ndenumerate(direction)
+                    if coefficient != 0.0
+                }
+            },
+            input_domain_shape=direction.shape,
+            input_codomain_shape=(),
+        )
+        super().__init__(
+            coefficients=coefficients,
+            input_shapes=(direction.shape, ()),
+        )
 
-        super().__init__(input_shapes=(self._direction.shape, ()))
+        self._direction = direction
 
     @property
     def direction(self) -> np.ndarray:

src/linpde_gp/linfuncops/diffops/_laplacian.py

@@ -12,6 +12,7 @@
 
 from linpde_gp import functions
 
+from ._coefficients import MultiIndex, PartialDerivativeCoefficients
 from ._lindiffop import LinearDifferentialOperator
 
 if TYPE_CHECKING:
@@ -29,13 +30,19 @@ class WeightedLaplacian(LinearDifferentialOperator):
     def __init__(self, weights: ArrayLike) -> None:
         weights = np.asarray(weights)
 
-        if weights.ndim > 1:
-            raise ValueError(
-                "The Laplacian operator only supports functions with input ndim of at "
-                "most 1."
-            )
+        coefficients = PartialDerivativeCoefficients(
+            {
+                (): {
+                    MultiIndex.from_index(domain_index, weights.shape, 2): coefficient
+                    for domain_index, coefficient in np.ndenumerate(weights)
+                    if coefficient != 0.0
+                }
+            },
+            input_domain_shape=weights.shape,
+            input_codomain_shape=(),
+        )
 
-        super().__init__(input_shapes=(weights.shape, ()))
+        super().__init__(coefficients=coefficients, input_shapes=(weights.shape, ()))
 
         self._weights = weights
 

src/linpde_gp/linfuncops/diffops/_lindiffop.py

@@ -10,21 +10,33 @@
 from linpde_gp.functions import JaxFunction, JaxLambdaFunction
 
 from .._linfuncop import LinearFunctionOperator
+from ._coefficients import PartialDerivativeCoefficients
 
 
 class LinearDifferentialOperator(LinearFunctionOperator):
+    """Linear differential operator that maps to functions with codomain R."""
+
     def __init__(
         self,
+        coefficients: PartialDerivativeCoefficients,
         input_shapes: tuple[ShapeLike, ShapeLike],
-        output_codomain_shape: ShapeLike = (),
     ) -> None:
-        input_shapes = tuple(input_shapes)
+        if coefficients.input_domain_shape != input_shapes[0]:
+            raise ValueError()
+        if coefficients.input_codomain_shape != input_shapes[1]:
+            raise ValueError()
 
         super().__init__(
             input_shapes=input_shapes,
-            output_shapes=(input_shapes[0], output_codomain_shape),
+            output_shapes=(input_shapes[0], ()),
         )
 
+        self._coefficients = coefficients
+
+    @property
+    def coefficients(self) -> PartialDerivativeCoefficients:
+        return self._coefficients
+
     @functools.singledispatchmethod
     def __call__(self, f, **kwargs):
         try:
@@ -72,25 +84,3 @@ def weak_form(
         self, basis: pn.functions.Function, /
     ) -> "linpde_gp.linfunctls.LinearFunctional":
         raise NotImplementedError()
-
-
-class LambdaLinearDifferentialOperator(LinearDifferentialOperator):
-    def __init__(
-        self,
-        jax_diffop_fn,
-        /,
-        input_shapes: tuple[ShapeLike, ShapeLike],
-        output_codomain_shape: ShapeLike = (),
-    ) -> None:
-        super().__init__(input_shapes, output_codomain_shape)
-
-        self._jax_diffop_fn = jax_diffop_fn
-
-    def _jax_fallback(self, f: Callable, /, **kwargs) -> Callable:
-        return self._jax_diffop_fn(f, **kwargs)
-
-    @functools.singledispatchmethod
-    def weak_form(
-        self, test_basis: pn.functions.Function, /
-    ) -> "linpde_gp.linfunctls.LinearFunctional":
-        raise NotImplementedError()