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 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,141 @@
+from collections.abc import Mapping
+from copy import deepcopy
+from typing import Dict, Iterator, Tuple
+
+import numpy as np
+from probnum.typing import ShapeType
+
+
+class PartialDerivativeCoefficients(
+    Mapping[ShapeType, Mapping[Tuple[int, ...], 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
+
+        {
+            (): {
+                (2, 0): 1.0,
+                (0, 2): 1.0,
+            }
+        }
+
+    Parameters
+    ----------
+    coefficient_dict :
+        Dictionary of coefficients.
+    """
+
+    def __init__(
+        self,
+        coefficient_dict: Dict[ShapeType, Dict[Tuple[int, ...], float]],
+        input_domain_shape: ShapeType,
+        input_codomain_shape: ShapeType,
+    ) -> None:
+        if len(input_domain_shape) > 1:
+            raise ValueError(
+                f"Input domain must be R or R^n, but got shape {input_domain_shape}."
+            )
+        input_size = int(np.prod(input_domain_shape))
+        self._num_entries = 0
+        for key in coefficient_dict.keys():
+            if len(key) != len(input_codomain_shape) or not all(
+                x < y for x, y in zip(key, input_codomain_shape)
+            ):
+                raise ValueError(
+                    f"Codomain index {key} does not match shape"
+                    "{input_codomain_shape}."
+                )
+            for sub_key in coefficient_dict[key].keys():
+                if len(sub_key) != input_size:
+                    raise ValueError(
+                        f"Multi-index {sub_key} does not match input domain shape "
+                        f"{input_domain_shape}."
+                    )
+                if any(x < 0 for x in sub_key):
+                    raise ValueError(
+                        f"Multi-index {sub_key} contains negative entries."
+                    )
+                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, key: ShapeType) -> Dict[Tuple[ShapeType, int], float]:
+        return self._coefficient_dict[key]
+
+    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_dic = deepcopy(self._coefficient_dict)
+            for key in dict(other).keys():
+                if key in new_dic.keys():
+                    for sub_key in dict(other)[key].keys():
+                        if sub_key in new_dic[key].keys():
+                            new_dic[key][sub_key] += dict(other)[key][sub_key]
+                        else:
+                            new_dic[key][sub_key] = dict(other)[key][sub_key]
+                else:
+                    new_dic[key] = deepcopy(dict(other)[key])
+            return PartialDerivativeCoefficients(
+                new_dic, 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_dic = deepcopy(self._coefficient_dict)
+            for key in scaled_dic.keys():
+                for sub_key in scaled_dic[key].keys():
+                    scaled_dic[key][sub_key] *= other
+            return PartialDerivativeCoefficients(
+                scaled_dic, 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 PartialDerivativeCoefficients
 from ._lindiffop import LinearDifferentialOperator
 
 
@@ -17,7 +18,10 @@ def __init__(
         if order < 0:
             raise ValueError(f"Order must be >= 0, but got {order}.")
 
-        super().__init__(input_shapes=((), ()))
+        super().__init__(
+            coefficients=PartialDerivativeCoefficients({(): {(order,): 1.0}}, (), ()),
+            input_shapes=((), ()),
+        )
 
         self._order = order
 

src/linpde_gp/linfuncops/diffops/_directional_derivative.py

@@ -6,14 +6,39 @@
 import probnum as pn
 from probnum.typing import ArrayLike, ShapeLike
 
-from ._lindiffop import LinearDifferentialOperator
+import linpde_gp  # pylint: disable=unused-import # for type hints
+
+from ._lindiffop import LinearDifferentialOperator, PartialDerivativeCoefficients
 
 
 class DirectionalDerivative(LinearDifferentialOperator):
     def __init__(self, direction: ArrayLike):
-        self._direction = np.asarray(direction)
+        direction = np.asarray(direction)
+        if direction.ndim > 1:
+            raise ValueError("Direction must be element of R^n.")
+
+        def get_one_hot(index: int) -> np.ndarray:
+            one_hot = np.zeros(direction.size, dtype=int)
+            one_hot[index] = 1
+            return tuple(one_hot)
+
+        coefficients = PartialDerivativeCoefficients(
+            {
+                (): {
+                    get_one_hot(domain_index): coefficient
+                    for domain_index, coefficient in enumerate(direction.reshape(-1))
+                    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,7 +12,7 @@
 
 from linpde_gp import functions
 
-from ._lindiffop import LinearDifferentialOperator
+from ._lindiffop import LinearDifferentialOperator, PartialDerivativeCoefficients
 
 if TYPE_CHECKING:
     import linpde_gp
@@ -35,7 +35,24 @@ def __init__(self, weights: ArrayLike) -> None:
                 "most 1."
             )
 
-        super().__init__(input_shapes=(weights.shape, ()))
+        def get_one_hot(index: int) -> np.ndarray:
+            one_hot = np.zeros(weights.size, dtype=int)
+            one_hot[index] = 2
+            return tuple(one_hot)
+
+        coefficients = PartialDerivativeCoefficients(
+            {
+                (): {
+                    get_one_hot(domain_index): coefficient
+                    for domain_index, coefficient in enumerate(weights.reshape(-1))
+                    if coefficient != 0.0
+                }
+            },
+            input_domain_shape=weights.shape,
+            input_codomain_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()