Merge branch 'main' into cetagostini/bsts_tech_debt_v0

Author: drbenvincent

causalpy/experiments/interrupted_time_series.py

@@ -629,12 +629,13 @@ def _compute_statistics(
     cumulative=True,
     relative=True,
     min_effect=None,
+    time_dim="obs_ind",
 ):
     """Compute all summary statistics from posterior draws."""
     stats = {}
 
     # Average effect over window
-    avg_effect = impact.mean(dim="obs_ind")
+    avg_effect = impact.mean(dim=time_dim)
     stats["avg"] = {
         "mean": float(avg_effect.mean(dim=["chain", "draw"]).values),
         "median": float(avg_effect.median(dim=["chain", "draw"]).values),
@@ -677,9 +678,9 @@ def _compute_statistics(
     # Cumulative effect
     if cumulative:
         # Use cumulative sum over window
-        cum_effect = impact.cumsum(dim="obs_ind")
+        cum_effect = impact.cumsum(dim=time_dim)
         # Take final value (cumulative over entire window)
-        cum_final = cum_effect.isel(obs_ind=-1)
+        cum_final = cum_effect.isel({time_dim: -1})
 
         stats["cum"] = {
             "mean": float(cum_final.mean(dim=["chain", "draw"]).values),
@@ -720,7 +721,7 @@ def _compute_statistics(
     # Relative effects
     if relative:
         epsilon = 1e-8  # Guard against division by zero
-        counterfactual_mean = counterfactual.mean(dim="obs_ind")
+        counterfactual_mean = counterfactual.mean(dim=time_dim)
         rel_avg = (avg_effect / (counterfactual_mean + epsilon)) * 100
 
         stats["avg"]["relative_mean"] = float(
@@ -746,7 +747,9 @@ def _compute_statistics(
 
         if cumulative:
             # Relative cumulative: (cumulative effect / cumulative counterfactual) * 100
-            counterfactual_cum = counterfactual.cumsum(dim="obs_ind").isel(obs_ind=-1)
+            counterfactual_cum = counterfactual.cumsum(dim=time_dim).isel(
+                {time_dim: -1}
+            )
             rel_cum = (cum_final / (counterfactual_cum + epsilon)) * 100
 
             stats["cum"]["relative_mean"] = float(
@@ -850,6 +853,7 @@ def _generate_prose(
     direction="increase",
     cumulative=True,
     relative=True,
+    prefix="Post-period",
 ):
     """Generate prose summary from statistics."""
     hdi_pct = int((1 - alpha) * 100)
@@ -883,7 +887,7 @@ def fmt_num(x, decimals=2):
         direction_text = "effect"
 
     prose_parts = [
-        f"Post-period ({window_str}), the average effect was {fmt_num(avg_mean)} "
+        f"{prefix} ({window_str}), the average effect was {fmt_num(avg_mean)} "
         f"({hdi_pct}% HDI [{fmt_num(avg_lower)}, {fmt_num(avg_upper)}]), "
         f"with a posterior probability of an {direction_text} of {fmt_num(p_val, 3)}."
     ]
@@ -1138,6 +1142,7 @@ def _generate_prose_ols(
     alpha=0.05,
     cumulative=True,
     relative=True,
+    prefix="Post-period",
 ):
     """Generate prose summary for OLS models."""
     ci_pct = int((1 - alpha) * 100)
@@ -1161,7 +1166,7 @@ def fmt_num(x, decimals=2):
     p_val = stats["avg"]["p_value"]
 
     prose_parts = [
-        f"Post-period ({window_str}), the average effect was {fmt_num(avg_mean)} "
+        f"{prefix} ({window_str}), the average effect was {fmt_num(avg_mean)} "
         f"({ci_pct}% CI [{fmt_num(avg_lower)}, {fmt_num(avg_upper)}]), "
         f"with a p-value of {fmt_num(p_val, 3)}."
     ]