Merge pull request #473 from pymc-labs/cetagostini/adding_bsts_to_causalpy

Implement Bayesian Structural Time Series (BSTS)

Author: cetagostini

causalpy/experiments/__init__.py

@@ -11,4 +11,24 @@
 #   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.
-"""Quasi-experimental designs for causal inference."""
+"""CausalPy experiment module"""
+
+from .diff_in_diff import DifferenceInDifferences
+from .instrumental_variable import InstrumentalVariable
+from .interrupted_time_series import InterruptedTimeSeries
+from .inverse_propensity_weighting import InversePropensityWeighting
+from .prepostnegd import PrePostNEGD
+from .regression_discontinuity import RegressionDiscontinuity
+from .regression_kink import RegressionKink
+from .synthetic_control import SyntheticControl
+
+__all__ = [
+    "DifferenceInDifferences",
+    "InstrumentalVariable",
+    "InversePropensityWeighting",
+    "PrePostNEGD",
+    "RegressionDiscontinuity",
+    "RegressionKink",
+    "SyntheticControl",
+    "InterruptedTimeSeries",
+]

causalpy/experiments/interrupted_time_series.py

@@ -62,23 +62,25 @@ def plot_xY(
     if plot_hdi_kwargs is None:
         plot_hdi_kwargs = {}
 
+    # Separate fill_kwargs for az.plot_hdi, as ax.plot doesn't accept them
+    line_kwargs = plot_hdi_kwargs.copy()
+    if "fill_kwargs" in line_kwargs:
+        del line_kwargs["fill_kwargs"]
+
     (h_line,) = ax.plot(
         x,
         Y.mean(dim=["chain", "draw"]),
         ls="-",
-        **plot_hdi_kwargs,
-        label=f"{label}",
+        **line_kwargs,  # Use kwargs without fill_kwargs
+        label=label,  # Use the provided label for the mean line
     )
     ax_hdi = az.plot_hdi(
         x,
         Y,
         hdi_prob=hdi_prob,
-        fill_kwargs={
-            "alpha": 0.25,
-            "label": " ",
-        },
-        smooth=False,
         ax=ax,
+        smooth=False,  # To prevent warning about resolution with few data points
+        # Pass original plot_hdi_kwargs which might include fill_kwargs for fill_between
         **plot_hdi_kwargs,
     )
     # Return handle to patch. We get a list of the children of the axis. Filter for just