Add get_default_holidays for modelling holiday effects.

PiperOrigin-RevId: 389269656

Author: g-luo

tensorflow_probability/python/sts/BUILD

@@ -167,6 +167,7 @@ py_library(
     srcs = ["holiday_effects.py"],
     srcs_version = "PY3",
     deps = [
+        # holidays dep,
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/experimental/util",
@@ -181,6 +182,8 @@ py_test(
     python_version = "PY3",
     shard_count = 6,
     srcs_version = "PY3",
+    # TODO(gcluo): Add the holidays package to the OSS testing setup.
+    tags = ["no-oss-ci"],
     deps = [
         # holidays dep,
         # numpy dep,

tensorflow_probability/python/sts/holiday_effects.py

@@ -15,13 +15,52 @@
 """Utilities for holiday regressors."""
 
 __all__ = [
+    'get_default_holidays',
     'create_holiday_regressors',
 ]
 
 # Defines expected holiday file fields.
 _HOLIDAY_FILE_FIELDS = frozenset({'geo', 'holiday', 'date'})
 
 
+# TODO(b/195771744): Add holidays as unofficial TFP dependency.
+def get_default_holidays(times, country):
+  """Creates default holidays for a specific country.
+
+  Args:
+    times: a Pandas `DatetimeIndex` that indexes time series data.
+    country: `str`, two-letter upper-case [ISO 3166-1 alpha-2 country code](
+      https://en.wikipedia.org/wiki/ISO_3166-1_alpha-2) for using holiday
+        regressors from a particular country.
+
+  Returns:
+    holidays: a Pandas `DataFrame` with default holidays relevant to the input
+      times. 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`
+  """
+  # pylint: disable=g-import-not-at-top
+  import pandas as pd
+  import holidays
+  # pylint: enable=g-import-not-at-top
+
+  years = range(times.min().year, times.max().year + 1)
+  holidays = holidays.CountryHoliday(country, years=years, expand=False)
+  holidays = pd.DataFrame(
+      [(country, holidays.get_list(date), date) for date in holidays],
+      columns=['geo', 'holiday', 'date'])
+  holidays = holidays.explode('holiday')
+  # Ensure that only holiday dates covered by times are used.
+  holidays = holidays[(holidays['date'] >= times.min())
+                      & (holidays['date'] <= times.max())]
+  holidays = holidays.reset_index(drop=True)
+  holidays['date'] = pd.to_datetime(holidays['date'])
+  holidays = holidays.sort_values('date')
+  holidays['date'] = holidays['date'].dt.strftime('%Y-%m-%d')
+  return holidays
+
+
 def create_holiday_regressors(times, holidays):
   """Creates a design matrix of holiday regressors for a given time series.
 

tensorflow_probability/python/sts/holiday_effects_test.py

@@ -24,6 +24,32 @@
 
 class HolidayEffectsTest(test_util.TestCase):
 
+  def test_get_default_holidays_invalid_country(self):
+    times = pd.to_datetime(['2012-12-25', '2013-01-01'])
+    country = 'AA'
+    with self.assertRaises(Exception):
+      holiday_effects.get_default_holidays(times, country)
+
+  def test_get_default_holidays_invalid_times(self):
+    times = ['2012-12-25', '2013-01-01']
+    country = 'US'
+    with self.assertRaises(Exception):
+      holiday_effects.get_default_holidays(times, country)
+
+  @parameterized.named_parameters(
+      ('united_states_holidays', 'US',
+       pd.DataFrame([['US', 'Christmas Day', '2012-12-25'],
+                     ['US', 'New Year\'s Day', '2013-01-01']],
+                    columns=HOLIDAY_FILE_FIELDS)),
+      ('egypt_holidays', 'EG',
+       pd.DataFrame([['EG', 'New Year\'s Day - Bank Holiday', '2013-01-01']],
+                    columns=HOLIDAY_FILE_FIELDS)))
+  def test_get_default_holidays(self, country, expected):
+    times = pd.date_range(
+        start='2012-12-25', end='2013-01-01', freq=pd.DateOffset(days=1))
+    holidays = holiday_effects.get_default_holidays(times, country)
+    self.assertTrue(holidays.equals(expected))
+
   @parameterized.named_parameters(
       ('date_wrong_order',
        pd.DataFrame([['US', 'TestHoliday', '12-20-2013']],