Preparation change for seasonality support in the GibbsSampler:

* Stop assuming Seasonal components are not supported in tests.
* Have a common function for deconstructing the list of STS component
* Stop requiring the STS component be in a certain order (not needed for seasonality support, but moves towards a world where build_model_for_gibbs_fitting is not required)
* Small typo fix in a test
* Small test improvement to share the batch value

PiperOrigin-RevId: 475990211

Author: kylepl

tensorflow_probability/python/experimental/sts_gibbs/BUILD

@@ -118,6 +118,7 @@ multi_substrate_py_test(
         "//tensorflow_probability/python/sts/components:local_linear_trend",
         "//tensorflow_probability/python/sts/components:regression",
         "//tensorflow_probability/python/sts/components:seasonal",
+        "//tensorflow_probability/python/sts/components:semilocal_linear_trend",
         "//tensorflow_probability/python/sts/components:sum",
         "//tensorflow_probability/python/sts/internal:missing_values_util",
 #         "//third_party/tensorflow/compiler/jit:xla_cpu_jit",  # DisableOnExport

tensorflow_probability/python/experimental/sts_gibbs/gibbs_sampler.py

@@ -395,10 +395,6 @@ def fit_with_gibbs_sampling(model,
   Returns:
     model: A `GibbsSamplerState` structure of posterior samples.
   """
-  if not hasattr(model, 'supports_gibbs_sampling'):
-    raise ValueError('This STS model does not support Gibbs sampling. Models '
-                     'for Gibbs sampling must be created using the '
-                     'method `build_model_for_gibbs_fitting`.')
   if not tf.nest.is_nested(num_chains):
     num_chains = [num_chains]
 
@@ -420,7 +416,8 @@ def fit_with_gibbs_sampling(model,
   # the slope_scale is always zero.
   initial_slope_scale = 0.
   initial_slope = 0.
-  if isinstance(model.components[0], local_linear_trend.LocalLinearTrend):
+  (_, level_component), _ = (_get_components_from_model(model))
+  if isinstance(level_component, local_linear_trend.LocalLinearTrend):
     initial_slope_scale = 1. * tf.ones(batch_shape, dtype=dtype)
     initial_slope = tf.zeros(level_slope_shape, dtype=dtype)
 
@@ -476,17 +473,22 @@ def model_parameter_samples_from_gibbs_samples(model, gibbs_samples):
    A set of posterior samples, that can be used with `make_state_space_model`
    or `sts.forecast`.
   """
-  if not hasattr(model, 'supports_gibbs_sampling'):
-    raise ValueError('This STS model does not support Gibbs sampling. Models '
-                     'for Gibbs sampling must be created using the '
-                     'method `build_model_for_gibbs_fitting`.')
+  # Make use of the indices in the model to avoid requiring a specific
+  # order of components.
+  ((level_component_index, level_component),
+   (regression_component_index, _)) = (
+       _get_components_from_model(model))
+  model_parameter_samples = (gibbs_samples.observation_noise_scale,)
 
-  if isinstance(model.components[0], local_linear_trend.LocalLinearTrend):
-    return (gibbs_samples.observation_noise_scale, gibbs_samples.level_scale,
-            gibbs_samples.slope_scale, gibbs_samples.weights)
-  else:
-    return (gibbs_samples.observation_noise_scale, gibbs_samples.level_scale,
-            gibbs_samples.weights)
+  for index in range(len(model.components)):
+    if index == level_component_index:
+      model_parameter_samples += (gibbs_samples.level_scale,)
+      if isinstance(level_component, local_linear_trend.LocalLinearTrend):
+        model_parameter_samples += (gibbs_samples.slope_scale,)
+    elif index == regression_component_index:
+      model_parameter_samples += (gibbs_samples.weights,)
+
+  return model_parameter_samples
 
 
 def one_step_predictive(model,
@@ -582,7 +584,7 @@ def one_step_predictive(model,
         tf.range(1., num_forecast_steps + 1., dtype=forecast_level.dtype))
 
   level_pred = tf.concat(
-      ([thinned_samples.level] if use_zero_step_prediction else [
+      ([thinned_samples.level] if use_zero_step_prediction else [  # pylint:disable=g-long-ternary
           thinned_samples.level[..., :1],  # t == 0
           (thinned_samples.level[..., :-1] + thinned_samples.slope[..., :-1]
           )  # 1 <= t < T. Constructs the next level from previous level
@@ -784,6 +786,56 @@ def _resample_scale(prior, observed_residuals, is_missing=None, seed=None):
       posterior, seed=seed)
 
 
+def _get_components_from_model(model):
+  """Returns the split-apart components from an STS model.
+
+  Args:
+    model: A `tf.sts.StructuralTimeSeries` to split apart.
+
+  Returns:
+    A tuple of (index and level component, index and regression component) or
+    an exception. The regression component may be None (along with its index -
+    (None, None)).
+
+    Each 'index and component' is a tuple of (index, component), where index
+    is the position in the model.
+  """
+  if not hasattr(model, 'supports_gibbs_sampling'):
+    raise ValueError('This STS model does not support Gibbs sampling. Models '
+                     'for Gibbs sampling must be created using the '
+                     'method `build_model_for_gibbs_fitting`.')
+
+  level_components = []
+  regression_components = []
+
+  for index, component in enumerate(model.components):
+    if (isinstance(component, local_level.LocalLevel) or
+        isinstance(component, local_linear_trend.LocalLinearTrend)):
+      level_components.append((index, component))
+    elif (isinstance(component, regression.LinearRegression) or
+          isinstance(component, SpikeAndSlabSparseLinearRegression)):
+      regression_components.append((index, component))
+    else:
+      raise NotImplementedError(
+          'Found unsupported model component for Gibbs Sampling: {}'.format(
+              component))
+
+  if len(level_components) != 1:
+    raise ValueError(
+        'Expected exactly one level component, found {} components.'.format(
+            len(level_components)))
+  level_component = level_components[0]
+
+  regression_component = (None, None)
+  if len(regression_components) > 1:
+    raise ValueError(
+        'Expected at most one regression component, found {} components.'
+        .format(len(regression_components)))
+  elif len(regression_components) == 1:
+    regression_component = regression_components[0]
+  return level_component, regression_component
+
+
 def _build_sampler_loop_body(model,
                              observed_time_series,
                              is_missing=None,
@@ -819,29 +871,12 @@ def _build_sampler_loop_body(model,
   """
   if JAX_MODE and experimental_use_dynamic_cholesky:
     raise ValueError('Dynamic Cholesky decomposition not supported in JAX')
-  level_component = model.components[0]
-  if not (isinstance(level_component, local_level.LocalLevel) or
-          isinstance(level_component, local_linear_trend.LocalLinearTrend)):
-    raise ValueError(
-        'Expected the first model component to be an instance of '
-        '`tfp.sts.LocalLevel` or `tfp.local_linear_trend.LocalLinearTrend`; '
-        'instead saw {}'.format(level_component))
+  (_, level_component), (_, regression_component) = (
+      _get_components_from_model(model))
   model_has_slope = isinstance(level_component,
                                local_linear_trend.LocalLinearTrend)
 
-  # TODO(kloveless): When we add support for more flexible models, remove
-  # this assumption.
-  regression_component = (None if len(model.components) != 2 else
-                          model.components[1])
-  if regression_component:
-    if not (isinstance(regression_component, regression.LinearRegression) or
-            isinstance(regression_component,
-                       SpikeAndSlabSparseLinearRegression)):
-      raise ValueError(
-          'Expected the second model component to be an instance of '
-          '`tfp.sts.LinearRegression` or '
-          '`SpikeAndSlabSparseLinearRegression`; '
-          'instead saw {}'.format(regression_component))
+  if regression_component is not None:
     model_has_spike_slab_regression = isinstance(
         regression_component, SpikeAndSlabSparseLinearRegression)
 

tensorflow_probability/python/experimental/sts_gibbs/gibbs_sampler_test.py

@@ -17,7 +17,6 @@
 from absl.testing import parameterized
 
 import numpy as np
-
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python.distributions import inverse_gamma
@@ -34,7 +33,7 @@
 from tensorflow_probability.python.sts.components import local_level
 from tensorflow_probability.python.sts.components import local_linear_trend
 from tensorflow_probability.python.sts.components import regression
-from tensorflow_probability.python.sts.components import seasonal
+from tensorflow_probability.python.sts.components import semilocal_linear_trend
 from tensorflow_probability.python.sts.components import sum as sum_lib
 from tensorflow_probability.python.sts.forecast import forecast
 from tensorflow_probability.python.sts.internal import missing_values_util
@@ -189,11 +188,12 @@ def test_forecasts_match_reference(self,
     if not tf.nest.is_nested(num_chains):
       num_results = num_results // num_chains
 
+    batch_shape = [3]
     model, observed_time_series, is_missing = self._build_test_model(
         num_timesteps=num_observed_steps + num_forecast_steps,
         true_slope_scale=0.5 if use_slope else None,
-        batch_shape=[3],
-        time_series_shift=time_series_shift)
+        time_series_shift=time_series_shift,
+        batch_shape=batch_shape)
 
     @tf.function(autograph=False)
     def do_sampling():
@@ -226,9 +226,9 @@ def reshape_chain_and_sample(x):
     predictive_mean, predictive_stddev = self.evaluate((
         predictive_dist.mean(), predictive_dist.stddev()))
     self.assertAllEqual(predictive_mean.shape,
-                        [3, num_observed_steps + num_forecast_steps])
+                        batch_shape + [num_observed_steps + num_forecast_steps])
     self.assertAllEqual(predictive_stddev.shape,
-                        [3, num_observed_steps + num_forecast_steps])
+                        batch_shape + [num_observed_steps + num_forecast_steps])
 
     # big tolerance, but makes sure the predictive mean initializes near
     # the initial time series value
@@ -457,31 +457,23 @@ def test_invalid_model_raises_error(self):
     ],
                             observed_time_series=observed_time_series)
 
-    with self.assertRaisesRegexp(ValueError, 'does not support Gibbs sampling'):
+    with self.assertRaisesRegex(ValueError, 'does not support Gibbs sampling'):
       gibbs_sampler.fit_with_gibbs_sampling(
           bad_model, observed_time_series, seed=test_util.test_seed())
 
     bad_model.supports_gibbs_sampling = True
-    with self.assertRaisesRegexp(
-        ValueError, 'Expected the first model component to be an instance of'):
-      gibbs_sampler.fit_with_gibbs_sampling(
-          bad_model, observed_time_series, seed=test_util.test_seed())
-
     bad_model_with_correct_params = sum_lib.Sum([
-        # A seasonal model with no drift has no parameters, so adding it
-        # won't break the check for correct params.
-        seasonal.Seasonal(
-            num_seasons=2,
-            allow_drift=False,
+        # An unsupported model component.
+        semilocal_linear_trend.SemiLocalLinearTrend(
             observed_time_series=observed_time_series),
         local_level.LocalLevel(observed_time_series=observed_time_series),
         regression.LinearRegression(design_matrix=tf.ones([5, 2]))
     ])
     bad_model_with_correct_params.supports_gibbs_sampling = True
 
-    with self.assertRaisesRegexp(ValueError,
-                                 'Expected the first model component to be an '
-                                 'instance of `tfp.sts.LocalLevel`'):
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        'Found unsupported model component for Gibbs Sampling'):
       gibbs_sampler.fit_with_gibbs_sampling(bad_model_with_correct_params,
                                             observed_time_series,
                                             seed=test_util.test_seed())
@@ -673,13 +665,13 @@ def do_sampling():
   @parameterized.named_parameters(
       {
           'testcase_name': 'Rank1Updates',
-          'use_dyanamic_cholesky': False,
+          'use_dynamic_cholesky': False,
       }, {
           'testcase_name': 'DynamicCholesky',
-          'use_dyanamic_cholesky': True,
+          'use_dynamic_cholesky': True,
       })
-  def test_sparse_regression_recovers_plausible_weights(
-      self, use_dyanamic_cholesky):
+  def test_sparse_regression_recovers_plausible_weights(self,
+                                                        use_dynamic_cholesky):
     true_weights = tf.constant([0., 0., 2., 0., -2.])
     model, observed_time_series, _ = self._build_test_model(
         num_timesteps=20,
@@ -698,9 +690,9 @@ def do_sampling():
           num_results=100,
           num_warmup_steps=100,
           seed=test_util.test_seed(sampler_type='stateless'),
-          experimental_use_dynamic_cholesky=use_dyanamic_cholesky)
+          experimental_use_dynamic_cholesky=use_dynamic_cholesky)
 
-    if JAX_MODE and use_dyanamic_cholesky:
+    if JAX_MODE and use_dynamic_cholesky:
       with self.assertRaises(ValueError):
         self.evaluate(do_sampling())
       return