diff --git a/ptgp/conditionals.py b/ptgp/conditionals.py index b3fb84c..aeba4ca 100644 --- a/ptgp/conditionals.py +++ b/ptgp/conditionals.py @@ -55,7 +55,7 @@ def conditional_unwhitened(A, Kmn, Knn, f, q_sqrt=None, full_cov=False): 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. + # 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. diff --git a/ptgp/gp/__init__.py b/ptgp/gp/__init__.py index c342fb8..8d69902 100644 --- a/ptgp/gp/__init__.py +++ b/ptgp/gp/__init__.py @@ -1,11 +1,15 @@ from ptgp.gp.svgp import SVGP, VariationalParams, init_variational_params from ptgp.gp.unapproximated import Unapproximated from ptgp.gp.vfe import VFE +from ptgp.gp.vgp import VGP, VGPParams, init_vgp_params __all__ = [ "Unapproximated", "VFE", "SVGP", + "VGP", "VariationalParams", "init_variational_params", + "VGPParams", + "init_vgp_params", ] diff --git a/ptgp/gp/svgp.py b/ptgp/gp/svgp.py index cdfa52b..84c9ed3 100644 --- a/ptgp/gp/svgp.py +++ b/ptgp/gp/svgp.py @@ -77,7 +77,7 @@ def init_variational_params(M, q_mu_init=None, q_sqrt_init=None): ``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 + ``pt.softplus`` to the diagonal, which 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. @@ -203,7 +203,7 @@ def extra_init(self): 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 + 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. @@ -265,7 +265,7 @@ 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 + iid-standard-normal noise. The caller owns the RNG; the function itself is deterministic. Parameters diff --git a/ptgp/gp/vfe.py b/ptgp/gp/vfe.py index 3c2550a..1d6f155 100644 --- a/ptgp/gp/vfe.py +++ b/ptgp/gp/vfe.py @@ -8,7 +8,7 @@ class VFE: """Variational Free Energy (SGPR) sparse Gaussian Process. - Uses Titsias' collapsed bound — inducing variables are analytically + Uses Titsias' collapsed bound; inducing variables are analytically integrated out. The observation model is Gaussian; parameterize the noise via ``sigma``. diff --git a/ptgp/gp/vgp.py b/ptgp/gp/vgp.py new file mode 100644 index 0000000..e091dbc --- /dev/null +++ b/ptgp/gp/vgp.py @@ -0,0 +1,225 @@ +import dataclasses + +import numpy as np +import pytensor.assumptions as pta +import pytensor.tensor as pt + +from ptgp.mean import Zero +from ptgp.objectives import vgp_elbo + + +@dataclasses.dataclass(frozen=True) +class VGPParams: + """Symbolic Opper-Archambeau variational parameters for VGP. + + Attributes + ---------- + alpha : TensorVariable + Mean weights, shape ``(N,)``. The posterior mean is ``mean(X) + K alpha``. + lam : TensorVariable + Posterior precision diagonal, shape ``(N,)``, strictly positive + (materialised as ``softplus(lambda_raw)``). + extra_vars : tuple of TensorVariable + Trainable leaves underlying ``alpha`` and ``lam``. + extra_init : tuple of ndarray + Initial values for ``extra_vars``, in matching order. + """ + + alpha: object + lam: object + extra_vars: tuple = () + extra_init: tuple = () + + +def init_vgp_params(N, alpha_init=None, lambda_init=None): + """Build symbolic Opper-Archambeau variational parameters sized to N. + + The posterior precision diagonal ``lam`` is stored as ``lambda_raw`` and + materialised via ``pt.softplus`` so it is strictly positive at every + optimizer step (matching the softplus convention in + :func:`ptgp.gp.init_variational_params`). ``alpha`` is a free length-``N`` + vector. + + Parameters + ---------- + N : int + Number of training points (the variational posterior is over the latent + function at these inputs). + alpha_init : ndarray of shape (N,), optional + Initial mean weights. Defaults to zeros. + lambda_init : ndarray of shape (N,), optional + Initial precision diagonal, strictly positive. Defaults to ones. + + Returns + ------- + VGPParams + + Examples + -------- + >>> vp = init_vgp_params(N=10) + >>> vgp = VGP(kernel=..., likelihood=..., variational_params=vp) + """ + if alpha_init is None: + alpha_init = np.zeros(N, dtype=np.float64) + else: + alpha_init = np.asarray(alpha_init, dtype=np.float64) + if alpha_init.shape != (N,): + raise ValueError(f"alpha_init must have shape ({N},); got {alpha_init.shape}.") + + if lambda_init is None: + lambda_init = np.ones(N, dtype=np.float64) + else: + lambda_init = np.asarray(lambda_init, dtype=np.float64) + if lambda_init.shape != (N,): + raise ValueError(f"lambda_init must have shape ({N},); got {lambda_init.shape}.") + if np.any(lambda_init <= 0): + raise ValueError("lambda_init must be strictly positive (it is a precision diagonal).") + + alpha = pt.vector("alpha", shape=(N,), dtype="float64") + lambda_raw = pt.vector("lambda_raw", shape=(N,), dtype="float64") + lam = pt.softplus(lambda_raw) + # Inverse softplus: lambda_raw such that softplus(lambda_raw) == lambda_init. + lambda_raw_init = np.log(np.expm1(lambda_init)) + return VGPParams( + alpha=alpha, + lam=lam, + extra_vars=(alpha, lambda_raw), + extra_init=(alpha_init, lambda_raw_init), + ) + + +class VGP: + """Full variational Gaussian Process (Opper-Archambeau). + + The variational posterior q(f) = N(mean(X) + K alpha, S) is placed directly + over the latent function at the N training inputs, with + S = (K^{-1} + diag(lambda))^{-1}. There are no inducing points; the free + parameters are ``alpha`` (N,) and the precision diagonal ``lambda`` (N,), + so the posterior costs O(N) parameters (Opper & Archambeau, 2009). + + All computation routes through the factor A = I + G K G, + ``G = diag(sqrt(lambda))``, whose eigenvalues are >= 1, so K itself is never + factorised or inverted and no jitter is required. + + Parameters + ---------- + kernel : Kernel + Covariance function. + mean : callable, optional + Mean function (default: ``Zero()``). + likelihood : Likelihood + Observation likelihood (Gaussian or any non-Gaussian likelihood). + variational_params : VGPParams + Symbolic variational parameters sized to N. Construct via + :func:`init_vgp_params` (recommended). ``vp.extra_vars`` / + ``vp.extra_init`` are surfaced for the optimizer. + """ + + default_objective = staticmethod(vgp_elbo) + predict_needs_data = True + + def __init__(self, kernel, mean=None, likelihood=None, variational_params=None): + """Store the kernel, mean, likelihood, and variational parameters.""" + if variational_params is None: + raise ValueError( + "VGP requires variational_params. Construct via " + "ptgp.gp.init_vgp_params(N) and pass as variational_params=..., " + "where N is the number of training points." + ) + self.kernel = kernel + self.mean = mean if mean is not None else Zero() + self.likelihood = likelihood + self.variational_params = variational_params + self.alpha = variational_params.alpha + self.lam = variational_params.lam + + @property + def extra_vars(self): + """Trainable variational leaves (alpha, lambda_raw).""" + return tuple(self.variational_params.extra_vars) + + @property + def extra_init(self): + """Initial values for the trainable variational leaves.""" + return tuple(self.variational_params.extra_init) + + def _build_A(self, X): + """Return ``(K, g, L)`` where K = kernel(X), g = sqrt(lambda), and L is + the Cholesky factor of A = I + G K G (G = diag(g)). + + A is symmetric positive definite (eigenvalues >= 1) for any PSD K, so no + jitter is needed and K is never factorised on its own. + """ + K = self.kernel(X) + n = K.shape[0] + g = pt.sqrt(self.lam) + A = pt.eye(n, dtype=K.dtype) + g[:, None] * K * g[None, :] + A = pta.assume(A, positive_definite=True, symmetric=True) + L = pt.linalg.cholesky(A) + return K, g, L + + def _train_marginals(self, X): + """Marginal mean and variance of q(f) at the training inputs. + + ``fmean = mean(X) + K alpha`` and ``fvar = diag(S)`` where + ``diag(S) = diag(K) - colsum((L^{-1} G K)**2)``. + + Returns + ------- + fmean : tensor, shape (N,) + fvar : tensor, shape (N,) + """ + K, g, L = self._build_A(X) + fmean = self.mean(X) + K @ self.alpha + B = pt.linalg.solve_triangular(L, g[:, None] * K, lower=True) + fvar = pt.diagonal(K) - pt.sum(B**2, axis=0) + return fmean, fvar + + def prior_kl(self, X): + """KL[q(f) || p(f)] in closed form. + + ``KL = 0.5 * (tr(A^{-1}) + alpha^T K alpha - N + log|A|)``, using + ``tr(K^{-1} S) = tr(A^{-1})`` and ``log|K| - log|S| = log|A|``. + + Returns + ------- + scalar + """ + K, g, L = self._build_A(X) + n = K.shape[0] + logdet_A = 2.0 * pt.sum(pt.log(pt.diagonal(L))) + Linv = pt.linalg.solve_triangular(L, pt.eye(n, dtype=K.dtype), lower=True) + tr_Ainv = pt.sum(Linv**2) + quad = self.alpha @ (K @ self.alpha) + return 0.5 * (tr_Ainv + quad - n + logdet_A) + + def predict_marginal(self, X_new, X_train, y_train=None, 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. ``y_train`` is accepted for API compatibility with + :func:`ptgp.optim.predict` but ignored: the posterior is carried by + ``alpha``/``lambda``, not recomputed from the targets. + + Parameters + ---------- + X_new : tensor, shape (N*, D) + X_train : tensor, shape (N, D) + y_train : tensor, optional + Ignored. + incl_lik : bool + If True, push through the likelihood's predictive mean/var. + + Returns + ------- + mean : tensor, shape (N*,) + var : tensor, shape (N*,) + """ + K, g, L = self._build_A(X_train) + Kfnew = self.kernel(X_train, X_new) # (N, N*) + fmean = self.mean(X_new) + Kfnew.T @ self.alpha + C = pt.linalg.solve_triangular(L, g[:, None] * Kfnew, lower=True) + fvar = self.kernel.diag(X_new) - pt.sum(C**2, axis=0) + if incl_lik: + return self.likelihood.predict_mean_and_var(fmean, fvar) + return fmean, fvar diff --git a/ptgp/idata.py b/ptgp/idata.py index c1954d5..d0248fb 100644 --- a/ptgp/idata.py +++ b/ptgp/idata.py @@ -23,18 +23,18 @@ def to_idata( Groups produced (empty groups are omitted): - - ``point_estimate`` — every optimized value in constrained space, on its + - ``point_estimate``: every optimized value in constrained space, on its own dims with no ``chain``/``draw`` axes (nothing was sampled). Includes PyMC RVs (``sigma``, ``eta``, …) *and* non-PyMC extras (``Z``, ``q_mu``, ``q_sqrt_flat``, …). - - ``unconstrained_point_estimate`` — the same values in value-var + - ``unconstrained_point_estimate``: the same values in value-var (transformed) space. - - ``optimizer_result`` — scalar fields from the scipy ``OptimizeResult`` + - ``optimizer_result``: scalar fields from the scipy ``OptimizeResult`` (``fun``, ``success``, ``status``, ``message``, ``nit``, ``nfev``, ``njev``) plus per-iteration trajectory DataArrays from ``history`` (``elbo``, ``trace_penalty``, …) on an ``iteration`` dimension. There is no ``sample_stats``: nothing was sampled. - - ``observed_data`` / ``constant_data`` — pulled from the PyMC model. + - ``observed_data`` / ``constant_data``: pulled from the PyMC model. Parameters ---------- @@ -103,7 +103,7 @@ def _collect_optimized(model, shared_params, shared_extras): PyMC RVs contribute to both dicts: the constrained dict uses ``rv.name`` and applies the backward transform; the unconstrained dict uses ``vv.name`` and - keeps the value-var-space value. Non-PyMC extras have no transform — they + keeps the value-var-space value. Non-PyMC extras have no transform, so they appear in both dicts under ``sv.name`` so callers can find them in either group. """ @@ -154,7 +154,7 @@ def _build_optimizer_result(result, history, phase_labels): populates ``elbo``, ``trace_penalty``, …). If a history field name collides with a result-scalar name, the trajectory - wins — it's strictly more informative than the terminal value. + wins, it's strictly more informative than the terminal value. """ data_vars = {} coords = {} diff --git a/ptgp/inducing.py b/ptgp/inducing.py index 4f9192f..53101b8 100644 --- a/ptgp/inducing.py +++ b/ptgp/inducing.py @@ -111,9 +111,9 @@ class KernelHealthDiagnostics: Per-data-point residual conditional variance ``Kff_diag - Q_diag``. Large values identify points poorly covered by ``Z``. total_variance : float - ``tr(Kff)`` — the kernel diagonal sum on ``X`` before conditioning. + ``tr(Kff)``, the kernel diagonal sum on ``X`` before conditioning. nystrom_residual : float - ``sum(d_final) = tr(Kff - Q)`` — total unexplained variance. + ``sum(d_final) = tr(Kff - Q)``, total unexplained variance. kuu_min_eigenvalue : float Smallest eigenvalue of ``K(Z, Z) + jitter * I``. kuu_max_eigenvalue : float @@ -159,7 +159,7 @@ class RandomSubsampleDiagnostics: The ``M`` argument. M_returned : int Rows in the returned ``Z``. Always equals ``M_requested`` for this - routine — kept for API symmetry with the other init routines. + routine, kept for API symmetry with the other init routines. N_candidates : int Rows in the input ``X``. n_unique : int @@ -223,7 +223,7 @@ class KMeansDiagnostics: dedup_tol : float The Euclidean-distance threshold used for deduplication. inertia : float - ``sum_i ||X[i] - centroid[label[i]]||^2`` — k-means' standard + ``sum_i ||X[i] - centroid[label[i]]||^2``, k-means' standard within-cluster sum of squares. Smaller is tighter. pairwise_min_distance : float Minimum Euclidean distance between returned centroids. ``nan`` if @@ -529,7 +529,7 @@ def greedy_variance_init( Implements the "ConditionalVariance" initialization of Burt et al. (2020), *Convergence of Sparse Variational Inference in GP Regression*. At each step, the next inducing point is the row of ``X`` with largest remaining - conditional variance given the already-selected points — equivalent to + conditional variance given the already-selected points, equivalent to running a partial pivoted Cholesky decomposition of ``K(X, X)`` with the standard max-diagonal pivot rule. Selected points are a **subset of X**; this is discrete subset selection, not continuous optimization. @@ -539,7 +539,7 @@ def greedy_variance_init( Recommended workflow -------------------- Burt et al. show that with a good greedy initialization, ``Z`` typically - does **not** need to be gradient-optimized during training — for most + does **not** need to be gradient-optimized during training. For most problems the frozen subset is within noise of jointly-optimized ``Z`` at a tiny fraction of the compute. The standard recipe: @@ -682,7 +682,7 @@ def compute_inducing_diagnostics( Computes the same kernel-derived metrics that :func:`greedy_variance_init` reports (Kuu eigenvalues, per-point residual conditional variance, total variance), but for a ``Z`` that - you already have — from k-means, manual placement, post-training, or + you already have, from k-means, manual placement, post-training, or elsewhere. No greedy selection is performed; no ``trace_curve`` history is available. diff --git a/ptgp/inducing_fourier.py b/ptgp/inducing_fourier.py index bac59d9..755d0f1 100644 --- a/ptgp/inducing_fourier.py +++ b/ptgp/inducing_fourier.py @@ -360,7 +360,7 @@ def Kuu_logdet(self, kernel): def Kuu_sqrt_solve(self, kernel, rhs): """Apply ``R^{-1}`` where ``R @ R.T = Kuu = diag(d) + U @ U.T``. ``rhs`` is ``(M, K)``. - Uses ``delta = -1 / (sqrt(1 + lam) * (1 + sqrt(1 + lam)))`` — algebraically + Uses ``delta = -1 / (sqrt(1 + lam) * (1 + sqrt(1 + lam)))``, algebraically equivalent to ``1/sqrt(1+lam) - 1`` divided by ``lam`` but with no division by ``lam`` or ``sqrt(lam)``, so it stays finite as ``lam -> 0``. """ diff --git a/ptgp/kernels/base.py b/ptgp/kernels/base.py index e51ae43..760236b 100644 --- a/ptgp/kernels/base.py +++ b/ptgp/kernels/base.py @@ -38,7 +38,7 @@ def __call__(self, X, Y=None): return self._eval(X, Y) def _eval(self, X, Y): - """Raw kernel matrix — subclasses implement.""" + """Raw kernel matrix; subclasses implement.""" raise NotImplementedError def diag(self, X): diff --git a/ptgp/kernels/categorical.py b/ptgp/kernels/categorical.py index ca2c848..e271703 100644 --- a/ptgp/kernels/categorical.py +++ b/ptgp/kernels/categorical.py @@ -10,7 +10,7 @@ class Overlap(Kernel): where D is the number of active categorical columns. The kernel returns the fraction of active columns whose levels match. With one active column, - this is simply ``1[x == y]``. + this is ``1[x == y]``. Values are expected to be non-negative integer level codes stored in a float matrix (the library-wide convention). Entries are cast to ``int64`` @@ -50,7 +50,7 @@ class LowRankCategorical(Kernel): kernel. Operates on a single categorical column. Level codes are expected to be non-negative integers stored in a float matrix (the library-wide convention). Codes are cast to ``int64`` inside ``_eval`` and ``diag``; - they are not range-checked — passing a code >= ``num_levels`` is user + they are not range-checked; passing a code >= ``num_levels`` is user error. Parameters diff --git a/ptgp/kernels/combination.py b/ptgp/kernels/combination.py index 62a6356..f9819fa 100644 --- a/ptgp/kernels/combination.py +++ b/ptgp/kernels/combination.py @@ -27,7 +27,7 @@ def _eval(self, X, Y): return self.k1(X, Y) + self.k2(X, Y) def diag(self, X): - """Return the per-point variance — sum of the two component diagonals.""" + """Per-point variance: sum of the two component diagonals.""" return self.k1.diag(X) + self.k2.diag(X) @@ -56,7 +56,7 @@ def _eval(self, X, Y): return self.k1 * self.k2(X, Y) def diag(self, X): - """Return the per-point variance — product of component diagonals, with scalar passthrough.""" + """Per-point variance: product of component diagonals, with scalar passthrough.""" k1_is_kernel = isinstance(self.k1, Kernel) k2_is_kernel = isinstance(self.k2, Kernel) if k1_is_kernel and k2_is_kernel: diff --git a/ptgp/kernels/nonstationary.py b/ptgp/kernels/nonstationary.py index 94d443b..f9526f2 100644 --- a/ptgp/kernels/nonstationary.py +++ b/ptgp/kernels/nonstationary.py @@ -136,7 +136,7 @@ class WarpedInput(Kernel): ---------- input_dim : int Number of columns of ``X`` this kernel expects. May differ from - ``kernel_func.input_dim`` — the warp may change dimensionality. + ``kernel_func.input_dim``; the warp may change dimensionality. kernel_func : Kernel Inner kernel applied to the warped inputs. warp_func : callable diff --git a/ptgp/likelihoods/__init__.py b/ptgp/likelihoods/__init__.py index f51989b..7eaf22d 100644 --- a/ptgp/likelihoods/__init__.py +++ b/ptgp/likelihoods/__init__.py @@ -1,8 +1,17 @@ from ptgp.likelihoods.base import Likelihood from ptgp.likelihoods.bernoulli import Bernoulli +from ptgp.likelihoods.composite import CompositeLikelihood from ptgp.likelihoods.gaussian import Gaussian from ptgp.likelihoods.negative_binomial import NegativeBinomial from ptgp.likelihoods.poisson import Poisson from ptgp.likelihoods.student_t import StudentT -__all__ = ["Likelihood", "Gaussian", "Bernoulli", "StudentT", "Poisson", "NegativeBinomial"] +__all__ = [ + "Likelihood", + "Gaussian", + "Bernoulli", + "StudentT", + "Poisson", + "NegativeBinomial", + "CompositeLikelihood", +] diff --git a/ptgp/likelihoods/base.py b/ptgp/likelihoods/base.py index 24771d6..4842018 100644 --- a/ptgp/likelihoods/base.py +++ b/ptgp/likelihoods/base.py @@ -60,7 +60,7 @@ def predict_mean_and_var(self, mu, var): return E_mean, E_var + E_mean_sq - E_mean**2 def predict_log_density(self, y, mu, var): - """log E_{q(f)}[p(y|f)] — predictive log-density at test points. + """log E_{q(f)}[p(y|f)], the predictive log-density at test points. Default: Gauss-Hermite quadrature in log-space for numerical stability. """ diff --git a/ptgp/likelihoods/composite.py b/ptgp/likelihoods/composite.py new file mode 100644 index 0000000..ea4a77a --- /dev/null +++ b/ptgp/likelihoods/composite.py @@ -0,0 +1,80 @@ +import numpy as np +import pytensor.tensor as pt + +from ptgp.likelihoods.base import Likelihood + + +class CompositeLikelihood(Likelihood): + """Apply a different likelihood to different observations. + + The data is partitioned into disjoint subsets, each governed by its own + sub-likelihood. Because the ELBO data term is a sum over independent points, + the combined ``variational_expectation`` is just the per-subset expectations + scattered back into a length-N vector, and similarly for prediction. This + composes directly with the VGP (and SVGP) objectives. + + Parameters + ---------- + likelihoods : list of Likelihood + One sub-likelihood per subset. + indices : list of array-like of int + Integer index arrays, one per sub-likelihood, that partition + ``range(N)`` (disjoint, exhaustive, non-empty). ``indices[k]`` selects + the observations governed by ``likelihoods[k]``. + + Notes + ----- + A sub-likelihood reads its own parameters internally. If a sub-likelihood + carries a per-point parameter (for example a vector ``sigma`` on a + heteroskedastic :class:`~ptgp.likelihoods.Gaussian`), that parameter must + already be aligned to its subset, since the combinator slices only ``y``, + ``mu``, and ``var``, not the sub-likelihood's internal tensors. + """ + + def __init__(self, likelihoods, indices): + """Validate that ``indices`` partition ``range(N)`` and store the parts.""" + if len(likelihoods) != len(indices): + raise ValueError( + f"likelihoods and indices must have equal length; got " + f"{len(likelihoods)} and {len(indices)}." + ) + idx_arrays = [np.asarray(idx, dtype=np.intp) for idx in indices] + for a in idx_arrays: + if a.ndim != 1 or a.size == 0: + raise ValueError("each entry of indices must be a non-empty 1D integer array.") + concat = np.concatenate(idx_arrays) + N = concat.size + if not np.array_equal(np.sort(concat), np.arange(N)): + raise ValueError( + "indices must partition range(N): disjoint, exhaustive, and covering 0..N-1." + ) + self.likelihoods = list(likelihoods) + self._indices = idx_arrays + self.num_data = N + + def variational_expectation(self, y, mu, var): + """Per-point E_{q(f)}[log p(y|f)], dispatched by subset. + + Returns a length-N vector; the VGP/SVGP objective sums it. + """ + out = pt.zeros_like(mu) + for idx, lik in zip(self._indices, self.likelihoods): + out = pt.set_subtensor(out[idx], lik.variational_expectation(y[idx], mu[idx], var[idx])) + return out + + def predict_mean_and_var(self, mu, var): + """Predictive mean and variance, dispatched by subset.""" + mean_out = pt.zeros_like(mu) + var_out = pt.zeros_like(var) + for idx, lik in zip(self._indices, self.likelihoods): + m_i, v_i = lik.predict_mean_and_var(mu[idx], var[idx]) + mean_out = pt.set_subtensor(mean_out[idx], m_i) + var_out = pt.set_subtensor(var_out[idx], v_i) + return mean_out, var_out + + def predict_log_density(self, y, mu, var): + """Predictive log-density at test points, dispatched by subset.""" + out = pt.zeros_like(mu) + for idx, lik in zip(self._indices, self.likelihoods): + out = pt.set_subtensor(out[idx], lik.predict_log_density(y[idx], mu[idx], var[idx])) + return out diff --git a/ptgp/likelihoods/gaussian.py b/ptgp/likelihoods/gaussian.py index ef2cb8a..dbdcca7 100644 --- a/ptgp/likelihoods/gaussian.py +++ b/ptgp/likelihoods/gaussian.py @@ -33,7 +33,7 @@ def sigma_at(self, X_train, X_new): Sigma is treated as data-dependent iff it has any free symbolic tensor input in its graph (covers both raw ``pt.matrix`` X and - ``pm.Data`` X — the latter is a SharedVariable but its type is + ``pm.Data`` X, the latter is a SharedVariable but its type is ``TensorType``). For data-independent sigma the call is a no-op. For data-dependent sigma, strict-mode ``graph_replace`` raises a clear error if ``X_train`` is not in sigma's graph. diff --git a/ptgp/linalg/operator.py b/ptgp/linalg/operator.py index ce97521..635d20a 100644 --- a/ptgp/linalg/operator.py +++ b/ptgp/linalg/operator.py @@ -1,4 +1,4 @@ -# LinearOperatorType — placeholder for future PyTensor Op implementation. +# LinearOperatorType: placeholder for a future PyTensor Op. # Will define a new PyTensor type representing a matrix implicitly via # its matrix-vector product, enabling CG solves and Lanczos log-det # without materialising the full N x N kernel matrix. diff --git a/ptgp/linalg/ops.py b/ptgp/linalg/ops.py index d0f3154..ffc1283 100644 --- a/ptgp/linalg/ops.py +++ b/ptgp/linalg/ops.py @@ -1,5 +1,5 @@ # Custom PyTensor Ops for lazy linear algebra. -# KernelLinearOp — represents K(X, X) as a LinearOperatorType -# LinearOpMatvec — computes A @ v without materialising A -# LinearOpSolve — CG-based solve: A x = b -# LinearOpLogdet — Lanczos-based log-determinant estimator +# KernelLinearOp: represents K(X, X) as a LinearOperatorType +# LinearOpMatvec: computes A @ v without materialising A +# LinearOpSolve: CG-based solve, A x = b +# LinearOpLogdet: Lanczos-based log-determinant estimator diff --git a/ptgp/objectives.py b/ptgp/objectives.py index 32c0f4f..12e3209 100644 --- a/ptgp/objectives.py +++ b/ptgp/objectives.py @@ -10,6 +10,7 @@ "CollapsedELBOTerms", ["elbo", "fit", "trace_penalty", "nystrom_residual"] ) FITCTerms = namedtuple("FITCTerms", ["fitc", "fit", "logdet"]) +VGPTerms = namedtuple("VGPTerms", ["elbo", "var_exp", "kl"]) # Diagonal jitter added to Kuu before Cholesky / inversion, to keep it PSD @@ -82,6 +83,37 @@ def elbo(svgp, X, y, n_data=None): return ELBOTerms(elbo=var_exp - kl, var_exp=var_exp, kl=kl) +def vgp_elbo(vgp, X, y, n_data=None): + """Opper-Archambeau variational GP evidence lower bound. + + ELBO = E_{q(f)}[log p(y|f)] - KL[q(f) || p(f)] + + The variational posterior q(f) = N(m + K alpha, S) is placed directly over + the latent function at the N training inputs (no inducing points), with + S = (K^{-1} + diag(lambda))^{-1}. The marginals and the KL are computed by + ``vgp`` through the well-conditioned factor A = I + G K G. + + Parameters + ---------- + vgp : VGP + Full variational GP model. + X : tensor, shape (N, D) + y : tensor, shape (N,) + n_data : int, optional + Accepted for signature parity with :func:`elbo`. VGP is full-batch, so + this is a no-op (minibatching is not meaningful when q is over all N + latent values). + + Returns + ------- + VGPTerms + """ + fmean, fvar = vgp._train_marginals(X) + var_exp = pt.sum(vgp.likelihood.variational_expectation(y, fmean, fvar)) + kl = vgp.prior_kl(X) + return VGPTerms(elbo=var_exp - kl, var_exp=var_exp, kl=kl) + + def collapsed_elbo(vfe, X, y): """VFE/SGPR collapsed ELBO (Titsias' bound), unified for scalar and callable sigma. @@ -153,7 +185,7 @@ def collapsed_elbo(vfe, X, y): def fitc_log_marginal_likelihood(vfe, X, y): """FITC (Fully Independent Training Conditional) approximate log marginal likelihood. - Unlike ``collapsed_elbo``, FITC is not a lower bound — it approximates the + Unlike ``collapsed_elbo``, FITC is not a lower bound; it approximates the log marginal likelihood using the true per-point diagonal rather than the Nystrom diagonal throughout. The FITC covariance is:: @@ -185,9 +217,9 @@ def fitc_log_marginal_likelihood(vfe, X, y): Returns ------- FITCTerms - ``fitc`` — FITC approximate log marginal likelihood (fit + logdet). - ``fit`` — quadratic term: ``-0.5 * (y^T K_fitc^{-1} y + N log 2π)``. - ``logdet`` — log-determinant term: ``-0.5 log|K_fitc|``. + ``fitc``: FITC approximate log marginal likelihood (fit + logdet). + ``fit``: quadratic term: ``-0.5 * (y^T K_fitc^{-1} y + N log 2π)``. + ``logdet``: log-determinant term: ``-0.5 log|K_fitc|``. """ sigma2 = vfe.likelihood.sigma**2 N = X.shape[0] @@ -294,20 +326,20 @@ def vfe_diagnostics(vfe, X, y): """Collapsed ELBO terms plus sigma and two normalised fit metrics. Returns a ``VFEDiagnostics`` namedtuple of symbolic TensorVariables, - suitable for use with :func:`ptgp.optim.compile_scipy_diagnostics`. + for use with :func:`ptgp.optim.compile_scipy_diagnostics`. Fields ------ elbo, fit, trace_penalty Direct from :func:`collapsed_elbo`. nystrom_residual - ``tr(Kff - Qff) / N`` — per-point Nyström approximation error. + ``tr(Kff - Qff) / N``, the per-point Nyström approximation error. sigma Likelihood noise (constrained space). fit_per_n - ``fit / N`` — scale-invariant data fit. + ``fit / N``, the per-point data fit. excess_fit_per_n - ``fit_per_n + 0.5 * log(2π σ²)`` — how much better than noise floor. + ``fit_per_n + 0.5 * log(2π σ²)``, how much better than the noise floor. Goes to zero when the model fits at the noise level only. """ terms = collapsed_elbo(vfe, X, y) diff --git a/ptgp/optim/api.py b/ptgp/optim/api.py index e82de13..3ccaf33 100644 --- a/ptgp/optim/api.py +++ b/ptgp/optim/api.py @@ -49,7 +49,7 @@ def fit( compile_kwargs=None, **scipy_kwargs, ): - """Batteries-included estimation method for GP models. + """Estimation method for GP models. Estimation uses reasonable defaults based on the provided approximation. For fine-grained control, use the low-level API. @@ -70,18 +70,18 @@ def fit( ``gp_model.default_objective``. method : str Optimization method passed to :func:`scipy.optimize.minimize`. Must - accept a Jacobian — ``fit`` always supplies one (``jac=True``). + accept a Jacobian; ``fit`` always supplies one (``jac=True``). Default ``"L-BFGS-B"``. init : str Strategy for the initial parameter vector. One of: - - ``"prior_median"`` (default) — median of each prior via + - ``"prior_median"`` (default): median of each prior via ``pm.icdf(rv, 0.5)``, with a 500-sample fallback when icdf is unimplemented and a per-RV fallback to PyMC's initial point for improper priors. - - ``"prior_draw"`` — one draw from each prior. Stochastic; pin with + - ``"prior_draw"``: one draw from each prior. Stochastic; pin with ``init_rng`` for reproducibility. - - ``"unconstrained_zero"`` — PyMC's ``model.initial_point()`` (0 in + - ``"unconstrained_zero"``: PyMC's ``model.initial_point()`` (0 in unconstrained space unless ``initval=`` was set on the RV). init_rng : int or numpy.random.Generator, optional Seed for the sampling-based portions of ``"prior_median"`` and @@ -93,7 +93,7 @@ def fit( **scipy_kwargs Additional keyword arguments forwarded to :func:`scipy.optimize.minimize` (e.g. ``tol``, ``options``, - ``bounds``). Do not pass ``jac`` — it is set internally. + ``bounds``). Do not pass ``jac``; it is set internally. Returns ------- diff --git a/ptgp/optim/training.py b/ptgp/optim/training.py index c85401d..2724113 100644 --- a/ptgp/optim/training.py +++ b/ptgp/optim/training.py @@ -1,7 +1,7 @@ """Training and prediction compilation for PTGP models with PyMC priors. Uses ``pytensor.shared`` variables so that training automatically updates -the parameters used by prediction — no model reconstruction needed. +the parameters used by prediction; no model reconstruction needed. """ import logging @@ -96,7 +96,7 @@ def _make_initial_point(model, init="prior_median", rng=None, n_median_samples=5 Some priors cannot be sampled or quantiled. Improper priors like ``pm.HalfFlat`` and ``pm.Flat`` raise ``NotImplementedError`` from both ``pm.icdf`` and ``pm.draw`` because they have no proper measure. When - every method fails for a given RV — or returns non-finite values — that + every method fails for a given RV, or returns non-finite values, that RV's value is taken from ``model.initial_point()`` instead. PyMC's initial point is 0 in unconstrained space for every parameter unless ``initval`` was set explicitly, which corresponds to: @@ -185,7 +185,7 @@ def _make_shared_params( If ``frozen_vars`` is given, value vars appearing as keys are initialized to their freeze values rather than the PyMC initial point. Without this, a frozen var's shared slot (and its scipy theta slice) would carry the - initial point — wrong for diagnostics and for ``unpack_to_shared``. + initial point, wrong for diagnostics and for ``unpack_to_shared``. Returns ------- @@ -275,7 +275,7 @@ def compile_training_step( model : pm.Model, optional PyMC model. Uses the enclosing ``with pm.Model()`` context if None. Every continuous free RV in the model is automatically - made into a trainable shared variable — you do not need to list + made into a trainable shared variable, so you do not need to list them. optimizer_fn : callable, optional Optimizer function (default: ``adam``). Must have signature @@ -320,7 +320,7 @@ def compile_training_step( ``(X_batch, y_batch) -> loss_value``. Updates shared parameters in place. shared_params : dict - ``{value_var: shared_var}`` — the shared variables holding the + ``{value_var: shared_var}``, the shared variables holding the unconstrained parameter values. Needed by ``compile_predict``. shared_extras : list Shared variables for ``extra_vars``. Needed by ``compile_predict``. @@ -440,7 +440,7 @@ def compile_scipy_objective( PTGP model whose hyperparameters are PyMC RVs. X_var : TensorVariable Symbolic input placeholder. ``X`` is passed to the compiled - function on each scipy iteration — typically the full training + function on each scipy iteration, typically the full training inputs for GP/VFE (batching is not used with quasi-Newton methods). y_var : TensorVariable Symbolic target placeholder, handled like ``X_var``. @@ -516,7 +516,7 @@ def compile_scipy_objective( shared_params : dict ``{value_var: shared_var}`` for every continuous PyMC value var. Needed by :func:`compile_predict` and :func:`get_trained_params`. - Not read by ``fun`` — present only for the predict handoff. + Not read by ``fun``; present only for the predict handoff. shared_extras : list Shared variables for ``extra_vars``, in the same order. Needed by :func:`compile_predict`. Not read by ``fun``. @@ -749,7 +749,7 @@ def tracked_minimize(fun, theta0, args, diag_fn=None, print_every=None, **scipy_ theta0 : ndarray Initial parameter vector. args : tuple - Extra arguments forwarded to both ``fun`` and ``diag_fn`` — typically + Extra arguments forwarded to both ``fun`` and ``diag_fn``, typically ``(X, y)``. diag_fn : callable, optional ``(theta, *args) -> namedtuple`` from @@ -924,17 +924,17 @@ def minimize_staged_vfe( the trace penalty instead of improving ``Z``. This function blocks that failure mode with a structured four-phase schedule: - 1. **Phase 1** — freeze ``sigma`` at ``sigma_init``; train all other + 1. **Phase 1**: freeze ``sigma`` at ``sigma_init``; train all other hyperparameters and (by default) ``Z`` together. - 2. **Phase 2a** — freeze all hyperparameters; train ``Z`` only. - 3. **Phase 2b** — freeze ``Z``; train all hyperparameters (sigma now free). - 4. **Phase 3** — joint fine-tuning of everything (nothing frozen). + 2. **Phase 2a**: freeze all hyperparameters; train ``Z`` only. + 3. **Phase 2b**: freeze ``Z``; train all hyperparameters (sigma now free). + 4. **Phase 3**: joint fine-tuning of everything (nothing frozen). Phases 2a/2b repeat ``phase2_cycles`` times. Pass ``phase2_cycles=0`` to skip straight from phase 1 to phase 3 when the ``Z`` initialisation is already good (e.g. from ``greedy_variance_init``). - ``sigma`` is auto-detected from ``gp_model.likelihood.sigma`` — no need to + ``sigma`` is auto-detected from ``gp_model.likelihood.sigma``; no need to pass a separate ``sigma_rv`` argument. All diagnostic history entries are :class:`ptgp.objectives.VFEDiagnostics` @@ -977,7 +977,7 @@ def minimize_staged_vfe( init : str Initialisation strategy for phase 1 hyperparameters. Same options as :func:`compile_scipy_objective`: ``"prior_median"`` (default), - ``"prior_draw"``, or ``"unconstrained_zero"`` (PyMC initial point — + ``"prior_draw"``, or ``"unconstrained_zero"`` (PyMC initial point, 0 in unconstrained space). Only affects phase 1; later phases inherit converged values from the previous phase. init_rng : int or numpy Generator, optional diff --git a/ptgp/utils.py b/ptgp/utils.py index d0dfc24..e62aede 100644 --- a/ptgp/utils.py +++ b/ptgp/utils.py @@ -12,7 +12,7 @@ def get_initial_params(model, init="prior_median", rng=None, n_median_samples=50 """Return constrained-space initial values for all free RVs in a PyMC model. Uses the same initialization strategies as ``compile_scipy_objective``. - Useful for building proxy kernels with concrete float values to pass to + Use it to build proxy kernels with concrete float values to pass to ``greedy_variance_init``. Parameters diff --git a/tests/test_composite_likelihood.py b/tests/test_composite_likelihood.py new file mode 100644 index 0000000..e618f5e --- /dev/null +++ b/tests/test_composite_likelihood.py @@ -0,0 +1,133 @@ +"""CompositeLikelihood tests: per-subset dispatch and end-to-end use with VGP.""" + +import numpy as np +import pymc as pm +import pytensor +import pytensor.tensor as pt +import pytest + +import ptgp as pg + + +def test_variational_expectation_equals_per_subset(): + """The combinator equals each sub-likelihood computed on its slice.""" + rng = np.random.default_rng(0) + N = 10 + mu = rng.standard_normal(N) + var = rng.uniform(0.1, 1.0, N) + y = rng.standard_normal(N) + idx0 = np.array([0, 2, 4, 6, 8]) # interleaved to exercise non-contiguous subsets + idx1 = np.array([1, 3, 5, 7, 9]) + + lik0 = pg.likelihoods.StudentT(nu=4.0, sigma=0.5) + lik1 = pg.likelihoods.Gaussian(0.3) + composite = pg.likelihoods.CompositeLikelihood([lik0, lik1], [idx0, idx1]) + + yv, muv, varv = pt.vector("y"), pt.vector("mu"), pt.vector("var") + comp_fn = pytensor.function([yv, muv, varv], composite.variational_expectation(yv, muv, varv)) + e0_fn = pytensor.function([yv, muv, varv], lik0.variational_expectation(yv, muv, varv)) + e1_fn = pytensor.function([yv, muv, varv], lik1.variational_expectation(yv, muv, varv)) + + out = comp_fn(y, mu, var) + expected = np.zeros(N) + expected[idx0] = e0_fn(y[idx0], mu[idx0], var[idx0]) + expected[idx1] = e1_fn(y[idx1], mu[idx1], var[idx1]) + np.testing.assert_allclose(out, expected, atol=1e-10) + + +def test_predict_mean_and_var_dispatch(): + """predict_mean_and_var is scattered per subset.""" + rng = np.random.default_rng(1) + N = 6 + mu = rng.standard_normal(N) + var = rng.uniform(0.1, 1.0, N) + idx0 = np.array([0, 1, 2]) + idx1 = np.array([3, 4, 5]) + + lik0 = pg.likelihoods.Poisson() + lik1 = pg.likelihoods.Gaussian(0.3) + composite = pg.likelihoods.CompositeLikelihood([lik0, lik1], [idx0, idx1]) + + muv, varv = pt.vector("mu"), pt.vector("var") + comp_fn = pytensor.function([muv, varv], list(composite.predict_mean_and_var(muv, varv))) + m0_fn = pytensor.function([muv, varv], list(lik0.predict_mean_and_var(muv, varv))) + m1_fn = pytensor.function([muv, varv], list(lik1.predict_mean_and_var(muv, varv))) + + sm, sv = comp_fn(mu, var) + em, ev = np.zeros(N), np.zeros(N) + m0, v0 = m0_fn(mu[idx0], var[idx0]) + m1, v1 = m1_fn(mu[idx1], var[idx1]) + em[idx0], ev[idx0] = m0, v0 + em[idx1], ev[idx1] = m1, v1 + np.testing.assert_allclose(sm, em, atol=1e-10) + np.testing.assert_allclose(sv, ev, atol=1e-10) + + +def test_heteroskedastic_sigma_subset_aligned(): + """A vector-sigma sub-likelihood works when sigma is aligned to its subset.""" + rng = np.random.default_rng(2) + N = 8 + mu = rng.standard_normal(N) + var = rng.uniform(0.1, 1.0, N) + y = rng.standard_normal(N) + idx0 = np.arange(0, 4) + idx1 = np.arange(4, 8) + + sigma0 = rng.uniform(0.2, 0.6, idx0.size) # length matches the subset + lik0 = pg.likelihoods.Gaussian(pt.as_tensor_variable(sigma0)) + lik1 = pg.likelihoods.Gaussian(0.3) + composite = pg.likelihoods.CompositeLikelihood([lik0, lik1], [idx0, idx1]) + + yv, muv, varv = pt.vector("y"), pt.vector("mu"), pt.vector("var") + comp_fn = pytensor.function([yv, muv, varv], composite.variational_expectation(yv, muv, varv)) + e0_fn = pytensor.function([yv, muv, varv], lik0.variational_expectation(yv, muv, varv)) + e1_fn = pytensor.function([yv, muv, varv], lik1.variational_expectation(yv, muv, varv)) + + out = comp_fn(y, mu, var) + expected = np.zeros(N) + expected[idx0] = e0_fn(y[idx0], mu[idx0], var[idx0]) + expected[idx1] = e1_fn(y[idx1], mu[idx1], var[idx1]) + np.testing.assert_allclose(out, expected, atol=1e-10) + + +def test_rejects_non_partition(): + """Index arrays that do not partition range(N) are rejected.""" + lik = pg.likelihoods.Gaussian(0.3) + with pytest.raises(ValueError): + pg.likelihoods.CompositeLikelihood([lik, lik], [np.array([0, 1]), np.array([1, 2])]) + + +def test_vgp_with_composite_likelihood_trains(): + """VGP with a composite likelihood (Gaussian + Student-t) trains.""" + rng = np.random.default_rng(3) + N = 40 + X = np.sort(rng.uniform(-3, 3, N))[:, None] + y = np.sin(X.ravel()) + 0.15 * rng.standard_normal(N) + idx0 = np.arange(0, 20) + idx1 = np.arange(20, N) + y[idx1[:3]] += 4.0 # outliers in the Student-t subset + + lik = pg.likelihoods.CompositeLikelihood( + [pg.likelihoods.Gaussian(0.2), pg.likelihoods.StudentT(nu=3.0, sigma=0.2)], + [idx0, idx1], + ) + vp = pg.gp.init_vgp_params(N) + with pm.Model() as model: + ls = pm.InverseGamma("ls", alpha=2.0, beta=1.0) + eta = pm.Exponential("eta", lam=1.0) + kernel = eta**2 * pg.kernels.Matern52(input_dim=1, ls=ls) + vgp = pg.gp.VGP(kernel=kernel, likelihood=lik, variational_params=vp) + + X_var, y_var = pt.matrix("X"), pt.vector("y") + train_step, _, _ = pg.optim.compile_training_step( + lambda gp, X, y: pg.objectives.vgp_elbo(gp, X, y).elbo, + vgp, + X_var, + y_var, + model=model, + extra_vars=vp.extra_vars, + extra_init=vp.extra_init, + learning_rate=1e-2, + ) + losses = [float(train_step(X, y)) for _ in range(300)] + assert losses[-1] < losses[0], "VGP + CompositeLikelihood loss should decrease" diff --git a/tests/test_vgp.py b/tests/test_vgp.py new file mode 100644 index 0000000..cbc1a59 --- /dev/null +++ b/tests/test_vgp.py @@ -0,0 +1,221 @@ +"""VGP (Opper-Archambeau full variational GP) tests. + +The decisive test is the exact-GP recovery: with a Gaussian likelihood, the VGP +at its closed-form optimum reproduces the exact ``Unapproximated`` GP +(predictions and ELBO), which validates the ELBO, KL, training marginals, and +predict math all at once. +""" + +import numpy as np +import pymc as pm +import pytensor +import pytensor.tensor as pt +import scipy.optimize + +import ptgp as pg + +from ptgp.mean import Constant, Zero +from ptgp.optim.training import compile_scipy_objective + + +def _eval_kernel(kernel, X): + """Evaluate a kernel Gram matrix numerically.""" + Xv = pt.matrix("X") + return pytensor.function([Xv], kernel(Xv))(X) + + +def _oa_optimal_params(K, y, m, sigma): + """Closed-form OA params for a Gaussian likelihood. + + The exact posterior over f has precision ``K^{-1} + sigma^{-2} I``, so + ``lambda = 1/sigma**2`` and ``alpha = (K + sigma**2 I)^{-1} (y - m)``. + """ + N = K.shape[0] + alpha = np.linalg.solve(K + sigma**2 * np.eye(N), y - m) + lam = (1.0 / sigma**2) * np.ones(N) + return alpha, lam + + +def _count_data(rng, n=80): + """1D Poisson data with a smoothly varying log-rate.""" + X = np.sort(rng.uniform(-3, 3, n))[:, None] + rate = np.exp(0.5 * np.sin(X[:, 0]) + 0.3) + y = rng.poisson(rate).astype(np.float64) + return X, y + + +def _binary_data(rng, n=120): + """1D Bernoulli data with the decision boundary near x=0.""" + from scipy.special import erf + + X = np.sort(rng.uniform(-3, 3, n))[:, None] + p = 0.5 * (1.0 + erf(X[:, 0] / np.sqrt(2.0))) + y = (rng.uniform(0, 1, n) < p).astype(np.float64) + return X, y + + +class TestVGPRecoversExactGP: + """Gaussian VGP at the closed-form optimum equals the exact GP.""" + + def _run(self, mean_c): + rng = np.random.default_rng(0) + N = 14 + X = np.sort(rng.uniform(-2, 2, N))[:, None] + y = np.sin(X.ravel()) + 0.1 * rng.standard_normal(N) + Xs = np.linspace(-2.5, 2.5, 9)[:, None] + sigma, ls, eta = 0.3, 0.7, 1.2 + + kernel = eta**2 * pg.kernels.Matern52(input_dim=1, ls=ls) + mean = Zero() if mean_c is None else Constant(float(mean_c)) + K = _eval_kernel(kernel, X) + m = np.zeros(N) if mean_c is None else mean_c * np.ones(N) + alpha, lam = _oa_optimal_params(K, y, m, sigma) + + vp = pg.gp.init_vgp_params(N, alpha_init=alpha, lambda_init=lam) + vgp = pg.gp.VGP( + kernel=kernel, + mean=mean, + likelihood=pg.likelihoods.Gaussian(sigma), + variational_params=vp, + ) + exact = pg.gp.Unapproximated(kernel=kernel, mean=mean, sigma=sigma) + + # ELBO at the optimum equals the exact marginal log-likelihood. + Xv, yv = pt.matrix("X"), pt.vector("y") + elbo_fn = pytensor.function( + [Xv, yv, *vp.extra_vars], pg.objectives.vgp_elbo(vgp, Xv, yv).elbo + ) + mll_fn = pytensor.function( + [Xv, yv], pg.objectives.marginal_log_likelihood(exact, Xv, yv).mll + ) + elbo_val = float(elbo_fn(X, y, *vp.extra_init)) + mll_val = float(mll_fn(X, y)) + np.testing.assert_allclose(elbo_val, mll_val, atol=1e-6) + + # Predictions match the exact GP. + Xn, Xt, yt = pt.matrix("Xn"), pt.matrix("Xt"), pt.vector("yt") + vm, vv = vgp.predict_marginal(Xn, Xt) + vgp_pred = pytensor.function([Xn, Xt, *vp.extra_vars], [vm, vv]) + em, ev = exact.predict_marginal(Xn, Xt, yt) + exact_pred = pytensor.function([Xn, Xt, yt], [em, ev]) + + vmean, vvar = vgp_pred(Xs, X, *vp.extra_init) + emean, evar = exact_pred(Xs, X, y) + np.testing.assert_allclose(vmean, emean, atol=1e-6) + np.testing.assert_allclose(vvar, evar, atol=1e-6) + + def test_zero_mean(self): + """VGP recovers the exact GP with a zero mean function.""" + self._run(mean_c=None) + + def test_nonzero_constant_mean(self): + """VGP recovers the exact GP with a non-zero constant mean (guards centering).""" + self._run(mean_c=1.5) + + +class TestVGPElboBound: + """The VGP ELBO is a lower bound on the exact marginal likelihood.""" + + def test_elbo_leq_mll(self): + """At arbitrary feasible variational params, ELBO <= MLL.""" + rng = np.random.default_rng(1) + N = 12 + X = np.sort(rng.uniform(0, 5, N))[:, None] + y = np.sin(X.ravel()) + 0.1 * rng.standard_normal(N) + sigma = 0.4 + kernel = pg.kernels.Matern52(input_dim=1, ls=1.0) + + # Feasible but non-optimal params. + vp = pg.gp.init_vgp_params( + N, alpha_init=0.1 * rng.standard_normal(N), lambda_init=2.0 * np.ones(N) + ) + vgp = pg.gp.VGP( + kernel=kernel, likelihood=pg.likelihoods.Gaussian(sigma), variational_params=vp + ) + exact = pg.gp.Unapproximated(kernel=kernel, sigma=sigma) + + Xv, yv = pt.matrix("X"), pt.vector("y") + elbo_val = float( + pytensor.function([Xv, yv, *vp.extra_vars], pg.objectives.vgp_elbo(vgp, Xv, yv).elbo)( + X, y, *vp.extra_init + ) + ) + mll_val = float( + pytensor.function([Xv, yv], pg.objectives.marginal_log_likelihood(exact, Xv, yv).mll)( + X, y + ) + ) + assert elbo_val <= mll_val + 1e-6 + + +class TestVGPNonGaussianSmoke: + """VGP trains end-to-end with non-Gaussian likelihoods.""" + + def _smoke(self, X, y, likelihood, n_steps=300, lr=1e-2): + N = X.shape[0] + vp = pg.gp.init_vgp_params(N) + with pm.Model() as model: + ls = pm.InverseGamma("ls", alpha=2.0, beta=1.0) + eta = pm.Exponential("eta", lam=1.0) + kernel = eta**2 * pg.kernels.Matern52(input_dim=1, ls=ls) + vgp = pg.gp.VGP(kernel=kernel, likelihood=likelihood, variational_params=vp) + + X_var, y_var = pt.matrix("X"), pt.vector("y") + train_step, _, _ = pg.optim.compile_training_step( + lambda gp, X, y: pg.objectives.vgp_elbo(gp, X, y).elbo, + vgp, + X_var, + y_var, + model=model, + extra_vars=vp.extra_vars, + extra_init=vp.extra_init, + learning_rate=lr, + ) + losses = [float(train_step(X, y)) for _ in range(n_steps)] + assert losses[-1] < losses[0], "VGP loss should decrease" + + def test_poisson(self): + """VGP + Poisson loss decreases.""" + rng = np.random.default_rng(2) + X, y = _count_data(rng, n=80) + self._smoke(X, y, pg.likelihoods.Poisson()) + + def test_bernoulli(self): + """VGP + Bernoulli loss decreases.""" + rng = np.random.default_rng(3) + X, y = _binary_data(rng, n=100) + self._smoke(X, y, pg.likelihoods.Bernoulli()) + + +class TestVGPGradient: + """Finite-difference gradient check of the VGP ELBO objective.""" + + def test_grad_matches_finite_difference(self): + """compile_scipy_objective's analytic gradient matches central differences.""" + rng = np.random.default_rng(4) + N = 8 + X = np.sort(rng.uniform(-2, 2, N))[:, None] + y = np.sin(X.ravel()) + 0.1 * rng.standard_normal(N) + + vp = pg.gp.init_vgp_params(N) + with pm.Model() as model: + ls = pm.InverseGamma("ls", alpha=2.0, beta=1.0) + eta = pm.Exponential("eta", lam=1.0) + kernel = eta**2 * pg.kernels.Matern52(input_dim=1, ls=ls) + vgp = pg.gp.VGP( + kernel=kernel, likelihood=pg.likelihoods.Poisson(), variational_params=vp + ) + + X_var, y_var = pt.matrix("X"), pt.vector("y") + fun, theta0, _, _, _ = compile_scipy_objective( + pg.objectives.vgp_elbo, vgp, X_var, y_var, model=model + ) + + def f(theta): + return fun(theta, X, y)[0] + + def g(theta): + return fun(theta, X, y)[1] + + err = scipy.optimize.check_grad(f, g, theta0, epsilon=1e-6) + assert err < 1e-3, f"gradient mismatch {err:.2e}"