replaces kernel_construction.py with covariance.py containing new extensible classes for combining operators and kernels into covariance matrices. user extensions for combining custom operators and kernels or for defining new combination rules can now be achieved by inheriting from covariance.TwoSided/OneSided and passing instances to the model constructor.

Author: julian-urban

fredipy/covariance.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, List
+if TYPE_CHECKING:
+    from .kernels import Kernel
+    from .constraints import LinearEquality
+
+from .operators import Integral, Identity, ConstMatrix, Derivative
+import numpy as np
+
+
+class TwoSided:
+    """
+    Construct the symmetric block matrix of covariances of all linear equality
+    constraints by sandwiching the kernel between all combinations of the
+    associated operators.
+
+    Rules for combining operators are defined by functions named with the
+    appropriate types, e.g. _Integral_Identity() to sandwich the kernel between
+    an integral operator and the identity operator.
+
+    New rules for user-defined operators can be defined by creating a class that
+    inherits from this class, adding the appropriate functions following the
+    naming scheme, and passing an instance to the model constructor.
+    """
+    def __init__(self):
+        pass
+
+    def __call__(
+            self,
+            kernel: Kernel,
+            constraints: List[LinearEquality]
+            ) -> np.ndarray:
+        rows = []
+        for c1 in constraints:
+            columns = []
+            for c2 in constraints:
+                try:
+                    entry = getattr(self, f"_{type(c1.op).__name__}_{type(c2.op).__name__}")(c1, kernel, c2)
+                except AttributeError:
+                    try:
+                        entry = getattr(self, f"_{type(c2.op).__name__}_{type(c1.op).__name__}")(c2, kernel, c1).T
+                    except AttributeError:
+                        raise NotImplementedError(
+                            f"No rule found to combine operators of types {type(c1.op).__name__} and {type(c2.op).__name__} with kernel.")
+                columns.append(entry)
+            rows.append(np.concatenate(columns, axis=1))
+        return np.concatenate(rows)
+
+    def _Integral_Integral(self, c1, k, c2):
+        if c1 == c2:
+            return c1.op.integrator.doubleIntegrationSymmetric(c1, k)
+        else:
+            return c1.op.integrator.doubleIntegration(c1, k, c2)
+
+    def _Identity_Identity(self, c1, k, c2):
+        return k(c1.x, c2.x)
+
+    def _Derivative_Derivative(self, c1, k, c2):
+        return k.d2K_dxdy(c1.x, c2.x)
+
+    def _Integral_Identity(self, c1, k, c2):
+        return c1.op.integrator.singleIntegration(c1, k, c2.x)
+
+    def _Integral_Derivative(self, c1, k, c2):
+        return c1.op.integrator.singleIntegration(c1, k.dK_dy, c2.x)
+
+    def _Identity_Derivative(self, c1, k, c2):
+        return k.dK_dy(c1.x, c2.x)
+
+    def _ConstMatrix_ConstMatrix(self, c1, k, c2):
+        return c1.op() @ k(c1.x, c2.x) @ c2.op().T
+
+
+class OneSided:
+    """
+    Construct the asymmetric block matrix of covariances of the prediction
+    points and all linear equality constraints by left-application of the
+    associated operators to the kernel.
+
+    Rules for combining operators with the kernel are defined by functions
+    named with the appropriate type, e.g. _Integral() to apply an integral
+    operator.
+
+    New rules for user-defined operators can be defined by creating a class that
+    inherits from this class, adding the appropriate functions following the
+    naming scheme, and passing an instance to the model constructor.
+    """
+    def __init__(self):
+        pass
+
+    def __call__(
+            self,
+            kernel: Kernel,
+            constraints: List[LinearEquality],
+            w: np.ndarray
+            ) -> np.ndarray:
+        entries = []
+        for c in constraints:
+            try:
+                entry = getattr(self, f"_{type(c.op).__name__}")(c, kernel, w)
+            except AttributeError:
+                raise NotImplementedError(
+                    f"No rule found to combine operator of type {type(c.op).__name__} with kernel.")
+            entries.append(entry)
+        return np.concatenate(entries)
+
+    def _Integral(self, c, k, w):
+        return c.op.integrator.singleIntegration(c, k, w)
+
+    def _Identity(self, c, k, w):
+        return k(c.x, w)
+
+    def _Derivative(self, c, k, w):
+        return k.dK_dx(c.x, w)
+
+    def _ConstMatrix(self, c, k, w):
+        return c.op()

fredipy/kernel_construction.py

@@ -7,7 +7,7 @@
 import numpy as np
 import scipy as sp
 
-from .kernel_construction import construct_OpKer, construct_OpKerOp
+from .covariance import TwoSided, OneSided
 from .util import make_column_vector
 
 
@@ -42,10 +42,16 @@ class GaussianProcess(Model):
     def __init__(
             self,
             kernel: Kernel,
-            constraints: List[LinearEquality]
+            constraints: List[LinearEquality],
+            OpKerOp: TwoSided =TwoSided(),
+            OpKer: OneSided =OneSided(),
             ) -> None:
         self.kernel = kernel
         self.constraints = constraints
+
+        self.OpKerOp = OpKerOp
+        self.OpKer = OpKer
+        
         self.y = np.concatenate([c.y for c in self.constraints])
         self.dy = np.concatenate([c.dy for c in self.constraints])
 
@@ -169,7 +175,7 @@ def log_likelihood_grad(self) -> List[float]:
         kernel_grad = self.kernel.params_gradient()
 
         for i in range(self.kernel.dim):
-            OpKer_grad = construct_OpKerOp(kernel_grad[i], self.constraints)
+            OpKer_grad = self.OpKerOp(kernel_grad[i], self.constraints)
             loglik_grad[i] = 0.5 * np.trace((alpha @ alpha.T - K_inv) @ OpKer_grad)
 
         return loglik_grad
@@ -194,7 +200,7 @@ def _empty_cache(self) -> None:
 
     def _maybe_prepare_posterior(self) -> None:
         if not self._posterior_cache:
-            OpKerOp = construct_OpKerOp(
+            OpKerOp = self.OpKerOp(
                 self.kernel, self.constraints)
             OpKerOp_cholesky = np.linalg.cholesky(
                 OpKerOp + self.dy**2 * np.eye(OpKerOp.shape[0]))
@@ -213,7 +219,7 @@ def _maybe_prepare_inference(
             w_pred: np.ndarray,
             ) -> None:
         if not self._inference_cache:
-            OpKer = construct_OpKer(
+            OpKer = self.OpKer(
                 self.kernel, self.constraints, w_pred)
             self._inference_cache = {
                 'OpKer': OpKer}