Gaussian Processes via JAX

.github/workflows/build.yml

@@ -14,10 +14,10 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v6
 
-      - name: Setup Python 3.9
+      - name: Setup Python 3.12
         uses: actions/setup-python@v6
         with:
-          python-version: 3.9
+          python-version: 3.12
           cache: 'pip'  # optional, adds caching for faster builds
 
       - name: Install pypa/build

.github/workflows/lint.yml

@@ -14,10 +14,10 @@ jobs:
     - name: Checkout code
       uses: actions/checkout@v6
 
-    - name: Setup Python 3.9
+    - name: Setup Python 3.12
       uses: actions/setup-python@v6
       with:
-        python-version: 3.9
+        python-version: 3.12
         cache: 'pip'  # optional, adds caching for faster builds
 
     - name: Install UncertainSCI

.github/workflows/test.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest, windows-latest, macOS-latest]
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.12", "3.13", "3.14"]
 
     steps:
     - name: Checkout code
@@ -43,7 +43,7 @@ jobs:
         coveralls --service=github
       env:
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        COVERALLS_FLAG_NAME: ${{ matrix.os }}
+        COVERALLS_FLAG_NAME: ${{ matrix.os }}-${{ matrix.python-version }}
         COVERALLS_PARALLEL: true
 
   coveralls:

.gitignore

@@ -68,8 +68,9 @@ instance/
 # Scrapy stuff:
 .scrapy
 
-# Sphinx documentation
+# Docs:
 docs/_build/
+docs/**/copied/
 docs/**/generated/
 
 # PyBuilder

.readthedocs.yaml

@@ -9,7 +9,7 @@ version: 2
 build:
   os: ubuntu-22.04
   tools:
-    python: "3.10"
+    python: "3.12"
   apt_packages:
     - pandoc
 

UncertainSCI/_equinox.py

@@ -0,0 +1,138 @@
+import abc
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import numpy.typing as npt
+
+
+class ComputedArray(eqx.Module, metaclass=abc.ABCMeta):
+    """
+    Computed array base class.
+
+    Primarily used to detect if PyTree nodes require (further) resolution to reach
+    underlying objects or data.
+    """
+    is_static: bool = eqx.field(static=True)
+
+    @abc.abstractmethod
+    def __call__(self) -> jax.Array:
+        raise NotImplementedError
+
+    @classmethod
+    @abc.abstractmethod
+    def from_array(cls, a: jax.Array, is_static: bool=False, **kwargs):
+        """
+        Convenience method for in-place replacement of types of
+        :class:`ComputedArray`, though not generally intended for use as a long-lived
+        instance; prefer direct instantiation when possible.
+
+        Args:
+            a (array):
+                The array data of this :class:`ComputedArray`.
+            is_static (bool, optional):
+                If static.
+            **kwargs:
+                Kwargs to hand off to :meth:`__init__` (generally depends
+                on child class).
+        """
+        raise NotImplementedError
+
+
+class PositiveDefinite(ComputedArray):
+    """
+    A positive-definite matrix.
+
+    Module contains some niceties that guarantee matrix remains positive definite
+    during training.
+
+    Attributes:
+        L_log_diag (array or None):
+            Log of diagonal of Cholesky factor of positive definite matrix.
+        L_off_diag (array or None):
+            Elements below main diagonal of Cholesky factor of positive definite.
+    """
+    L_log_diag: jax.Array
+    L_off_diag: jax.Array
+
+    def __init__(
+        self,
+        *,  # TODO: make this NOT kwargs-only.
+        D: jax.Array | npt.NDArray | None = None,
+        L_log_diag: jax.Array | npt.NDArray | None = None,
+        L_off_diag: jax.Array | npt.NDArray | None = None,
+        is_static: bool = False,
+    ):
+        """
+        Create a positive definite matrix.
+
+        All args must be passed as kwargs to avoid name collisions with parent
+        methods.
+
+        Args:
+            D (array or None):
+                Positive definite matrix; defined only if ``L_log_diag is None`` and
+                ``L_off_diag is None``.
+            L_log_diag (array or None):
+                Log of diagonal of Cholesky factor of positive definite matrix;
+                defined only if ``D is None``.
+            L_off_diag (array or None):
+                Elements below main diagonal of Cholesky factor of positive definite
+                matrix; defined only if ``D is None``.
+            is_static (bool):
+                If the matrix is static.
+        """
+        if not (
+            (D is not None) ^
+            (L_log_diag is not None and L_off_diag is not None)
+        ):
+            raise ValueError('D xor (L_log_diag and L_off_diag) must be defined!')
+
+        if D is not None:
+            if D.ndim != 2:
+                raise ValueError(f'Expected ndim = 2 for matrix, got D.ndim = {D.ndim}!')
+            if D.shape[0] != D.shape[1]:
+                raise ValueError(f'Expected square matrix, got D.shape = {D.shape}!')
+
+            n = D.shape[0]
+            L = jnp.linalg.cholesky(D)
+            self.L_log_diag = jnp.log(jnp.diag(L))
+            self.L_off_diag = L[jnp.tril_indices(n, k=-1)]
+
+        elif L_log_diag is not None and L_off_diag is not None:
+            if L_log_diag.ndim != 1:
+                raise ValueError(
+                    f'Expected ndim = 1 for L_log_diag, got L_log_diag.ndim = {L_log_diag.ndim}!'
+                )
+            if L_off_diag.ndim != 1:
+                raise ValueError(
+                    f'Expected ndim = 1 for L_off_diag, got L_off_diag.ndim = {L_off_diag.ndim}!'
+                )
+
+
+            n = L_log_diag.shape[0]
+            if L_off_diag.shape[0] != (n) * (n - 1) / 2:
+                raise ValueError(
+                    f'Expected L_off_diag.shape[0] = (n) * (n - 1) / 2 = '
+                    f'{(n) * (n - 1) / 2} (numel below main diagonal), got '
+                    f'L_off_diag.shape[0] = {L_off_diag.shape[0]}!'
+                )
+
+            self.L_log_diag = jnp.asarray(L_log_diag)
+            self.L_off_diag = jnp.asarray(L_off_diag)
+
+        else:
+            raise ValueError('Got unexpected combination of parameters; check caller!')
+
+        super().__init__(is_static=is_static)
+
+    def __call__(self) -> jax.Array:
+        n = self.L_log_diag.shape[0]
+        L = jnp.diag(jnp.exp(self.L_log_diag))
+        L = L.at[jnp.tril_indices(n, k=-1)].set(self.L_off_diag)
+        D = L @ L.T
+
+        return jax.lax.stop_gradient(D) if self.is_static else D
+
+    @classmethod
+    def from_array(cls, a, is_static: bool = False, **kwargs):
+        return cls(D=a, is_static=is_static, **kwargs)

UncertainSCI/_linalg.py

@@ -0,0 +1,80 @@
+import jax
+import jax.numpy as jnp
+
+
+@jax.jit
+def d_foerstner(a: jax.Array, b: jax.Array) -> jax.Array:
+    r"""
+    Compute the Förstner distance between covariance matrices ``a`` and ``b``.
+
+    The Förstner metric is defined as
+
+    .. math::
+
+        d_\text{Förstner}(A, B) = \sqrt{ \sum_i \log^2 \lambda_i }
+
+    for :math:`\lambda_i` eigenvalues of the generalized eigenvalue problem
+
+    .. math::
+
+        \operatorname{det}\left( A - \lambda B \right) = 0.
+
+    See [A Metric on Covariance Matrices](https://doi.org/10.1007/978-3-662-05296-9_31)
+    for more details.
+
+    Args:
+        a (array):
+            Matrix ``a``.
+        b (array):
+            Matric ``b``.
+
+    Returns:
+        A scalar array (``ndim == 0`` and ``shape == ()``) of the Förstner distance
+        between ``a`` and ``b`` of complex float.
+    """
+    eig_vals, _ = eigh_generalized(a, b)
+    return jnp.sqrt(jnp.sum(jnp.log(eig_vals) ** 2))
+
+
+@jax.jit
+def eigh_generalized(
+    a: jax.Array,
+    b: jax.Array
+) -> tuple[jax.Array, jax.Array]:
+    r"""
+    Solves the generalized eigenvalue problem :math:`A v = \lambda B v`
+    for symmetric/Hermitian matrices where :math:`B` is positive-definite
+    in a JAX-compatible way.
+
+    This exists because :func:`jax.scipy.linalg.eigh` does not support
+    ``b is not None`` (in particular, it supports only the standard eigenvalue
+    problem).
+
+    Args:
+        a (array):
+            Symmetric or Hermitian matrix of shape (N, N).
+        b (array):
+            Symmetric positive-definite matrix of shape (N, N).
+
+    Returns:
+        eigenvalues (array):
+            Array of shape (N,) sorted in ascending order.
+        eigenvectors (array):
+            Matrix of shape (N, N) of columns of eigenvectors.
+    """
+    l = jnp.linalg.cholesky(b)
+    x = jax.scipy.linalg.solve_triangular(l, a, lower=True)
+    c = jax.scipy.linalg.solve_triangular(l, x.conj().T, lower=True)
+    evals, w = jnp.linalg.eigh(c)
+    evecs = jax.scipy.linalg.solve_triangular(l.conj().T, w, lower=False)
+
+    return evals, evecs
+
+
+@jax.jit
+def solve_cholesky(
+    l: jax.Array,
+    y: jax.Array
+):
+    w = jax.scipy.linalg.solve_triangular(l, y, lower=True)
+    return jax.scipy.linalg.solve_triangular(l.conj().T, w)