Merge pull request #7 from YuminosukeSato/feat/dynamic-regression

feat: add dynamic regression (time-varying coefficients)

Author: YuminosukeSato

python/causal_impact/data.py

@@ -185,9 +185,7 @@ def _parse_periods(
         return pre_start, pre_end, post_start, post_end
 
     @staticmethod
-    def _validate_periods(
-        pre_start, pre_end, post_start, post_end, time_index
-    ) -> None:
+    def _validate_periods(pre_start, pre_end, post_start, post_end, time_index) -> None:
         idx_min = time_index.min()
         idx_max = time_index.max()
 

python/causal_impact/main.py

@@ -25,6 +25,7 @@
     "standardize_data": True,
     "prior_level_sd": 0.01,
     "expected_model_size": 2,
+    "dynamic_regression": False,
     "nseasons": None,
     "season_duration": None,
 }
@@ -100,6 +101,7 @@ def _run_sampler(self, prepared: PreparedData, args: dict):
             expected_model_size=float(args["expected_model_size"]),
             nseasons=args["nseasons"],
             season_duration=args["season_duration"],
+            dynamic_regression=bool(args.get("dynamic_regression", False)),
         )
 
     def _compute_results(self, prepared: PreparedData, samples) -> CausalImpactResults:

python/causal_impact/options.py

@@ -19,6 +19,7 @@ class ModelOptions:
     standardize_data: bool = True
     prior_level_sd: float = 0.01
     expected_model_size: int = 1
+    dynamic_regression: bool = False
     nseasons: int | None = None
     season_duration: int | None = None
 
@@ -38,6 +39,12 @@ def __post_init__(self) -> None:
         if self.expected_model_size <= 0:
             msg = f"expected_model_size must be > 0, got {self.expected_model_size}"
             raise ValueError(msg)
+        if not isinstance(self.dynamic_regression, bool):
+            msg = (
+                "dynamic_regression must be a bool, "
+                f"got {type(self.dynamic_regression).__name__}"
+            )
+            raise ValueError(msg)
         if self.nseasons is not None:
             if not isinstance(self.nseasons, int):
                 msg = f"nseasons must be an integer, got {self.nseasons}"

python/causal_impact/plot.py

@@ -57,12 +57,18 @@ def plot(
     def _plot_original(ax, y, time_index, post_index, results):
         ax.plot(time_index, y, color="black", linewidth=1, label="Observed")
         ax.plot(
-            post_index, results.predictions_mean, color="blue",
-            linestyle="--", label="Counterfactual",
+            post_index,
+            results.predictions_mean,
+            color="blue",
+            linestyle="--",
+            label="Counterfactual",
         )
         ax.fill_between(
-            post_index, results.predictions_lower, results.predictions_upper,
-            alpha=0.2, color="blue",
+            post_index,
+            results.predictions_lower,
+            results.predictions_upper,
+            alpha=0.2,
+            color="blue",
         )
         ax.set_ylabel("Response")
         ax.legend(loc="upper left", fontsize=8)