Add utility to STS to build a default model from a Pandas series or dataframe.

davmre · tensorflower-gardener · commit 4e176ba0f313 · 2021-07-08T19:53:43.000-07:00
Q: Should this support other model components, like regression against user-provided covariates?
A: I think it's best to start simple; we can always add options later. An alternative would be to make it easy to extend the default model, e.g., implement an  `add_component` method for `Sum` instances.

Q: What should we name this?
A: I've been swinging between 'auto' and 'default' prefixes, and we could also try to get at the fact that what we're specifically automating is constructing seasonal components from a list of dates. For example:
 - `tfp.sts.auto_build_model`
 - `tfp.sts.build_default_model`
 - `tfp.sts.default_model`
 - `tfp.sts.model_from_dates`
Thoughts welcome.
PiperOrigin-RevId: 383763071
diff --git a/tensorflow_probability/python/sts/BUILD b/tensorflow_probability/python/sts/BUILD
@@ -29,6 +29,7 @@ py_library(
     srcs_version = "PY3",
     deps = [
         ":decomposition",
+        ":default_model",
         ":fitting",
         ":forecast",
         ":regularization",
@@ -66,6 +67,38 @@ py_test(
     ],
 )
 
+py_library(
+    name = "default_model",
+    srcs = ["default_model.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":regularization",
+        ":structural_time_series",
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/distributions:inverse_gamma",
+        "//tensorflow_probability/python/distributions:normal",
+        "//tensorflow_probability/python/experimental/sts_gibbs",
+        "//tensorflow_probability/python/experimental/util",
+        "//tensorflow_probability/python/sts/components",
+        "//tensorflow_probability/python/sts/internal",
+    ],
+)
+
+py_test(
+    name = "default_model_test",
+    size = "medium",
+    srcs = ["default_model_test.py"],
+    deps = [
+        # absl/testing:parameterized dep,
+        # numpy dep,
+        # pandas dep,
+        # tensorflow dep,
+        "//tensorflow_probability",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)
+
 py_library(
     name = "fitting",
     srcs = ["fitting.py"],
diff --git a/tensorflow_probability/python/sts/default_model.py b/tensorflow_probability/python/sts/default_model.py
@@ -0,0 +1,174 @@
+# Copyright 2021 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.
+# ============================================================================
+"""Utilities for automatically building StructuralTimeSeries models."""
+
+import tensorflow.compat.v2 as tf
+
+from tensorflow_probability.python.sts import components as sts_components
+from tensorflow_probability.python.sts import structural_time_series
+from tensorflow_probability.python.sts.internal import seasonality_util
+from tensorflow_probability.python.sts.internal import util as sts_util
+
+__all__ = [
+    'build_default_model',
+]
+
+
+# TODO(davmre): before exposing publicly, consider simplifying this function
+# (e.g., not exposing prior specification args) and/or renaming it to something
+# like `auto_build_model`.
+def build_default_model(observed_time_series,
+                        base_component=sts_components.LocalLinearTrend,
+                        observation_noise_scale_prior=None,
+                        drift_scale_prior=None,
+                        allow_seasonal_effect_drift=True,
+                        name=None):
+  """Builds a model with seasonality from a Pandas Series or DataFrame.
+
+  Returns a model of the form
+  `tfp.sts.Sum([base_component] + seasonal_components)`, where
+  `seasonal_components` are automatically selected using the frequency from the
+  `DatetimeIndex` of the provided `pd.Series` or `pd.DataFrame`. If the index
+  does not have a set frequency, one will be inferred from the index dates, and
+
+  Args:
+    observed_time_series: Instance of `pd.Series` or `pd.DataFrame` containing
+      one or more time series indexed by a `pd.DatetimeIndex`.
+    base_component: Optional subclass of `tfp.sts.StructuralTimeSeries`
+      specifying the model used for residual variation in the series not
+      explained by seasonal or other effects. May also be an *instance* of such
+      a class with specific priors set; if not provided, such an instance will
+      be constructed with heuristic default priors.
+      Default value: `tfp.sts.LocalLinearTrend`.
+    observation_noise_scale_prior: Optional `tfd.Distribution` instance
+      specifying a prior on `observation_noise_scale`. If `None`, a heuristic
+      default prior is constructed based on the provided `observed_time_series`.
+      Default value: `None`.
+    drift_scale_prior: Optional `tfd.Distribution` instance
+      specifying a prior on the `drift_scale` parameter of Seasonal components.
+      If `None`, a heuristic default prior is constructed based on the provided
+      `observed_time_series`.
+      Default value: `None`.
+    allow_seasonal_effect_drift: optional Python `bool` specifying whether the
+      seasonal effects can drift over time.  Setting this to `False`
+      removes the `drift_scale` parameter from the model. This is
+      mathematically equivalent to `drift_scale_prior = tfd.Deterministic(0.)`,
+      but removing drift directly is preferred because it avoids the use of a
+      degenerate prior.
+      Default value: `True`.
+    name: Python `str` name for ops created by this function.
+      Default value: `None` (i.e., 'build_default_model').
+  Returns:
+    model: instance of `tfp.sts.Sum` representing a model for the given data.
+
+  #### Example
+
+  Consider a series of eleven data points, covering a period of two weeks
+  with three missing days.
+
+  ```python
+  import pandas as pd
+  import tensorflow as tf
+  import tensorflow_probability as tfp
+
+  series = pd.Series(
+    [100., 27., 92., 66., 51., 126., 113., 95., 48., 20., 59.,],
+    index=pd.to_datetime(['2020-01-01', '2020-01-02', '2020-01-04',
+                          '2020-01-05', '2020-01-06', '2020-01-07',
+                          '2020-01-10', '2020-01-11', '2020-01-12',
+                          '2020-01-13', '2020-01-14']))
+  ```
+
+  Before calling `build_default_model`, we must regularize the series to follow
+  a fixed frequency (here, daily observations):
+
+  ```python
+  series = tfp.sts.regularize_series(series)
+  # len(series) ==> 14
+  ```
+
+  The default model will combine a LocalLinearTrend baseline with a Seasonal
+  component to capture day-of-week effects. We can then fit this model to our
+  observed data. Here we'll use variational inference:
+
+  ```python
+  model = tfp.sts.build_default_model(series)
+  # len(model.components) == 2
+
+  # Fit the model using variational inference.
+  surrogate_posterior = tfp.sts.build_factored_surrogate_posterior(model)
+  losses = tfp.vi.fit_surrogate_posterior(
+    target_log_prob_fn=model.joint_log_prob(series),
+    surrogate_posterior=surrogate_posterior,
+    optimizer=tf.optimizers.Adam(0.1),
+    num_steps=1000,
+    convergence_criterion=(
+      tfp.optimizer.convergence_criteria.SuccessiveGradientsAreUncorrelated(
+        window_size=20, min_num_steps=50)),
+    jit_compile=True)
+  parameter_samples = surrogate_posterior.sample(50)
+  ```
+
+  Finally, use the fitted parameters to forecast the next week of data:
+
+  ```python
+  forecast_dist = tfp.sts.forecast(model,
+                                   observed_time_series=series,
+                                   parameter_samples=parameter_samples,
+                                   num_steps_forecast=7)
+  # Strip trailing unit dimension from LinearGaussianStateSpaceModel events.
+  forecast_mean = forecast_dist.mean()[..., 0]
+  forecast_stddev = forecast_dist.stddev()[..., 0]
+
+  forecast = pd.DataFrame(
+      {'mean': forecast_mean,
+       'lower_bound': forecast_mean - 2. * forecast_stddev,
+       'upper_bound': forecast_mean + 2. * forecast_stddev}
+      index=pd.date_range(start=series.index[-1] + series.index.freq,
+                          periods=7,
+                          freq=series.index.freq))
+  ```
+
+  """
+  with tf.name_scope(name or 'build_default_model'):
+    frequency = getattr(observed_time_series.index, 'freq', None)
+    if frequency is None:
+      raise ValueError('Provided series has no set frequency. Consider '
+                       'using `tfp.sts.regularize_series` to infer a frequency '
+                       'and build a regularly spaced series.')
+    observed_time_series = sts_util.canonicalize_observed_time_series_with_mask(
+        observed_time_series)
+
+    if not isinstance(base_component,
+                      structural_time_series.StructuralTimeSeries):
+      # Build a component of the given type using default priors.
+      base_component = base_component(observed_time_series=observed_time_series)
+
+    components = [base_component]
+    seasonal_structure = seasonality_util.create_seasonal_structure(
+        frequency=frequency,
+        num_steps=int(observed_time_series.time_series.shape[-2]))
+    for season_type, season in seasonal_structure.items():
+      components.append(
+          sts_components.Seasonal(num_seasons=season.num,
+                                  num_steps_per_season=season.duration,
+                                  drift_scale_prior=drift_scale_prior,
+                                  allow_drift=allow_seasonal_effect_drift,
+                                  observed_time_series=observed_time_series,
+                                  name=str(season_type)))
+    return sts_components.Sum(
+        components,
+        observed_time_series=observed_time_series,
+        observation_noise_scale_prior=observation_noise_scale_prior)
diff --git a/tensorflow_probability/python/sts/default_model_test.py b/tensorflow_probability/python/sts/default_model_test.py
@@ -0,0 +1,137 @@
+# Copyright 2021 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 automatically building StructuralTimeSeries models."""
+
+# Dependency imports
+
+import pandas as pd
+import tensorflow.compat.v2 as tf
+import tensorflow_probability as tfp
+
+from tensorflow_probability.python.internal import test_util
+from tensorflow_probability.python.sts import default_model
+
+tfb = tfp.bijectors
+tfd = tfp.distributions
+
+
+class DefaultModelTests(test_util.TestCase):
+
+  def _build_test_series(self, shape, freq, start='2020-01-01 00:00:00'):
+    values = self.evaluate(tf.random.stateless_normal(
+        shape, seed=test_util.test_seed(sampler_type='stateless')))
+    index = pd.date_range('2020-01-01 00:00:00',
+                          periods=shape[0],
+                          freq=freq)
+    if len(shape) > 1:
+      num_columns = shape[1]
+      return pd.DataFrame(values,
+                          columns=['series{}'.format(i)
+                                   for i in range(num_columns)],
+                          index=index)
+    else:
+      return pd.Series(values, index=index)
+
+  def test_has_expected_seasonality(self):
+    model = default_model.build_default_model(
+        self._build_test_series(shape=[168 * 2], freq=pd.DateOffset(hours=1)))
+
+    self.assertIsInstance(model, tfp.sts.Sum)
+    self.assertLen(model.components, 3)
+    self.assertIsInstance(model.components[0], tfp.sts.LocalLinearTrend)
+    self.assertIsInstance(model.components[1], tfp.sts.Seasonal)
+    self.assertContainsSubsequence(model.components[1].name, 'HOUR_OF_DAY')
+    self.assertIsInstance(model.components[2], tfp.sts.Seasonal)
+    self.assertContainsSubsequence(model.components[2].name, 'DAY_OF_WEEK')
+
+  def test_explicit_base_component_and_priors(self):
+    series = self._build_test_series(shape=[48], freq=pd.DateOffset(hours=1))
+    model = default_model.build_default_model(
+        series,
+        base_component=tfp.sts.SemiLocalLinearTrend(
+            level_scale_prior=tfd.Exponential(5.),
+            slope_scale_prior=tfd.Exponential(0.1),
+            slope_mean_prior=tfd.Normal(0., 100.),
+            constrain_ar_coef_positive=True,
+            constrain_ar_coef_stationary=True,
+            observed_time_series=series),
+        observation_noise_scale_prior=tfd.Exponential(3.),
+        drift_scale_prior=tfd.Exponential(1.))
+    self.assertLen(model.components, 2)
+
+    param_by_name = lambda n: [p for p in model.parameters if n in p.name][0]
+    self.assertAllClose(param_by_name('level_scale').prior.rate, 5.)
+    self.assertAllClose(param_by_name('slope_mean').prior.scale, 100.)
+    self.assertAllClose(param_by_name('slope_scale').prior.rate, 0.1)
+    self.assertAllClose(param_by_name('drift_scale').prior.rate, 1.)
+    self.assertAllClose(param_by_name('observation_noise_scale').prior.rate, 3.)
+
+  def test_creates_batch_model_from_multiple_series(self):
+    model = default_model.build_default_model(
+        self._build_test_series(shape=[48, 3], freq=pd.DateOffset(hours=1)))
+    self.assertAllEqual(model.batch_shape, [3])
+
+  def test_docstring_fitting_example(self):
+    # Construct a series of eleven data points, covering a period of two weeks
+    # with three missing days.
+    series = pd.Series(
+        [100., 27., 92., 66., 51., 126., 113., 95., 48., 20., 59.,],
+        index=pd.to_datetime(['2020-01-01', '2020-01-02', '2020-01-04',
+                              '2020-01-05', '2020-01-06', '2020-01-07',
+                              '2020-01-10', '2020-01-11', '2020-01-12',
+                              '2020-01-13', '2020-01-14']))
+    series = tfp.sts.regularize_series(series)
+    self.assertLen(series, 14)
+
+    # Default model captures day-of-week effects with a LocalLinearTrend
+    # baseline.
+    model = default_model.build_default_model(series)
+    self.assertLen(model.components, 2)
+
+    # Fit the model using variational inference.
+    surrogate_posterior = tfp.sts.build_factored_surrogate_posterior(model)
+    _ = tfp.vi.fit_surrogate_posterior(
+        target_log_prob_fn=model.joint_log_prob(series),
+        surrogate_posterior=surrogate_posterior,
+        optimizer=tf.optimizers.Adam(0.1),
+        num_steps=1000,
+        convergence_criterion=(tfp.optimizer.convergence_criteria.
+                               SuccessiveGradientsAreUncorrelated(
+                                   window_size=15, min_num_steps=50)),
+        jit_compile=True)
+
+    # Forecast the next week.
+    parameter_samples = surrogate_posterior.sample(50)
+    forecast_dist = tfp.sts.forecast(model,
+                                     observed_time_series=series,
+                                     parameter_samples=parameter_samples,
+                                     num_steps_forecast=7)
+    # Strip trailing unit dimension from LinearGaussianStateSpaceModel events.
+    self.evaluate(
+        [v.initializer for v in surrogate_posterior.trainable_variables])
+    forecast_mean, forecast_stddev = self.evaluate(
+        (forecast_dist.mean()[..., 0], forecast_dist.stddev()[..., 0]))
+
+    pd.DataFrame(
+        {'mean': forecast_mean,
+         'lower_bound': forecast_mean - 2. * forecast_stddev,
+         'upper_bound': forecast_mean + 2. * forecast_stddev},
+        index=pd.date_range(start=series.index[-1] + series.index.freq,
+                            periods=7,
+                            freq=series.index.freq))
+
+
+if __name__ == '__main__':
+  tf.test.main()
