Add seasonal-effects component for structural time series models.

PiperOrigin-RevId: 213710428
ixkael · Sep 19, 2018 · d8fecd0 · d8fecd0
1 parent 795bacd
commit d8fecd0
Show file tree

Hide file tree

Showing 5 changed files with 757 additions and 7 deletions.
diff --git a/tensorflow_probability/python/sts/BUILD b/tensorflow_probability/python/sts/BUILD
@@ -31,6 +31,7 @@ py_library(
     srcs_version = "PY2AND3",
     deps = [
         ":local_linear_trend",
+        ":seasonal",
         ":structural_time_series",
         ":sum",
     ],
@@ -58,6 +59,28 @@ py_test(
     ],
 )
 
+py_library(
+    name = "seasonal",
+    srcs = ["seasonal.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/sts/internal",
+    ],
+)
+
+py_test(
+    name = "seasonal_test",
+    size = "medium",
+    srcs = ["seasonal_test.py"],
+    deps = [
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability",
+    ],
+)
+
 py_library(
     name = "structural_time_series",
     srcs = ["structural_time_series.py"],

diff --git a/tensorflow_probability/python/sts/__init__.py b/tensorflow_probability/python/sts/__init__.py
@@ -20,6 +20,8 @@
 
 from tensorflow_probability.python.sts.local_linear_trend import LocalLinearTrend
 from tensorflow_probability.python.sts.local_linear_trend import LocalLinearTrendStateSpaceModel
+from tensorflow_probability.python.sts.seasonal import Seasonal
+from tensorflow_probability.python.sts.seasonal import SeasonalStateSpaceModel
 from tensorflow_probability.python.sts.structural_time_series import StructuralTimeSeries
 from tensorflow_probability.python.sts.sum import AdditiveStateSpaceModel
 from tensorflow_probability.python.sts.sum import Sum
@@ -30,6 +32,8 @@
     'AdditiveStateSpaceModel',
     'LocalLinearTrend',
     'LocalLinearTrendStateSpaceModel',
+    'Seasonal',
+    'SeasonalStateSpaceModel',
     'StructuralTimeSeries',
     'Sum',
 ]