First piece of Seasonality support in the GibbsSampler - support resampling drift scale in Seasonality components.

This includes:
* Splitting out some sampling logic for scale into a separate file to be shared with seasonality
* A convenience method for accessing parameters by name.
* Comparison of expectation generated by sampling the mean/variance from a similar estimation

PiperOrigin-RevId: 475877493

Author: kylepl

tensorflow_probability/python/experimental/sts_gibbs/BUILD

@@ -63,6 +63,7 @@ multi_substrate_py_library(
     srcs = ["gibbs_sampler.py"],
     deps = [
         ":dynamic_spike_and_slab",
+        ":sample_parameters",
         ":spike_and_slab",
         # numpy dep,
         # tensorflow dep,
@@ -204,3 +205,32 @@ multi_substrate_py_test(
 #         "//third_party/tensorflow/compiler/jit:xla_cpu_jit",  # DisableOnExport
     ],
 )
+
+multi_substrate_py_library(
+    name = "sample_parameters",
+    srcs = ["sample_parameters.py"],
+    deps = [
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/bijectors:invert",
+        "//tensorflow_probability/python/bijectors:square",
+        "//tensorflow_probability/python/distributions:distribution",
+        "//tensorflow_probability/python/distributions:inverse_gamma",
+        "//tensorflow_probability/python/distributions:transformed_distribution",
+        "//tensorflow_probability/python/internal:dtype_util",
+        "//tensorflow_probability/python/internal:prefer_static",
+    ],
+)
+
+multi_substrate_py_test(
+    name = "sample_parameters_test",
+    size = "medium",
+    srcs = ["sample_parameters_test.py"],
+    deps = [
+        ":sample_parameters",
+        # absl/testing:parameterized dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/distributions:inverse_gamma",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)

tensorflow_probability/python/experimental/sts_gibbs/gibbs_sampler.py

@@ -76,6 +76,7 @@ class is somewhat general, in that we assume that any seasonal/holiday variation
 from tensorflow_probability.python.distributions import normal_conjugate_posteriors
 from tensorflow_probability.python.experimental.distributions import mvn_precision_factor_linop as mvnpflo
 from tensorflow_probability.python.experimental.sts_gibbs import dynamic_spike_and_slab
+from tensorflow_probability.python.experimental.sts_gibbs import sample_parameters
 from tensorflow_probability.python.experimental.sts_gibbs import spike_and_slab
 from tensorflow_probability.python.internal import distribution_util as dist_util
 from tensorflow_probability.python.internal import dtype_util
@@ -777,27 +778,10 @@ def _resample_scale(prior, observed_residuals, is_missing=None, seed=None):
   Returns:
     sampled_scale: A `Tensor` sample from the posterior `p(scale | x)`.
   """
-  dtype = observed_residuals.dtype
-
-  if is_missing is not None:
-    num_missing = tf.reduce_sum(tf.cast(is_missing, dtype), axis=-1)
-  num_observations = prefer_static.shape(observed_residuals)[-1]
-  if is_missing is not None:
-    observed_residuals = tf.where(is_missing, tf.zeros_like(observed_residuals),
-                                  observed_residuals)
-    num_observations -= num_missing
-
-  variance_posterior = type(prior)(
-      concentration=prior.concentration + tf.cast(num_observations / 2., dtype),
-      scale=prior.scale +
-      tf.reduce_sum(tf.square(observed_residuals), axis=-1) / 2.)
-  new_scale = tf.sqrt(variance_posterior.sample(seed=seed))
-
-  # Support truncated priors.
-  if hasattr(prior, 'upper_bound') and prior.upper_bound is not None:
-    new_scale = tf.minimum(new_scale, prior.upper_bound)
-
-  return new_scale
+  posterior = sample_parameters.normal_scale_posterior_inverse_gamma_conjugate(
+      prior, observed_residuals, is_missing)
+  return sample_parameters.sample_with_optional_upper_bound(
+      posterior, seed=seed)
 
 
 def _build_sampler_loop_body(model,

tensorflow_probability/python/experimental/sts_gibbs/gibbs_sampler_test.py

@@ -106,7 +106,8 @@ def _build_test_model(self,
     observation_noise_variance_prior = prior_class(
         concentration=tf.cast(0.01, dtype),
         scale=tf.cast(0.01 * 0.01, dtype))
-    observation_noise_variance_prior.upper_bound = 100.0
+    observation_noise_variance_prior.upper_bound = tf.constant(
+        100.0, dtype=dtype)
 
     observed_time_series = missing_values_util.MaskedTimeSeries(
         time_series[..., tf.newaxis], is_missing)

tensorflow_probability/python/experimental/sts_gibbs/sample_parameters.py

@@ -0,0 +1,89 @@
+# Copyright 2022 The TensorFlow Probability Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Functions for sampling parameters useful in Gibbs Sampling."""
+
+import tensorflow.compat.v2 as tf
+
+from tensorflow_probability.python.bijectors import invert
+from tensorflow_probability.python.bijectors import square
+from tensorflow_probability.python.distributions import inverse_gamma
+from tensorflow_probability.python.distributions import transformed_distribution
+from tensorflow_probability.python.internal import prefer_static
+
+
+def normal_scale_posterior_inverse_gamma_conjugate(variance_prior,
+                                                   observations,
+                                                   is_missing=None):
+  """Returns the conditional posterior of Normal scale given observations.
+
+  We assume the conjugate InverseGamma->Normal model:
+
+  ```
+  scale ~ Sqrt(InverseGamma(variance_prior.concentration, variance_prior.scale))
+  for i in [1, ..., num_observations]:
+    x[i] ~ Normal(0, scale)
+  ```
+
+  and return a sample from `p(scale | x)`.
+
+  Args:
+    variance_prior: Variance prior distribution as a `tfd.InverseGamma`
+      instance. Note that the prior is given on the variance, but the value
+      returned is a sample of the scale.
+    observations: Float `Tensor` of shape `[..., num_observations]`, specifying
+      the centered observations `(x)`.
+    is_missing: Optional `bool` `Tensor` of shape `[..., num_observations]`. A
+      `True` value indicates that the corresponding observation is missing.
+
+  Returns:
+    sampled_scale: A `tfd.Distribution` of the conditional posterior
+      of the inverse gamma scale.
+  """
+  dtype = observations.dtype
+  num_observations = prefer_static.shape(observations)[-1]
+  if is_missing is not None:
+    num_missing = tf.reduce_sum(tf.cast(is_missing, dtype), axis=-1)
+    observations = tf.where(is_missing, tf.zeros_like(observations),
+                            observations)
+    num_observations -= num_missing
+
+  variance_posterior = inverse_gamma.InverseGamma(
+      concentration=variance_prior.concentration +
+      tf.cast(num_observations / 2, dtype),
+      scale=variance_prior.scale +
+      tf.reduce_sum(tf.square(observations), axis=-1) / 2.)
+  scale_posterior = transformed_distribution.TransformedDistribution(
+      bijector=invert.Invert(square.Square()), distribution=variance_posterior)
+
+  if hasattr(variance_prior,
+             'upper_bound') and variance_prior.upper_bound is not None:
+    variance_posterior.upper_bound = variance_prior.upper_bound
+    # TODO(kloveless): This should have sqrt applied, but it is not for
+    # temporary backwards compatibility.
+    scale_posterior.upper_bound = variance_prior.upper_bound
+
+  return scale_posterior
+
+
+# TODO(kloveless): This seems like this should be a function on the
+# distribution itself.
+def sample_with_optional_upper_bound(distribution, sample_shape=(), seed=None):
+  """Samples from the given distribution with an optional upper bound."""
+  sample = distribution.sample(sample_shape=sample_shape, seed=seed)
+  if hasattr(distribution,
+             'upper_bound') and distribution.upper_bound is not None:
+    sample = tf.minimum(sample, distribution.upper_bound)
+
+  return sample

tensorflow_probability/python/experimental/sts_gibbs/sample_parameters_test.py

@@ -0,0 +1,115 @@
+# Copyright 2022 The TensorFlow Probability Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Tests for sample_parameters."""
+
+import tensorflow as tf
+from tensorflow_probability.python.bijectors import invert
+from tensorflow_probability.python.bijectors import square
+from tensorflow_probability.python.distributions import inverse_gamma
+from tensorflow_probability.python.distributions import transformed_distribution
+from tensorflow_probability.python.experimental.sts_gibbs import sample_parameters
+from tensorflow_probability.python.internal import samplers
+from tensorflow_probability.python.internal import test_util
+
+
+@test_util.test_all_tf_execution_regimes
+class NormalScalePosteriorInverseGammaConjugate(test_util.TestCase):
+
+  def testNoObservations(self):
+    distribution = inverse_gamma.InverseGamma(16., 4.)
+    posterior_distribution = sample_parameters.normal_scale_posterior_inverse_gamma_conjugate(
+        distribution, observations=tf.constant([], dtype=tf.float32))
+    self.assertIsInstance(posterior_distribution,
+                          transformed_distribution.TransformedDistribution)
+    self.assertIsInstance(posterior_distribution.bijector, invert.Invert)
+    self.assertIsInstance(posterior_distribution.bijector.bijector,
+                          square.Square)
+    self.assertIsInstance(posterior_distribution.distribution,
+                          inverse_gamma.InverseGamma)
+    self.assertAllEqual(distribution.concentration,
+                        posterior_distribution.distribution.concentration)
+    self.assertAllEqual(distribution.scale,
+                        posterior_distribution.distribution.scale)
+
+  def testSingleObservation(self):
+    concentration = 16.
+    scale = 4.
+    distribution = inverse_gamma.InverseGamma(
+        concentration=concentration, scale=scale)
+    posterior_distribution = sample_parameters.normal_scale_posterior_inverse_gamma_conjugate(
+        distribution, observations=tf.constant([10.]))
+    self.assertAllEqual(
+        concentration + 0.5,  # Add half the number of observations
+        posterior_distribution.distribution.concentration)
+    self.assertAllEqual(
+        scale + 50,  # Add half the square of observations sum.
+        posterior_distribution.distribution.scale)
+
+  def testTwoObservations(self):
+    concentration = 16.
+    scale = 4.
+    distribution = inverse_gamma.InverseGamma(
+        concentration=concentration, scale=scale)
+    posterior_distribution = sample_parameters.normal_scale_posterior_inverse_gamma_conjugate(
+        distribution, observations=tf.constant([10., 4.]))
+    self.assertAllEqual(
+        concentration + 1,  # Add half the number of observations
+        posterior_distribution.distribution.concentration)
+    self.assertAllEqual(
+        scale + 58,  # Add half the square of observations sum.
+        posterior_distribution.distribution.scale)
+
+  def testUpperBoundPropagatedFromPrior(self):
+    # If no upper bound is provided, expect there to be none.
+    distribution = inverse_gamma.InverseGamma(16., 4.)
+    posterior_distribution = sample_parameters.normal_scale_posterior_inverse_gamma_conjugate(
+        distribution, observations=tf.constant([], dtype=tf.float32))
+    self.assertFalse(hasattr(posterior_distribution, 'upper_bound'))
+    self.assertFalse(
+        hasattr(posterior_distribution.distribution, 'upper_bound'))
+
+    distribution = inverse_gamma.InverseGamma(16., 4.)
+    distribution.upper_bound = 16.
+    posterior_distribution = sample_parameters.normal_scale_posterior_inverse_gamma_conjugate(
+        distribution, observations=tf.constant([], dtype=tf.float32))
+    self.assertAllEqual(posterior_distribution.distribution.upper_bound, 16.)
+    self.assertAllEqual(
+        posterior_distribution.upper_bound,
+        # TODO(kloveless): This should have sqrt applied, but it is not for
+        # temporary backwards compatibility.
+        16.)
+
+
+@test_util.test_all_tf_execution_regimes
+class SampleWithOptionalUpperBoundTest(test_util.TestCase):
+
+  def testBasic(self):
+    distribution = inverse_gamma.InverseGamma(16., 4.)
+    seed = samplers.sanitize_seed((0, 1))
+    unbounded_result = sample_parameters.sample_with_optional_upper_bound(
+        distribution, seed=seed)
+    # Whatever the result, we want to get that value again minus a fixed offset.
+    # Since we fix the seed, we know it is equal just because of the upper
+    # bound.
+    new_target_result = unbounded_result - 0.5
+    distribution.upper_bound = new_target_result
+    self.assertAllEqual(
+        new_target_result,
+        sample_parameters.sample_with_optional_upper_bound(
+            distribution, seed=seed))
+
+
+if __name__ == '__main__':
+  test_util.main()

tensorflow_probability/python/sts/components/BUILD

@@ -291,13 +291,18 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:chain",
+        "//tensorflow_probability/python/bijectors:invert",
         "//tensorflow_probability/python/bijectors:scale",
         "//tensorflow_probability/python/bijectors:softplus",
+        "//tensorflow_probability/python/bijectors:square",
+        "//tensorflow_probability/python/distributions:inverse_gamma",
         "//tensorflow_probability/python/distributions:linear_gaussian_ssm",
         "//tensorflow_probability/python/distributions:lognormal",
         "//tensorflow_probability/python/distributions:mvn_diag",
         "//tensorflow_probability/python/distributions:mvn_tril",
         "//tensorflow_probability/python/distributions:normal",
+        "//tensorflow_probability/python/distributions:transformed_distribution",
+        "//tensorflow_probability/python/experimental/sts_gibbs:sample_parameters",
         "//tensorflow_probability/python/internal:docstring_util",
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:prefer_static",
@@ -309,7 +314,7 @@ multi_substrate_py_test(
     name = "seasonal_test",
     size = "medium",
     srcs = ["seasonal_test.py"],
-    shard_count = 6,
+    shard_count = 12,
     deps = [
         ":seasonal",
         # numpy dep,