Add gp.posterior_predictive for constructing a posterior predictive from a GP.

This is equivalent (and a wrapper) to calling GaussianProcessRegressionModel.

PiperOrigin-RevId: 384605668

Author: srvasude

tensorflow_probability/python/distributions/gaussian_process.py

@@ -596,6 +596,53 @@ def _mode(self, index_points=None):
   def _default_event_space_bijector(self):
     return identity_bijector.Identity(validate_args=self.validate_args)
 
+  def posterior_predictive(self, observations, predictive_index_points=None):
+    """Return the posterior predictive distribution associated with this distribution.
+
+    Given `predictive_index_points` and `observations`, return the posterior
+    predictive distribution on the `predictive_index_points` conditioned on
+    `index_points` and `observations` associated to `index_points`.
+
+    This is equivalent to using the `GaussianProcessRegressionModel` class.
+
+    Args:
+      observations: `float` `Tensor` representing collection, or batch of
+        collections, of observations corresponding to
+        `self.index_points`. Shape has the form `[b1, ..., bB, e]`, which
+        must be broadcastable with the batch and example shapes of
+        `self.index_points`. The batch shape `[b1, ..., bB]` must be
+        broadcastable with the shapes of all other batched parameters
+      predictive_index_points: `float` `Tensor` representing finite collection,
+        or batch of collections, of points in the index set over which the GP
+        is defined.
+        Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the
+        number of feature dimensions and must equal `kernel.feature_ndims` and
+        `e` is the number (size) of predictive index points in each batch.
+        The batch shape must be broadcastable with this distributions
+        `batch_shape`.
+        Default value: `None`.
+
+    Returns:
+      gprm: An instance of `Distribution` that represents the posterior
+        predictive.
+    """
+    from tensorflow_probability.python.distributions import gaussian_process_regression_model as gprm  # pylint:disable=g-import-not-at-top
+    if self.index_points is None:
+      raise ValueError(
+          'Expected that `self.index_points` is not `None`. Using '
+          '`self.index_points=None` is equivalent to using a `GaussianProcess` '
+          'prior, which this class encapsulates.')
+    return gprm.GaussianProcessRegressionModel(
+        kernel=self.kernel,
+        observation_index_points=self.index_points,
+        observations=observations,
+        index_points=predictive_index_points,
+        observation_noise_variance=self.observation_noise_variance,
+        mean_fn=self.mean_fn,
+        jitter=self.jitter,
+        validate_args=self.validate_args,
+        allow_nan_stats=self.allow_nan_stats)
+
 
 def _assert_kl_compatible(marginal, other):
   if ((isinstance(marginal, normal.Normal) and

tensorflow_probability/python/distributions/gaussian_process_test.py

@@ -257,6 +257,52 @@ def test_marginal_fn(
         np.eye(10),
         gp.get_marginal_distribution().covariance())
 
+  def testGPPosteriorPredictive(self):
+    amplitude = np.float64(.5)
+    length_scale = np.float64(2.)
+    jitter = np.float64(1e-4)
+    observation_noise_variance = np.float64(3e-3)
+    kernel = psd_kernels.ExponentiatedQuadratic(amplitude, length_scale)
+
+    index_points = np.random.uniform(-1., 1., 10)[..., np.newaxis]
+
+    gp = tfd.GaussianProcess(
+        kernel,
+        index_points,
+        observation_noise_variance=observation_noise_variance,
+        jitter=jitter,
+        validate_args=True)
+
+    predictive_index_points = np.random.uniform(1., 2., 10)[..., np.newaxis]
+    observations = np.linspace(1., 10., 10)
+
+    expected_gprm = tfd.GaussianProcessRegressionModel(
+        kernel=kernel,
+        observation_index_points=index_points,
+        observations=observations,
+        observation_noise_variance=observation_noise_variance,
+        jitter=jitter,
+        index_points=predictive_index_points,
+        validate_args=True)
+
+    actual_gprm = gp.posterior_predictive(
+        predictive_index_points=predictive_index_points,
+        observations=observations)
+
+    samples = self.evaluate(actual_gprm.sample(10, seed=test_util.test_seed()))
+
+    self.assertAllClose(
+        self.evaluate(expected_gprm.log_prob(samples)),
+        self.evaluate(actual_gprm.log_prob(samples)))
+
+    self.assertAllClose(
+        self.evaluate(expected_gprm.mean()),
+        self.evaluate(actual_gprm.mean()))
+
+    self.assertAllClose(
+        self.evaluate(expected_gprm.covariance()),
+        self.evaluate(actual_gprm.covariance()))
+
 
 @test_util.test_all_tf_execution_regimes
 class GaussianProcessStaticTest(_GaussianProcessTest, test_util.TestCase):