Add VGP model and CompositeLikelihood

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.

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",
 ]

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

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``.
 

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

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 = {}

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.
 

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``.
         """