Create STS module for modeling holiday effects.

PiperOrigin-RevId: 389190872

Author: g-luo

tensorflow_probability/python/sts/BUILD

@@ -32,6 +32,7 @@ py_library(
         ":default_model",
         ":fitting",
         ":forecast",
+        ":holiday_effects",
         ":regularization",
         ":structural_time_series",
         "//tensorflow_probability/python/internal:all_util",
@@ -161,6 +162,34 @@ py_test(
     ],
 )
 
+py_library(
+    name = "holiday_effects",
+    srcs = ["holiday_effects.py"],
+    srcs_version = "PY3",
+    deps = [
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/experimental/util",
+        "//tensorflow_probability/python/sts/internal",
+    ],
+)
+
+py_test(
+    name = "holiday_effects_test",
+    size = "medium",
+    srcs = ["holiday_effects_test.py"],
+    python_version = "PY3",
+    shard_count = 6,
+    srcs_version = "PY3",
+    deps = [
+        # holidays dep,
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)
+
 py_library(
     name = "regularization",
     srcs = ["regularization.py"],

tensorflow_probability/python/sts/holiday_effects.py

@@ -0,0 +1,127 @@
+# 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 holiday regressors."""
+
+__all__ = [
+    'create_holiday_regressors',
+]
+
+# Defines expected holiday file fields.
+_HOLIDAY_FILE_FIELDS = frozenset({'geo', 'holiday', 'date'})
+
+
+def create_holiday_regressors(times, holidays):
+  """Creates a design matrix of holiday regressors for a given time series.
+
+  Args:
+    times: a Pandas `DatetimeIndex` that indexes time series data.
+    holidays: a Pandas `DataFrame` containing the dates of holidays. The
+      `DataFrame` should have the following columns:
+      * `geo`: `str`, two-letter upper-case country code
+      * `holiday`: `str`, the name of the holiday
+      * `date`: `str`, dates in the form of `YYYY-MM-DD`
+
+  Returns:
+    holiday_regressors: a Pandas `DataFrame` where the columns are the names of
+      holidays. This matrix of one hot encodings is shape
+      (N, H), where N is the length of `times` and H is the number of unique
+      holiday names in `holidays.holiday`.
+  """
+  # pylint: disable=g-import-not-at-top
+  import pandas as pd
+  # pylint: enable=g-import-not-at-top
+
+  # TODO(b/195346554): Expand fixed holidays.
+  _check_times(times)
+  _check_holidays(holidays)
+
+  holidays = holidays.sort_values('date')
+
+  holiday_types = list(holidays.holiday.unique())
+  holiday_regressors = pd.DataFrame()
+  for holiday in holiday_types:
+    holiday_dates = holidays.loc[holidays.holiday == holiday]
+    holiday_dates = pd.to_datetime(
+        list(holiday_dates.date), errors='raise', format='%Y-%m-%d')
+    holiday_regressors.loc[:, holiday] = _match_dates(times, holiday_dates)
+
+  # Remove all regressors with only zeros.
+  holiday_regressors = (
+      holiday_regressors.loc[:, (holiday_regressors != 0).any(axis=0)])
+  return holiday_regressors
+
+
+def _check_times(times):
+  """Checks that times are in the correct format.
+
+  Args:
+    times: a Pandas `DatetimeIndex` that indexes time series data.
+
+  Raises:
+    ValueError: if times is not a Pandas `DatetimeIndex` or does not have a
+    frequency.
+  """
+  # pylint: disable=g-import-not-at-top
+  import pandas as pd
+  # pylint: enable=g-import-not-at-top
+  if not isinstance(times, pd.core.indexes.datetimes.DatetimeIndex):
+    raise ValueError('Times is not a Pandas DatetimeIndex.')
+  if not times.freq:
+    raise ValueError('Times does not have a frequency.')
+
+
+def _check_holidays(holidays):
+  """Checks that holiday files are in the correct format.
+
+  Args:
+    holidays: a Pandas `DataFrame` containing the dates of holidays.
+
+  Raises:
+    ValueError: if the holidays column names are improperly formatted.
+  """
+  all_column_names = _HOLIDAY_FILE_FIELDS.issubset(holidays.columns)
+  if not all_column_names:
+    raise ValueError(
+        'Holidays column names must contain: {0}.'.format(_HOLIDAY_FILE_FIELDS))
+
+
+def _match_dates(times, dates):
+  """Creates a 0-1 dummy variable that marks every instance of dates that also occurs in times.
+
+  Args:
+    times: a Pandas `DatetimeIndex` for observed data.
+    dates: a Pandas `DatetimeIndex` for relevant dates of a single holiday.
+
+  Returns:
+    regressor: a list the same length as `times`, with a 1 where there is a date
+    match, and otherwise 0.
+  """
+  regressor = [0] * len(times)
+  rounded_times = times.floor('d')
+  rounded_dates = dates.floor('d')
+
+  # TODO(b/195347492): Approximate to the nearest prior day
+  # for greater than daily granularity.
+  # TODO(b/195347492): Add _MIN_HOLIDAY_OCCURRENCES.
+  date_intersection = rounded_times.intersection(rounded_dates).unique()
+  for date in date_intersection:
+    date_slice = rounded_times.get_loc(date)
+    regressor_slice = regressor[date_slice]
+    if isinstance(regressor_slice, int):
+      replacement = 1
+    else:
+      replacement = [1] * len(regressor[date_slice])
+    regressor[date_slice] = replacement
+  return regressor

tensorflow_probability/python/sts/holiday_effects_test.py

@@ -0,0 +1,96 @@
+# 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 holiday_effects."""
+from absl.testing import parameterized
+import pandas as pd
+import tensorflow as tf
+from tensorflow_probability.python.internal import test_util
+from tensorflow_probability.python.sts import holiday_effects
+
+HOLIDAY_FILE_FIELDS = ['geo', 'holiday', 'date']
+
+
+class HolidayEffectsTest(test_util.TestCase):
+
+  @parameterized.named_parameters(
+      ('date_wrong_order',
+       pd.DataFrame([['US', 'TestHoliday', '12-20-2013']],
+                    columns=HOLIDAY_FILE_FIELDS)),
+      ('date_invalid',
+       pd.DataFrame([['US', 'TestHoliday', '12-00-2013']],
+                    columns=HOLIDAY_FILE_FIELDS)),
+      ('bad_column_names',
+       pd.DataFrame([['US', 'TestHoliday', '2013-12-20']],
+                    columns=['geo', 'wrong', 'date'])))
+  def test_holidays_raise_error(self, holidays):
+    times = pd.date_range(
+        start='2013-12-20', end='2015-12-20', freq=pd.DateOffset(years=1))
+    with self.assertRaises(ValueError):
+      holiday_effects.create_holiday_regressors(times, holidays)
+
+  @parameterized.named_parameters(
+      ('data_wrong_format', pd.Series(['2013-12-20'])),
+      ('data_no_frequency', pd.DatetimeIndex(['2013-12-20'])))
+  def test_times_raise_error(self, times):
+    holidays = pd.DataFrame([['US', 'TestHoliday', '2013-12-20']],
+                            columns=HOLIDAY_FILE_FIELDS)
+    with self.assertRaises(ValueError):
+      holiday_effects.create_holiday_regressors(times, holidays)
+
+  @parameterized.named_parameters(
+      ('holiday_daily', pd.DateOffset(days=1), '2012-01-01', '2012-12-31',
+       [0] * 359 + [1] + [0] * 6),
+      ('holiday_hourly', pd.DateOffset(hours=1), '2012-01-01',
+       '2012-12-31 23:00:00', [0] * 359 * 24 + [1] * 24 + [0] * 6 * 24),
+      # Note that expected should be `[0] * 51 + [1] + [0]` if
+      # _match_dates supports rounding timestamps to the nearest prior day
+      ('holiday_weekly', pd.DateOffset(weeks=1), '2012-01-01', '2012-12-31',
+       [0] * 51 + [0] + [0]))
+  def test_match_dates_by_frequency(self, freq, start, end, expected):
+    holiday_dates = pd.to_datetime(['2012-12-25'])
+    index = pd.date_range(start, end, freq=freq)
+    matched_dates = holiday_effects._match_dates(index, holiday_dates)
+    self.assertEqual(matched_dates, expected)
+
+  @parameterized.named_parameters(
+      ('holiday_disjoint', '2011-01-01', '2011-12-31', [0] * 365),
+      ('holiday_intersection', '2011-02-01', '2012-01-31',
+       [0] * 334 + [1] * 31),
+      ('holiday_subset', '2012-01-01', '2012-01-31', [1] * 31))
+  def test_match_dates_by_overlap(self, start, end, expected):
+    holiday_dates = pd.date_range(
+        '2012-01-01', '2012-12-31', freq=pd.DateOffset(days=1))
+    index = pd.date_range(start, end, freq=pd.DateOffset(days=1))
+    matched_dates = holiday_effects._match_dates(index, holiday_dates)
+    self.assertEqual(matched_dates, expected)
+
+  @parameterized.named_parameters(
+      ('diagonal_pattern', [('H1', 0), ('H2', 1)], [[1, 0], [0, 1]]),
+      ('row_pattern', [('H1', 0), ('H2', 0)], [[1, 1], [0, 0]]),
+      ('column_pattern', [('H1', 0), ('H1', 1)], [[1], [1]]))
+  def test_create_holiday_regressors(self, holiday_patterns, expected):
+    times = pd.date_range(
+        '2011-01-01', '2012-01-01', freq=pd.DateOffset(years=1))
+    holidays_list = []
+    for name, date_index in holiday_patterns:
+      holidays_list.append(['US', name, times[date_index]])
+    holidays = pd.DataFrame(holidays_list, columns=HOLIDAY_FILE_FIELDS)
+    holiday_regressors = holiday_effects.create_holiday_regressors(
+        times, holidays)
+    self.assertEqual(holiday_regressors.values.tolist(), expected)
+
+
+if __name__ == '__main__':
+  tf.test.main()