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/.gitignore b/.gitignore index 1de96dd..cdf1e3f 100644 --- a/.gitignore +++ b/.gitignore @@ -68,8 +68,9 @@ instance/ # Scrapy stuff: .scrapy -# Sphinx documentation +# Docs: docs/_build/ +docs/**/copied/ docs/**/generated/ # PyBuilder 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 diff --git a/UncertainSCI/_equinox.py b/UncertainSCI/_equinox.py new file mode 100644 index 0000000..9ddf27c --- /dev/null +++ b/UncertainSCI/_equinox.py @@ -0,0 +1,138 @@ +import abc +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy.typing as npt + + +class ComputedArray(eqx.Module, metaclass=abc.ABCMeta): + """ + Computed array base class. + + Primarily used to detect if PyTree nodes require (further) resolution to reach + underlying objects or data. + """ + is_static: bool = eqx.field(static=True) + + @abc.abstractmethod + def __call__(self) -> jax.Array: + raise NotImplementedError + + @classmethod + @abc.abstractmethod + def from_array(cls, a: jax.Array, is_static: bool=False, **kwargs): + """ + Convenience method for in-place replacement of types of + :class:`ComputedArray`, though not generally intended for use as a long-lived + instance; prefer direct instantiation when possible. + + Args: + a (array): + The array data of this :class:`ComputedArray`. + is_static (bool, optional): + If static. + **kwargs: + Kwargs to hand off to :meth:`__init__` (generally depends + on child class). + """ + raise NotImplementedError + + +class PositiveDefinite(ComputedArray): + """ + A positive-definite matrix. + + Module contains some niceties that guarantee matrix remains positive definite + during training. + + Attributes: + L_log_diag (array or None): + Log of diagonal of Cholesky factor of positive definite matrix. + L_off_diag (array or None): + Elements below main diagonal of Cholesky factor of positive definite. + """ + L_log_diag: jax.Array + L_off_diag: jax.Array + + def __init__( + self, + *, # TODO: make this NOT kwargs-only. + D: jax.Array | npt.NDArray | None = None, + L_log_diag: jax.Array | npt.NDArray | None = None, + L_off_diag: jax.Array | npt.NDArray | None = None, + is_static: bool = False, + ): + """ + Create a positive definite matrix. + + All args must be passed as kwargs to avoid name collisions with parent + methods. + + Args: + D (array or None): + Positive definite matrix; defined only if ``L_log_diag is None`` and + ``L_off_diag is None``. + L_log_diag (array or None): + Log of diagonal of Cholesky factor of positive definite matrix; + defined only if ``D is None``. + L_off_diag (array or None): + Elements below main diagonal of Cholesky factor of positive definite + matrix; defined only if ``D is None``. + is_static (bool): + If the matrix is static. + """ + if not ( + (D is not None) ^ + (L_log_diag is not None and L_off_diag is not None) + ): + raise ValueError('D xor (L_log_diag and L_off_diag) must be defined!') + + if D is not None: + if D.ndim != 2: + raise ValueError(f'Expected ndim = 2 for matrix, got D.ndim = {D.ndim}!') + if D.shape[0] != D.shape[1]: + raise ValueError(f'Expected square matrix, got D.shape = {D.shape}!') + + n = D.shape[0] + L = jnp.linalg.cholesky(D) + self.L_log_diag = jnp.log(jnp.diag(L)) + self.L_off_diag = L[jnp.tril_indices(n, k=-1)] + + elif L_log_diag is not None and L_off_diag is not None: + if L_log_diag.ndim != 1: + raise ValueError( + f'Expected ndim = 1 for L_log_diag, got L_log_diag.ndim = {L_log_diag.ndim}!' + ) + if L_off_diag.ndim != 1: + raise ValueError( + f'Expected ndim = 1 for L_off_diag, got L_off_diag.ndim = {L_off_diag.ndim}!' + ) + + + n = L_log_diag.shape[0] + if L_off_diag.shape[0] != (n) * (n - 1) / 2: + raise ValueError( + f'Expected L_off_diag.shape[0] = (n) * (n - 1) / 2 = ' + f'{(n) * (n - 1) / 2} (numel below main diagonal), got ' + f'L_off_diag.shape[0] = {L_off_diag.shape[0]}!' + ) + + self.L_log_diag = jnp.asarray(L_log_diag) + self.L_off_diag = jnp.asarray(L_off_diag) + + else: + raise ValueError('Got unexpected combination of parameters; check caller!') + + super().__init__(is_static=is_static) + + def __call__(self) -> jax.Array: + n = self.L_log_diag.shape[0] + L = jnp.diag(jnp.exp(self.L_log_diag)) + L = L.at[jnp.tril_indices(n, k=-1)].set(self.L_off_diag) + D = L @ L.T + + return jax.lax.stop_gradient(D) if self.is_static else D + + @classmethod + def from_array(cls, a, is_static: bool = False, **kwargs): + return cls(D=a, is_static=is_static, **kwargs) 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..d9478e1 --- /dev/null +++ b/UncertainSCI/gp/__init__.py @@ -0,0 +1,630 @@ +""" +Build Gaussian processes. +""" + +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np +import numpy.typing as npt +import scipy.spatial as spatial +import optax + +from . import kernel +from . import mean +from . import vis + + +DEFAULT_OPTIM = optax.adam +DEFAULT_OPTIM_KWARGS = { + 'learning_rate': 1e-2 +} + +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 + k: kernel.Kernel + + nugget: float + + train_x: jax.Array + train_y: jax.Array + train_s: jax.Array + train_cov_factor: kernel.CovarianceFactor + + 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: + """ + 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: + 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 + ): + """ + 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) + # 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( + (self.mu, self.k), + self.train_x, + self.train_y, + self.train_s + ) + + 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 + ): + """ + 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) + + if nx != ny: + raise ValueError + + s_shape = jnp.shape(s) + if s_shape == (): + return + elif s_shape not in { + (ny, self.cdim), + (ny * self.cdim,), + (ny * self.cdim, ny * self.cdim), + }: + raise ValueError + + 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 + + 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]: + """ + 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 + + n, c = y.shape + if c != self.cdim: + raise ValueError + + return n, c + + def get_train_cov_factor( + self, + x: jax.Array | npt.NDArray, + y: jax.Array | npt.NDArray, + s: jax.Array | npt.NDArray | float | int + ) -> 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) + # 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) + + # 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_train_cov_factor( + p: tuple[mean.Mean, kernel.Kernel], + x: jax.Array, + y: jax.Array, + s: jax.Array + ) -> kernel.CovarianceFactor: + _, k = p + return k.covariance_factor(x, s) + + def tune( + self, + n: int = 1000, + optim = DEFAULT_OPTIM, + optim_kwargs: dict = DEFAULT_OPTIM_KWARGS, + ) -> jax.Array: + """ + 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_hyperparameters( + optim, + optim_state, + (self.mu, self.k), + self.train_x, + self.train_y, + self.train_s + ) + losses.append(loss) + + self.train_cov_factor = GaussianProcess._get_train_cov_factor( + (self.mu, self.k), + self.train_x, + self.train_y, + self.train_s + ) + + return jnp.asarray(losses) + + @staticmethod + @jax.jit(static_argnames=('optim',)) + def _step_hyperparameters( + 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_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_hyperparameters( + self, + x: jax.Array | npt.NDArray, + y: jax.Array | npt.NDArray, + s: jax.Array | npt.NDArray | float | int + ) -> jax.Array: + """ + 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) + # 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) + + @staticmethod + @jax.jit + def _loss_hyperparameters( + p: tuple[mean.Mean, kernel.Kernel], + x: jax.Array, + y: jax.Array, + s: jax.Array + ) -> jax.Array: + mu, k = p + 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) / 2 + + cov_factor.log_sqrt_det() + ) + + def get_sample_point( + self, + 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 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 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 (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.') + + x = jnp.asarray(x) + if not jnp.issubdtype(x.dtype, jnp.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 not ((hull is not None) ^ (ranges is not None)): + raise ValueError + + 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) + + 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) & + (x_cand <= ranges[1, 0] + 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.train_x, + self.train_cov_factor, + x + ) + + if inside_domain(x_cand): + x = x_cand + else: + return jnp.asarray(x) + + return jnp.asarray(x) + + @staticmethod + @jax.jit(static_argnames=('optim',)) + def _step_sample_point( + optim, + optim_state, + k: kernel.Kernel, + train_x: jax.Array, + train_cov_factor: kernel.CovarianceFactor, + x: jax.Array + ) -> tuple[jax.Array, jax.Array, jax.Array]: + 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 + + 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: jax.Array | npt.NDArray, p: int = 1) -> jax.Array: + """ + 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``. + + 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) + + L = jnp.linalg.cholesky( + self.prior_covariance(x, x) + self.nugget * jnp.eye(n * self.cdim) + ) + 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) -> jax.Array: + """ + 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) + 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) -> jax.Array: + """ + 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) + + return ( + self.k(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) -> jax.Array: + """ + 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``. + + 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) + + L = jnp.linalg.cholesky( + self.posterior_covariance(x, x) + self.nugget * jnp.eye(n * self.cdim) + ) + 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/UncertainSCI/gp/kernel.py b/UncertainSCI/gp/kernel.py new file mode 100644 index 0000000..177351c --- /dev/null +++ b/UncertainSCI/gp/kernel.py @@ -0,0 +1,421 @@ +""" +Kernels for Gaussian processes. +""" + +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 +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. + * 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 (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: + 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): + """ + 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. + """ + size: int = eqx.field(static=True) + + @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 CholeskyFactor(CovarianceFactor): + """ + 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 + 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) + + def log_sqrt_det(self) -> jax.Array: + return jnp.sum(jnp.log(jnp.diag(self.L))) + + def cholesky(self) -> jax.Array: + return self.L + + +class KroneckerScalarNoiseFactor(CovarianceFactor): + """ + 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 + 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, + ): + """ + 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]: + 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) + # FIXME: Ensure integer inputs have an inexact dtype when traced by JAX. + 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 + + def log_sqrt_det(self) -> jax.Array: + return 0.5 * jnp.sum(jnp.log(self.eigenvalues)) + + 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): + """ + 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 + ) -> jax.Array: + """ + Evaluate this kernel. + + Args: + 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 + ) -> 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 (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 CholeskyFactor(jnp.linalg.cholesky(covariance + s)) + + +class Gaussian(Kernel): + r""" + A Gaussian (square exponential) scalar-valued covariance kernel. + + This produces a kernel according to the function + + .. math:: + + k(x, y) = \exp{\left[ - (x - y)^\top D \;(x - y) \right]}. + + :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 = (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 + + def __init__( + self, + dim: int, + cdim: int, + 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) + + 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 + ) + ) + + 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""" + 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. + """ + k: Kernel + C: _equinox.PositiveDefinite + + def __init__( + self, + dim: int, + cdim: int, + k: Kernel, + 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.') + + 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()) + + def covariance_factor( + self, + x: jax.Array | npt.NDArray, + 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: + 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 new file mode 100644 index 0000000..e65d7cf --- /dev/null +++ b/UncertainSCI/gp/mean.py @@ -0,0 +1,110 @@ +""" +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 = (n, d)`` for some number of coordinates ``n`` and + dimension of the domain ``d``. + + 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)): + # FIXME: Ensure integer inputs have an inexact dtype when traced by JAX. + 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/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 new file mode 100644 index 0000000..a31bc31 --- /dev/null +++ b/demos/gp-1d.ipynb @@ -0,0 +1,361 @@ +{ + "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, + "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 jax.numpy as jnp\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "import UncertainSCI.gp as gp\n", + "\n", + "\n", + "DIM = 1 \n", + "CDIM = 1\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", + "# 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", + " \"\"\"\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 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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "598efb92", + "metadata": {}, + "outputs": [], + "source": [ + "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", + "\n", + "def f_noisy(x):\n", + " y = f_hidden(x)\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 Realization')\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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ee181026", + "metadata": {}, + "outputs": [], + "source": [ + "mu = gp.mean.Affine(\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=DIM,\n", + " cdim=CDIM,\n", + " D=jnp.ones((1, 1)),\n", + ")\n", + "g = gp.GaussianProcess(dim=DIM, cdim=CDIM, mu=mu, k=k, seed=GP_SEED, nugget=1e-3)\n" + ] + }, + { + "cell_type": "markdown", + "id": "prior-heading", + "metadata": {}, + "source": [ + "## Inspect the prior\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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4a5952f3", + "metadata": {}, + "outputs": [], + "source": [ + "fig, (ax1, ax2) = plot_in_subplots(2, 1, dpi=FIGDPI)\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" + ] + }, + { + "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, + "id": "ee0fc8ac", + "metadata": {}, + "outputs": [], + "source": [ + "N_INIT = 10\n", + "\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" + ] + }, + { + "cell_type": "markdown", + "id": "2cf76bf6", + "metadata": {}, + "source": [ + "## 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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6ca89cdd", + "metadata": {}, + "outputs": [], + "source": [ + "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', 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" + ] + }, + { + "cell_type": "markdown", + "id": "9f7a058a", + "metadata": {}, + "source": [ + "## 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." + ] + }, + { + "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": [ + "## 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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "885d3e14", + "metadata": {}, + "outputs": [], + "source": [ + "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', 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" + ] + }, + { + "cell_type": "markdown", + "id": "4c34f19e", + "metadata": {}, + "source": [ + "## 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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "818fb48c", + "metadata": {}, + "outputs": [], + "source": [ + "N_TRAIN = 10\n", + "N_SAMPLE = 10\n", + "\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, DIM)\n", + "\n", + " # Iteratively optimize posterior variance:\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", + "\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", + " # Condition with new data:\n", + " g.condition(train_x, train_y, train_s)\n", + " g.tune()\n", + "\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" + ] + }, + { + "cell_type": "markdown", + "id": "result", + "metadata": {}, + "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": { + "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 +} diff --git a/docs/Makefile b/docs/Makefile index d4bb2cb..f7b5ea6 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -14,7 +14,13 @@ 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/conf.py b/docs/conf.py index 3cbea56..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'] @@ -56,9 +56,12 @@ autosummary_generate = True autosummary_imported_members = True +autoclass_content = 'both' # Path for bibtex files -bibtex_bibfiles = ['references.bib'] +bibtex_bibfiles = [ + 'references.bib', +] bibtex_default_style = 'unsrt' bibtex_encoding = 'latin' @@ -69,6 +72,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 = [ @@ -78,9 +82,7 @@ 'dollarmath', 'html_image', ] - myst_heading_anchors = 4 - nb_execution_mode = 'auto' 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 diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst index e76ede3..48c4b76 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,4 @@ Tutorials models parameters template - + 1-d Gaussian Process diff --git a/pyproject.toml b/pyproject.toml index 4b64a17..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", @@ -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", ] diff --git a/tests/test_gp.py b/tests/test_gp.py new file mode 100644 index 0000000..92deda0 --- /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.KroneckerScalarNoiseFactor) + 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.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( + 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)