From 4f954e815ce38f5f6ee93f9cfd6442ce89b65ddc Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Wed, 29 Jul 2026 16:30:22 -0600 Subject: [PATCH 01/22] gp: Gaussian processes via jax * pending commit on MWE * tutorial works (all features, all dimensions), but it's kinda messy * want to clean up a few things, then happy to commit --- UncertainSCI/_equinox.py | 117 ++++++++++++++ UncertainSCI/_linalg.py | 80 ++++++++++ UncertainSCI/gp/__init__.py | 297 ++++++++++++++++++++++++++++++++++++ UncertainSCI/gp/kernel.py | 131 ++++++++++++++++ UncertainSCI/gp/mean.py | 108 +++++++++++++ pyproject.toml | 3 + 6 files changed, 736 insertions(+) create mode 100644 UncertainSCI/_equinox.py create mode 100644 UncertainSCI/_linalg.py create mode 100644 UncertainSCI/gp/__init__.py create mode 100644 UncertainSCI/gp/kernel.py create mode 100644 UncertainSCI/gp/mean.py diff --git a/UncertainSCI/_equinox.py b/UncertainSCI/_equinox.py new file mode 100644 index 0000000..c7bc6b0 --- /dev/null +++ b/UncertainSCI/_equinox.py @@ -0,0 +1,117 @@ +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy.typing as npt + + +class PositiveDefinite(eqx.Module): + """ + A positive-definite matrix. + + Module contains some niceties that guarantee matrix remains positive definite + during training. + + Attributes: + n (int): + Dimension of the (square) matrix, inferred from initialization. + shape (tuple of int): + Shape of the matrix (always like ``(n,) * 2``). + ndim (int): + Number of dimensions of the matrix (always 2). + """ + n: int = eqx.field(static=True) + L_log_diag: jax.Array + L_off_diag: jax.Array + is_static: bool = eqx.field(static=True) + + def __init__( + self, + *, + 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 xor( + 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}!') + + self.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(self.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}!' + ) + + self.n = L_log_diag.shape[0] + + if L_off_diag.shape[0] != (self.n) * (self.n - 1) / 2: + raise ValueError( + f'Expected L_off_diag.shape[0] = (n) * (n - 1) / 2 = ' + f'{(self.n) * (self.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!') + + self.is_static = is_static + super().__init__() + + def __call__(self) -> jax.Array: + L = jnp.diag(jnp.exp(self.L_log_diag)) + L = L.at[jnp.tril_indices(self.n, k=-1)].set(self.L_off_diag) + D = L @ L.T + + return jax.lax.stop_gradient(D) if self.is_static else D + + @property + def shape(self): + return (self.n, self.n) + + @property + def ndim(self): + return 2 + + +def xor(a: bool, b: bool) -> bool: + return (a or b) and not (a and b) diff --git a/UncertainSCI/_linalg.py b/UncertainSCI/_linalg.py new file mode 100644 index 0000000..e57c46c --- /dev/null +++ b/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) diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py new file mode 100644 index 0000000..589ba7a --- /dev/null +++ b/UncertainSCI/gp/__init__.py @@ -0,0 +1,297 @@ +""" +Build Gaussian processes. +""" + +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np +import numpy.typing as npt +import optax + +from .. import _linalg +from . import kernel +from . import mean + + +DEFAULT_OPTIM = optax.adam +DEFAULT_OPTIM_KWARGS = { + 'learning_rate': 1e-2 +} + +class GaussianProcess: + dim: int + cdim: int + mu: mean.Mean + k: kernel.Kernel + + nugget: float + + x_train: jax.Array + y_train: jax.Array + s_train: jax.Array + L_train: jax.Array + + seed: int + key: jax.Array + + def __init__( + self, + dim: int, + cdim: int, + mu: mean.Mean, + k: kernel.Kernel, + nugget: float = 1e-3, + seed: int | None = None + ) -> None: + if mu.dim != dim: + raise ValueError + if k.dim != dim: + raise ValueError + if mu.cdim != cdim: + raise ValueError + if k.cdim != cdim: + raise ValueError + + self.dim = dim + self.cdim = cdim + self.mu = mu + self.k = k + + self.nugget = nugget + + if seed is None: + self.seed = np.random.randint(0x00000000, 0xffffffff) + else: + self.seed = seed + self.key = jax.random.key(self.seed) + + def condition( + self, + x: jax.Array | npt.NDArray, + y: jax.Array | npt.NDArray, + s: jax.Array | npt.NDArray | float | int + ): + _, ny = self.validate_shapes(x, y, s) + + self.x_train = jnp.asarray(x) + self.y_train = jnp.asarray(y) + if isinstance(s, (float, int)): + self.s_train = s * jnp.eye(ny * self.cdim) + else: + self.s_train = jnp.asarray(s) + + self.L_train = GaussianProcess._get_L_train( + (self.mu, self.k), + self.x_train, + self.y_train, + self.s_train + ) + + def add_point(self, x, y, s): + raise NotImplementedError + + def validate_shapes( + self, + x: jax.Array | npt.NDArray, + y: jax.Array | npt.NDArray, + s: jax.Array | npt.NDArray | float | int + ) -> tuple[int, int]: + nx, _ = self.validate_input_shape(x) + ny, _ = self.validate_output_shape(y) + + if nx != ny: + raise ValueError + + if not isinstance(s, (float, int)): + if s.shape != (ny, self.cdim): + raise ValueError + + return nx, ny + + def validate_input_shape(self, x: jax.Array | npt.NDArray) -> tuple[int, int]: + if x.ndim != 2: + raise ValueError + + n, d = x.shape + if x.shape[1] != self.dim: + raise ValueError + + return n, d + + def validate_output_shape(self, y: jax.Array | npt.NDArray) -> tuple[int, int]: + if y.ndim != 2: + raise ValueError + + n, c = y.shape + if c != self.cdim: + raise ValueError + + return n, c + + def get_L_train( + self, + x: jax.Array | npt.NDArray, + y: jax.Array | npt.NDArray, + s: jax.Array | npt.NDArray | float | int + ) -> jax.Array: + _, ny = self.validate_shapes(x, y, s) + + x = jnp.asarray(x) + y = jnp.asarray(y) + if isinstance(s, (float, int)): + s = s * jnp.eye(ny * self.cdim) + else: + s = jnp.asarray(s) + + return GaussianProcess._get_L_train((self.mu, self.k), x, y, s) + + # TODO: may want to change args for this function, but it's nice to maintain the same signature as + # other core/mathematical routines. + @staticmethod + @jax.jit + def _get_L_train( + p: tuple[mean.Mean, kernel.Kernel], + x: jax.Array, + y: jax.Array, + s: jax.Array + ) -> jax.Array: + mu, k = p + return jnp.linalg.cholesky(k(x, x) + s) + + def tune( + self, + n: int = 1000, + optim = DEFAULT_OPTIM, + optim_kwargs: dict = DEFAULT_OPTIM_KWARGS, + ): + """ + Args: + optim (callable): + Optax or Optax-style optimizer. + optim_kwargs (dict): + kwargs for ``optim``. + """ + optim = optim(**optim_kwargs) + optim_state = optim.init(eqx.filter((self.mu, self.k), eqx.is_inexact_array)) + + losses = [] + for i in range(n): + optim_state, (self.mu, self.k), loss = GaussianProcess._step( + optim, + optim_state, + (self.mu, self.k), + self.x_train, + self.y_train, + self.s_train + ) + losses.append(loss) + + # Recompute the Cholesky factor of observation matrix after tuning: + self.L_train = GaussianProcess._get_L_train( + (self.mu, self.k), + self.x_train, + self.y_train, + self.s_train + ) + + return losses + + @staticmethod + @jax.jit(static_argnames=('optim',)) + def _step( + optim, + optim_state, + p: tuple[mean.Mean, kernel.Kernel], + x: jax.Array, + y: jax.Array, + s: jax.Array + ) -> tuple[jax.Array, tuple[mean.Mean, kernel.Kernel], jax.Array]: + loss, grads = eqx.filter_value_and_grad(GaussianProcess._loss)(p, x, y, s) + updates, optim_state = optim.update(grads, optim_state, p) + p = eqx.apply_updates(p, updates) + return optim_state, p, loss + + def loss(self, + x: jax.Array | npt.NDArray, + y: jax.Array | npt.NDArray, + s: jax.Array | npt.NDArray | float | int + ) -> jax.Array: + _, ny = self.validate_shapes(x, y, s) + + x = jnp.asarray(x) + y = jnp.asarray(y) + if isinstance(s, (float, int)): + s = s * jnp.eye(ny * self.cdim) + else: + s = jnp.asarray(s) + + return GaussianProcess._loss((self.mu, self.k), x, y, s) + + @staticmethod + @jax.jit + def _loss( + p: tuple[mean.Mean, kernel.Kernel], + x: jax.Array, + y: jax.Array, + s: jax.Array + ) -> jax.Array: + mu, k = p + d = (y - mu(x)).flatten() + L_train = GaussianProcess._get_L_train((mu, k), x, y, s) + return ( + d.T @ _linalg.solve_cholesky(L_train, d) + + jnp.sum(jnp.log(jnp.diag(L_train))) + ) + + def prior_mean(self, x): + self.validate_input_shape(x) + + return self.mu(x) + + def prior_covariance(self, u, v): + self.validate_input_shape(u) + self.validate_input_shape(v) + + return self.k(u, v) + + def prior_realization(self, x, p=1): + self.key, key = jax.random.split(self.key) + n, _ = self.validate_input_shape(x) + + L = jnp.linalg.cholesky( + self.prior_covariance(x, x) + self.nugget * jnp.eye(n * self.cdim) + ) + return self.prior_mean(x) + (L @ jax.random.normal(key, (n, p))) + + def posterior_mean(self, x): + self.validate_input_shape(x) + + return ( + self.mu(x) + + self.k(x, self.x_train) @ _linalg.solve_cholesky( + self.L_train, + self.y_train - self.mu(self.x_train) + ) + ) + + def posterior_covariance(self, u, v): + self.validate_input_shape(u) + self.validate_input_shape(v) + + return ( + self.k(u, v) - + self.k(u, self.x_train) @ _linalg.solve_cholesky( + self.L_train, + self.k(self.x_train, v) + ) + ) + + def posterior_realization(self, x, p=1): + self.key, key = jax.random.split(self.key) + n, _ = self.validate_input_shape(x) + + L = jnp.linalg.cholesky( + self.posterior_covariance(x, x) + self.nugget * jnp.eye(n) + ) + return self.posterior_mean(x) + (L @ jax.random.normal(key, (n, p))) diff --git a/UncertainSCI/gp/kernel.py b/UncertainSCI/gp/kernel.py new file mode 100644 index 0000000..e0b0e43 --- /dev/null +++ b/UncertainSCI/gp/kernel.py @@ -0,0 +1,131 @@ +""" +Kernels for Gaussian processes. +""" + +import abc +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy.typing as npt + +from .. import _equinox + + +class Kernel(eqx.Module, metaclass=abc.ABCMeta): + """ + Kernel base class. + + Attributes: + dim (int): + Dimension of the domain. + cdim (int): + Dimension of the codomain. + """ + dim: int = eqx.field(static=True) + cdim: int = eqx.field(static=True) + + @abc.abstractmethod + def __call__( + self, + x: jax.Array | npt.NDArray, + y: jax.Array | npt.NDArray + ): + """ + Evaluate this kernel. + + Args: + x (array): + Array of coordinates shaped as ``(num_x, dim)``. May be passed as 1-d + array in the case ``dim == 1``. + y (array): + Array of coordinates shaped as ``(num_y, dim)``. May be passed as 1-d + array in the case ``dim == 1``. + """ + raise NotImplementedError + + +class Gaussian(Kernel): + r""" + A Gaussian (square exponential) covariance kernel. + + This produces a kernel according to the function + + .. math:: + + k(x, y) = \exp{\left( - y^\top D x \right)}. + + Attributes: + D (array): + Matrix defining metric on coordinates. + + ``D`` is a matrix defining some metric on the coordinates :math:`x` and :math:`y` + according to + + .. math:: + + \left[ \operatorname{dist}(x, y) \right]^2 = y^\top D x. + """ + D: _equinox.PositiveDefinite + + def __init__( + self, + dim: int, + cdim: int, + D: jax.Array | npt.NDArray, + D_is_static: bool = False + ): + super().__init__(dim=dim, cdim=cdim) + self.D = _equinox.PositiveDefinite(D=D, is_static=D_is_static) + + def __call__( + self, + x: jax.Array | npt.NDArray, + y: jax.Array | npt.NDArray + ): + d = x.reshape((-1, 1, self.dim)) - y.reshape((1, -1, self.dim)) + return jnp.exp(-1 * jnp.einsum( + 'ijk,kk,ijk->ij', d, self.D(), d + ) + ) + + +class Kronecker(Kernel): + r""" + A Kronecker-structure matrix-valued covariance kernel. + + This produces a kernel according to the function + + .. math:: + + k_\text{Kronecker}(x, y) = k(x, y) \otimes C, + + where :math:`k` is a scalar kernel (e.g., a square exponential kernel) and + :math:`\otimes` is the Kronecker product. + + Attributes: + k (Kernel): + Covariance kernel on coordindates. + C (array): + Covariance on outputs. + """ + k: Kernel + C: _equinox.PositiveDefinite + + def __init__( + self, + dim: int, + cdim: int, + k: Kernel, + C: jax.Array | npt.NDArray, + C_is_static: bool = False + ): + super().__init__(dim=dim, cdim=cdim) + self.k = k + self.C = _equinox.PositiveDefinite(D=C, is_static=C_is_static) + + def __call__( + self, + x: jax.Array | npt.NDArray, + y: jax.Array | npt.NDArray + ) -> jax.Array: + return jnp.kron(self.k(x, y), self.C()) diff --git a/UncertainSCI/gp/mean.py b/UncertainSCI/gp/mean.py new file mode 100644 index 0000000..4d2c6f4 --- /dev/null +++ b/UncertainSCI/gp/mean.py @@ -0,0 +1,108 @@ +""" +Mean functions for Gaussian processes. +""" + +import abc +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy.typing as npt + + +class Mean(eqx.Module, metaclass=abc.ABCMeta): + """ + Mean function base class. + + Attributes: + dim (int): + Dimension of the domain. + cdim (int): + Dimension of the codomain. + """ + dim: int = eqx.field(static=True) + cdim: int = eqx.field(static=True) + + @abc.abstractmethod + def __call__( + self, + x: jax.Array | npt.NDArray + ) -> jax.Array: + """ + Evaluate the mean. + + Arguments: + x (array): + x values of evaluation. + """ + raise NotImplementedError + + def validate_input(self, x: jax.Array | npt.NDArray): + if x.ndim != 2: + raise ValueError + if x.shape[1] != self.dim: + raise ValueError + + +class Affine(Mean): + a: jax.Array + a_is_static: bool = eqx.field(static=True) + b: jax.Array + b_is_static: bool = eqx.field(static=True) + + def __init__( + self, + dim: int, + cdim: int, + a: jax.Array | npt.NDArray, + b: jax.Array | npt.NDArray | float | int, + a_is_static: bool = False, + b_is_static: bool = False + ): + """ + Like: + + .. math:: + + y = x @ A.T + b + + where ``x.shape = (number, dim)``. + + Args: + dim (int): + Dimension of the domain. + cdim (int): + Dimension of the codomain. + a (array): + Matrix :math:`A`. + b (array): + Intercept :math:`b`. + a_is_static (bool): + If ``a`` is static. + b_is_static (bool): + If ``b`` is static. + """ + super().__init__(dim, cdim) + + a = jnp.asarray(a) + if isinstance(b, (float, int)): + b = b * jnp.ones((cdim,)) + else: + b = jnp.asarray(b) + + if a.shape != (cdim, dim): + raise ValueError + if b.shape != (cdim,): + raise ValueError + + self.a = a + self.a_is_static = a_is_static + self.b = b + self.b_is_static = b_is_static + + def __call__( + self, + x: jax.Array | npt.NDArray + ) -> jax.Array: + a = jax.lax.stop_gradient(self.a) if self.a_is_static else self.a + b = jax.lax.stop_gradient(self.b) if self.b_is_static else self.b + return jnp.einsum('id,cd->ic', x.reshape(-1, self.dim), a) + b diff --git a/pyproject.toml b/pyproject.toml index 4b64a17..1f9044c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,8 +35,11 @@ classifiers = [ "Operating System :: MacOS", ] dependencies = [ + "equinox", + "jax", "matplotlib>=3.1.3", "numpy>=1.25,!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*", + "optax", "packaging", "scipy>=1.4.1", ] From edf39139cd4c0a2b60b03ce73b26911380bba258 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 30 Jul 2026 12:55:22 -0400 Subject: [PATCH 02/22] gp: MWE for gps via JAX * small changes to gp and _equinox modules * bug fixes * some changes to facilitate real use-cases * still want to implement dynamic selection of next sampling coordinate --- UncertainSCI/_equinox.py | 77 ++++--- UncertainSCI/gp/__init__.py | 8 +- gp-jax-tuning.ipynb | 428 ++++++++++++++++++++++++++++++++++++ 3 files changed, 483 insertions(+), 30 deletions(-) create mode 100644 gp-jax-tuning.ipynb diff --git a/UncertainSCI/_equinox.py b/UncertainSCI/_equinox.py index c7bc6b0..c9d61c1 100644 --- a/UncertainSCI/_equinox.py +++ b/UncertainSCI/_equinox.py @@ -1,10 +1,44 @@ +import abc import equinox as eqx import jax import jax.numpy as jnp import numpy.typing as npt -class PositiveDefinite(eqx.Module): +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. @@ -12,21 +46,17 @@ class PositiveDefinite(eqx.Module): during training. Attributes: - n (int): - Dimension of the (square) matrix, inferred from initialization. - shape (tuple of int): - Shape of the matrix (always like ``(n,) * 2``). - ndim (int): - Number of dimensions of the matrix (always 2). + 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. """ - n: int = eqx.field(static=True) L_log_diag: jax.Array L_off_diag: jax.Array - is_static: bool = eqx.field(static=True) 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, @@ -63,11 +93,10 @@ def __init__( if D.shape[0] != D.shape[1]: raise ValueError(f'Expected square matrix, got D.shape = {D.shape}!') - self.n = D.shape[0] - + 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(self.n, k=-1)] + 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: @@ -79,12 +108,12 @@ def __init__( f'Expected ndim = 1 for L_off_diag, got L_off_diag.ndim = {L_off_diag.ndim}!' ) - self.n = L_log_diag.shape[0] - if L_off_diag.shape[0] != (self.n) * (self.n - 1) / 2: + 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'{(self.n) * (self.n - 1) / 2} (numel below main diagonal), got ' + f'{(n) * (n - 1) / 2} (numel below main diagonal), got ' f'L_off_diag.shape[0] = {L_off_diag.shape[0]}!' ) @@ -94,23 +123,19 @@ def __init__( else: raise ValueError('Got unexpected combination of parameters; check caller!') - self.is_static = is_static - super().__init__() + 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(self.n, k=-1)].set(self.L_off_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 - @property - def shape(self): - return (self.n, self.n) - - @property - def ndim(self): - return 2 + @classmethod + def from_array(cls, a, is_static: bool = False, **kwargs): + return cls(D=a, is_static=is_static, **kwargs) def xor(a: bool, b: bool) -> bool: diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index 589ba7a..b6dde54 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -195,7 +195,7 @@ def tune( self.s_train ) - return losses + return jnp.asarray(losses) @staticmethod @jax.jit(static_argnames=('optim',)) @@ -262,7 +262,7 @@ def prior_realization(self, x, p=1): L = jnp.linalg.cholesky( self.prior_covariance(x, x) + self.nugget * jnp.eye(n * self.cdim) ) - return self.prior_mean(x) + (L @ jax.random.normal(key, (n, p))) + return self.prior_mean(x) + (L @ jax.random.normal(key, (n * self.cdim, p))) def posterior_mean(self, x): self.validate_input_shape(x) @@ -292,6 +292,6 @@ def posterior_realization(self, x, p=1): n, _ = self.validate_input_shape(x) L = jnp.linalg.cholesky( - self.posterior_covariance(x, x) + self.nugget * jnp.eye(n) + self.posterior_covariance(x, x) + self.nugget * jnp.eye(n * self.cdim) ) - return self.posterior_mean(x) + (L @ jax.random.normal(key, (n, p))) + return self.posterior_mean(x) + (L @ jax.random.normal(key, (n * self.cdim, p))) diff --git a/gp-jax-tuning.ipynb b/gp-jax-tuning.ipynb new file mode 100644 index 0000000..e04af89 --- /dev/null +++ b/gp-jax-tuning.ipynb @@ -0,0 +1,428 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "8a631981", + "metadata": {}, + "outputs": [], + "source": [ + "import jax\n", + "\n", + "# Un/comment this for double/single precision:\n", + "# jax.config.update('jax_enable_x64', True)\n", + "\n", + "import copy\n", + "import equinox as eqx\n", + "import jax.numpy as jnp\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "import UncertainSCI._equinox as _eqx\n", + "import UncertainSCI.gp as gp\n", + "\n", + "\n", + "D = 1 \n", + "C = 1\n", + "\n", + "np.random.seed(0x01234567)\n", + "seed = 0xdeadbeef # Used in gp.GaussianProcess initialization.\n", + "\n", + "FIGSIZE = (7, 7 / 1.6)\n", + "FIGDPI = None\n", + "\n", + "def print_loss_here(g): # Simple helper for this notebook; does not generalize.\n", + " print(\n", + " f'Loss at D = {jnp.squeeze(g.k.D()):.4e} (length scale {1 / jnp.squeeze(g.k.D())**2:.4e}): '\n", + " f'{g.loss(x_train, y_train, s_train):.4e}'\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "598efb92", + "metadata": {}, + "outputs": [], + "source": [ + "x_plot = jnp.linspace(0, 2 * jnp.pi, 1000).reshape((-1, D))\n", + "\n", + "def f_hidden(x):\n", + " return 2 * jnp.cos(x) * jnp.sin(4 * x)\n", + "\n", + "def f_noisy(x):\n", + " y = f_hidden(x)\n", + " return (\n", + " y,\n", + " 1e-3 * jnp.ones_like(y)\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ee181026", + "metadata": {}, + "outputs": [], + "source": [ + "mu = gp.mean.Affine(\n", + " dim=D,\n", + " cdim=C,\n", + " a=0. * jnp.ones((C, D)),\n", + " b=0.,\n", + " a_is_static=True,\n", + " b_is_static=True\n", + ")\n", + "k = gp.kernel.Gaussian(\n", + " dim=D,\n", + " cdim=C,\n", + " D=jnp.ones((1, 1)),\n", + ")\n", + "g = gp.GaussianProcess(dim=D, cdim=C, mu=mu, k=k, seed=seed, nugget=1e-3)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2ae4ac3b", + "metadata": {}, + "outputs": [], + "source": [ + "def gp_stats(g: gp.GaussianProcess, x, which='posterior', p=100):\n", + " if which == 'posterior':\n", + " m = jnp.squeeze(g.posterior_mean(x))\n", + " c = jnp.squeeze(jnp.diag(g.posterior_covariance(x, x)))\n", + " r = jnp.squeeze(g.posterior_realization(x, p))\n", + " elif which == 'prior':\n", + " m = jnp.squeeze(g.prior_mean(x))\n", + " c = jnp.squeeze(jnp.diag(g.prior_covariance(x, x)))\n", + " r = jnp.squeeze(g.prior_realization(x, p))\n", + " else:\n", + " raise ValueError\n", + "\n", + " return m, c, r\n", + "\n", + "\n", + "def plot_distribution_mean(ax, g: gp.GaussianProcess, x, which='posterior', p=100):\n", + " m, c, r = gp_stats(g, x, which, p)\n", + " ax.fill_between(\n", + " x.squeeze(),\n", + " (m - c),\n", + " (m + c),\n", + " color='#00000040',\n", + " edgecolor='none'\n", + " )\n", + " ax.plot(\n", + " x.squeeze(),\n", + " r,\n", + " color='#00000010'\n", + " )\n", + " ax.plot(\n", + " x.squeeze(),\n", + " f_hidden(x).squeeze(),\n", + " )\n", + " if which == 'posterior':\n", + " ax.errorbar(\n", + " g.x_train.squeeze(),\n", + " g.y_train.squeeze(),\n", + " yerr=g.s_train.squeeze(),\n", + " capsize=2,\n", + " linestyle='none'\n", + " )\n", + " ax.set_title('Distribution')\n", + "\n", + "\n", + "def plot_distribution_variance(ax, g: gp.GaussianProcess, x, which='posterior', p=100, colorlast=True):\n", + " m, c, r = gp_stats(g, x, which, p)\n", + " ax.plot(\n", + " x.squeeze(),\n", + " c\n", + " )\n", + " if which == 'posterior':\n", + " if colorlast:\n", + " ax.vlines(g.x_train.squeeze()[:-1], 0, 1, color='tab:green')\n", + " ax.vlines(g.x_train.squeeze()[-1], 0, 1, color='tab:red')\n", + " else:\n", + " ax.vlines(g.x_train.squeeze(), 0, 1, color='tab:green')\n", + " ax.set_title('Variance')\n", + "\n", + "\n", + "def plot_loss_landscape(\n", + " ax,\n", + " g: gp.GaussianProcess,\n", + " p_name: tuple[str, ...],\n", + " p_range: jax.Array,\n", + " x: jax.Array | None = None,\n", + " ymin: jax.Array | None = None,\n", + " ymax: jax.Array | None = None,\n", + " ytop: float = 1e2\n", + " ):\n", + " \"\"\"\n", + " Plot the loss landscape with respect to the parameter in Gaussian process ``g``\n", + " at location ``p_name``.\n", + "\n", + " Performs recursive ``getattr`` until ``p_name`` is exhausted.\n", + "\n", + " Args:\n", + " ax (matplotlib.axes.Axes):\n", + " The axes to plot into.\n", + " g (gp.GaussianProcess):\n", + " The Gaussian process to evaluate.\n", + " p_name (tuple of str):\n", + " Attr name chain leading to the parameter (:class:`jax.Array` or :class:`_eqx.ComputedArray`)\n", + " to be sampled.\n", + " p_range (array):\n", + " Range of parameter values to test.\n", + " x (array or float, optional):\n", + " x values at which to place vlines, default ``None``.\n", + " ymin (array or float, optional):\n", + " y range for vlines, must match shape of x, default ``None``.\n", + " ymax (array or float, optional):\n", + " y range for vlines, must match shape of x, default ``None``.\n", + " ytop (float, optional):\n", + " Top of y-axis to plot, default 1e2.\n", + " \"\"\"\n", + " p_parent = None\n", + " g = copy.deepcopy(g)\n", + " p = g\n", + " for name in p_name:\n", + " p_parent = p\n", + " p = getattr(p, name)\n", + "\n", + " if p_parent is None:\n", + " raise ValueError(\n", + " f'Did not find parent of target from keys {p_name}! '\n", + " f'This usually happens when keys {p_name} has zero length.'\n", + " )\n", + "\n", + " if isinstance(p, _eqx.ComputedArray):\n", + " numel = p().size\n", + " shape = p().shape\n", + " elif isinstance(p, jax.Array):\n", + " numel = p.size\n", + " shape = p.shape\n", + " else:\n", + " raise ValueError(\n", + " f'Resolution of target from keys {p_name} led to unsupported type {type(p)}!'\n", + " )\n", + "\n", + " if numel > 1:\n", + " raise NotImplementedError('Loss landscape for non-scalar parameter is not supported!')\n", + "\n", + " losses = jnp.empty(p_range.shape)\n", + " for i, p_value in enumerate(p_range):\n", + " if isinstance(p, _eqx.ComputedArray):\n", + " p_value = type(p).from_array(p_value * jnp.ones(shape))\n", + " else:\n", + " p_value = p_value * jnp.ones(shape)\n", + "\n", + " setattr(p_parent, p_name[-1], p_value)\n", + " losses = losses.at[i].set(g.loss(g.x_train, g.y_train, g.s_train))\n", + "\n", + " ax.plot(p_range, losses)\n", + " if x is not None:\n", + " if ymin is None:\n", + " ymin = min(losses)\n", + " if ymax is None:\n", + " ymax = max(losses)\n", + " ax.vlines(x, ymin, ymax)\n", + " ax.set_xscale('symlog')\n", + " ax.set_yscale('symlog')\n", + " ax.set_ylim(top=ytop)\n", + " ax.set_title('Loss Landscape')\n", + " ax.set_xlabel(p_name[-1])\n", + "\n", + "\n", + "def plot_in_subplots(nrows=1, ncols=1, **kwargs):\n", + " \"\"\"\n", + " Returns ``(fig, axes)`` with plot scaled correctly for ``(nrows, ncols)``.\n", + " \"\"\"\n", + " if 'figsize' in kwargs:\n", + " raise ValueError(\"kwargs had 'figsize' key: that's the whole point of this function!\")\n", + " return plt.subplots(\n", + " nrows,\n", + " ncols,\n", + " figsize=(ncols * FIGSIZE[0], nrows * FIGSIZE[1]),\n", + " **kwargs\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4a5952f3", + "metadata": {}, + "outputs": [], + "source": [ + "fig, (ax1, ax2) = plot_in_subplots(2, 1, dpi=FIGDPI)\n", + "plot_distribution_mean(ax1, g, x_plot, 'prior')\n", + "plot_distribution_variance(ax2, g, x_plot, 'prior')\n", + "plt.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ee0fc8ac", + "metadata": {}, + "outputs": [], + "source": [ + "N_INIT = 10\n", + "\n", + "x_train = jnp.linspace(0, 2 * jnp.pi, N_INIT).reshape(-1, D)\n", + "y_train, s_train = f_noisy(x_train) \n", + "\n", + "g.condition(x_train, y_train, s_train)\n" + ] + }, + { + "cell_type": "markdown", + "id": "2cf76bf6", + "metadata": {}, + "source": [ + "Plot what that looks like:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6ca89cdd", + "metadata": {}, + "outputs": [], + "source": [ + "fig, (ax1, ax2, ax3) = plot_in_subplots(3, 1, dpi=FIGDPI)\n", + "plot_distribution_mean(ax1, g, x_plot, 'prior')\n", + "plot_distribution_variance(ax2, g, x_plot, 'prior')\n", + "plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", + "plt.show()\n", + "\n", + "print_loss_here(g)\n" + ] + }, + { + "cell_type": "markdown", + "id": "9f7a058a", + "metadata": {}, + "source": [ + "Tune:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "59c5a570", + "metadata": {}, + "outputs": [], + "source": [ + "losses = g.tune()\n", + "\n", + "plt.figure(figsize=FIGSIZE)\n", + "plt.plot(losses)\n", + "plt.xlabel('Step')\n", + "plt.ylabel('Loss')\n", + "plt.show()\n" + ] + }, + { + "cell_type": "markdown", + "id": "cf45781f", + "metadata": {}, + "source": [ + "Plot what things look like after tuning:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "885d3e14", + "metadata": {}, + "outputs": [], + "source": [ + "fig, (ax1, ax2, ax3) = plot_in_subplots(3, 1, dpi=FIGDPI)\n", + "plot_distribution_mean(ax1, g, x_plot, 'posterior')\n", + "plot_distribution_variance(ax2, g, x_plot, 'posterior', colorlast=False)\n", + "plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", + "plt.show()\n", + "\n", + "print_loss_here(g)\n" + ] + }, + { + "cell_type": "markdown", + "id": "4c34f19e", + "metadata": {}, + "source": [ + "Iteratively sample points to choose for tuning, then tune:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "818fb48c", + "metadata": {}, + "outputs": [], + "source": [ + "N_TRAIN = 10\n", + "N_SAMPLE = 20\n", + "\n", + "x_sample = jnp.linspace(0, 2 * jnp.pi, N_SAMPLE).reshape(-1, D)\n", + "\n", + "# def covariance(...): ...\n", + "\n", + "# def optimize_x_over_covariance(...): ...\n", + "\n", + "for i in range(N_TRAIN):\n", + "\n", + " # ...do those things here...\n", + "\n", + " _, c, _ = gp_stats(g, x_sample, p=0) # Greedy sampling on dense mesh.\n", + " _x = x_sample[jnp.argmax(c)].reshape(-1, D)\n", + " _y, _s = f_noisy(_x)\n", + "\n", + " x_train = jnp.concat((x_train, _x))\n", + " y_train = jnp.concat((y_train, _y))\n", + " s_train = jnp.concat((s_train, _s))\n", + "\n", + " g.condition(x_train, y_train, s_train)\n", + " g.tune()\n", + "\n", + " fig, (ax1, ax2, ax3) = plot_in_subplots(3, 1, dpi=FIGDPI)\n", + " plot_distribution_mean(ax1, g, x_plot, 'posterior')\n", + " plot_distribution_variance(ax2, g, x_plot, 'posterior')\n", + " plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", + " plt.show()\n", + " print_loss_here(g)\n", + " print('\\n' * 3)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5e2afcd6", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "uncertainsci", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.15" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 9f0d554fd1748f2f8c09e746c6a0c1a0056ed7c9 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Mon, 17 Aug 2026 17:31:14 -0600 Subject: [PATCH 03/22] gp: adaptive sampling coordinate selection * still want to change some symbol names, but works --- UncertainSCI/gp/__init__.py | 123 ++++++++++++++++++++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index b6dde54..7b2ba07 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -7,6 +7,7 @@ import jax.numpy as jnp import numpy as np import numpy.typing as npt +import scipy.spatial as spatial import optax from .. import _linalg @@ -244,6 +245,128 @@ def _loss( jnp.sum(jnp.log(jnp.diag(L_train))) ) + def get_sample_point( + self, + x: jax.Array, + hull: spatial.ConvexHull | None = None, + ranges: jax.Array | npt.NDArray | None = None, + tol: float = 1e-6, + n: int = 1000, + optim = DEFAULT_OPTIM, + optim_kwargs: dict = DEFAULT_OPTIM_KWARGS, + ): + """ + Args: + x (jax.Array): + Candidate coordinate from which to start optimization. + hull (scipy.spatial.ConvexHull, optional): + Convex hull of coordinate domain, + must be supplied if ``ranges`` is not. + ranges (jax.Array or array-like, optional): + Array of shape (2, dim) of corners of prism of coordinate domain. + Must be supplied if ``hull`` is not. + tol (float): + Tolerance of hull inclusion criterion. + n (int): + Number of steps to take in optimization. + optim (...): + ... + optim_kwargs (...): + ... + """ + if not hasattr(self, 'x_train') or not hasattr(self, 'L_train'): + raise ValueError('Gaussian process must be conditioned before sampling.') + + x = jnp.asarray(x) + if not np.issubdtype(x.dtype, np.inexact): + x = x.astype(jnp.float32) + if x.ndim > 1: + raise ValueError + if x.shape != (self.dim,): + raise ValueError + x = x.reshape((1, self.dim)) + + if hull is None and ranges is None: + raise ValueError + + if hull is None: + ranges = np.asarray(ranges) + if ranges.ndim != 2: + raise ValueError + if ranges.shape != (2, self.dim): + raise ValueError + + hull = spatial.ConvexHull( + np.stack( + np.meshgrid( + ranges.T, + indexing='ij' + ), + axis=-1 + ).reshape((-1, self.dim)) + ) + + A = hull.equations[:, :-1] + b = hull.equations[:, -1] + + def inside_domain(x_cand: jax.Array): + return jnp.all(jnp.dot(x_cand, A.T) + b <= tol) + + if not inside_domain(x): + raise ValueError + + optim = optim(**optim_kwargs) + optim_state = optim.init(x) + + for i in range(n): + optim_state, x_cand, _variance = GaussianProcess._step_sample_point( + optim, + optim_state, + self.k, + self.x_train, + self.L_train, + x + ) + + if inside_domain(x_cand): + x = x_cand + else: + return x + + return x + + @staticmethod + @jax.jit(static_argnames=('optim',)) + def _step_sample_point( + optim, + optim_state, + k: kernel.Kernel, + x_train: jax.Array, + L_train: jax.Array, + x: jax.Array + ) -> tuple[jax.Array, jax.Array, jax.Array]: + def loss_fn(x): + return -GaussianProcess._posterior_variance(k, x_train, L_train, x) + + loss, grads = jax.value_and_grad(loss_fn)(x) + updates, optim_state = optim.update(grads, optim_state, x) + x = eqx.apply_updates(x, updates) + return optim_state, x, -loss + + @staticmethod + @jax.jit + def _posterior_variance( + k: kernel.Kernel, + x_train: jax.Array, + L_train: jax.Array, + x: jax.Array + ) -> jax.Array: + covariance = ( + k(x, x) - + k(x, x_train) @ _linalg.solve_cholesky(L_train, k(x_train, x)) + ) + return jnp.sum(jnp.diag(covariance)) + def prior_mean(self, x): self.validate_input_shape(x) From 7120116245c862df91b47dd55ca0060d7ca13405 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Wed, 19 Aug 2026 15:33:33 -0600 Subject: [PATCH 04/22] gp: efficient Kronecker-structure kernel solves * more coming for docs * more coming for different noise structures, too --- UncertainSCI/_equinox.py | 10 +- UncertainSCI/gp/__init__.py | 184 ++++++++++++++++---------------- UncertainSCI/gp/kernel.py | 205 +++++++++++++++++++++++++++++++++++- tests/test_gp.py | 157 +++++++++++++++++++++++++++ 4 files changed, 457 insertions(+), 99 deletions(-) create mode 100644 tests/test_gp.py diff --git a/UncertainSCI/_equinox.py b/UncertainSCI/_equinox.py index c9d61c1..9ddf27c 100644 --- a/UncertainSCI/_equinox.py +++ b/UncertainSCI/_equinox.py @@ -81,9 +81,9 @@ def __init__( is_static (bool): If the matrix is static. """ - if not xor( - D is not None, - L_log_diag is not None and L_off_diag is not None + 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!') @@ -136,7 +136,3 @@ def __call__(self) -> jax.Array: @classmethod def from_array(cls, a, is_static: bool = False, **kwargs): return cls(D=a, is_static=is_static, **kwargs) - - -def xor(a: bool, b: bool) -> bool: - return (a or b) and not (a and b) diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index 7b2ba07..172e08e 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -10,7 +10,6 @@ import scipy.spatial as spatial import optax -from .. import _linalg from . import kernel from . import mean @@ -28,10 +27,10 @@ class GaussianProcess: nugget: float - x_train: jax.Array - y_train: jax.Array - s_train: jax.Array - L_train: jax.Array + train_x: jax.Array + train_y: jax.Array + train_s: jax.Array + train_cov_factor: kernel.CovarianceFactor seed: int key: jax.Array @@ -75,18 +74,15 @@ def condition( ): _, ny = self.validate_shapes(x, y, s) - self.x_train = jnp.asarray(x) - self.y_train = jnp.asarray(y) - if isinstance(s, (float, int)): - self.s_train = s * jnp.eye(ny * self.cdim) - else: - self.s_train = jnp.asarray(s) + self.train_x = jnp.asarray(x) + self.train_y = jnp.asarray(y) + self.train_s = jnp.asarray(s) - self.L_train = GaussianProcess._get_L_train( + self.train_cov_factor = GaussianProcess._get_train_cov_factor( (self.mu, self.k), - self.x_train, - self.y_train, - self.s_train + self.train_x, + self.train_y, + self.train_s ) def add_point(self, x, y, s): @@ -104,9 +100,16 @@ def validate_shapes( if nx != ny: raise ValueError - if not isinstance(s, (float, int)): - if s.shape != (ny, self.cdim): - raise ValueError + s_shape = jnp.shape(s) + if s_shape == (): + return nx, ny + + if s_shape not in { + (ny, self.cdim), + (ny * self.cdim,), + (ny * self.cdim, ny * self.cdim), + }: + raise ValueError return nx, ny @@ -130,35 +133,32 @@ def validate_output_shape(self, y: jax.Array | npt.NDArray) -> tuple[int, int]: return n, c - def get_L_train( + def get_train_cov_factor( self, x: jax.Array | npt.NDArray, y: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int - ) -> jax.Array: - _, ny = self.validate_shapes(x, y, s) + ) -> kernel.CovarianceFactor: + _, _ = self.validate_shapes(x, y, s) x = jnp.asarray(x) y = jnp.asarray(y) - if isinstance(s, (float, int)): - s = s * jnp.eye(ny * self.cdim) - else: - s = jnp.asarray(s) + s = jnp.asarray(s) - return GaussianProcess._get_L_train((self.mu, self.k), x, y, s) + return GaussianProcess._get_train_cov_factor((self.mu, self.k), x, y, s) # TODO: may want to change args for this function, but it's nice to maintain the same signature as # other core/mathematical routines. @staticmethod @jax.jit - def _get_L_train( + def _get_train_cov_factor( p: tuple[mean.Mean, kernel.Kernel], x: jax.Array, y: jax.Array, s: jax.Array - ) -> jax.Array: - mu, k = p - return jnp.linalg.cholesky(k(x, x) + s) + ) -> kernel.CovarianceFactor: + _, k = p + return k.covariance_factor(x, s) def tune( self, @@ -182,18 +182,17 @@ def tune( optim, optim_state, (self.mu, self.k), - self.x_train, - self.y_train, - self.s_train + self.train_x, + self.train_y, + self.train_s ) losses.append(loss) - # Recompute the Cholesky factor of observation matrix after tuning: - self.L_train = GaussianProcess._get_L_train( + self.train_cov_factor = GaussianProcess._get_train_cov_factor( (self.mu, self.k), - self.x_train, - self.y_train, - self.s_train + self.train_x, + self.train_y, + self.train_s ) return jnp.asarray(losses) @@ -213,19 +212,17 @@ def _step( p = eqx.apply_updates(p, updates) return optim_state, p, loss - def loss(self, + def loss( + self, x: jax.Array | npt.NDArray, y: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int ) -> jax.Array: - _, ny = self.validate_shapes(x, y, s) + _, _ = self.validate_shapes(x, y, s) x = jnp.asarray(x) y = jnp.asarray(y) - if isinstance(s, (float, int)): - s = s * jnp.eye(ny * self.cdim) - else: - s = jnp.asarray(s) + s = jnp.asarray(s) return GaussianProcess._loss((self.mu, self.k), x, y, s) @@ -238,11 +235,11 @@ def _loss( s: jax.Array ) -> jax.Array: mu, k = p - d = (y - mu(x)).flatten() - L_train = GaussianProcess._get_L_train((mu, k), x, y, s) + d = (y - mu(x)).reshape(-1) + cov_factor = GaussianProcess._get_train_cov_factor((mu, k), x, y, s) return ( - d.T @ _linalg.solve_cholesky(L_train, d) + - jnp.sum(jnp.log(jnp.diag(L_train))) + d.T @ cov_factor.solve(d) + + cov_factor.log_sqrt_det() ) def get_sample_point( @@ -274,11 +271,11 @@ def get_sample_point( optim_kwargs (...): ... """ - if not hasattr(self, 'x_train') or not hasattr(self, 'L_train'): + if not hasattr(self, 'train_x') or not hasattr(self, 'train_cov_factor'): raise ValueError('Gaussian process must be conditioned before sampling.') x = jnp.asarray(x) - if not np.issubdtype(x.dtype, np.inexact): + if not jnp.issubdtype(x.dtype, jnp.inexact): x = x.astype(jnp.float32) if x.ndim > 1: raise ValueError @@ -286,31 +283,40 @@ def get_sample_point( raise ValueError x = x.reshape((1, self.dim)) - if hull is None and ranges is None: + if not ((hull is not None) ^ (ranges is not None)): raise ValueError - if hull is None: - ranges = np.asarray(ranges) - if ranges.ndim != 2: - raise ValueError - if ranges.shape != (2, self.dim): - raise ValueError - - hull = spatial.ConvexHull( - np.stack( - np.meshgrid( - ranges.T, - indexing='ij' - ), - axis=-1 - ).reshape((-1, self.dim)) - ) - - A = hull.equations[:, :-1] - b = hull.equations[:, -1] + if self.dim > 1: + if hull is None: + ranges = jnp.asarray(ranges) + if ranges.ndim != 2: + raise ValueError + if ranges.shape != (2, self.dim): + raise ValueError + + hull = spatial.ConvexHull( + jnp.stack( + jnp.meshgrid( + *ranges.T, + indexing='ij' + ), + axis=-1 + ).reshape((-1, self.dim)) + ) + + A = hull.equations[:, :-1] + b = hull.equations[:, -1] + def inside_domain(x_cand: jax.Array): + return jnp.all(jnp.dot(x_cand, A.T) + b <= tol) - def inside_domain(x_cand: jax.Array): - return jnp.all(jnp.dot(x_cand, A.T) + b <= tol) + else: + if ranges is None: + raise ValueError + def inside_domain(x_cand: jax.Array): + return jnp.all( + (x_cand >= ranges[0, 0] - tol) & + (x_cand <= ranges[1, 0] + tol) + ) if not inside_domain(x): raise ValueError @@ -323,8 +329,8 @@ def inside_domain(x_cand: jax.Array): optim, optim_state, self.k, - self.x_train, - self.L_train, + self.train_x, + self.train_cov_factor, x ) @@ -341,12 +347,15 @@ def _step_sample_point( optim, optim_state, k: kernel.Kernel, - x_train: jax.Array, - L_train: jax.Array, + train_x: jax.Array, + train_cov_factor: kernel.CovarianceFactor, x: jax.Array ) -> tuple[jax.Array, jax.Array, jax.Array]: def loss_fn(x): - return -GaussianProcess._posterior_variance(k, x_train, L_train, x) + return ( + -1 * + GaussianProcess._posterior_variance(k, train_x, train_cov_factor, x) + ) loss, grads = jax.value_and_grad(loss_fn)(x) updates, optim_state = optim.update(grads, optim_state, x) @@ -357,13 +366,13 @@ def loss_fn(x): @jax.jit def _posterior_variance( k: kernel.Kernel, - x_train: jax.Array, - L_train: jax.Array, + train_x: jax.Array, + train_cov_factor: kernel.CovarianceFactor, x: jax.Array ) -> jax.Array: covariance = ( k(x, x) - - k(x, x_train) @ _linalg.solve_cholesky(L_train, k(x_train, x)) + k(x, train_x) @ train_cov_factor.solve(k(train_x, x)) ) return jnp.sum(jnp.diag(covariance)) @@ -390,13 +399,9 @@ def prior_realization(self, x, p=1): def posterior_mean(self, x): self.validate_input_shape(x) - return ( - self.mu(x) + - self.k(x, self.x_train) @ _linalg.solve_cholesky( - self.L_train, - self.y_train - self.mu(self.x_train) - ) - ) + d = (self.train_y - self.mu(self.train_x)).reshape(-1) + correction = self.k(x, self.train_x) @ self.train_cov_factor.solve(d) + return self.mu(x) + correction.reshape((-1, self.cdim)) def posterior_covariance(self, u, v): self.validate_input_shape(u) @@ -404,10 +409,7 @@ def posterior_covariance(self, u, v): return ( self.k(u, v) - - self.k(u, self.x_train) @ _linalg.solve_cholesky( - self.L_train, - self.k(self.x_train, v) - ) + self.k(u, self.train_x) @ self.train_cov_factor.solve(self.k(self.train_x, v)) ) def posterior_realization(self, x, p=1): diff --git a/UncertainSCI/gp/kernel.py b/UncertainSCI/gp/kernel.py index e0b0e43..cb7b962 100644 --- a/UncertainSCI/gp/kernel.py +++ b/UncertainSCI/gp/kernel.py @@ -9,6 +9,171 @@ import numpy.typing as npt from .. import _equinox +from .. import _linalg + + +def noise_as_matrix( + s: jax.Array | npt.NDArray | float | int, + matrix_size: int, + n: int, + cdim: int, +) -> jax.Array: + """ + Given ``s``, infer and return noise matrix. + + ``s`` can be: + * Scalar, interpreted as isotropic noise. + * A vector of length ``matrix_size``, + interpreted as noise per-observation per-channel. + * A matrix of shape ``(matrix_size, matrix_size)``, returned as-is. + * Of shape ``(n, cdim)`` such that ``n * cdim == matrix_size``, + interpreted as noise per-observation per-channel. + + Args: + s (array or scalar of noise): + Noise of observations. + """ + s = jnp.asarray(s) + + if s.ndim == 0: + return s * jnp.eye(matrix_size) + + if s.ndim == 1: + if s.shape != (matrix_size,): + raise ValueError + return jnp.diag(s) + + if s.ndim != 2: + raise ValueError + + if s.shape == (matrix_size, matrix_size): + return s + + if s.shape == (n, cdim) and n * cdim == matrix_size: + return jnp.diag(s.reshape(-1)) + + raise ValueError + + +class CovarianceFactor(eqx.Module, metaclass=abc.ABCMeta): + """ + Factorization of a covariance matrix used for solves and log determinants. + + Subclasses may store a Cholesky factor, an eigendecomposition, or another + structured representation. + """ + size: int = eqx.field(static=True) + + @abc.abstractmethod + def solve(self, y: jax.Array | npt.NDArray) -> jax.Array: + raise NotImplementedError + + @property + @abc.abstractmethod + def log_sqrt_det(self) -> jax.Array: + raise NotImplementedError + + @property + @abc.abstractmethod + def cholesky(self) -> jax.Array: + raise NotImplementedError + + +class Cholesky(CovarianceFactor): + """ + Dense lower Cholesky factorization of a noisy covariance matrix. + """ + L: jax.Array + + def __init__(self, L: jax.Array | npt.NDArray): + L = jnp.asarray(L) + if L.ndim != 2: + raise ValueError + if L.shape[0] != L.shape[1]: + raise ValueError + + self.size = L.shape[0] + self.L = L + + def solve(self, y: jax.Array | npt.NDArray) -> jax.Array: + return _linalg.solve_cholesky(self.L, y) + + @property + def log_sqrt_det(self) -> jax.Array: + return jnp.sum(jnp.log(jnp.diag(self.L))) + + @property + def cholesky(self) -> jax.Array: + return self.L + + +class KroneckerScalarNoise(CovarianceFactor): + """ + Eigendecomposition of ``left`` ⊗ ``right`` plus scalar isotropic noise. + """ + left_eigenvalues: jax.Array + left_eigenvectors: jax.Array + right_eigenvalues: jax.Array + right_eigenvectors: jax.Array + noise: jax.Array + left_size: int = eqx.field(static=True) + right_size: int = eqx.field(static=True) + + def __init__( + self, + left: jax.Array | npt.NDArray, + right: jax.Array | npt.NDArray, + noise: jax.Array | npt.NDArray | float | int, + ): + left = jnp.asarray(left) + right = jnp.asarray(right) + if left.ndim != 2 or left.shape[0] != left.shape[1]: + raise ValueError + if right.ndim != 2 or right.shape[0] != right.shape[1]: + raise ValueError + + self.size = left.shape[0] * right.shape[0] + self.left_size = left.shape[0] + self.right_size = right.shape[0] + self.left_eigenvalues, self.left_eigenvectors = jnp.linalg.eigh(left) + self.right_eigenvalues, self.right_eigenvectors = jnp.linalg.eigh(right) + self.noise = jnp.asarray(noise) + + @property + def eigenvalues(self) -> jax.Array: + return ( + self.left_eigenvalues.reshape((-1, 1)) * + self.right_eigenvalues.reshape((1, -1)) + + self.noise + ) + + def _solve_vector(self, y: jax.Array) -> jax.Array: + y = y.reshape((self.left_size, self.right_size)) + y = self.left_eigenvectors.T @ y @ self.right_eigenvectors + y = y / self.eigenvalues + y = self.left_eigenvectors @ y @ self.right_eigenvectors.T + return y.reshape((self.size,)) + + def solve(self, y: jax.Array | npt.NDArray) -> jax.Array: + y = jnp.asarray(y) + + if y.ndim == 1: + return self._solve_vector(y) + + if y.ndim == 2: + return jax.vmap(self._solve_vector, in_axes=1, out_axes=1)(y) + + raise ValueError + + @property + def log_sqrt_det(self) -> jax.Array: + return 0.5 * jnp.sum(jnp.log(self.eigenvalues)) + + @property + def cholesky(self) -> jax.Array: + q = jnp.kron(self.left_eigenvectors, self.right_eigenvectors) + covariance = (q * self.eigenvalues.reshape(-1)) @ q.T + return jnp.linalg.cholesky(covariance) class Kernel(eqx.Module, metaclass=abc.ABCMeta): @@ -29,7 +194,7 @@ def __call__( self, x: jax.Array | npt.NDArray, y: jax.Array | npt.NDArray - ): + ) -> jax.Array: """ Evaluate this kernel. @@ -43,6 +208,28 @@ def __call__( """ raise NotImplementedError + def covariance_factor( + self, + x: jax.Array | npt.NDArray, + s: jax.Array | npt.NDArray | float | int + ) -> CovarianceFactor: + """ + Factor the covariance of observed values at ``x``. + + The default implementation materializes a dense covariance matrix. + Structured kernels can override this to avoid dense factorization. + + Args: + x (array): + Array of coordinates shaped as ``(num_x, dim)``. + s (array): + Array of noise associated with each observation. + """ + x = jnp.asarray(x) + covariance = self(x, x) + s = noise_as_matrix(s, covariance.shape[0], x.shape[0], self.cdim) + return Cholesky(jnp.linalg.cholesky(covariance + s)) + class Gaussian(Kernel): r""" @@ -119,6 +306,9 @@ def __init__( C: jax.Array | npt.NDArray, C_is_static: bool = False ): + if k.cdim != 1: + raise ValueError('Kronecker coordinate kernel must be scalar-valued.') + super().__init__(dim=dim, cdim=cdim) self.k = k self.C = _equinox.PositiveDefinite(D=C, is_static=C_is_static) @@ -129,3 +319,16 @@ def __call__( y: jax.Array | npt.NDArray ) -> jax.Array: return jnp.kron(self.k(x, y), self.C()) + + def covariance_factor( + self, + x: jax.Array | npt.NDArray, + s: jax.Array | npt.NDArray | float | int + ) -> CovarianceFactor: + x = jnp.asarray(x) + s = jnp.asarray(s) + + if s.ndim == 0: + return KroneckerScalarNoise(self.k(x, x), self.C(), s) + + return super().covariance_factor(x, s) diff --git a/tests/test_gp.py b/tests/test_gp.py new file mode 100644 index 0000000..9227062 --- /dev/null +++ b/tests/test_gp.py @@ -0,0 +1,157 @@ +import unittest + +import jax.numpy as jnp + +import UncertainSCI.gp as gp + + +class GaussianProcessTestCase(unittest.TestCase): + """ + Performs basic tests for Gaussian process routines. + """ + + def setUp(self): + self.longMessage = True + + def test_get_sample_point_optimizes_posterior_variance(self): + """Test posterior-variance sampling over an interval.""" + mu = gp.mean.Affine( + dim=1, + cdim=1, + a=jnp.zeros((1, 1)), + b=0.0, + a_is_static=True, + b_is_static=True + ) + k = gp.kernel.Gaussian( + dim=1, + cdim=1, + D=jnp.eye(1), + D_is_static=True + ) + g = gp.GaussianProcess(dim=1, cdim=1, mu=mu, k=k, seed=0) + g.condition( + jnp.array([[0.0]]), + jnp.array([[0.0]]), + 1e-8 + ) + + x_start = jnp.array([0.25]) + x_sample = g.get_sample_point( + x_start, + ranges=jnp.array([[0.0], [2.0]]), + n=100, + optim_kwargs={'learning_rate': 0.05} + ) + + x_start = x_start.reshape((1, 1)) + start_variance = g.posterior_covariance(x_start, x_start)[0, 0] + sample_variance = g.posterior_covariance(x_sample, x_sample)[0, 0] + + self.assertEqual(x_sample.shape, (1, 1)) + self.assertGreaterEqual(float(x_sample[0, 0]), 1.5) + self.assertLessEqual(float(x_sample[0, 0]), 2.0 + 1e-6) + self.assertGreater(float(sample_variance), float(start_variance)) + + def test_kronecker_scalar_noise_factor_matches_dense_solve(self): + """Test Kronecker structured solves with scalar observation noise.""" + x = jnp.array([[0.0], [0.5], [1.0]]) + base = gp.kernel.Gaussian( + dim=1, + cdim=1, + D=jnp.eye(1), + D_is_static=True + ) + c = jnp.array([[2.0, 0.3], [0.3, 1.5]]) + k = gp.kernel.Kronecker( + dim=1, + cdim=2, + k=base, + C=c, + C_is_static=True + ) + + s = 1e-4 + factor = k.covariance_factor(x, s) + dense_covariance = k(x, x) + s * jnp.eye(x.shape[0] * k.cdim) + rhs = jnp.arange(x.shape[0] * k.cdim, dtype=float) + + dense_cholesky = jnp.linalg.cholesky(dense_covariance) + dense_log_sqrt_det = jnp.sum(jnp.log(jnp.diag(dense_cholesky))) + + self.assertIsInstance(factor, gp.kernel.KroneckerScalarNoise) + self.assertTrue( + jnp.allclose( + factor.solve(rhs), + jnp.linalg.solve(dense_covariance, rhs), + rtol=1e-4, + atol=1e-4 + ) + ) + self.assertAlmostEqual( + float(factor.log_sqrt_det), + float(dense_log_sqrt_det), + places=4 + ) + + def test_kronecker_posterior_handles_vector_outputs(self): + """Test posterior calculations for vector-valued Kronecker kernels.""" + mu = gp.mean.Affine( + dim=1, + cdim=2, + a=jnp.zeros((2, 1)), + b=jnp.zeros(2), + a_is_static=True, + b_is_static=True + ) + base = gp.kernel.Gaussian( + dim=1, + cdim=1, + D=jnp.eye(1), + D_is_static=True + ) + c = jnp.array([[2.0, 0.3], [0.3, 1.5]]) + k = gp.kernel.Kronecker( + dim=1, + cdim=2, + k=base, + C=c, + C_is_static=True + ) + g = gp.GaussianProcess(dim=1, cdim=2, mu=mu, k=k, seed=0) + x_train = jnp.array([[0.0], [0.5], [1.0]]) + y_train = jnp.array([[0.0, 1.0], [0.5, -0.5], [1.0, 0.25]]) + x_test = jnp.array([[0.25], [0.75]]) + + g.condition(x_train, y_train, 1e-4) + + dense_covariance = k(x_train, x_train) + 1e-4 * jnp.eye( + x_train.shape[0] * k.cdim + ) + expected_mean = ( + k(x_test, x_train) @ + jnp.linalg.solve(dense_covariance, y_train.reshape(-1)) + ).reshape((x_test.shape[0], k.cdim)) + expected_covariance = ( + k(x_test, x_test) - + k(x_test, x_train) @ + jnp.linalg.solve(dense_covariance, k(x_train, x_test)) + ) + + self.assertIsInstance(g.train_cov_factor, gp.kernel.KroneckerScalarNoise) + self.assertEqual(g.posterior_mean(x_test).shape, (2, 2)) + self.assertEqual(g.posterior_covariance(x_test, x_test).shape, (4, 4)) + self.assertTrue( + jnp.allclose(g.posterior_mean(x_test), expected_mean, atol=1e-4) + ) + self.assertTrue( + jnp.allclose( + g.posterior_covariance(x_test, x_test), + expected_covariance, + atol=1e-4 + ) + ) + + +if __name__ == "__main__": + unittest.main(verbosity=2) From fc908c383ac4e911c05a06dfe678f0705717019e Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 20 Aug 2026 10:14:02 -0600 Subject: [PATCH 05/22] gp: cast ranges to array in posterior variance optimization --- UncertainSCI/gp/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index 172e08e..0d4a69f 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -312,6 +312,8 @@ def inside_domain(x_cand: jax.Array): else: if ranges is None: raise ValueError + + ranges = jnp.asarray(ranges) def inside_domain(x_cand: jax.Array): return jnp.all( (x_cand >= ranges[0, 0] - tol) & From adabd1f40797e5ac879bbac16f972d2e713eaa63 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 20 Aug 2026 10:14:18 -0600 Subject: [PATCH 06/22] gp.kernel: revert @property decorator in CovarianceFactor --- UncertainSCI/gp/kernel.py | 6 ------ tests/test_gp.py | 2 +- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/UncertainSCI/gp/kernel.py b/UncertainSCI/gp/kernel.py index cb7b962..b600919 100644 --- a/UncertainSCI/gp/kernel.py +++ b/UncertainSCI/gp/kernel.py @@ -68,12 +68,10 @@ class CovarianceFactor(eqx.Module, metaclass=abc.ABCMeta): def solve(self, y: jax.Array | npt.NDArray) -> jax.Array: raise NotImplementedError - @property @abc.abstractmethod def log_sqrt_det(self) -> jax.Array: raise NotImplementedError - @property @abc.abstractmethod def cholesky(self) -> jax.Array: raise NotImplementedError @@ -98,11 +96,9 @@ def __init__(self, L: jax.Array | npt.NDArray): def solve(self, y: jax.Array | npt.NDArray) -> jax.Array: return _linalg.solve_cholesky(self.L, y) - @property def log_sqrt_det(self) -> jax.Array: return jnp.sum(jnp.log(jnp.diag(self.L))) - @property def cholesky(self) -> jax.Array: return self.L @@ -165,11 +161,9 @@ def solve(self, y: jax.Array | npt.NDArray) -> jax.Array: raise ValueError - @property def log_sqrt_det(self) -> jax.Array: return 0.5 * jnp.sum(jnp.log(self.eigenvalues)) - @property def cholesky(self) -> jax.Array: q = jnp.kron(self.left_eigenvectors, self.right_eigenvectors) covariance = (q * self.eigenvalues.reshape(-1)) @ q.T diff --git a/tests/test_gp.py b/tests/test_gp.py index 9227062..096cbb8 100644 --- a/tests/test_gp.py +++ b/tests/test_gp.py @@ -89,7 +89,7 @@ def test_kronecker_scalar_noise_factor_matches_dense_solve(self): ) ) self.assertAlmostEqual( - float(factor.log_sqrt_det), + float(factor.log_sqrt_det()), float(dense_log_sqrt_det), places=4 ) From 82ffc52802409afae6d54998abe2a7095faa1713 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 20 Aug 2026 10:14:38 -0600 Subject: [PATCH 07/22] gp: add 1-d demo notebook --- gp-jax-tuning.ipynb => demos/gp-1d.ipynb | 56 ++++++++++++------------ 1 file changed, 29 insertions(+), 27 deletions(-) rename gp-jax-tuning.ipynb => demos/gp-1d.ipynb (88%) diff --git a/gp-jax-tuning.ipynb b/demos/gp-1d.ipynb similarity index 88% rename from gp-jax-tuning.ipynb rename to demos/gp-1d.ipynb index e04af89..c7eeb3f 100644 --- a/gp-jax-tuning.ipynb +++ b/demos/gp-1d.ipynb @@ -33,8 +33,8 @@ "\n", "def print_loss_here(g): # Simple helper for this notebook; does not generalize.\n", " print(\n", - " f'Loss at D = {jnp.squeeze(g.k.D()):.4e} (length scale {1 / jnp.squeeze(g.k.D())**2:.4e}): '\n", - " f'{g.loss(x_train, y_train, s_train):.4e}'\n", + " f'Loss at D = {jnp.squeeze(g.k.D()):.4e} (length scale {1 / jnp.sqrt(jnp.squeeze(g.k.D())):.4e}): '\n", + " f'{g.loss(train_x, train_y, train_s):.4e}'\n", " )\n" ] }, @@ -53,7 +53,7 @@ "def f_noisy(x):\n", " y = f_hidden(x)\n", " return (\n", - " y,\n", + " y + 1e-3 * np.random.randn(*y.shape),\n", " 1e-3 * jnp.ones_like(y)\n", " )\n" ] @@ -78,7 +78,8 @@ " cdim=C,\n", " D=jnp.ones((1, 1)),\n", ")\n", - "g = gp.GaussianProcess(dim=D, cdim=C, mu=mu, k=k, seed=seed, nugget=1e-3)\n" + "g = gp.GaussianProcess(dim=D, cdim=C, mu=mu, k=k, seed=seed, nugget=1e-3)\n", + "# g = gp.GaussianProcess(dim=D, cdim=C, mu=mu, k=k, seed=seed, nugget=1e-4)\n" ] }, { @@ -123,9 +124,9 @@ " )\n", " if which == 'posterior':\n", " ax.errorbar(\n", - " g.x_train.squeeze(),\n", - " g.y_train.squeeze(),\n", - " yerr=g.s_train.squeeze(),\n", + " g.train_x.squeeze(),\n", + " g.train_y.squeeze(),\n", + " yerr=g.train_s.squeeze(),\n", " capsize=2,\n", " linestyle='none'\n", " )\n", @@ -140,10 +141,10 @@ " )\n", " if which == 'posterior':\n", " if colorlast:\n", - " ax.vlines(g.x_train.squeeze()[:-1], 0, 1, color='tab:green')\n", - " ax.vlines(g.x_train.squeeze()[-1], 0, 1, color='tab:red')\n", + " ax.vlines(g.train_x.squeeze()[:-1], 0, 1, color='tab:green')\n", + " ax.vlines(g.train_x.squeeze()[-1], 0, 1, color='tab:red')\n", " else:\n", - " ax.vlines(g.x_train.squeeze(), 0, 1, color='tab:green')\n", + " ax.vlines(g.train_x.squeeze(), 0, 1, color='tab:green')\n", " ax.set_title('Variance')\n", "\n", "\n", @@ -217,7 +218,7 @@ " p_value = p_value * jnp.ones(shape)\n", "\n", " setattr(p_parent, p_name[-1], p_value)\n", - " losses = losses.at[i].set(g.loss(g.x_train, g.y_train, g.s_train))\n", + " losses = losses.at[i].set(g.loss(g.train_x, g.train_y, g.train_s))\n", "\n", " ax.plot(p_range, losses)\n", " if x is not None:\n", @@ -269,10 +270,10 @@ "source": [ "N_INIT = 10\n", "\n", - "x_train = jnp.linspace(0, 2 * jnp.pi, N_INIT).reshape(-1, D)\n", - "y_train, s_train = f_noisy(x_train) \n", + "train_x = jnp.linspace(0, 2 * jnp.pi, N_INIT).reshape(-1, D)\n", + "train_y, train_s = f_noisy(train_x) \n", "\n", - "g.condition(x_train, y_train, s_train)\n" + "g.condition(train_x, train_y, train_s)\n" ] }, { @@ -291,8 +292,8 @@ "outputs": [], "source": [ "fig, (ax1, ax2, ax3) = plot_in_subplots(3, 1, dpi=FIGDPI)\n", - "plot_distribution_mean(ax1, g, x_plot, 'prior')\n", - "plot_distribution_variance(ax2, g, x_plot, 'prior')\n", + "plot_distribution_mean(ax1, g, x_plot, 'posterior')\n", + "plot_distribution_variance(ax2, g, x_plot, 'posterior', colorlast=False)\n", "plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", "plt.show()\n", "\n", @@ -367,25 +368,26 @@ "\n", "x_sample = jnp.linspace(0, 2 * jnp.pi, N_SAMPLE).reshape(-1, D)\n", "\n", - "# def covariance(...): ...\n", - "\n", - "# def optimize_x_over_covariance(...): ...\n", - "\n", "for i in range(N_TRAIN):\n", + " # Greedily sample posterior variance:\n", + " _, c, _ = gp_stats(g, x_sample, p=0)\n", + " _x = x_sample[jnp.argmax(c)].reshape(-1, D)\n", "\n", - " # ...do those things here...\n", + " # Iteratively optimize posterior variance:\n", + " _x = g.get_sample_point(_x.reshape((1,)), ranges=[[0.], [2 * jnp.pi]])\n", "\n", - " _, c, _ = gp_stats(g, x_sample, p=0) # Greedy sampling on dense mesh.\n", - " _x = x_sample[jnp.argmax(c)].reshape(-1, D)\n", + " # Observe that point:\n", " _y, _s = f_noisy(_x)\n", "\n", - " x_train = jnp.concat((x_train, _x))\n", - " y_train = jnp.concat((y_train, _y))\n", - " s_train = jnp.concat((s_train, _s))\n", + " train_x = jnp.concat((train_x, _x))\n", + " train_y = jnp.concat((train_y, _y))\n", + " train_s = jnp.concat((train_s, _s))\n", "\n", - " g.condition(x_train, y_train, s_train)\n", + " # Condition with new data:\n", + " g.condition(train_x, train_y, train_s)\n", " g.tune()\n", "\n", + " # Plot:\n", " fig, (ax1, ax2, ax3) = plot_in_subplots(3, 1, dpi=FIGDPI)\n", " plot_distribution_mean(ax1, g, x_plot, 'posterior')\n", " plot_distribution_variance(ax2, g, x_plot, 'posterior')\n", From 6a1a2110049eece1010d108b969d0442d1f39387 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Sat, 29 Aug 2026 13:40:06 -0600 Subject: [PATCH 08/22] gp: clean up public namespace and method names --- UncertainSCI/gp/__init__.py | 33 ++++++++++----------------------- demos/gp-1d.ipynb | 4 ++-- 2 files changed, 12 insertions(+), 25 deletions(-) diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index 0d4a69f..3c60496 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -199,7 +199,7 @@ def tune( @staticmethod @jax.jit(static_argnames=('optim',)) - def _step( + def _step_hyperparameters( optim, optim_state, p: tuple[mean.Mean, kernel.Kernel], @@ -207,12 +207,12 @@ def _step( y: jax.Array, s: jax.Array ) -> tuple[jax.Array, tuple[mean.Mean, kernel.Kernel], jax.Array]: - loss, grads = eqx.filter_value_and_grad(GaussianProcess._loss)(p, x, y, s) + loss, grads = eqx.filter_value_and_grad(GaussianProcess._loss_hyperparameters)(p, x, y, s) updates, optim_state = optim.update(grads, optim_state, p) p = eqx.apply_updates(p, updates) return optim_state, p, loss - def loss( + def loss_hyperparameters( self, x: jax.Array | npt.NDArray, y: jax.Array | npt.NDArray, @@ -224,11 +224,11 @@ def loss( y = jnp.asarray(y) s = jnp.asarray(s) - return GaussianProcess._loss((self.mu, self.k), x, y, s) + return GaussianProcess._loss_hyperparameters((self.mu, self.k), x, y, s) @staticmethod @jax.jit - def _loss( + def _loss_hyperparameters( p: tuple[mean.Mean, kernel.Kernel], x: jax.Array, y: jax.Array, @@ -353,31 +353,18 @@ def _step_sample_point( train_cov_factor: kernel.CovarianceFactor, x: jax.Array ) -> tuple[jax.Array, jax.Array, jax.Array]: - def loss_fn(x): - return ( - -1 * - GaussianProcess._posterior_variance(k, train_x, train_cov_factor, x) + def loss_fn(x): # Posterior variance: + covariance = ( + k(x, x) - + k(x, train_x) @ train_cov_factor.solve(k(train_x, x)) ) + return -1 * jnp.sum(jnp.diag(covariance)) loss, grads = jax.value_and_grad(loss_fn)(x) updates, optim_state = optim.update(grads, optim_state, x) x = eqx.apply_updates(x, updates) return optim_state, x, -loss - @staticmethod - @jax.jit - def _posterior_variance( - k: kernel.Kernel, - train_x: jax.Array, - train_cov_factor: kernel.CovarianceFactor, - x: jax.Array - ) -> jax.Array: - covariance = ( - k(x, x) - - k(x, train_x) @ train_cov_factor.solve(k(train_x, x)) - ) - return jnp.sum(jnp.diag(covariance)) - def prior_mean(self, x): self.validate_input_shape(x) diff --git a/demos/gp-1d.ipynb b/demos/gp-1d.ipynb index c7eeb3f..b1f72a4 100644 --- a/demos/gp-1d.ipynb +++ b/demos/gp-1d.ipynb @@ -34,7 +34,7 @@ "def print_loss_here(g): # Simple helper for this notebook; does not generalize.\n", " print(\n", " f'Loss at D = {jnp.squeeze(g.k.D()):.4e} (length scale {1 / jnp.sqrt(jnp.squeeze(g.k.D())):.4e}): '\n", - " f'{g.loss(train_x, train_y, train_s):.4e}'\n", + " f'{g.loss_hyperparameters(train_x, train_y, train_s):.4e}'\n", " )\n" ] }, @@ -218,7 +218,7 @@ " p_value = p_value * jnp.ones(shape)\n", "\n", " setattr(p_parent, p_name[-1], p_value)\n", - " losses = losses.at[i].set(g.loss(g.train_x, g.train_y, g.train_s))\n", + " losses = losses.at[i].set(g.loss_hyperparameters(g.train_x, g.train_y, g.train_s))\n", "\n", " ax.plot(p_range, losses)\n", " if x is not None:\n", From b2bfb192d298b6b978e5c614e34fd4d46fd59869 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Sat, 29 Aug 2026 13:42:33 -0600 Subject: [PATCH 09/22] gp.kernel: make CovarianceFactor(s) private classes inside kernel/gp namespace * may break this out further later * for now, I don't want any _CF* objects getting confused with covariance kernels --- UncertainSCI/gp/__init__.py | 10 +++++----- UncertainSCI/gp/kernel.py | 14 +++++++------- tests/test_gp.py | 4 ++-- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index 3c60496..0c212f3 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -30,7 +30,7 @@ class GaussianProcess: train_x: jax.Array train_y: jax.Array train_s: jax.Array - train_cov_factor: kernel.CovarianceFactor + train_cov_factor: kernel._CF seed: int key: jax.Array @@ -138,7 +138,7 @@ def get_train_cov_factor( x: jax.Array | npt.NDArray, y: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int - ) -> kernel.CovarianceFactor: + ) -> kernel._CF: _, _ = self.validate_shapes(x, y, s) x = jnp.asarray(x) @@ -156,7 +156,7 @@ def _get_train_cov_factor( x: jax.Array, y: jax.Array, s: jax.Array - ) -> kernel.CovarianceFactor: + ) -> kernel._CF: _, k = p return k.covariance_factor(x, s) @@ -178,7 +178,7 @@ def tune( losses = [] for i in range(n): - optim_state, (self.mu, self.k), loss = GaussianProcess._step( + optim_state, (self.mu, self.k), loss = GaussianProcess._step_hyperparameters( optim, optim_state, (self.mu, self.k), @@ -350,7 +350,7 @@ def _step_sample_point( optim_state, k: kernel.Kernel, train_x: jax.Array, - train_cov_factor: kernel.CovarianceFactor, + train_cov_factor: kernel._CF, x: jax.Array ) -> tuple[jax.Array, jax.Array, jax.Array]: def loss_fn(x): # Posterior variance: diff --git a/UncertainSCI/gp/kernel.py b/UncertainSCI/gp/kernel.py index b600919..514836c 100644 --- a/UncertainSCI/gp/kernel.py +++ b/UncertainSCI/gp/kernel.py @@ -55,7 +55,7 @@ def noise_as_matrix( raise ValueError -class CovarianceFactor(eqx.Module, metaclass=abc.ABCMeta): +class _CF(eqx.Module, metaclass=abc.ABCMeta): """ Factorization of a covariance matrix used for solves and log determinants. @@ -77,7 +77,7 @@ def cholesky(self) -> jax.Array: raise NotImplementedError -class Cholesky(CovarianceFactor): +class _CF_Cholesky(_CF): """ Dense lower Cholesky factorization of a noisy covariance matrix. """ @@ -103,7 +103,7 @@ def cholesky(self) -> jax.Array: return self.L -class KroneckerScalarNoise(CovarianceFactor): +class _CF_Kronecker_ScalarNoise(_CF): """ Eigendecomposition of ``left`` ⊗ ``right`` plus scalar isotropic noise. """ @@ -206,7 +206,7 @@ def covariance_factor( self, x: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int - ) -> CovarianceFactor: + ) -> _CF: """ Factor the covariance of observed values at ``x``. @@ -222,7 +222,7 @@ def covariance_factor( x = jnp.asarray(x) covariance = self(x, x) s = noise_as_matrix(s, covariance.shape[0], x.shape[0], self.cdim) - return Cholesky(jnp.linalg.cholesky(covariance + s)) + return _CF_Cholesky(jnp.linalg.cholesky(covariance + s)) class Gaussian(Kernel): @@ -318,11 +318,11 @@ def covariance_factor( self, x: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int - ) -> CovarianceFactor: + ) -> _CF: x = jnp.asarray(x) s = jnp.asarray(s) if s.ndim == 0: - return KroneckerScalarNoise(self.k(x, x), self.C(), s) + return _CF_Kronecker_ScalarNoise(self.k(x, x), self.C(), s) return super().covariance_factor(x, s) diff --git a/tests/test_gp.py b/tests/test_gp.py index 096cbb8..aca1054 100644 --- a/tests/test_gp.py +++ b/tests/test_gp.py @@ -79,7 +79,7 @@ def test_kronecker_scalar_noise_factor_matches_dense_solve(self): dense_cholesky = jnp.linalg.cholesky(dense_covariance) dense_log_sqrt_det = jnp.sum(jnp.log(jnp.diag(dense_cholesky))) - self.assertIsInstance(factor, gp.kernel.KroneckerScalarNoise) + self.assertIsInstance(factor, gp.kernel._CF_Kronecker_ScalarNoise) self.assertTrue( jnp.allclose( factor.solve(rhs), @@ -138,7 +138,7 @@ def test_kronecker_posterior_handles_vector_outputs(self): jnp.linalg.solve(dense_covariance, k(x_train, x_test)) ) - self.assertIsInstance(g.train_cov_factor, gp.kernel.KroneckerScalarNoise) + self.assertIsInstance(g.train_cov_factor, gp.kernel._CF_Kronecker_ScalarNoise) self.assertEqual(g.posterior_mean(x_test).shape, (2, 2)) self.assertEqual(g.posterior_covariance(x_test, x_test).shape, (4, 4)) self.assertTrue( From 374dd855dc7be789407ce467000e1fcb9dfd5246 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Sat, 29 Aug 2026 15:28:50 -0600 Subject: [PATCH 10/22] gp: missing factor of 1/2 on quadractic term in log marginal likelihood --- UncertainSCI/gp/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index 0c212f3..7cb6e5e 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -238,7 +238,7 @@ def _loss_hyperparameters( d = (y - mu(x)).reshape(-1) cov_factor = GaussianProcess._get_train_cov_factor((mu, k), x, y, s) return ( - d.T @ cov_factor.solve(d) + + d.T @ cov_factor.solve(d) / 2 + cov_factor.log_sqrt_det() ) From 58aae91aa5d33de4d1a525b96648d0b7bc259702 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Sun, 30 Aug 2026 12:16:34 -0600 Subject: [PATCH 11/22] gp: add/correct docstrings across gp submodule --- UncertainSCI/gp/__init__.py | 249 ++++++++++++++++++++++++++++++++---- UncertainSCI/gp/kernel.py | 170 ++++++++++++++++++------ UncertainSCI/gp/mean.py | 5 +- tests/test_gp.py | 4 +- 4 files changed, 358 insertions(+), 70 deletions(-) diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index 7cb6e5e..a3a0b3d 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -20,6 +20,22 @@ } class GaussianProcess: + """ + A Gaussian process emulator. + + Builds a distribution on functions like + + .. math:: + + f \\sim \\mathcal{GP}(\\mu, k) + + where :math:`f` is an unknown function, :math:`\\mu` is a mean function and + :math:`k` is a positive definite covariance kernel. + + See `Gaussian Processes on Wikipedia `_ + or `Kanagawa et al. (2018) `_ for more + information. + """ dim: int cdim: int mu: mean.Mean @@ -30,7 +46,7 @@ class GaussianProcess: train_x: jax.Array train_y: jax.Array train_s: jax.Array - train_cov_factor: kernel._CF + train_cov_factor: kernel.CovarianceFactor seed: int key: jax.Array @@ -44,6 +60,22 @@ def __init__( nugget: float = 1e-3, seed: int | None = None ) -> None: + """ + Args: + dim (int): + Dimension of the domain of unknown function :math:`f`. + cdim (int): + Dimension of the codomain of unknown function :math:`f`. + mu (mean.Mean): + Mean function of Gaussian process. + k (kernel.Kernel): + Covariance kernel of Gaussian process. + nugget (float): + Numerical nugget added in some methods for stability. + seed (int, optional): + Seed to initialize JAX PRNG. Default initialized from + :func:`numpy.random.randint`. + """ if mu.dim != dim: raise ValueError if k.dim != dim: @@ -72,7 +104,28 @@ def condition( y: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int ): - _, ny = self.validate_shapes(x, y, s) + """ + Condition this Gaussian process from data ``(x, y, s)`` where ``s`` is the + variance of the observation pair ``(x, y)``. + + Args: + x (jax.Array or numpy.ndarray): + x coordinates. Must have shape ``(n, d)`` where ``d`` is the dimension + of the domain. + y (jax.Array or numpy.ndarray): + y coordinates. Must have shape ``(n, c)`` where ``c`` is the dimension + of the codomain. + s (jax.Array, numpy.ndarray, float, or int): + Covariances of ``(x, y)``. If array-type, must have shape: + + * ``()`` + * ``(n, c)`` + * ``(n * c)`` + * ``(n * c, n * c)`` + + where ``c`` is the dimension of the codomain. + """ + self.validate_shapes(x, y, s) self.train_x = jnp.asarray(x) self.train_y = jnp.asarray(y) @@ -93,7 +146,27 @@ def validate_shapes( x: jax.Array | npt.NDArray, y: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int - ) -> tuple[int, int]: + ): + """ + Validate shapes of arguments to various methods. + + Args: + x (jax.Array or numpy.ndarray): + x coordinates. Must have shape ``(n, d)`` where ``d`` is the dimension + of the domain. + y (jax.Array or numpy.ndarray): + y coordinates. Must have shape ``(n, c)`` where ``c`` is the dimension + of the codomain. + s (jax.Array, numpy.ndarray, float, or int): + Covariances of ``(x, y)``. If array-type, must have shape: + + * ``()`` + * ``(n, c)`` + * ``(n * c)`` + * ``(n * c, n * c)`` + + where ``c`` is the dimension of the codomain. + """ nx, _ = self.validate_input_shape(x) ny, _ = self.validate_output_shape(y) @@ -102,18 +175,27 @@ def validate_shapes( s_shape = jnp.shape(s) if s_shape == (): - return nx, ny - - if s_shape not in { + return + elif s_shape not in { (ny, self.cdim), (ny * self.cdim,), (ny * self.cdim, ny * self.cdim), }: raise ValueError - return nx, ny - def validate_input_shape(self, x: jax.Array | npt.NDArray) -> tuple[int, int]: + """ + Validate shapes of input-type arguments to various methods (e.g., x values). + + Args: + x (jax.Array or numpy.ndarray): + x coordinates. Must have shape ``(n, d)`` where ``d`` is the dimension + of the domain. + + Returns: + tuple[int, int]: + A tuple of sizes like ``(n, d)``. + """ if x.ndim != 2: raise ValueError @@ -124,6 +206,18 @@ def validate_input_shape(self, x: jax.Array | npt.NDArray) -> tuple[int, int]: return n, d def validate_output_shape(self, y: jax.Array | npt.NDArray) -> tuple[int, int]: + """ + Validate shapes of output-type arguments to various methods (e.g., y values). + + Args: + y (jax.Array or numpy.ndarray): + y coordinates. Must have shape ``(n, c)`` where ``c`` is the dimension + of the codomain. + + Returns: + tuple[int, int]: + A tuple of sizes like ``(n, c)``. + """ if y.ndim != 2: raise ValueError @@ -138,8 +232,28 @@ def get_train_cov_factor( x: jax.Array | npt.NDArray, y: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int - ) -> kernel._CF: - _, _ = self.validate_shapes(x, y, s) + ) -> kernel.CovarianceFactor: + """ + Get training covariance factor associated with the data ``(x, y, s)``. + + Args: + x (jax.Array or numpy.ndarray): + x coordinates. Must have shape ``(n, d)`` where ``d`` is the dimension + of the domain. + y (jax.Array or numpy.ndarray): + y coordinates. Must have shape ``(n, c)`` where ``c`` is the dimension + of the codomain. + s (jax.Array, numpy.ndarray, float, or int): + Covariances of ``(x, y)``. If array-type, must have shape: + + * ``()`` + * ``(n, c)`` + * ``(n * c)`` + * ``(n * c, n * c)`` + + where ``c`` is the dimension of the codomain. + """ + self.validate_shapes(x, y, s) x = jnp.asarray(x) y = jnp.asarray(y) @@ -156,7 +270,7 @@ def _get_train_cov_factor( x: jax.Array, y: jax.Array, s: jax.Array - ) -> kernel._CF: + ) -> kernel.CovarianceFactor: _, k = p return k.covariance_factor(x, s) @@ -165,7 +279,7 @@ def tune( n: int = 1000, optim = DEFAULT_OPTIM, optim_kwargs: dict = DEFAULT_OPTIM_KWARGS, - ): + ) -> jax.Array: """ Args: optim (callable): @@ -218,7 +332,30 @@ def loss_hyperparameters( y: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int ) -> jax.Array: - _, _ = self.validate_shapes(x, y, s) + """ + Compute the loss of Gaussian process' prior given data ``(x, y, s)``. + + This is the negative log marginal likelihood of the Gaussian process given data + ``(x, y, s)``, discarding constants. + + Args: + x (jax.Array or numpy.ndarray): + x coordinates. Must have shape ``(n, d)`` where ``d`` is the dimension + of the domain. + y (jax.Array or numpy.ndarray): + y coordinates. Must have shape ``(n, c)`` where ``c`` is the dimension + of the codomain. + s (jax.Array, numpy.ndarray, float, or int): + Covariances of ``(x, y)``. If array-type, must have shape: + + * ``()`` + * ``(n, c)`` + * ``(n * c)`` + * ``(n * c, n * c)`` + + where ``c`` is the dimension of the codomain. + """ + self.validate_shapes(x, y, s) x = jnp.asarray(x) y = jnp.asarray(y) @@ -244,32 +381,34 @@ def _loss_hyperparameters( def get_sample_point( self, - x: jax.Array, + x: jax.Array | npt.NDArray, hull: spatial.ConvexHull | None = None, ranges: jax.Array | npt.NDArray | None = None, tol: float = 1e-6, n: int = 1000, optim = DEFAULT_OPTIM, optim_kwargs: dict = DEFAULT_OPTIM_KWARGS, - ): + ) -> jax.Array: """ + Optimize a coordinate ``x`` to maximize posterior variance. + Args: - x (jax.Array): + x (jax.Array or numpy.ndarray): Candidate coordinate from which to start optimization. hull (scipy.spatial.ConvexHull, optional): Convex hull of coordinate domain, must be supplied if ``ranges`` is not. - ranges (jax.Array or array-like, optional): + ranges (jax.Array or numpy.ndarray, optional): Array of shape (2, dim) of corners of prism of coordinate domain. Must be supplied if ``hull`` is not. tol (float): Tolerance of hull inclusion criterion. n (int): Number of steps to take in optimization. - optim (...): - ... - optim_kwargs (...): - ... + optim (callable): + Optax or Optax-style optimizer. + optim_kwargs (dict): + kwargs for ``optim``. """ if not hasattr(self, 'train_x') or not hasattr(self, 'train_cov_factor'): raise ValueError('Gaussian process must be conditioned before sampling.') @@ -339,9 +478,9 @@ def inside_domain(x_cand: jax.Array): if inside_domain(x_cand): x = x_cand else: - return x + return jnp.asarray(x) - return x + return jnp.asarray(x) @staticmethod @jax.jit(static_argnames=('optim',)) @@ -350,7 +489,7 @@ def _step_sample_point( optim_state, k: kernel.Kernel, train_x: jax.Array, - train_cov_factor: kernel._CF, + train_cov_factor: kernel.CovarianceFactor, x: jax.Array ) -> tuple[jax.Array, jax.Array, jax.Array]: def loss_fn(x): # Posterior variance: @@ -366,17 +505,46 @@ def loss_fn(x): # Posterior variance: return optim_state, x, -loss def prior_mean(self, x): + """ + Evaluate the prior mean at ``x``. + + Args: + x (jax.Array or numpy.ndarray): + x coordinates. Must have shape ``(n, d)`` where ``d`` is the dimension + of the domain. + """ self.validate_input_shape(x) return self.mu(x) def prior_covariance(self, u, v): + """ + Evaluate the prior covariance as :math:`k(\\texttt{u}, \\texttt{v})`. + + Args: + u (jax.Array or numpy.ndarray): + u coordinates. Must have shape ``(nu, d)`` where ``d`` is the + dimension of the domain. + v (jax.Array or numpy.ndarray): + v coordinates. Must have shape ``(nv, d)`` where ``d`` is the + dimension of the domain. + """ self.validate_input_shape(u) self.validate_input_shape(v) return self.k(u, v) - def prior_realization(self, x, p=1): + def prior_realization(self, x: jax.Array | npt.NDArray, p: int = 1): + """ + Produce ``p`` realizations of the prior at locations ``x``. + + Args: + x (jax.Array or numpy.ndarray): + x coordinates. Must have shape ``(n, d)`` where ``d`` is the dimension + of the domain. + p (int, optional): + Number of realizations to generate, default ``1``. + """ self.key, key = jax.random.split(self.key) n, _ = self.validate_input_shape(x) @@ -386,6 +554,14 @@ def prior_realization(self, x, p=1): return self.prior_mean(x) + (L @ jax.random.normal(key, (n * self.cdim, p))) def posterior_mean(self, x): + """ + Evaluate the posterior mean at ``x``. + + Args: + x (jax.Array or numpy.ndarray): + x coordinates. Must have shape ``(n, d)`` where ``d`` is the dimension + of the domain. + """ self.validate_input_shape(x) d = (self.train_y - self.mu(self.train_x)).reshape(-1) @@ -393,6 +569,17 @@ def posterior_mean(self, x): return self.mu(x) + correction.reshape((-1, self.cdim)) def posterior_covariance(self, u, v): + """ + Evaluate the posterior covariance at ``(u, v)``. + + Args: + u (jax.Array or numpy.ndarray): + u coordinates. Must have shape ``(nu, d)`` where ``d`` is the + dimension of the domain. + v (jax.Array or numpy.ndarray): + v coordinates. Must have shape ``(nv, d)`` where ``d`` is the + dimension of the domain. + """ self.validate_input_shape(u) self.validate_input_shape(v) @@ -401,7 +588,17 @@ def posterior_covariance(self, u, v): self.k(u, self.train_x) @ self.train_cov_factor.solve(self.k(self.train_x, v)) ) - def posterior_realization(self, x, p=1): + def posterior_realization(self, x: jax.Array | npt.NDArray, p: int = 1): + """ + Produce ``p`` realizations of the posterior at locations ``x``. + + Args: + x (jax.Array or numpy.ndarray): + x coordinates. Must have shape ``(n, d)`` where ``d`` is the dimension + of the domain. + p (int, optional): + Number of realizations to generate, default ``1``. + """ self.key, key = jax.random.split(self.key) n, _ = self.validate_input_shape(x) diff --git a/UncertainSCI/gp/kernel.py b/UncertainSCI/gp/kernel.py index 514836c..d940fef 100644 --- a/UncertainSCI/gp/kernel.py +++ b/UncertainSCI/gp/kernel.py @@ -3,9 +3,11 @@ """ import abc +from typing import Any import equinox as eqx import jax import jax.numpy as jnp +from numpy._typing import NDArray import numpy.typing as npt from .. import _equinox @@ -22,15 +24,17 @@ def noise_as_matrix( Given ``s``, infer and return noise matrix. ``s`` can be: + * Scalar, interpreted as isotropic noise. * A vector of length ``matrix_size``, - interpreted as noise per-observation per-channel. + interpreted as noise per-observation per-channel. * A matrix of shape ``(matrix_size, matrix_size)``, returned as-is. - * Of shape ``(n, cdim)`` such that ``n * cdim == matrix_size``, - interpreted as noise per-observation per-channel. + * A matrix of shape ``(n, c)`` such that ``n * c == matrix_size`` + where ``c`` is the dimension of the codomain, + interpreted as noise per-observation per-channel. Args: - s (array or scalar of noise): + s (jax.Array, numpy.ndarray, float, or int): Noise of observations. """ s = jnp.asarray(s) @@ -55,9 +59,24 @@ def noise_as_matrix( raise ValueError -class _CF(eqx.Module, metaclass=abc.ABCMeta): +class CovarianceFactor(eqx.Module, metaclass=abc.ABCMeta): """ - Factorization of a covariance matrix used for solves and log determinants. + Covariance factor base class. Stores factorization of a covariance matrix + :math:`\\tilde{K}` used for solves and log determinants. + + In particular, this class stores factorization methods and data to solve + + .. math:: + + x = \\tilde{K}^{-1} y + + where, for instance, + + .. math:: + + \\tilde{K} = K_{XX} + \\sigma^2 I + + or :math:`\\tilde{K}` has some other more complicated structure. Subclasses may store a Cholesky factor, an eigendecomposition, or another structured representation. @@ -66,24 +85,44 @@ class _CF(eqx.Module, metaclass=abc.ABCMeta): @abc.abstractmethod def solve(self, y: jax.Array | npt.NDArray) -> jax.Array: + """ + Compute :math:`x = \\tilde{K} \\backslash y`. + """ raise NotImplementedError @abc.abstractmethod def log_sqrt_det(self) -> jax.Array: + """ + Compute :math:`\\log{ \\sqrt{ \\det \\tilde{K} } }`. + """ raise NotImplementedError @abc.abstractmethod def cholesky(self) -> jax.Array: + """ + Compute (or return) the Cholesky factor :math:`L` of :math:`\\tilde{K}` such + that + + .. math:: + + L L^\\top = \\tilde{K}. + """ raise NotImplementedError -class _CF_Cholesky(_CF): +class CholeskyFactor(CovarianceFactor): """ - Dense lower Cholesky factorization of a noisy covariance matrix. + Dense lower Cholesky factorization of a noisy covariance matrix + :math:`\\tilde{K} = K + S` for covariance matrix :math:`K` and noise :math:`S`. """ L: jax.Array def __init__(self, L: jax.Array | npt.NDArray): + """ + Args: + L (array): + The Cholesky factor of noisy covariance matrix. + """ L = jnp.asarray(L) if L.ndim != 2: raise ValueError @@ -103,9 +142,11 @@ def cholesky(self) -> jax.Array: return self.L -class _CF_Kronecker_ScalarNoise(_CF): +class KroneckerScalarNoiseFactor(CovarianceFactor): """ - Eigendecomposition of ``left`` ⊗ ``right`` plus scalar isotropic noise. + Eigendecomposition of noisy covariance matrix + :math:`\\tilde{K} = L \\otimes R + s I`, where :math:`L` is ``left``, + :math:`R` is ``right``, and :math:`s` is ``noise``. """ left_eigenvalues: jax.Array left_eigenvectors: jax.Array @@ -121,6 +162,15 @@ def __init__( right: jax.Array | npt.NDArray, noise: jax.Array | npt.NDArray | float | int, ): + """ + Args: + left (jax.Array or numpy.ndarray): + ``left`` matrix in Kronecker-structure covariance matrix. + right (jax.Array or numpy.ndarray): + ``right`` matrix in Kronecker-structure covariance matrix. + noise (jax.Array, numpy.ndarray, float, or int): + ``noise`` in noisy covariance matrix. + """ left = jnp.asarray(left) right = jnp.asarray(right) if left.ndim != 2 or left.shape[0] != left.shape[1]: @@ -173,7 +223,7 @@ def cholesky(self) -> jax.Array: class Kernel(eqx.Module, metaclass=abc.ABCMeta): """ Kernel base class. - + Attributes: dim (int): Dimension of the domain. @@ -193,20 +243,26 @@ def __call__( Evaluate this kernel. Args: - x (array): - Array of coordinates shaped as ``(num_x, dim)``. May be passed as 1-d - array in the case ``dim == 1``. - y (array): - Array of coordinates shaped as ``(num_y, dim)``. May be passed as 1-d - array in the case ``dim == 1``. + x (jax.Array or numpy.ndarray): + Array of coordinates shaped as ``(nx, d)``. May be passed as 1-d + array in the case dimension of domain ``d == 1``. + y (jax.Array or numpy.ndarray): + Array of coordinates shaped as ``(ny, d)``. May be passed as 1-d + array in the case dimension of domain ``d == 1``. + + Returns: + jax.Array: + Array of shape (c * nx, c * ny) of covariance where ``c`` is the + dimension of the codomain. """ raise NotImplementedError + @abc.abstractmethod def covariance_factor( self, x: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int - ) -> _CF: + ) -> CovarianceFactor: """ Factor the covariance of observed values at ``x``. @@ -214,37 +270,44 @@ def covariance_factor( Structured kernels can override this to avoid dense factorization. Args: - x (array): - Array of coordinates shaped as ``(num_x, dim)``. - s (array): - Array of noise associated with each observation. + x (jax.Array or numpy.ndarray): + Array of coordinates shaped as ``(n, d)`` where ``d`` is the dimension + of the domain. + s (jax.Array, numpy.ndarray, float, or int): + Noise associated with each coordinate. If array-type, must have shape: + + * ``()`` + * ``(n, c)`` + * ``(n * c)`` + * ``(n * c, n * c)`` + + where ``c`` is the dimension of the codomain. """ x = jnp.asarray(x) covariance = self(x, x) s = noise_as_matrix(s, covariance.shape[0], x.shape[0], self.cdim) - return _CF_Cholesky(jnp.linalg.cholesky(covariance + s)) + return CholeskyFactor(jnp.linalg.cholesky(covariance + s)) class Gaussian(Kernel): r""" - A Gaussian (square exponential) covariance kernel. + A Gaussian (square exponential) scalar-valued covariance kernel. This produces a kernel according to the function .. math:: - k(x, y) = \exp{\left( - y^\top D x \right)}. - - Attributes: - D (array): - Matrix defining metric on coordinates. + k(x, y) = \exp{\left[ - (x - y)^\top D \;(x - y) \right]}. - ``D`` is a matrix defining some metric on the coordinates :math:`x` and :math:`y` - according to + :math:`D` is probably best viewed as a matrix defining a metric on the coordinates + :math:`x` and :math:`y`. Specifically, .. math:: - \left[ \operatorname{dist}(x, y) \right]^2 = y^\top D x. + \left[ \operatorname{dist}(x, y) \right]^2 = (x - y)^\top D \;(x - y). + + If the coordinate space is isotropic, a reasonable choice is :math:`D = I_d` where + :math:`d` is the dimension of the domain. """ D: _equinox.PositiveDefinite @@ -255,6 +318,20 @@ def __init__( D: jax.Array | npt.NDArray, D_is_static: bool = False ): + """ + Args: + dim (int): + Dimension of the domain of this kernel. + cdim (int): + Dimension of the codomain of this kernel. Must be 1. + D (jax.Array or numpy.ndarray): + Matrix defining metric on coordinates. + D_is_static (bool, optional): + If ``D`` is a tunable hyperparameter, default ``False``. + """ + if cdim != 1: + raise ValueError + super().__init__(dim=dim, cdim=cdim) self.D = _equinox.PositiveDefinite(D=D, is_static=D_is_static) @@ -269,6 +346,12 @@ def __call__( ) ) + def covariance_factor( + self, + x: jax.Array | npt.NDArray, + s: jax.Array | npt.NDArray | float | int + ) -> CovarianceFactor: + return super().covariance_factor(x, s) class Kronecker(Kernel): r""" @@ -282,12 +365,6 @@ class Kronecker(Kernel): where :math:`k` is a scalar kernel (e.g., a square exponential kernel) and :math:`\otimes` is the Kronecker product. - - Attributes: - k (Kernel): - Covariance kernel on coordindates. - C (array): - Covariance on outputs. """ k: Kernel C: _equinox.PositiveDefinite @@ -300,6 +377,19 @@ def __init__( C: jax.Array | npt.NDArray, C_is_static: bool = False ): + """ + Args: + dim (int): + Dimension of the domain of this kernel. + cdim (int): + Dimension of the codomain of this kernel. + k (Kernel): + Scalar covariance kernel. + C (jax.Array or numpy.ndarray): + Output covariance matrix in Kronecker kernel. + C_is_static (bool, optional): + If ``C`` is a tunable hyperparameter, default ``False``. + """ if k.cdim != 1: raise ValueError('Kronecker coordinate kernel must be scalar-valued.') @@ -318,11 +408,11 @@ def covariance_factor( self, x: jax.Array | npt.NDArray, s: jax.Array | npt.NDArray | float | int - ) -> _CF: + ) -> CovarianceFactor: x = jnp.asarray(x) s = jnp.asarray(s) if s.ndim == 0: - return _CF_Kronecker_ScalarNoise(self.k(x, x), self.C(), s) + return KroneckerScalarNoiseFactor(self.k(x, x), self.C(), s) return super().covariance_factor(x, s) diff --git a/UncertainSCI/gp/mean.py b/UncertainSCI/gp/mean.py index 4d2c6f4..d78bbaa 100644 --- a/UncertainSCI/gp/mean.py +++ b/UncertainSCI/gp/mean.py @@ -63,9 +63,10 @@ def __init__( .. math:: - y = x @ A.T + b + y = x A.T + b - where ``x.shape = (number, dim)``. + where ``x.shape = (n, d)`` for some number of coordinates ``n`` and + dimension of the domain ``d``. Args: dim (int): diff --git a/tests/test_gp.py b/tests/test_gp.py index aca1054..92deda0 100644 --- a/tests/test_gp.py +++ b/tests/test_gp.py @@ -79,7 +79,7 @@ def test_kronecker_scalar_noise_factor_matches_dense_solve(self): dense_cholesky = jnp.linalg.cholesky(dense_covariance) dense_log_sqrt_det = jnp.sum(jnp.log(jnp.diag(dense_cholesky))) - self.assertIsInstance(factor, gp.kernel._CF_Kronecker_ScalarNoise) + self.assertIsInstance(factor, gp.kernel.KroneckerScalarNoiseFactor) self.assertTrue( jnp.allclose( factor.solve(rhs), @@ -138,7 +138,7 @@ def test_kronecker_posterior_handles_vector_outputs(self): jnp.linalg.solve(dense_covariance, k(x_train, x_test)) ) - self.assertIsInstance(g.train_cov_factor, gp.kernel._CF_Kronecker_ScalarNoise) + self.assertIsInstance(g.train_cov_factor, gp.kernel.KroneckerScalarNoiseFactor) self.assertEqual(g.posterior_mean(x_test).shape, (2, 2)) self.assertEqual(g.posterior_covariance(x_test, x_test).shape, (4, 4)) self.assertTrue( From 9b297fe05a04bbcf3a3e5594af67972f669a6f47 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Mon, 31 Aug 2026 10:32:12 -0600 Subject: [PATCH 12/22] gp: mark out jax inexact problematic dtype statements * will fix later --- UncertainSCI/gp/__init__.py | 3 +++ UncertainSCI/gp/kernel.py | 3 +++ UncertainSCI/gp/mean.py | 1 + 3 files changed, 7 insertions(+) diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index a3a0b3d..f411271 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -129,6 +129,7 @@ def condition( self.train_x = jnp.asarray(x) self.train_y = jnp.asarray(y) + # FIXME: Ensure integer inputs have an inexact dtype when traced by JAX. self.train_s = jnp.asarray(s) self.train_cov_factor = GaussianProcess._get_train_cov_factor( @@ -257,6 +258,7 @@ def get_train_cov_factor( x = jnp.asarray(x) y = jnp.asarray(y) + # FIXME: Ensure integer inputs have an inexact dtype when traced by JAX. s = jnp.asarray(s) return GaussianProcess._get_train_cov_factor((self.mu, self.k), x, y, s) @@ -359,6 +361,7 @@ def loss_hyperparameters( x = jnp.asarray(x) y = jnp.asarray(y) + # FIXME: Ensure integer inputs have an inexact dtype when traced by JAX. s = jnp.asarray(s) return GaussianProcess._loss_hyperparameters((self.mu, self.k), x, y, s) diff --git a/UncertainSCI/gp/kernel.py b/UncertainSCI/gp/kernel.py index d940fef..177351c 100644 --- a/UncertainSCI/gp/kernel.py +++ b/UncertainSCI/gp/kernel.py @@ -37,6 +37,7 @@ def noise_as_matrix( s (jax.Array, numpy.ndarray, float, or int): Noise of observations. """ + # FIXME: Ensure integer inputs have an inexact dtype when traced by JAX. s = jnp.asarray(s) if s.ndim == 0: @@ -183,6 +184,7 @@ def __init__( self.right_size = right.shape[0] self.left_eigenvalues, self.left_eigenvectors = jnp.linalg.eigh(left) self.right_eigenvalues, self.right_eigenvectors = jnp.linalg.eigh(right) + # FIXME: Ensure integer inputs have an inexact dtype when traced by JAX. self.noise = jnp.asarray(noise) @property @@ -410,6 +412,7 @@ def covariance_factor( s: jax.Array | npt.NDArray | float | int ) -> CovarianceFactor: x = jnp.asarray(x) + # FIXME: Ensure integer inputs have an inexact dtype when traced by JAX. s = jnp.asarray(s) if s.ndim == 0: diff --git a/UncertainSCI/gp/mean.py b/UncertainSCI/gp/mean.py index d78bbaa..e65d7cf 100644 --- a/UncertainSCI/gp/mean.py +++ b/UncertainSCI/gp/mean.py @@ -86,6 +86,7 @@ def __init__( a = jnp.asarray(a) if isinstance(b, (float, int)): + # FIXME: Ensure integer inputs have an inexact dtype when traced by JAX. b = b * jnp.ones((cdim,)) else: b = jnp.asarray(b) From 0e5396dacc560e106c5f1bdf54b04ca3dcdf3a5e Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 3 Sep 2026 08:26:39 -0600 Subject: [PATCH 13/22] gp: break out 1-d visualization routines from demo --- UncertainSCI/gp/__init__.py | 1 + UncertainSCI/gp/vis.py | 268 ++++++++++++++++++++++++++++++++++++ demos/gp-1d.ipynb | 171 ++--------------------- 3 files changed, 283 insertions(+), 157 deletions(-) create mode 100644 UncertainSCI/gp/vis.py diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index f411271..2602698 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -12,6 +12,7 @@ from . import kernel from . import mean +from . import vis DEFAULT_OPTIM = optax.adam diff --git a/UncertainSCI/gp/vis.py b/UncertainSCI/gp/vis.py new file mode 100644 index 0000000..bc3ef3d --- /dev/null +++ b/UncertainSCI/gp/vis.py @@ -0,0 +1,268 @@ +""" +Visualize Gaussian processes. +""" + +from __future__ import annotations + +import copy +import jax +import jax.numpy as jnp +import matplotlib.axes as axes +import numpy.typing as npt +import typing + +from .. import _equinox as _eqx + +if typing.TYPE_CHECKING: + from . import GaussianProcess + + +def stats( + g: GaussianProcess, + x: jax.Array | npt.NDArray, + which: str = 'posterior', + p: int = 100 + ): + """ + Evaluate summary statistics and realizations of a Gaussian process. + + Args: + g (gp.GaussianProcess): + The Gaussian process to evaluate. + x (jax.Array or numpy.ndarray): + Coordinates at which to evaluate the process. Must have shape ``(n, d)``, + where ``d`` is the dimension of the domain. + which ({'posterior', 'prior'}, optional): + Distribution to evaluate, default ``'posterior'``. + p (int, optional): + Number of realizations to generate, default ``100``. + + Returns: + tuple of jax.Array: + The squeezed mean, variance, and realizations, respectively. + + Raises: + ValueError: + If ``which`` is neither ``'posterior'`` nor ``'prior'``. + """ + if which == 'posterior': + m = jnp.squeeze(g.posterior_mean(x)) + v = jnp.squeeze(jnp.diag(g.posterior_covariance(x, x))) + r = jnp.squeeze(g.posterior_realization(x, p)) + elif which == 'prior': + m = jnp.squeeze(g.prior_mean(x)) + v = jnp.squeeze(jnp.diag(g.prior_covariance(x, x))) + r = jnp.squeeze(g.prior_realization(x, p)) + else: + raise ValueError + + return m, v, r + + +def plot_distribution_mean( + ax: axes.Axes, + g: GaussianProcess, + x: jax.Array | npt.NDArray, + f_true: typing.Callable | None, + which: str = 'posterior', + p: int = 100 + ): + """ + Plot a Gaussian process distribution and (optionally) a reference function. + + Plots the process mean with a shaded band of one marginal variance on either + side, sampled realizations, and (optionally) the values of ``f_true``. For + the posterior distribution, the training observations and their uncertainties + are also shown. + + Args: + ax (matplotlib.axes.Axes): + The axes to plot into. + g (gp.GaussianProcess): + The Gaussian process to evaluate. + x (jax.Array or numpy.ndarray): + Coordinates at which to evaluate the process. Must have shape ``(n, d)``, + where ``d`` is the dimension of the domain. + f_true (callable or None): + Reference function evaluated at ``x`` and plotted for comparison. + which ({'posterior', 'prior'}, optional): + Distribution to plot, default ``'posterior'``. + p (int, optional): + Number of process realizations to plot, default ``100``. + + Raises: + ValueError: + If ``which`` is neither ``'posterior'`` nor ``'prior'``. + """ + m, c, r = stats(g, x, which, p) + ax.fill_between( + x.squeeze(), + (m - c), + (m + c), + color='#00000040', + edgecolor='none' + ) + ax.plot( + x.squeeze(), + r, + color='#00000010' + ) + if f_true is not None: + ax.plot( + x.squeeze(), + f_true(x).squeeze(), + ) + if which == 'posterior': + ax.errorbar( + g.train_x.squeeze(), + g.train_y.squeeze(), + yerr=g.train_s.squeeze(), + capsize=2, + linestyle='none' + ) + ax.set_title('Distribution') + + +def plot_distribution_variance( + ax: axes.Axes, + g: GaussianProcess, + x: jax.Array | npt.NDArray, + which: str = 'posterior', + p: int = 100, + colorlast: bool = True + ): + """ + Plot the marginal variance of a Gaussian process distribution. + + For the posterior distribution, vertical lines mark the training coordinates. + When ``colorlast`` is true, the final coordinate is highlighted separately. + + Args: + ax (matplotlib.axes.Axes): + The axes to plot into. + g (gp.GaussianProcess): + The Gaussian process to evaluate. + x (jax.Array or numpy.ndarray): + Coordinates at which to evaluate the process. Must have shape ``(n, d)``, + where ``d`` is the dimension of the domain. + which ({'posterior', 'prior'}, optional): + Distribution whose variance is plotted, default ``'posterior'``. + p (int, optional): + Number of process realizations generated while evaluating the distribution, + default ``100``. + colorlast (bool, optional): + Whether to highlight the final posterior training coordinate in red while + drawing the preceding coordinates in green, default ``True``. If false, + all training coordinates are drawn in green. + + Raises: + ValueError: + If ``which`` is neither ``'posterior'`` nor ``'prior'``. + """ + m, c, r = stats(g, x, which, p) + ax.plot( + x.squeeze(), + c + ) + if which == 'posterior': + if colorlast: + ax.vlines(g.train_x.squeeze()[:-1], 0, 1, color='tab:green') + ax.vlines(g.train_x.squeeze()[-1], 0, 1, color='tab:red') + else: + ax.vlines(g.train_x.squeeze(), 0, 1, color='tab:green') + ax.set_title('Variance') + + +def plot_loss_landscape( + ax: axes.Axes, + g: GaussianProcess, + p_name: tuple[str, ...], + p_range: jax.Array | npt.NDArray, + x: jax.Array | npt.NDArray | None = None, + ymin: jax.Array | npt.NDArray | float | int | None = None, + ymax: jax.Array | npt.NDArray | float | int | None = None, + ytop: float = 1e2 + ): + """ + Plot a Gaussian process loss landscape for a scalar parameter. + + The parameter is resolved by following the attribute path in ``p_name`` on a + deep copy of ``g``. Its value is replaced by each entry in ``p_range``, and the + hyperparameter loss on the process training data is plotted. Both plot axes use + a symmetric logarithmic scale. + + Args: + ax (matplotlib.axes.Axes): + The axes to plot into. + g (gp.GaussianProcess): + The Gaussian process to evaluate. + p_name (tuple of str): + Attribute name chain leading to the :class:`jax.Array` or + :class:`UncertainSCI._equinox.ComputedArray` parameter to vary. + p_range (jax.Array or numpy.ndarray): + One-dimensional sequence of parameter values at which to evaluate the + loss. + x (array or float, optional): + Coordinates at which to draw vertical reference lines, default ``None``. + ymin (array or float, optional): + Lower endpoints of the reference lines. Must be compatible with ``x``. + By default, the minimum evaluated loss is used. + ymax (array or float, optional): + Upper endpoints of the reference lines. Must be compatible with ``x``. + By default, the maximum evaluated loss is used. + ytop (float, optional): + Upper limit of the y-axis, default ``1e2``. + + Raises: + AttributeError: + If an attribute in ``p_name`` does not exist. + ValueError: + If ``p_name`` is empty or resolves to an unsupported parameter type. + NotImplementedError: + If the resolved parameter is not scalar. + """ + p_parent = None + g = copy.deepcopy(g) + p = g + for name in p_name: + p_parent = p + p = getattr(p, name) + + if p_parent is None: + raise ValueError( + f'Did not find parent of target from keys {p_name}! ' + f'This usually happens when keys {p_name} has zero length.' + ) + + if isinstance(p, _eqx.ComputedArray): + numel = p().size + shape = p().shape + elif isinstance(p, jax.Array): + numel = p.size + shape = p.shape + else: + raise ValueError( + f'Resolution of target from keys {p_name} led to unsupported type {type(p)}!' + ) + + if numel > 1: + raise NotImplementedError('Loss landscape for non-scalar parameter is not supported!') + + losses = jnp.empty(p_range.shape) + for i, p_value in enumerate(p_range): + if isinstance(p, _eqx.ComputedArray): + p_value = type(p).from_array(p_value * jnp.ones(shape)) + else: + p_value = p_value * jnp.ones(shape) + + setattr(p_parent, p_name[-1], p_value) + losses = losses.at[i].set(g.loss_hyperparameters(g.train_x, g.train_y, g.train_s)) + + ax.plot(p_range, losses) + if x is not None: + ax.vlines(x, min(losses) if ymin is None else ymin, max(losses) if ymax is None else ymax) + ax.set_xscale('symlog') + ax.set_yscale('symlog') + ax.set_ylim(top=ytop) + ax.set_title('Loss Landscape') + ax.set_xlabel(p_name[-1]) diff --git a/demos/gp-1d.ipynb b/demos/gp-1d.ipynb index b1f72a4..48ba1df 100644 --- a/demos/gp-1d.ipynb +++ b/demos/gp-1d.ipynb @@ -89,149 +89,6 @@ "metadata": {}, "outputs": [], "source": [ - "def gp_stats(g: gp.GaussianProcess, x, which='posterior', p=100):\n", - " if which == 'posterior':\n", - " m = jnp.squeeze(g.posterior_mean(x))\n", - " c = jnp.squeeze(jnp.diag(g.posterior_covariance(x, x)))\n", - " r = jnp.squeeze(g.posterior_realization(x, p))\n", - " elif which == 'prior':\n", - " m = jnp.squeeze(g.prior_mean(x))\n", - " c = jnp.squeeze(jnp.diag(g.prior_covariance(x, x)))\n", - " r = jnp.squeeze(g.prior_realization(x, p))\n", - " else:\n", - " raise ValueError\n", - "\n", - " return m, c, r\n", - "\n", - "\n", - "def plot_distribution_mean(ax, g: gp.GaussianProcess, x, which='posterior', p=100):\n", - " m, c, r = gp_stats(g, x, which, p)\n", - " ax.fill_between(\n", - " x.squeeze(),\n", - " (m - c),\n", - " (m + c),\n", - " color='#00000040',\n", - " edgecolor='none'\n", - " )\n", - " ax.plot(\n", - " x.squeeze(),\n", - " r,\n", - " color='#00000010'\n", - " )\n", - " ax.plot(\n", - " x.squeeze(),\n", - " f_hidden(x).squeeze(),\n", - " )\n", - " if which == 'posterior':\n", - " ax.errorbar(\n", - " g.train_x.squeeze(),\n", - " g.train_y.squeeze(),\n", - " yerr=g.train_s.squeeze(),\n", - " capsize=2,\n", - " linestyle='none'\n", - " )\n", - " ax.set_title('Distribution')\n", - "\n", - "\n", - "def plot_distribution_variance(ax, g: gp.GaussianProcess, x, which='posterior', p=100, colorlast=True):\n", - " m, c, r = gp_stats(g, x, which, p)\n", - " ax.plot(\n", - " x.squeeze(),\n", - " c\n", - " )\n", - " if which == 'posterior':\n", - " if colorlast:\n", - " ax.vlines(g.train_x.squeeze()[:-1], 0, 1, color='tab:green')\n", - " ax.vlines(g.train_x.squeeze()[-1], 0, 1, color='tab:red')\n", - " else:\n", - " ax.vlines(g.train_x.squeeze(), 0, 1, color='tab:green')\n", - " ax.set_title('Variance')\n", - "\n", - "\n", - "def plot_loss_landscape(\n", - " ax,\n", - " g: gp.GaussianProcess,\n", - " p_name: tuple[str, ...],\n", - " p_range: jax.Array,\n", - " x: jax.Array | None = None,\n", - " ymin: jax.Array | None = None,\n", - " ymax: jax.Array | None = None,\n", - " ytop: float = 1e2\n", - " ):\n", - " \"\"\"\n", - " Plot the loss landscape with respect to the parameter in Gaussian process ``g``\n", - " at location ``p_name``.\n", - "\n", - " Performs recursive ``getattr`` until ``p_name`` is exhausted.\n", - "\n", - " Args:\n", - " ax (matplotlib.axes.Axes):\n", - " The axes to plot into.\n", - " g (gp.GaussianProcess):\n", - " The Gaussian process to evaluate.\n", - " p_name (tuple of str):\n", - " Attr name chain leading to the parameter (:class:`jax.Array` or :class:`_eqx.ComputedArray`)\n", - " to be sampled.\n", - " p_range (array):\n", - " Range of parameter values to test.\n", - " x (array or float, optional):\n", - " x values at which to place vlines, default ``None``.\n", - " ymin (array or float, optional):\n", - " y range for vlines, must match shape of x, default ``None``.\n", - " ymax (array or float, optional):\n", - " y range for vlines, must match shape of x, default ``None``.\n", - " ytop (float, optional):\n", - " Top of y-axis to plot, default 1e2.\n", - " \"\"\"\n", - " p_parent = None\n", - " g = copy.deepcopy(g)\n", - " p = g\n", - " for name in p_name:\n", - " p_parent = p\n", - " p = getattr(p, name)\n", - "\n", - " if p_parent is None:\n", - " raise ValueError(\n", - " f'Did not find parent of target from keys {p_name}! '\n", - " f'This usually happens when keys {p_name} has zero length.'\n", - " )\n", - "\n", - " if isinstance(p, _eqx.ComputedArray):\n", - " numel = p().size\n", - " shape = p().shape\n", - " elif isinstance(p, jax.Array):\n", - " numel = p.size\n", - " shape = p.shape\n", - " else:\n", - " raise ValueError(\n", - " f'Resolution of target from keys {p_name} led to unsupported type {type(p)}!'\n", - " )\n", - "\n", - " if numel > 1:\n", - " raise NotImplementedError('Loss landscape for non-scalar parameter is not supported!')\n", - "\n", - " losses = jnp.empty(p_range.shape)\n", - " for i, p_value in enumerate(p_range):\n", - " if isinstance(p, _eqx.ComputedArray):\n", - " p_value = type(p).from_array(p_value * jnp.ones(shape))\n", - " else:\n", - " p_value = p_value * jnp.ones(shape)\n", - "\n", - " setattr(p_parent, p_name[-1], p_value)\n", - " losses = losses.at[i].set(g.loss_hyperparameters(g.train_x, g.train_y, g.train_s))\n", - "\n", - " ax.plot(p_range, losses)\n", - " if x is not None:\n", - " if ymin is None:\n", - " ymin = min(losses)\n", - " if ymax is None:\n", - " ymax = max(losses)\n", - " ax.vlines(x, ymin, ymax)\n", - " ax.set_xscale('symlog')\n", - " ax.set_yscale('symlog')\n", - " ax.set_ylim(top=ytop)\n", - " ax.set_title('Loss Landscape')\n", - " ax.set_xlabel(p_name[-1])\n", "\n", "\n", "def plot_in_subplots(nrows=1, ncols=1, **kwargs):\n", @@ -256,8 +113,8 @@ "outputs": [], "source": [ "fig, (ax1, ax2) = plot_in_subplots(2, 1, dpi=FIGDPI)\n", - "plot_distribution_mean(ax1, g, x_plot, 'prior')\n", - "plot_distribution_variance(ax2, g, x_plot, 'prior')\n", + "gp.vis.plot_distribution_mean(ax1, g, x_plot, f_hidden, 'prior')\n", + "gp.vis.plot_distribution_variance(ax2, g, x_plot, 'prior')\n", "plt.show()\n" ] }, @@ -292,9 +149,9 @@ "outputs": [], "source": [ "fig, (ax1, ax2, ax3) = plot_in_subplots(3, 1, dpi=FIGDPI)\n", - "plot_distribution_mean(ax1, g, x_plot, 'posterior')\n", - "plot_distribution_variance(ax2, g, x_plot, 'posterior', colorlast=False)\n", - "plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", + "gp.vis.plot_distribution_mean(ax1, g, x_plot, f_hidden, 'posterior')\n", + "gp.vis.plot_distribution_variance(ax2, g, x_plot, 'posterior', colorlast=False)\n", + "gp.vis.plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", "plt.show()\n", "\n", "print_loss_here(g)\n" @@ -340,9 +197,9 @@ "outputs": [], "source": [ "fig, (ax1, ax2, ax3) = plot_in_subplots(3, 1, dpi=FIGDPI)\n", - "plot_distribution_mean(ax1, g, x_plot, 'posterior')\n", - "plot_distribution_variance(ax2, g, x_plot, 'posterior', colorlast=False)\n", - "plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", + "gp.vis.plot_distribution_mean(ax1, g, x_plot, f_hidden, 'posterior')\n", + "gp.vis.plot_distribution_variance(ax2, g, x_plot, 'posterior', colorlast=False)\n", + "gp.vis.plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", "plt.show()\n", "\n", "print_loss_here(g)\n" @@ -370,8 +227,8 @@ "\n", "for i in range(N_TRAIN):\n", " # Greedily sample posterior variance:\n", - " _, c, _ = gp_stats(g, x_sample, p=0)\n", - " _x = x_sample[jnp.argmax(c)].reshape(-1, D)\n", + " _v = jnp.diag(g.posterior_covariance(x_sample, x_sample))\n", + " _x = x_sample[jnp.argmax(_v)].reshape(-1, D)\n", "\n", " # Iteratively optimize posterior variance:\n", " _x = g.get_sample_point(_x.reshape((1,)), ranges=[[0.], [2 * jnp.pi]])\n", @@ -389,9 +246,9 @@ "\n", " # Plot:\n", " fig, (ax1, ax2, ax3) = plot_in_subplots(3, 1, dpi=FIGDPI)\n", - " plot_distribution_mean(ax1, g, x_plot, 'posterior')\n", - " plot_distribution_variance(ax2, g, x_plot, 'posterior')\n", - " plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", + " gp.vis.plot_distribution_mean(ax1, g, x_plot, f_hidden, 'posterior')\n", + " gp.vis.plot_distribution_variance(ax2, g, x_plot, 'posterior')\n", + " gp.vis.plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", " plt.show()\n", " print_loss_here(g)\n", " print('\\n' * 3)\n" @@ -422,7 +279,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.15" + "version": "3.11.16" } }, "nbformat": 4, From f150cb6bea90e2c1cbb04739bcf99ca7fc4c023f Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 3 Sep 2026 08:31:42 -0600 Subject: [PATCH 14/22] gp: add docs and support scripts * will want to convert formalism/gp.md to .rst eventually * (cross)linking in .md is more terse than .rst --- docs/clean.sh | 8 +++ docs/conf.py | 2 + docs/dev_docs/api/index.rst | 10 ++++ docs/formalism/gp.md | 103 ++++++++++++++++++++++++++++++++++++ docs/formalism/index.rst | 1 + 5 files changed, 124 insertions(+) create mode 100755 docs/clean.sh create mode 100644 docs/formalism/gp.md diff --git a/docs/clean.sh b/docs/clean.sh new file mode 100755 index 0000000..78b77a5 --- /dev/null +++ b/docs/clean.sh @@ -0,0 +1,8 @@ +#!/bin/bash +set -e + +# Clean up build products not located in _build/: +rm -rf dev_docs/api/generated/ + +# Clean up build products with `make clean`: +make clean diff --git a/docs/conf.py b/docs/conf.py index 3cbea56..cd2cd37 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -56,6 +56,7 @@ autosummary_generate = True autosummary_imported_members = True +autoclass_content = 'both' # Path for bibtex files bibtex_bibfiles = ['references.bib'] @@ -69,6 +70,7 @@ 'numpy': ('https://numpy.org/doc/stable/', None), 'scipy': ('https://docs.scipy.org/doc/scipy/', None), 'matplotlib': ('https://matplotlib.org/stable/', None), + 'jax': ('https://docs.jax.dev/en/latest/', None), } myst_enable_extensions = [ diff --git a/docs/dev_docs/api/index.rst b/docs/dev_docs/api/index.rst index 2269188..2439591 100644 --- a/docs/dev_docs/api/index.rst +++ b/docs/dev_docs/api/index.rst @@ -26,3 +26,13 @@ API Reference transformations ttr vis + + +.. rubric:: GP Modules + +.. autosummary:: + :toctree: generated + :recursive: + :signatures: none + + gp diff --git a/docs/formalism/gp.md b/docs/formalism/gp.md new file mode 100644 index 0000000..6b6503c --- /dev/null +++ b/docs/formalism/gp.md @@ -0,0 +1,103 @@ +# Gaussian Processes + +For the purposes here, we define a Gaussian process as follows: + +Let $\Omega$ some domain, $k: \Omega \times \Omega \rightarrow \mathbb{R}$ a +positive-definite kernel, and $\mu: \Omega \rightarrow \mathbb{R}$ any function. +Then a random function $f: \Omega \rightarrow \mathbb{R}$ is said to be a +*Gaussian process* (Kanagawa et al. 2018) if for any set +$X = (x_1, \dots, x_n) \subset \Omega$ for any $n \in \mathbb{N}$ + +$$ +f_X \sim \mathcal{N}^n(\mu_X, k_{XX}) +$$ + +where $\mathcal{N}^n$ is the $n$-multivariate normal distribution, and + +$$ +\begin{aligned} + f_X &= (f(x_1), \dots, f(x_n))^\top \in \mathbb{R}^n, \\ + \mu_X &= (\mu(x_1), \dots, \mu(x_n))^\top \in \mathbb{R}^n, \text{ and} \\ + (k_{XX})_{ij} &= k(x_i, x_j). +\end{aligned} +$$ + +Notationally, we write + +$$ +f \sim \mathcal{GP}(\mu, k). +$$ + +A *positive-definite kernel* $k$ is any symmetric function such that for any +$X = (x_1, \dots, x_n) \subset \Omega$ then for all $u \in \mathbb{R}^n$ + +$$ +u^\top k_{XX} \,u \geq 0. +$$ + +Note that this is closely analogous to the standard definition of a +positive-(semi-)definite operator. + +## Mathematical Notes on Implementation + +### Efficient Solves of GPs with Kronecker-structure Kernels + +Let a Gaussian process $\mathcal{GP}(\mu, k)$ such that $f \sim \mathcal{GP}$ where +$f: X \rightarrow Y$ where $X = \mathbb{R}^d$ and $Y = \mathbb{R}^c$. In such cases, +$k$ is a matrix-valued kernel, i.e., $k(x_i, x_j) = M \in \mathbb{R}^{c \times c}$ +positive-definite for $x_i, x_j \in X$. + +In some cases, it is appropriate to model $k$ with Kronecker structure, + +$$ +k(x_i, x_j) = k^0(x_i, x_j) \otimes C +$$ + +for a scalar kernel $k^0$, where $x_i, x_j \in X$, and output covariance +$C \in \mathbb{R}^{c \times c}$ positive-definite. + +For notational convenience, let + +$$ +K = K^0 \otimes C +$$ + +where $K_{ij} = k(x_i, x_j)$ and $K^0_{ij} = k^0(x_i, x_j)$ for $x_i, x_j \in X$. + +Let the noisy observation matrix $M = K + \operatorname{diag}(s)$. When computing the +posterior of a Gaussian process with this structure, it is necessary to compute + +$$ +\begin{aligned} + M^{-1} &= (K + \operatorname{diag}(s))^{-1} \\ + &= (K^0 \otimes C + \operatorname{diag}(s))^{-1}. +\end{aligned} +$$ + +Depending on the structure of $s$, this computation can be factored to accelerate +solves involving this inverse. + +1. $s = \sigma^2 \mathbb{1}$ (totally isotropic noise) + + Here + + $$ + M = K^0 \otimes C + \sigma^2 I. + $$ + + Note that $K^0$ and $C$ are positive definite, thus they admit diagonlizations + + $$ + K^0 = Q \Lambda Q^\top \qquad \text{and} \qquad C = U \Gamma U^\top. + $$ + + Thus + + $$ + \begin{aligned} + M^{-1} &= (Q \Lambda Q^\top \otimes U \Gamma U^\top + \sigma^2 I)^{-1} \\ + &= \left( (Q \otimes U)(\Lambda \otimes \Gamma + \sigma^2 I)(Q \otimes U)^\top \right)^{-1} + \end{aligned} + $$ + + where the second line follows from $\sigma^2 I$ diagonal. diff --git a/docs/formalism/index.rst b/docs/formalism/index.rst index 70a45d7..bfb9109 100644 --- a/docs/formalism/index.rst +++ b/docs/formalism/index.rst @@ -5,3 +5,4 @@ Mathematical Formalism :maxdepth: 2 pce + gp From dc61b8dc5b95a3b496e56d8e5f85e4c1ff1ab96f Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 3 Sep 2026 09:21:15 -0600 Subject: [PATCH 15/22] gp: commentary in 1-d demo notebook --- demos/gp-1d.ipynb | 148 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 109 insertions(+), 39 deletions(-) diff --git a/demos/gp-1d.ipynb b/demos/gp-1d.ipynb index 48ba1df..e6ddfbc 100644 --- a/demos/gp-1d.ipynb +++ b/demos/gp-1d.ipynb @@ -1,5 +1,15 @@ { "cells": [ + { + "cell_type": "markdown", + "id": "introduction", + "metadata": {}, + "source": [ + "# 1-D Gaussian Process Adaptive Sampling\n", + "\n", + "This example builds and tunes a scalar-valued Gaussian process over a one-dimensional domain. We begin with a small set of noisy observations, tune the covariance length scale, and then add observations where the posterior variance is largest." + ] + }, { "cell_type": "code", "execution_count": null, @@ -12,21 +22,21 @@ "# Un/comment this for double/single precision:\n", "# jax.config.update('jax_enable_x64', True)\n", "\n", - "import copy\n", - "import equinox as eqx\n", "import jax.numpy as jnp\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "\n", - "import UncertainSCI._equinox as _eqx\n", "import UncertainSCI.gp as gp\n", "\n", "\n", "D = 1 \n", "C = 1\n", + "SEED = 0x01234567\n", + "GP_SEED = 0xDEADBEEF\n", + "rng = np.random.default_rng(SEED)\n", "\n", - "np.random.seed(0x01234567)\n", - "seed = 0xdeadbeef # Used in gp.GaussianProcess initialization.\n", + "DOMAIN = (0, 2 * jnp.pi)\n", + "NOISE_VARIANCE = 1e-2\n", "\n", "FIGSIZE = (7, 7 / 1.6)\n", "FIGDPI = None\n", @@ -35,9 +45,32 @@ " print(\n", " f'Loss at D = {jnp.squeeze(g.k.D()):.4e} (length scale {1 / jnp.sqrt(jnp.squeeze(g.k.D())):.4e}): '\n", " f'{g.loss_hyperparameters(train_x, train_y, train_s):.4e}'\n", + " )\n", + "\n", + "def plot_in_subplots(nrows=1, ncols=1, **kwargs):\n", + " \"\"\"\n", + " Returns ``(fig, axes)`` with plot scaled correctly for ``(nrows, ncols)``.\n", + " \"\"\"\n", + " if 'figsize' in kwargs:\n", + " raise ValueError(\"kwargs had 'figsize' key: that's the whole point of this function!\")\n", + " return plt.subplots(\n", + " nrows,\n", + " ncols,\n", + " figsize=(ncols * FIGSIZE[0], nrows * FIGSIZE[1]),\n", + " **kwargs\n", " )\n" ] }, + { + "cell_type": "markdown", + "id": "true-function-heading", + "metadata": {}, + "source": [ + "## True function and noisy observations\n", + "\n", + "The hidden function combines two oscillatory factors and is evaluated on the interval from zero to $2\\pi$. Observations have independent Gaussian noise with a known, constant variance. `GaussianProcess.condition` expects this variance—not its square root." + ] + }, { "cell_type": "code", "execution_count": null, @@ -45,17 +78,48 @@ "metadata": {}, "outputs": [], "source": [ - "x_plot = jnp.linspace(0, 2 * jnp.pi, 1000).reshape((-1, D))\n", + "x_plot = jnp.linspace(*DOMAIN, 1000).reshape((-1, D))\n", "\n", "def f_hidden(x):\n", " return 2 * jnp.cos(x) * jnp.sin(4 * x)\n", "\n", "def f_noisy(x):\n", " y = f_hidden(x)\n", - " return (\n", - " y + 1e-3 * np.random.randn(*y.shape),\n", - " 1e-3 * jnp.ones_like(y)\n", - " )\n" + " noise = jnp.sqrt(NOISE_VARIANCE) * jnp.asarray(\n", + " rng.normal(size=y.shape)\n", + " )\n", + " return y + noise, NOISE_VARIANCE * jnp.ones_like(y)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6bc298f7", + "metadata": {}, + "outputs": [], + "source": [ + "truth = f_hidden(x_plot)\n", + "noisy, noise = f_noisy(x_plot)\n", + "\n", + "fig, (ax1, ax2) = plot_in_subplots(2, 1, sharex=True, sharey=True)\n", + "\n", + "ax1.plot(truth)\n", + "ax1.set_title('Truth')\n", + "\n", + "ax2.plot(noisy)\n", + "ax2.set_title('Noisy View')\n", + "\n", + "plt.show()\n" + ] + }, + { + "cell_type": "markdown", + "id": "define-gp-heading", + "metadata": {}, + "source": [ + "## Define the Gaussian process\n", + "\n", + "The prior has a fixed zero affine mean and a Gaussian covariance kernel. The kernel parameter $D$ determines how quickly correlation decays with distance: its associated length scale is $1 / \\sqrt{D}$. We leave $D$ tunable so that it can be learned from the observations." ] }, { @@ -78,31 +142,17 @@ " cdim=C,\n", " D=jnp.ones((1, 1)),\n", ")\n", - "g = gp.GaussianProcess(dim=D, cdim=C, mu=mu, k=k, seed=seed, nugget=1e-3)\n", - "# g = gp.GaussianProcess(dim=D, cdim=C, mu=mu, k=k, seed=seed, nugget=1e-4)\n" + "g = gp.GaussianProcess(dim=D, cdim=C, mu=mu, k=k, seed=GP_SEED, nugget=1e-3)\n" ] }, { - "cell_type": "code", - "execution_count": null, - "id": "2ae4ac3b", + "cell_type": "markdown", + "id": "prior-heading", "metadata": {}, - "outputs": [], "source": [ + "## Inspect the prior\n", "\n", - "\n", - "def plot_in_subplots(nrows=1, ncols=1, **kwargs):\n", - " \"\"\"\n", - " Returns ``(fig, axes)`` with plot scaled correctly for ``(nrows, ncols)``.\n", - " \"\"\"\n", - " if 'figsize' in kwargs:\n", - " raise ValueError(\"kwargs had 'figsize' key: that's the whole point of this function!\")\n", - " return plt.subplots(\n", - " nrows,\n", - " ncols,\n", - " figsize=(ncols * FIGSIZE[0], nrows * FIGSIZE[1]),\n", - " **kwargs\n", - " )\n" + "Before conditioning, the process reflects only the chosen mean and covariance. The upper panel shows prior realizations and their mean, with the hidden function included for reference; the lower panel shows the marginal prior variance. At this point the GP has not observed the hidden function." ] }, { @@ -118,6 +168,16 @@ "plt.show()\n" ] }, + { + "cell_type": "markdown", + "id": "condition-heading", + "metadata": {}, + "source": [ + "## Condition on initial observations\n", + "\n", + "We begin with ten evenly spaced observations. Each call to `f_noisy` returns both an observed value and its noise variance, which are supplied to `condition` along with the observation coordinates." + ] + }, { "cell_type": "code", "execution_count": null, @@ -138,7 +198,9 @@ "id": "2cf76bf6", "metadata": {}, "source": [ - "Plot what that looks like:" + "## Initial posterior\n", + "\n", + "Conditioning updates the distribution while retaining the initial kernel parameter. The first two panels show the posterior distribution and marginal variance. The final panel evaluates the negative log marginal likelihood over possible values of $D$, showing where hyperparameter tuning should move the current kernel." ] }, { @@ -162,7 +224,9 @@ "id": "9f7a058a", "metadata": {}, "source": [ - "Tune:" + "## Tune the kernel hyperparameter\n", + "\n", + "`tune` minimizes the negative log marginal likelihood with respect to the non-static GP parameters. Here that means learning $D$, and therefore the covariance length scale, from the initial observations. The optimization trace below shows how the loss changes with each step." ] }, { @@ -186,7 +250,9 @@ "id": "cf45781f", "metadata": {}, "source": [ - "Plot what things look like after tuning:" + "## Posterior after tuning\n", + "\n", + "Replotting the same diagnostics makes the effect of tuning visible. The learned length scale changes both the posterior interpolation and how rapidly uncertainty grows between observations, while the loss landscape locates the tuned parameter relative to nearby alternatives." ] }, { @@ -210,7 +276,9 @@ "id": "4c34f19e", "metadata": {}, "source": [ - "Iteratively sample points to choose for tuning, then tune:" + "## Adaptively add observations\n", + "\n", + "At each iteration, we first select the best point from a coarse candidate grid using marginal posterior variance. Starting from that candidate, `get_sample_point` directly maximizes posterior variance within the domain. We observe the hidden function there, condition on all available data, and tune the kernel again. In each variance plot, the newly added coordinate is highlighted separately from the earlier observations." ] }, { @@ -231,7 +299,7 @@ " _x = x_sample[jnp.argmax(_v)].reshape(-1, D)\n", "\n", " # Iteratively optimize posterior variance:\n", - " _x = g.get_sample_point(_x.reshape((1,)), ranges=[[0.], [2 * jnp.pi]])\n", + " _x = g.get_sample_point(_x.reshape((1,)), ranges=np.atleast_2d(DOMAIN).T.tolist())\n", "\n", " # Observe that point:\n", " _y, _s = f_noisy(_x)\n", @@ -255,12 +323,14 @@ ] }, { - "cell_type": "code", - "execution_count": null, - "id": "5e2afcd6", + "cell_type": "markdown", + "id": "result", "metadata": {}, - "outputs": [], - "source": [] + "source": [ + "## Result\n", + "\n", + "Adaptive sampling directs new observations toward regions where the current model is least certain. Reconditioning reduces uncertainty near each new observation, while repeated tuning updates the covariance length scale as more information becomes available." + ] } ], "metadata": { From cd64b8e4b88355ff276ffcb2997743cd6201a369 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 10 Sep 2026 13:44:09 -0600 Subject: [PATCH 16/22] docs: add clean target to Makefile --- docs/Makefile | 4 ++++ docs/clean.sh | 8 -------- 2 files changed, 4 insertions(+), 8 deletions(-) delete mode 100755 docs/clean.sh diff --git a/docs/Makefile b/docs/Makefile index d4bb2cb..ebc7539 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -14,6 +14,10 @@ help: .PHONY: help Makefile +clean: + rm -rf ${SOURCEDIR}/dev_docs/api/generated/ + @$(SPHINXBUILD) -M clean "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + # Catch-all target: route all unknown targets to Sphinx using the new # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). %: Makefile diff --git a/docs/clean.sh b/docs/clean.sh deleted file mode 100755 index 78b77a5..0000000 --- a/docs/clean.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/bin/bash -set -e - -# Clean up build products not located in _build/: -rm -rf dev_docs/api/generated/ - -# Clean up build products with `make clean`: -make clean From bb2a0092adfafe83831c7c816faa99f8d5d1520c Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 10 Sep 2026 14:38:21 -0600 Subject: [PATCH 17/22] gp: get ready to add 1-d example notebook to docs * need to reorganize tutorial docs * reverted python 3.9 breaking changes on rebase * need to find solution for copying demos/sources from their permanent home into docs at build --- .gitignore | 2 +- UncertainSCI/gp/__init__.py | 30 +++++++++++++++++++++++------ demos/gp-1d.ipynb | 38 ++++++++++++++++++++----------------- docs/conf.py | 22 ++++++++++----------- docs/tutorials/index.rst | 8 ++++---- 5 files changed, 61 insertions(+), 39 deletions(-) diff --git a/.gitignore b/.gitignore index 1de96dd..165109f 100644 --- a/.gitignore +++ b/.gitignore @@ -68,7 +68,7 @@ instance/ # Scrapy stuff: .scrapy -# Sphinx documentation +# Docs: docs/_build/ docs/**/generated/ diff --git a/UncertainSCI/gp/__init__.py b/UncertainSCI/gp/__init__.py index 2602698..d9478e1 100644 --- a/UncertainSCI/gp/__init__.py +++ b/UncertainSCI/gp/__init__.py @@ -538,7 +538,7 @@ def prior_covariance(self, u, v): return self.k(u, v) - def prior_realization(self, x: jax.Array | npt.NDArray, p: int = 1): + def prior_realization(self, x: jax.Array | npt.NDArray, p: int = 1) -> jax.Array: """ Produce ``p`` realizations of the prior at locations ``x``. @@ -548,6 +548,12 @@ def prior_realization(self, x: jax.Array | npt.NDArray, p: int = 1): of the domain. p (int, optional): Number of realizations to generate, default ``1``. + + Returns: + jax.Array: + Array of shape ``(n, c, p)`` where the first dimension corresponds to + sampling coordinate, the second dimension corresponds to the output + channel, and the third dimension stores the number of realizations. """ self.key, key = jax.random.split(self.key) n, _ = self.validate_input_shape(x) @@ -555,9 +561,12 @@ def prior_realization(self, x: jax.Array | npt.NDArray, p: int = 1): L = jnp.linalg.cholesky( self.prior_covariance(x, x) + self.nugget * jnp.eye(n * self.cdim) ) - return self.prior_mean(x) + (L @ jax.random.normal(key, (n * self.cdim, p))) + return ( + self.prior_mean(x).reshape(n, self.cdim, 1) + + (L @ jax.random.normal(key, (n * self.cdim, p))).reshape(n, self.cdim, p) + ) - def posterior_mean(self, x): + def posterior_mean(self, x) -> jax.Array: """ Evaluate the posterior mean at ``x``. @@ -572,7 +581,7 @@ def posterior_mean(self, x): correction = self.k(x, self.train_x) @ self.train_cov_factor.solve(d) return self.mu(x) + correction.reshape((-1, self.cdim)) - def posterior_covariance(self, u, v): + def posterior_covariance(self, u, v) -> jax.Array: """ Evaluate the posterior covariance at ``(u, v)``. @@ -592,7 +601,7 @@ def posterior_covariance(self, u, v): self.k(u, self.train_x) @ self.train_cov_factor.solve(self.k(self.train_x, v)) ) - def posterior_realization(self, x: jax.Array | npt.NDArray, p: int = 1): + def posterior_realization(self, x: jax.Array | npt.NDArray, p: int = 1) -> jax.Array: """ Produce ``p`` realizations of the posterior at locations ``x``. @@ -602,6 +611,12 @@ def posterior_realization(self, x: jax.Array | npt.NDArray, p: int = 1): of the domain. p (int, optional): Number of realizations to generate, default ``1``. + + Returns: + jax.Array: + Array of shape ``(n, c, p)`` where the first dimension corresponds to + sampling coordinate, the second dimension corresponds to the output + channel, and the third dimension stores the number of realizations. """ self.key, key = jax.random.split(self.key) n, _ = self.validate_input_shape(x) @@ -609,4 +624,7 @@ def posterior_realization(self, x: jax.Array | npt.NDArray, p: int = 1): L = jnp.linalg.cholesky( self.posterior_covariance(x, x) + self.nugget * jnp.eye(n * self.cdim) ) - return self.posterior_mean(x) + (L @ jax.random.normal(key, (n * self.cdim, p))) + return ( + self.posterior_mean(x).reshape(n, self.cdim, 1) + + (L @ jax.random.normal(key, (n * self.cdim, p))).reshape(n, self.cdim, p) + ) diff --git a/demos/gp-1d.ipynb b/demos/gp-1d.ipynb index e6ddfbc..e39e552 100644 --- a/demos/gp-1d.ipynb +++ b/demos/gp-1d.ipynb @@ -31,20 +31,23 @@ "\n", "D = 1 \n", "C = 1\n", - "SEED = 0x01234567\n", - "GP_SEED = 0xDEADBEEF\n", - "rng = np.random.default_rng(SEED)\n", - "\n", "DOMAIN = (0, 2 * jnp.pi)\n", "NOISE_VARIANCE = 1e-2\n", "\n", + "SEED = 0x01234567\n", + "GP_SEED = 0xdeadbeef\n", + "rng = np.random.default_rng(SEED)\n", + "\n", + "N_PLOT = 1000\n", "FIGSIZE = (7, 7 / 1.6)\n", "FIGDPI = None\n", "\n", - "def print_loss_here(g): # Simple helper for this notebook; does not generalize.\n", + "# Simple helper for this notebook; does not generalize.\n", + "def print_loss_here(g, i: int | None = None):\n", " print(\n", " f'Loss at D = {jnp.squeeze(g.k.D()):.4e} (length scale {1 / jnp.sqrt(jnp.squeeze(g.k.D())):.4e}): '\n", " f'{g.loss_hyperparameters(train_x, train_y, train_s):.4e}'\n", + " + str('' if i is None else f' at iteration {i}')\n", " )\n", "\n", "def plot_in_subplots(nrows=1, ncols=1, **kwargs):\n", @@ -68,7 +71,7 @@ "source": [ "## True function and noisy observations\n", "\n", - "The hidden function combines two oscillatory factors and is evaluated on the interval from zero to $2\\pi$. Observations have independent Gaussian noise with a known, constant variance. `GaussianProcess.condition` expects this variance—not its square root." + "The hidden function is evaluated on the interval from $0$ to $2\\pi$. Observations have independent Gaussian noise with a known, constant variance. `GaussianProcess.condition` expects this variance—not its square root." ] }, { @@ -107,7 +110,7 @@ "ax1.set_title('Truth')\n", "\n", "ax2.plot(noisy)\n", - "ax2.set_title('Noisy View')\n", + "ax2.set_title('Noisy Realization')\n", "\n", "plt.show()\n" ] @@ -289,7 +292,7 @@ "outputs": [], "source": [ "N_TRAIN = 10\n", - "N_SAMPLE = 20\n", + "N_SAMPLE = 10\n", "\n", "x_sample = jnp.linspace(0, 2 * jnp.pi, N_SAMPLE).reshape(-1, D)\n", "\n", @@ -312,14 +315,15 @@ " g.condition(train_x, train_y, train_s)\n", " g.tune()\n", "\n", - " # Plot:\n", - " fig, (ax1, ax2, ax3) = plot_in_subplots(3, 1, dpi=FIGDPI)\n", - " gp.vis.plot_distribution_mean(ax1, g, x_plot, f_hidden, 'posterior')\n", - " gp.vis.plot_distribution_variance(ax2, g, x_plot, 'posterior')\n", - " gp.vis.plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", - " plt.show()\n", - " print_loss_here(g)\n", - " print('\\n' * 3)\n" + " # Plot on first and last:\n", + " if i == 0 or i == N_TRAIN - 1:\n", + " fig, (ax1, ax2, ax3) = plot_in_subplots(3, 1, dpi=FIGDPI)\n", + " gp.vis.plot_distribution_mean(ax1, g, x_plot, f_hidden, 'posterior')\n", + " gp.vis.plot_distribution_variance(ax2, g, x_plot, 'posterior')\n", + " gp.vis.plot_loss_landscape(ax3, g, ('k', 'D'), jnp.logspace(-2, 2, 1000))\n", + " plt.show()\n", + " print_loss_here(g, i)\n", + " print('\\n' * 3)\n" ] }, { @@ -349,7 +353,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.16" + "version": "3.11.15" } }, "nbformat": 4, diff --git a/docs/conf.py b/docs/conf.py index cd2cd37..4179684 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -28,14 +28,14 @@ # -- General configuration --------------------------------------------------- extensions = [ - 'sphinx.ext.autodoc', - 'sphinx.ext.autosummary', - 'sphinx.ext.intersphinx', - 'sphinx.ext.napoleon', - 'sphinx.ext.viewcode', - 'myst_nb', - 'sphinx_copybutton', - 'sphinxcontrib.bibtex', + 'sphinx.ext.autodoc', + 'sphinx.ext.autosummary', + 'sphinx.ext.intersphinx', + 'sphinx.ext.napoleon', + 'sphinx.ext.viewcode', + 'myst_nb', + 'sphinx_copybutton', + 'sphinxcontrib.bibtex', ] templates_path = ['_templates'] @@ -59,7 +59,9 @@ autoclass_content = 'both' # Path for bibtex files -bibtex_bibfiles = ['references.bib'] +bibtex_bibfiles = [ + 'references.bib', +] bibtex_default_style = 'unsrt' bibtex_encoding = 'latin' @@ -80,9 +82,7 @@ 'dollarmath', 'html_image', ] - myst_heading_anchors = 4 - nb_execution_mode = 'auto' diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst index e76ede3..db4f56d 100644 --- a/docs/tutorials/index.rst +++ b/docs/tutorials/index.rst @@ -1,9 +1,10 @@ Tutorials -======================================== +========= + +These are some helpful tutorials to familiarize yourself with UncertainSCI. .. toctree:: - :maxdepth: 2 - :caption: Contents: + :maxdepth: 1 notebooks/build_pce demos/index @@ -11,4 +12,3 @@ Tutorials models parameters template - From 4cd17660afaddcdbd23c86687f59c7aaedc14c14 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 10 Sep 2026 14:38:45 -0600 Subject: [PATCH 18/22] docs: recursively delete all occurences of 'generated' folder(s) --- docs/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/Makefile b/docs/Makefile index ebc7539..b8eb905 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -15,7 +15,7 @@ help: .PHONY: help Makefile clean: - rm -rf ${SOURCEDIR}/dev_docs/api/generated/ + find "$(SOURCEDIR)" -type d -name generated -prune -exec rm -rf {} + @$(SPHINXBUILD) -M clean "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) # Catch-all target: route all unknown targets to Sphinx using the new From 9b1056d259b058d568b46029d3d49929447f90b8 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 10 Sep 2026 14:57:09 -0600 Subject: [PATCH 19/22] gp: cleanup some variable names in example notebook --- demos/gp-1d.ipynb | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/demos/gp-1d.ipynb b/demos/gp-1d.ipynb index e39e552..a31bc31 100644 --- a/demos/gp-1d.ipynb +++ b/demos/gp-1d.ipynb @@ -29,8 +29,8 @@ "import UncertainSCI.gp as gp\n", "\n", "\n", - "D = 1 \n", - "C = 1\n", + "DIM = 1 \n", + "CDIM = 1\n", "DOMAIN = (0, 2 * jnp.pi)\n", "NOISE_VARIANCE = 1e-2\n", "\n", @@ -81,7 +81,7 @@ "metadata": {}, "outputs": [], "source": [ - "x_plot = jnp.linspace(*DOMAIN, 1000).reshape((-1, D))\n", + "x_plot = jnp.linspace(*DOMAIN, 1000).reshape((-1, DIM))\n", "\n", "def f_hidden(x):\n", " return 2 * jnp.cos(x) * jnp.sin(4 * x)\n", @@ -133,19 +133,19 @@ "outputs": [], "source": [ "mu = gp.mean.Affine(\n", - " dim=D,\n", - " cdim=C,\n", - " a=0. * jnp.ones((C, D)),\n", + " dim=DIM,\n", + " cdim=CDIM,\n", + " a=0. * jnp.ones((CDIM, DIM)),\n", " b=0.,\n", " a_is_static=True,\n", " b_is_static=True\n", ")\n", "k = gp.kernel.Gaussian(\n", - " dim=D,\n", - " cdim=C,\n", + " dim=DIM,\n", + " cdim=CDIM,\n", " D=jnp.ones((1, 1)),\n", ")\n", - "g = gp.GaussianProcess(dim=D, cdim=C, mu=mu, k=k, seed=GP_SEED, nugget=1e-3)\n" + "g = gp.GaussianProcess(dim=DIM, cdim=CDIM, mu=mu, k=k, seed=GP_SEED, nugget=1e-3)\n" ] }, { @@ -190,7 +190,7 @@ "source": [ "N_INIT = 10\n", "\n", - "train_x = jnp.linspace(0, 2 * jnp.pi, N_INIT).reshape(-1, D)\n", + "train_x = jnp.linspace(0, 2 * jnp.pi, N_INIT).reshape(-1, DIM)\n", "train_y, train_s = f_noisy(train_x) \n", "\n", "g.condition(train_x, train_y, train_s)\n" @@ -294,12 +294,12 @@ "N_TRAIN = 10\n", "N_SAMPLE = 10\n", "\n", - "x_sample = jnp.linspace(0, 2 * jnp.pi, N_SAMPLE).reshape(-1, D)\n", + "x_sample = jnp.linspace(0, 2 * jnp.pi, N_SAMPLE).reshape(-1, DIM)\n", "\n", "for i in range(N_TRAIN):\n", " # Greedily sample posterior variance:\n", " _v = jnp.diag(g.posterior_covariance(x_sample, x_sample))\n", - " _x = x_sample[jnp.argmax(_v)].reshape(-1, D)\n", + " _x = x_sample[jnp.argmax(_v)].reshape(-1, DIM)\n", "\n", " # Iteratively optimize posterior variance:\n", " _x = g.get_sample_point(_x.reshape((1,)), ranges=np.atleast_2d(DOMAIN).T.tolist())\n", From b1c4d2f86726ea524ccae1562c35357cd31e6371 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Wed, 30 Sep 2026 20:33:03 -0600 Subject: [PATCH 20/22] docs: just copy stuff in the Makefile --- .gitignore | 1 + docs/Makefile | 2 ++ docs/tutorials/index.rst | 1 + 3 files changed, 4 insertions(+) diff --git a/.gitignore b/.gitignore index 165109f..cdf1e3f 100644 --- a/.gitignore +++ b/.gitignore @@ -70,6 +70,7 @@ instance/ # Docs: docs/_build/ +docs/**/copied/ docs/**/generated/ # PyBuilder diff --git a/docs/Makefile b/docs/Makefile index b8eb905..f7b5ea6 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -15,10 +15,12 @@ help: .PHONY: help Makefile clean: + find "$(SOURCEDIR)" -type d -name copied -prune -exec rm -rf {} + find "$(SOURCEDIR)" -type d -name generated -prune -exec rm -rf {} + @$(SPHINXBUILD) -M clean "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) # Catch-all target: route all unknown targets to Sphinx using the new # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). %: Makefile + cp -r ../demos tutorials/copied @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst index db4f56d..48c4b76 100644 --- a/docs/tutorials/index.rst +++ b/docs/tutorials/index.rst @@ -12,3 +12,4 @@ These are some helpful tutorials to familiarize yourself with UncertainSCI. models parameters template + 1-d Gaussian Process From ef4261e2478a3db81fc8d6794f8a7a866e0c0724 Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Wed, 30 Sep 2026 20:47:31 -0600 Subject: [PATCH 21/22] project: bump python support to 3.12-3.14 (inclusive) --- .github/workflows/build.yml | 4 ++-- .github/workflows/lint.yml | 4 ++-- .github/workflows/test.yml | 4 ++-- pyproject.toml | 8 ++++---- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 47977ee..0d4ba55 100644 --- a/.github/workflows/build.yml +++ b/.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 diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 5191298..cb459df 100644 --- a/.github/workflows/lint.yml +++ b/.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 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index f343d13..3d2f286 100644 --- a/.github/workflows/test.yml +++ b/.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: diff --git a/pyproject.toml b/pyproject.toml index 1f9044c..d41dcbc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "UncertainSCI" version = "1.0.1.1" description = "A Non-invasive Uncertainty Quantification tool" readme = { file = "README.md", content-type = "text/markdown" } -requires-python = ">=3.9,<3.12" +requires-python = ">=3.12,<3.15" license = { text = "MIT" } authors = [ { name = "UncertainSCI Developers", email = "uncertainsci@sci.utah.edu" }, @@ -16,9 +16,9 @@ keywords = ["Uncertainty Quantification", "Simulation"] classifiers = [ "License :: OSI Approved :: MIT License", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", "Development Status :: 4 - Beta", "Intended Audience :: Developers", "Intended Audience :: Education", From e8a0654b26bb3125e3677a1c015bc823b7eb933a Mon Sep 17 00:00:00 2001 From: Asher Merrill Date: Thu, 1 Oct 2026 14:53:20 -0600 Subject: [PATCH 22/22] project: bump python to 3.12 in readthedocs config --- .readthedocs.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.readthedocs.yaml b/.readthedocs.yaml index f7d9006..0065108 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -9,7 +9,7 @@ version: 2 build: os: ubuntu-22.04 tools: - python: "3.10" + python: "3.12" apt_packages: - pandoc