ptgp 0.1.0__py3-none-any.whl

ptgp/__init__.py

@@ -0,0 +1,62 @@
+import logging
+import sys
+
+from ptgp import (
+    gp,
+    idata,
+    inducing,
+    kernels,
+    likelihoods,
+    mean,
+    objectives,
+    optim,
+    rewrites,  # noqa: F401  (registers PyTensor rewrites at import)
+    utils,
+)
+from ptgp.inducing_fourier import FourierFeatures1D
+from ptgp.optim.api import FitResult, fit, predict
+
+
+class _StdoutHandler(logging.StreamHandler):
+    """Handler that resolves sys.stdout at emit time, not at init time."""
+
+    def __init__(self):
+        super().__init__()
+
+    @property
+    def stream(self):
+        return sys.stdout
+
+    @stream.setter
+    def stream(self, _):
+        pass
+
+
+_logger = logging.getLogger("ptgp")
+if not _logger.handlers:
+    _handler = _StdoutHandler()
+    _handler.setFormatter(logging.Formatter("%(message)s"))
+    _logger.addHandler(_handler)
+    _logger.setLevel(logging.INFO)
+
+__all__ = [
+    "FourierFeatures1D",
+    "FitResult",
+    "fit",
+    "gp",
+    "idata",
+    "inducing",
+    "kernels",
+    "likelihoods",
+    "mean",
+    "objectives",
+    "optim",
+    "predict",
+    "utils",
+]
+
+
+try:
+    from ptgp._version import __version__
+except ImportError:
+    __version__ = "0.0.0+unknown"

ptgp/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

ptgp/conditionals.py

@@ -0,0 +1,72 @@
+import pytensor.assumptions as pta
+import pytensor.tensor as pt
+
+# Diagonal jitter added to Kmm before Cholesky / inversion, to keep it PSD
+# under floating-point noise. Matches GPflow / GPJax / PyMC defaults of 1e-6.
+_DEFAULT_JITTER = 1e-6
+
+
+def conditional_whitened(A_white, Knn, f, q_sqrt=None, full_cov=False):
+    """Posterior conditional in whitened parameterisation.
+
+    A_white : (M, N), satisfies A_white = R^{-1} @ Kmn where R @ R.T = Kmm.
+    Knn     : (N,) prior diagonal if ``full_cov=False``, else (N, N) prior covariance.
+    f       : (M,) variational mean (whitened: prior on v is N(0, I)).
+    q_sqrt  : (M, K), optional.
+    full_cov: if True, return full (N, N) covariance; else (N,) marginals.
+    """
+    fmean = A_white.T @ f
+    if full_cov:
+        fvar = Knn - A_white.T @ A_white
+    else:
+        fvar = Knn - pt.sum(A_white**2, axis=0)
+    if q_sqrt is not None:
+        B = A_white.T @ q_sqrt
+        if full_cov:
+            fvar = fvar + B @ B.T
+        else:
+            fvar = fvar + pt.sum(B**2, axis=1)
+    return fmean, fvar
+
+
+def conditional_unwhitened(A, Kmn, Knn, f, q_sqrt=None, full_cov=False):
+    """Posterior conditional in unwhitened parameterisation.
+
+    A    : (M, N), satisfies A = Kmm^{-1} @ Kmn.
+    Kmn  : (M, N) cross-covariance.
+    Knn  : (N,) prior diagonal if ``full_cov=False``, else (N, N) prior covariance.
+    f    : (M,) variational mean (prior on u is N(0, Kmm)).
+    full_cov: if True, return full (N, N) covariance; else (N,) marginals.
+    """
+    fmean = A.T @ f
+    if full_cov:
+        fvar = Knn - A.T @ Kmn
+    else:
+        fvar = Knn - pt.sum(A * Kmn, axis=0)
+    if q_sqrt is not None:
+        B = A.T @ q_sqrt
+        if full_cov:
+            fvar = fvar + B @ B.T
+        else:
+            fvar = fvar + pt.sum(B**2, axis=1)
+    return fmean, fvar
+
+
+def base_conditional(Kmn, Kmm, Knn, f, q_sqrt=None, white=False, full_cov=False):
+    """Back-compat wrapper. Materialises Kmm; use the helpers above directly
+    when you have a structured Kuu_solve / Kuu_sqrt_solve."""
+    # Add jitter to keep Kmm PSD under float noise — matches GPflow / PyMC default.
+    # Re-annotate after the addition: PyTensor canonicalizes ``Kmm + c·I`` into a
+    # ``set_subtensor`` on the diagonal, which our PSD-inference rules don't see
+    # through. The mathematical identity (PSD + c·I PSD ⇒ PSD) is sound.
+    Kmm = pta.assume(
+        Kmm + _DEFAULT_JITTER * pt.eye(Kmm.shape[-1], dtype=Kmm.dtype),
+        positive_definite=True,
+        symmetric=True,
+    )
+    if white:
+        L = pt.linalg.cholesky(Kmm)
+        A_white = pt.linalg.solve(L, Kmn)
+        return conditional_whitened(A_white, Knn, f, q_sqrt, full_cov=full_cov)
+    A = pt.linalg.inv(Kmm) @ Kmn
+    return conditional_unwhitened(A, Kmn, Knn, f, q_sqrt, full_cov=full_cov)

ptgp/gp/__init__.py

@@ -0,0 +1,11 @@
+from ptgp.gp.svgp import SVGP, VariationalParams, init_variational_params
+from ptgp.gp.unapproximated import Unapproximated
+from ptgp.gp.vfe import VFE
+
+__all__ = [
+    "Unapproximated",
+    "VFE",
+    "SVGP",
+    "VariationalParams",
+    "init_variational_params",
+]

ptgp/gp/svgp.py

@@ -0,0 +1,299 @@
+import dataclasses
+
+import numpy as np
+import pytensor.assumptions as pta
+import pytensor.tensor as pt
+
+from ptgp.conditionals import conditional_unwhitened, conditional_whitened
+from ptgp.kl import gauss_kl, gauss_kl_structured
+from ptgp.mean import Zero
+from ptgp.objectives import elbo
+
+
+def _softplus_lower_triangular(flat, M):
+    """Build M×M lower-triangular matrix from a flat vector, with softplus on the diagonal.
+
+    The flat vector has length M·(M+1)/2 in row-major lower-triangular layout
+    (matching ``np.tril_indices(M)``). Diagonal entries are passed through
+    ``pt.softplus`` so the result is a lower-triangular Cholesky factor with
+    strictly positive diagonal at every optimizer step.
+    """
+    rows, cols = np.tril_indices(M)
+    L = pt.set_subtensor(pt.zeros((M, M))[rows, cols], flat)
+    diag_idx = np.arange(M)
+    L = pt.set_subtensor(L[diag_idx, diag_idx], pt.softplus(L[diag_idx, diag_idx]))
+    return L
+
+
+def _matrix_to_softplus_flat_init(L_init, M):
+    """Inverse of ``_softplus_lower_triangular`` at init time.
+
+    Given an M×M lower-triangular Cholesky factor with strictly positive
+    diagonal, return the flat vector ``flat`` such that
+    ``_softplus_lower_triangular(flat, M)`` evaluates to ``L_init``.
+    """
+    if L_init.shape != (M, M):
+        raise ValueError(f"L_init must have shape ({M}, {M}); got {L_init.shape}.")
+    diag_vals = np.diag(L_init)
+    if np.any(diag_vals <= 0):
+        raise ValueError(
+            "L_init must have strictly positive diagonal (it represents a "
+            "Cholesky factor under softplus parameterization)."
+        )
+    rows, cols = np.tril_indices(M)
+    flat = L_init[rows, cols].astype(np.float64)
+    diag_positions = np.cumsum(np.arange(1, M + 1)) - 1
+    flat[diag_positions] = np.log(np.expm1(diag_vals.astype(np.float64)))
+    return flat
+
+
+@dataclasses.dataclass(frozen=True)
+class VariationalParams:
+    """Symbolic variational parameters for SVGP.
+
+    Attributes
+    ----------
+    q_mu : TensorVariable
+        Symbolic variational mean, shape ``(M,)``. Pass to ``SVGP`` via
+        ``variational_params=...``.
+    q_sqrt : TensorVariable
+        Symbolic Cholesky factor of the variational covariance, shape ``(M, M)``,
+        annotated ``lower_triangular=True``. Pass to ``SVGP`` via
+        ``variational_params=...``.
+    extra_vars : tuple of TensorVariable
+        Trainable leaves underlying ``q_mu`` and ``q_sqrt``.
+    extra_init : tuple of ndarray
+        Initial values for ``extra_vars``, in matching order.
+    """
+
+    q_mu: object
+    q_sqrt: object
+    extra_vars: tuple = ()
+    extra_init: tuple = ()
+
+
+def init_variational_params(M, q_mu_init=None, q_sqrt_init=None):
+    """Build symbolic variational parameters with GPJax-style parameterization.
+
+    ``q_sqrt`` is stored as a flat vector of length ``M·(M+1)/2`` and
+    materialised by filling a lower-triangular matrix and applying
+    ``pt.softplus`` to the diagonal — guarantees a true Cholesky factor
+    (lower-triangular with strictly positive diagonal) at every optimizer step.
+    ``q_mu`` is stored as a length-``M`` flat vector.
+
+    Parameters
+    ----------
+    M : int
+        Number of inducing points.
+    q_mu_init : ndarray of shape (M,), optional
+        Initial variational mean. Defaults to zeros.
+    q_sqrt_init : ndarray of shape (M, M), optional
+        Initial Cholesky factor. Must be lower-triangular with positive diagonal.
+        Defaults to the identity matrix.
+
+    Returns
+    -------
+    VariationalParams
+
+    Examples
+    --------
+    >>> vp = init_variational_params(M=8)
+    >>> svgp = SVGP(kernel=..., likelihood=..., inducing_variable=..., variational_params=vp)
+    >>> train_step, _, _ = compile_training_step(
+    ...     elbo,
+    ...     svgp,
+    ...     X,
+    ...     y,
+    ...     model=model,
+    ...     extra_vars=vp.extra_vars,
+    ...     extra_init=vp.extra_init,
+    ... )
+    """
+    if q_mu_init is None:
+        q_mu_init = np.zeros(M, dtype=np.float64)
+    else:
+        q_mu_init = np.asarray(q_mu_init, dtype=np.float64)
+        if q_mu_init.shape != (M,):
+            raise ValueError(f"q_mu_init must have shape ({M},); got {q_mu_init.shape}.")
+    if q_sqrt_init is None:
+        q_sqrt_init = np.eye(M, dtype=np.float64)
+    else:
+        q_sqrt_init = np.asarray(q_sqrt_init, dtype=np.float64)
+
+    n_lower = M * (M + 1) // 2
+    q_mu = pt.vector("q_mu", shape=(M,), dtype="float64")
+    q_sqrt_flat = pt.vector("q_sqrt_flat", shape=(n_lower,), dtype="float64")
+    q_sqrt = pta.assume(
+        _softplus_lower_triangular(q_sqrt_flat, M),
+        lower_triangular=True,
+    )
+    flat_init = _matrix_to_softplus_flat_init(q_sqrt_init, M)
+    return VariationalParams(
+        q_mu=q_mu,
+        q_sqrt=q_sqrt,
+        extra_vars=(q_mu, q_sqrt_flat),
+        extra_init=(q_mu_init, flat_init),
+    )
+
+
+class SVGP:
+    """Stochastic Variational Gaussian Process.
+
+    Parameters
+    ----------
+    kernel : Kernel
+        Covariance function.
+    mean : callable, optional
+        Mean function (default: Zero()).
+    likelihood : Likelihood
+        Observation likelihood.
+    inducing_variable : InducingVariables
+        Inducing point locations (or structured inducing variables).
+    variational_params : VariationalParams
+        Symbolic variational parameters and their backing trainable
+        leaves. Construct via :func:`init_variational_params` (recommended)
+        or build a :class:`VariationalParams` directly for custom
+        parameterizations. ``vp.extra_vars`` / ``vp.extra_init`` should be
+        passed to ``compile_training_step``.
+    whiten : bool
+        If True, use whitened variational parameterization (default True).
+    """
+
+    default_objective = staticmethod(elbo)
+    predict_needs_data = False
+
+    def __init__(
+        self,
+        kernel,
+        mean=None,
+        likelihood=None,
+        inducing_variable=None,
+        variational_params=None,
+        whiten=True,
+    ):
+        if variational_params is None:
+            raise ValueError(
+                "SVGP requires variational_params. Construct via "
+                "ptgp.gp.init_variational_params(M) and pass as "
+                "variational_params=..., or build a VariationalParams directly."
+            )
+        self.kernel = kernel
+        self.mean = mean if mean is not None else Zero()
+        self.likelihood = likelihood
+        self.inducing_variable = inducing_variable
+        self.whiten = whiten
+        self.variational_params = variational_params
+        self.q_mu = variational_params.q_mu
+        self.q_sqrt = variational_params.q_sqrt
+
+    @property
+    def extra_vars(self):
+        return (
+            *self.variational_params.extra_vars,
+            *self.inducing_variable.extra_vars,
+        )
+
+    @property
+    def extra_init(self):
+        return (
+            *self.variational_params.extra_init,
+            *self.inducing_variable.extra_init,
+        )
+
+    def predict_marginal(self, X, incl_lik=False):
+        """Posterior marginal mean and variance at each point in X.
+
+        Returns the per-point posterior — correlations between test
+        points are discarded. Use ``predict_joint`` for the full (N, N)
+        covariance or ``predict_f_samples`` to draw smooth function samples.
+
+        Parameters
+        ----------
+        X : tensor, shape (N, D)
+        incl_lik : bool
+            If True, push through the likelihood's predictive mean/var
+            (observation-space uncertainty).
+
+        Returns
+        -------
+        mean : tensor, shape (N,)
+        var : tensor, shape (N,)
+        """
+        ind, kernel = self.inducing_variable, self.kernel
+        Kmn = ind.K_uf(kernel, X)
+        Knn_diag = kernel.diag(X)
+        if self.whiten:
+            A_w = ind.Kuu_sqrt_solve(kernel, Kmn)
+            fmean, fvar = conditional_whitened(A_w, Knn_diag, self.q_mu, self.q_sqrt)
+        else:
+            A = ind.Kuu_solve(kernel, Kmn)
+            fmean, fvar = conditional_unwhitened(A, Kmn, Knn_diag, self.q_mu, self.q_sqrt)
+        fmean = fmean + self.mean(X)
+        if incl_lik:
+            return self.likelihood.predict_mean_and_var(fmean, fvar)
+        return fmean, fvar
+
+    def predict_joint(self, X):
+        """Posterior joint mean and full covariance of the latent f at X.
+
+        The diagonal of the returned covariance equals the variance from
+        ``predict_marginal``; the off-diagonals capture correlations
+        between test points, which are needed to draw smooth samples.
+
+        Parameters
+        ----------
+        X : tensor, shape (N, D)
+
+        Returns
+        -------
+        mean : tensor, shape (N,)
+        cov : tensor, shape (N, N)
+        """
+        ind, kernel = self.inducing_variable, self.kernel
+        Kmn = ind.K_uf(kernel, X)
+        Knn = kernel(X)
+        if self.whiten:
+            A_w = ind.Kuu_sqrt_solve(kernel, Kmn)
+            fmean, fcov = conditional_whitened(A_w, Knn, self.q_mu, self.q_sqrt, full_cov=True)
+        else:
+            A = ind.Kuu_solve(kernel, Kmn)
+            fmean, fcov = conditional_unwhitened(A, Kmn, Knn, self.q_mu, self.q_sqrt, full_cov=True)
+        fmean = fmean + self.mean(X)
+        return fmean, fcov
+
+    def predict_f_samples(self, X, epsilon, jitter=1e-6):
+        """Draw samples of the latent f at X from the joint posterior.
+
+        Samples are produced via a Cholesky transform of caller-supplied
+        iid-standard-normal noise. The caller owns the RNG — the function
+        itself is deterministic.
+
+        Parameters
+        ----------
+        X : tensor, shape (N, D)
+        epsilon : tensor, shape (S, N)
+            iid N(0, 1) draws. ``S`` is the number of samples.
+        jitter : float
+            Added to the diagonal of the posterior covariance before
+            Cholesky, for numerical stability (default 1e-6).
+
+        Returns
+        -------
+        samples : tensor, shape (S, N)
+        """
+        fmean, fcov = self.predict_joint(X)
+        N = X.shape[0]
+        L = pt.linalg.cholesky(fcov + jitter * pt.eye(N))
+        return fmean[None, :] + epsilon @ L.T
+
+    def prior_kl(self):
+        """KL divergence KL[q(u) || p(u)]."""
+        if self.whiten:
+            return gauss_kl(self.q_mu, self.q_sqrt, K=None)
+        ind, kernel = self.inducing_variable, self.kernel
+        return gauss_kl_structured(
+            self.q_mu,
+            self.q_sqrt,
+            K_solve=lambda rhs: ind.Kuu_solve(kernel, rhs),
+            K_logdet=ind.Kuu_logdet(kernel),
+        )

ptgp/gp/unapproximated.py

@@ -0,0 +1,69 @@
+import pytensor.tensor as pt
+
+from ptgp.likelihoods import Gaussian
+from ptgp.mean import Zero
+from ptgp.objectives import marginal_log_likelihood
+
+
+class Unapproximated:
+    """Exact (unapproximated) Gaussian process.
+
+    The observation model is Gaussian; parameterize the noise via ``sigma``.
+
+    Parameters
+    ----------
+    kernel : Kernel
+        Covariance function.
+    mean : callable, optional
+        Mean function (default: ``Zero()``).
+    sigma : tensor or PyMC random variable
+        Observation noise standard deviation.
+    """
+
+    extra_vars = ()
+    extra_init = ()
+    default_objective = staticmethod(marginal_log_likelihood)
+    predict_needs_data = True
+
+    def __init__(self, kernel, mean=None, sigma=None):
+        """Store the kernel and mean; build a Gaussian likelihood from sigma."""
+        self.kernel = kernel
+        self.mean = mean if mean is not None else Zero()
+        self.likelihood = Gaussian(sigma)
+
+    def predict_marginal(self, X_new, X_train, y_train, incl_lik=False):
+        """Posterior marginal mean and variance at each point in X_new.
+
+        Returns the per-point posterior; correlations between test points
+        are discarded.
+
+        Parameters
+        ----------
+        X_new : tensor, shape (N*, D)
+        X_train : tensor, shape (N, D)
+        y_train : tensor, shape (N,)
+        incl_lik : bool
+            If True, include likelihood noise in the predictions.
+
+        Returns
+        -------
+        mean : tensor, shape (N*,)
+        var : tensor, shape (N*,)
+        """
+        Knn = self.kernel(X_train)
+        sigma = self.likelihood.sigma
+        sigma2_train = sigma**2 * pt.ones(X_train.shape[0])  # (N,); scalar broadcasts
+        Knn_noisy = Knn + pt.diag(sigma2_train)
+        Kns = self.kernel(X_train, X_new)  # (N, N*)
+        Kss_diag = self.kernel.diag(X_new)
+
+        Knn_inv = pt.linalg.inv(Knn_noisy)
+
+        mu_train = self.mean(X_train)
+        fmean = self.mean(X_new) + Kns.T @ Knn_inv @ (y_train - mu_train)
+        fvar = Kss_diag - pt.sum(Kns * (Knn_inv @ Kns), axis=0)
+
+        if incl_lik:
+            sigma_new = self.likelihood.sigma_at(X_train, X_new)
+            return fmean, fvar + sigma_new**2
+        return fmean, fvar

ptgp/gp/vfe.py

@@ -0,0 +1,92 @@
+import pytensor.tensor as pt
+
+from ptgp.likelihoods import Gaussian
+from ptgp.mean import Zero
+from ptgp.objectives import collapsed_elbo
+
+
+class VFE:
+    """Variational Free Energy (SGPR) sparse Gaussian Process.
+
+    Uses Titsias' collapsed bound — inducing variables are analytically
+    integrated out. The observation model is Gaussian; parameterize the noise
+    via ``sigma``.
+
+    Parameters
+    ----------
+    kernel : Kernel
+        Covariance function.
+    mean : callable, optional
+        Mean function (default: ``Zero()``).
+    sigma : tensor or PyMC random variable
+        Observation noise standard deviation.
+    inducing_variable : InducingVariables
+        Inducing point locations.
+    """
+
+    default_objective = staticmethod(collapsed_elbo)
+    predict_needs_data = True
+
+    def __init__(self, kernel, mean=None, sigma=None, inducing_variable=None):
+        """Store the kernel, mean, and inducing variable; build a Gaussian likelihood from sigma."""
+        if not hasattr(inducing_variable, "Z"):
+            raise TypeError(
+                f"VFE requires inducing variables with a .Z attribute "
+                f"(got {type(inducing_variable).__name__}). "
+                f"Use SVGP for structured inducing variables like FourierFeatures1D."
+            )
+        self.kernel = kernel
+        self.mean = mean if mean is not None else Zero()
+        self.likelihood = Gaussian(sigma)
+        self.inducing_variable = inducing_variable
+
+    @property
+    def extra_vars(self):
+        return tuple(self.inducing_variable.extra_vars)
+
+    @property
+    def extra_init(self):
+        return tuple(self.inducing_variable.extra_init)
+
+    def predict_marginal(self, X_new, X_train, y_train, incl_lik=False):
+        """Posterior marginal mean and variance at each point in X_new.
+
+        Returns the per-point posterior; correlations between test points
+        are discarded.
+
+        Parameters
+        ----------
+        X_new : tensor, shape (N*, D)
+        X_train : tensor, shape (N, D)
+        y_train : tensor, shape (N,)
+        incl_lik : bool
+            If True, include likelihood noise in the predictions.
+
+        Returns
+        -------
+        mean : tensor, shape (N*,)
+        var : tensor, shape (N*,)
+        """
+        Z = self.inducing_variable.Z
+        sigma = self.likelihood.sigma
+        sigma2_vec = sigma**2 * pt.ones(X_train.shape[0])  # (N,); scalar broadcasts
+
+        Kuu = self.kernel(Z)  # (M, M)
+        Kuf = self.kernel(Z, X_train)  # (M, N)
+        Kus = self.kernel(Z, X_new)  # (M, N*)
+        Kss_diag = self.kernel.diag(X_new)
+
+        diff = y_train - self.mean(X_train)
+        Kuf_laminv = Kuf / sigma2_vec[None, :]  # each column ÷ σᵢ²
+        Sigma = Kuu + Kuf_laminv @ Kuf.T
+        Sigma_inv = pt.linalg.inv(Sigma)
+        alpha = Sigma_inv @ Kuf_laminv @ diff
+
+        fmean = self.mean(X_new) + Kus.T @ alpha
+        Kuu_inv = pt.linalg.inv(Kuu)
+        fvar = Kss_diag - pt.sum(Kus * ((Kuu_inv - Sigma_inv) @ Kus), axis=0)
+
+        if incl_lik:
+            sigma_new = self.likelihood.sigma_at(X_train, X_new)
+            return fmean, fvar + sigma_new**2
+        return fmean, fvar