Custom likelihood variables

ptgp/gp/svgp.py

@@ -230,7 +230,7 @@ def predict_marginal(self, X, incl_lik=False):
             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 self.likelihood.at(X).predict_mean_and_var(fmean, fvar)
         return fmean, fvar
 
     def predict_joint(self, X):

ptgp/gp/unapproximated.py

@@ -18,18 +18,22 @@ class Unapproximated:
         Mean function (default: ``Zero()``).
     sigma : tensor or PyMC random variable
         Observation noise standard deviation.
+    x : tensor, optional
+        The design matrix ``sigma`` was built against, for heteroskedastic
+        noise. Pass it so sigma can be re-rooted onto the test inputs at predict
+        time.
     """
 
     extra_vars = ()
     extra_init = ()
     default_objective = staticmethod(marginal_log_likelihood)
     predict_needs_data = True
 
-    def __init__(self, kernel, mean=None, sigma=None):
+    def __init__(self, kernel, mean=None, sigma=None, x=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)
+        self.likelihood = Gaussian(sigma, x=x)
 
     def predict_marginal(self, X_new, X_train, y_train, incl_lik=False):
         """Posterior marginal mean and variance at each point in X_new.
@@ -64,6 +68,5 @@ def predict_marginal(self, X_new, X_train, y_train, incl_lik=False):
         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 self.likelihood.at(X_new).predict_mean_and_var(fmean, fvar)
         return fmean, fvar

ptgp/gp/vfe.py

@@ -22,12 +22,16 @@ class VFE:
         Observation noise standard deviation.
     inducing_variable : InducingVariables
         Inducing point locations.
+    x : tensor, optional
+        The design matrix ``sigma`` was built against, for heteroskedastic
+        noise. Pass it so sigma can be re-rooted onto the test inputs at predict
+        time.
     """
 
     default_objective = staticmethod(collapsed_elbo)
     predict_needs_data = True
 
-    def __init__(self, kernel, mean=None, sigma=None, inducing_variable=None):
+    def __init__(self, kernel, mean=None, sigma=None, inducing_variable=None, x=None):
         """Store the kernel, mean, and inducing variable; build a Gaussian likelihood from sigma."""
         if not hasattr(inducing_variable, "Z"):
             raise TypeError(
@@ -37,7 +41,7 @@ def __init__(self, kernel, mean=None, sigma=None, inducing_variable=None):
             )
         self.kernel = kernel
         self.mean = mean if mean is not None else Zero()
-        self.likelihood = Gaussian(sigma)
+        self.likelihood = Gaussian(sigma, x=x)
         self.inducing_variable = inducing_variable
 
     @property
@@ -87,6 +91,5 @@ def predict_marginal(self, X_new, X_train, y_train, incl_lik=False):
         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 self.likelihood.at(X_new).predict_mean_and_var(fmean, fvar)
         return fmean, fvar

ptgp/likelihoods/__init__.py

@@ -1,8 +1,40 @@
-from ptgp.likelihoods.base import Likelihood
+from ptgp.likelihoods.base import (
+    GPData,
+    LikelihoodOp,
+    LikelihoodType,
+    LikelihoodVariable,
+    at,
+    gp_data,
+    op_of,
+    param,
+    predict_log_density,
+    predict_mean_and_var,
+    variational_expectation,
+)
 from ptgp.likelihoods.bernoulli import Bernoulli
 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__ = [
+    # Family helpers — build a LikelihoodVariable.
+    "Gaussian",
+    "Bernoulli",
+    "StudentT",
+    "Poisson",
+    "NegativeBinomial",
+    # Graph types.
+    "LikelihoodOp",
+    "LikelihoodType",
+    "LikelihoodVariable",
+    "GPData",
+    "gp_data",
+    # Purely functional API — operate on a likelihood node via ``owner.op``.
+    "op_of",
+    "param",
+    "at",
+    "variational_expectation",
+    "predict_mean_and_var",
+    "predict_log_density",
+]