Fix an embarassing off-by-one error in STS anomaly detection predictive distributions.

I'd understood `tfp.sts.one_step_predictive` to return a distribution over timesteps `1:T + 1`, so the previous code went to a lot of effort to send it a series padded with an extra initial step in order to get results for steps `0:T` (matching the Gibbs predictive dist). But it turns out that I was wrong, and this padding is not only unnecessary but actually screws things up.

This change also fixes a dtype issue with the Gibbs predictive distribution that came up in testing the above fix.

PiperOrigin-RevId: 388789203

Author: davmre

tensorflow_probability/python/experimental/sts_gibbs/gibbs_sampler.py

@@ -340,9 +340,9 @@ def one_step_predictive(model,
       distribution of each timestep given previous timesteps.
   """
   dtype = dtype_util.common_dtype([
-      posterior_samples.level_scale.dtype,
-      posterior_samples.observation_noise_scale.dtype,
-      posterior_samples.level.dtype,
+      posterior_samples.level_scale,
+      posterior_samples.observation_noise_scale,
+      posterior_samples.level,
       original_mean,
       original_scale], dtype_hint=tf.float32)
   num_observed_steps = prefer_static.shape(posterior_samples.level)[-1]

tensorflow_probability/python/experimental/sts_gibbs/gibbs_sampler_test.py

@@ -202,7 +202,7 @@ def reshape_chain_and_sample(x):
 
   @parameterized.named_parameters(
       {'testcase_name': 'float32_xla', 'dtype': tf.float32, 'use_xla': True},
-      {'testcase_name': 'float16', 'dtype': tf.float16, 'use_xla': False})
+      {'testcase_name': 'float64', 'dtype': tf.float64, 'use_xla': False})
   def test_end_to_end_prediction_works_and_is_deterministic(
       self, dtype, use_xla):
     if not tf.executing_eagerly():
@@ -211,7 +211,8 @@ def test_end_to_end_prediction_works_and_is_deterministic(
     model, observed_time_series, is_missing = self._build_test_model(
         num_timesteps=5,
         batch_shape=[3],
-        prior_class=gibbs_sampler.XLACompilableInverseGamma)
+        prior_class=gibbs_sampler.XLACompilableInverseGamma,
+        dtype=dtype)
 
     @tf.function(jit_compile=use_xla)
     def do_sampling(observed_time_series, is_missing):

tensorflow_probability/python/sts/anomaly_detection/anomaly_detection_lib.py

@@ -24,7 +24,6 @@
 from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.sts import regularization
 from tensorflow_probability.python.sts.forecast import one_step_predictive
-from tensorflow_probability.python.sts.internal import missing_values_util
 from tensorflow_probability.python.sts.internal import seasonality_util
 from tensorflow_probability.python.sts.internal import util as sts_util
 
@@ -177,24 +176,8 @@ def _detect_anomalies_inner(observed_time_series,
       predictive_dist = gibbs_sampler.one_step_predictive(model,
                                                           posterior_samples)
     else:
-      # Build the filtering predictive distribution.
-      # First, pad the time series with an unobserved point at time -1, so that
-      # we get a 'prediction' (which will just be the prior) at time 0, matching
-      # the shape of the input (and of the Gibbs predictive distribution).
-      is_missing = observed_time_series.is_missing
-      if is_missing is None:
-        is_missing = tf.zeros(tf.shape(observed_time_series.time_series)[:-1],
-                              dtype=tf.bool)
-      initial_value = observed_time_series.time_series[..., 0:1, :]
-      padded_series = missing_values_util.MaskedTimeSeries(
-          time_series=tf.concat(
-              [initial_value,
-               observed_time_series.time_series[..., :-1, :]], axis=-2),
-          is_missing=tf.concat(
-              [tf.ones(tf.shape(initial_value)[:-1], dtype=tf.bool),
-               is_missing[..., :-1]], axis=-1))
       predictive_dist = one_step_predictive(model,
-                                            padded_series,
+                                            observed_time_series,
                                             timesteps_are_event_shape=False,
                                             parameter_samples=parameter_samples)
 

tensorflow_probability/python/sts/anomaly_detection/anomaly_detection_test.py

@@ -128,6 +128,33 @@ def test_constant_series(self):
                         tf.ones_like(predictions.mean),
                         atol=0.1)
 
+  @parameterized.named_parameters(('', False),
+                                  ('_gibbs_predictive', True))
+  def test_predictions_align_with_series(self, use_gibbs_predictive_dist):
+    np.random.seed(0)
+    # Simulate data with very clear daily and hourly effects, so that an
+    # off-by-one error will almost certainly lead to out-of-bounds predictions.
+    daily_effects = [100., 0., 20., -50., -100., -20., 70.]
+    hourly_effects = [
+        20., 0., 10., -10., 0., -20., -10., -30., -15., -5., -10., 0.] * 2
+    effects = [daily_effects[(t // 24) % 7] + hourly_effects[t % 24]
+               for t in range(24 * 7 * 2)]
+    series = pd.Series(effects + np.random.randn(len(effects)),
+                       index=pd.date_range('2020-01-01',
+                                           periods=len(effects),
+                                           freq=pd.DateOffset(hours=1)))
+    predictions = anomaly_detection.detect_anomalies(
+        series,
+        seed=test_util.test_seed(sampler_type='stateless'),
+        num_samples=100,
+        num_warmup_steps=50,
+        use_gibbs_predictive_dist=use_gibbs_predictive_dist,
+        jit_compile=False)
+    # An off-by-one error in the predictive distribution would generate
+    # anomalies at most steps.
+    num_anomalies = tf.reduce_sum(tf.cast(predictions.is_anomaly, tf.int32))
+    self.assertLessEqual(self.evaluate(num_anomalies), 5)
+
 
 if __name__ == '__main__':
   tf.test.main()