diff --git a/causalpy/experiments/base.py b/causalpy/experiments/base.py index 74e8e0ec7..be44df6b9 100644 --- a/causalpy/experiments/base.py +++ b/causalpy/experiments/base.py @@ -27,6 +27,7 @@ import pandas as pd from sklearn.base import RegressorMixin +from causalpy.plot_utils import ResponseType from causalpy.pymc_models import PyMCModel from causalpy.reporting import EffectSummary from causalpy.skl_models import create_causalpy_compatible_class @@ -156,6 +157,7 @@ def effect_summary( treated_unit: str | None = None, period: Literal["intervention", "post", "comparison"] | None = None, prefix: str = "Post-period", + response_type: ResponseType = "expectation", **kwargs: Any, ) -> EffectSummary: """ @@ -192,6 +194,19 @@ def effect_summary( prefix : str, optional Prefix for prose generation (e.g., "During intervention", "Post-intervention"). Defaults to "Post-period". + response_type : {"expectation", "prediction"}, default="expectation" + Response type to compute effect sizes (ITS/SC only, ignored for DiD/RD/RKink): + + - ``"expectation"``: Effect size HDI based on model expectation (μ). + Excludes observation noise, focusing on the systematic causal effect. + - ``"prediction"``: Effect size HDI based on posterior predictive (ŷ). + Includes observation noise, showing full predictive uncertainty. + + Note: This parameter only affects experiments where the causal effect is + calculated as the difference between observed and predicted values + (ITS, Synthetic Control). For experiments where the effect is a model + coefficient (DiD, RD, RKink), the HDI is always computed from the + posterior of the coefficient and this parameter is ignored. Returns ------- diff --git a/causalpy/experiments/diff_in_diff.py b/causalpy/experiments/diff_in_diff.py index 2d1f4c2f2..f76913d84 100644 --- a/causalpy/experiments/diff_in_diff.py +++ b/causalpy/experiments/diff_in_diff.py @@ -30,7 +30,12 @@ DataException, FormulaException, ) -from causalpy.plot_utils import plot_xY +from causalpy.plot_utils import ( + ResponseType, + _log_response_type_info_once, + add_hdi_annotation, + plot_xY, +) from causalpy.pymc_models import LinearRegression, PyMCModel from causalpy.reporting import ( EffectSummary, @@ -335,21 +340,54 @@ def _causal_impact_summary_stat(self, round_to: int | None = None) -> str: return f"Causal impact = {convert_to_string(self.causal_impact, round_to=round_to)}" def _bayesian_plot( - self, round_to: int | None = None, **kwargs: dict + self, + round_to: int | None = None, + response_type: ResponseType = "expectation", + show_hdi_annotation: bool = False, + **kwargs: dict, ) -> tuple[plt.Figure, plt.Axes]: """ Plot the results - :param round_to: - Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. + Parameters + ---------- + round_to : int, optional + Number of decimals used to round results. Defaults to 2. + Use None to return raw numbers. + response_type : {"expectation", "prediction"}, default="expectation" + The response type to display in the HDI band: + + - ``"expectation"``: HDI of the model expectation (μ). This shows + uncertainty from model parameters only, excluding observation noise. + Results in narrower intervals that represent the uncertainty in + the expected value of the outcome. + - ``"prediction"``: HDI of the posterior predictive (ŷ). This includes + observation noise (σ) in addition to parameter uncertainty, resulting + in wider intervals that represent the full predictive uncertainty + for new observations. + show_hdi_annotation : bool, default=False + Whether to display a text annotation at the bottom of the figure + explaining what the HDI represents. Set to False to hide the annotation. + **kwargs : dict + Additional keyword arguments. + + Returns + ------- + tuple[plt.Figure, plt.Axes] + The matplotlib figure and axes. """ + # Log HDI type info once per session + _log_response_type_info_once() + + # Select the variable name based on response_type + var_name = "mu" if response_type == "expectation" else "y_hat" def _plot_causal_impact_arrow(results, ax): """ draw a vertical arrow between `y_pred_counterfactual` and `y_pred_counterfactual` """ - # Calculate y values to plot the arrow between + # Calculate y values to plot the arrow between - always use mu for arrow position y_pred_treatment = ( results.y_pred_treatment["posterior_predictive"] .mu.isel({"obs_ind": 1}) @@ -409,7 +447,7 @@ def _plot_causal_impact_arrow(results, ax): time_points = self.x_pred_control[self.time_variable_name].values h_line, h_patch = plot_xY( time_points, - self.y_pred_control["posterior_predictive"].mu.isel(treated_units=0), + self.y_pred_control["posterior_predictive"][var_name].isel(treated_units=0), ax=ax, plot_hdi_kwargs={"color": "C0"}, label="Control group", @@ -421,7 +459,9 @@ def _plot_causal_impact_arrow(results, ax): time_points = self.x_pred_control[self.time_variable_name].values h_line, h_patch = plot_xY( time_points, - self.y_pred_treatment["posterior_predictive"].mu.isel(treated_units=0), + self.y_pred_treatment["posterior_predictive"][var_name].isel( + treated_units=0 + ), ax=ax, plot_hdi_kwargs={"color": "C1"}, label="Treatment group", @@ -436,7 +476,7 @@ def _plot_causal_impact_arrow(results, ax): y_pred_cf = az.extract( self.y_pred_counterfactual, group="posterior_predictive", - var_names="mu", + var_names=var_name, ) # Select single unit data for plotting y_pred_cf_single = y_pred_cf.isel(treated_units=0) @@ -459,7 +499,7 @@ def _plot_causal_impact_arrow(results, ax): else: h_line, h_patch = plot_xY( time_points, - self.y_pred_counterfactual.posterior_predictive.mu.isel( + self.y_pred_counterfactual.posterior_predictive[var_name].isel( treated_units=0 ), ax=ax, @@ -482,6 +522,11 @@ def _plot_causal_impact_arrow(results, ax): labels=labels, fontsize=LEGEND_FONT_SIZE, ) + + # Add HDI type annotation to the title + if show_hdi_annotation: + add_hdi_annotation(ax, response_type) + return fig, ax def _ols_plot( diff --git a/causalpy/experiments/interrupted_time_series.py b/causalpy/experiments/interrupted_time_series.py index 1d5aeb449..c208cdc1f 100644 --- a/causalpy/experiments/interrupted_time_series.py +++ b/causalpy/experiments/interrupted_time_series.py @@ -27,7 +27,14 @@ from causalpy.custom_exceptions import BadIndexException from causalpy.date_utils import _combine_datetime_indices, format_date_axes -from causalpy.plot_utils import get_hdi_to_df, plot_xY +from causalpy.plot_utils import ( + ResponseType, + _log_response_type_effect_summary_once, + _log_response_type_info_once, + add_hdi_annotation, + get_hdi_to_df, + plot_xY, +) from causalpy.pymc_models import LinearRegression, PyMCModel from causalpy.reporting import EffectSummary from causalpy.utils import round_num @@ -236,7 +243,12 @@ def algorithm(self) -> None: elif isinstance(self.model, RegressorMixin): self.post_pred = self.model.predict(X=self.post_X) - # calculate impact - all PyMC models now use 2D data with treated_units + # Calculate impact - all PyMC models now use 2D data with treated_units. + # TODO: REFACTOR TARGET - Currently, stored impacts use the model expectation + # (mu) by default. When users request response_type="prediction" in plot(), the + # y_hat-based impact is calculated on-the-fly in _bayesian_plot(). This works + # but is not ideal: consider storing both mu and y_hat based impacts, or + # refactoring to always calculate on-demand. See calculate_impact() for details. if isinstance(self.model, PyMCModel): self.pre_impact = self.model.calculate_impact(self.pre_y, self.pre_pred) self.post_impact = self.model.calculate_impact(self.post_y, self.post_pred) @@ -601,26 +613,62 @@ def summary(self, round_to: int | None = None) -> None: self.print_coefficients(round_to) def _bayesian_plot( - self, round_to: int | None = 2, **kwargs: dict + self, + round_to: int | None = 2, + response_type: ResponseType = "expectation", + show_hdi_annotation: bool = False, + **kwargs: dict, ) -> tuple[plt.Figure, list[plt.Axes]]: """ Plot the results - :param round_to: - Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. + Parameters + ---------- + round_to : int, optional + Number of decimals used to round results. Defaults to 2. + Use None to return raw numbers. + response_type : {"expectation", "prediction"}, default="expectation" + The response type to display in the HDI band: + + - ``"expectation"``: HDI of the model expectation (μ). This shows + uncertainty from model parameters only, excluding observation noise. + Results in narrower intervals that represent the uncertainty in + the expected value of the outcome. + - ``"prediction"``: HDI of the posterior predictive (ŷ). This includes + observation noise (σ) in addition to parameter uncertainty, resulting + in wider intervals that represent the full predictive uncertainty + for new observations. + show_hdi_annotation : bool, default=False + Whether to display a text annotation at the bottom of the figure + explaining what the HDI represents. Set to False to hide the annotation. + **kwargs : dict + Additional keyword arguments. + + Returns + ------- + tuple[plt.Figure, list[plt.Axes]] + The matplotlib figure and axes. """ + # Log HDI type info once per session + _log_response_type_info_once() + counterfactual_label = "Counterfactual" + # Select the variable name based on response_type + var_name = "mu" if response_type == "expectation" else "y_hat" + fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8)) # TOP PLOT -------------------------------------------------- # pre-intervention period - pre_mu = self.pre_pred["posterior_predictive"].mu - pre_mu_plot = ( - pre_mu.isel(treated_units=0) if "treated_units" in pre_mu.dims else pre_mu + pre_pred_var = self.pre_pred["posterior_predictive"][var_name] + pre_pred_plot = ( + pre_pred_var.isel(treated_units=0) + if "treated_units" in pre_pred_var.dims + else pre_pred_var ) h_line, h_patch = plot_xY( self.datapre.index, - pre_mu_plot, + pre_pred_plot, ax=ax[0], plot_hdi_kwargs={"color": "C0"}, ) @@ -639,15 +687,15 @@ def _bayesian_plot( labels.append("Observations") # post intervention period - post_mu = self.post_pred["posterior_predictive"].mu - post_mu_plot = ( - post_mu.isel(treated_units=0) - if "treated_units" in post_mu.dims - else post_mu + post_pred_var = self.post_pred["posterior_predictive"][var_name] + post_pred_plot = ( + post_pred_var.isel(treated_units=0) + if "treated_units" in post_pred_var.dims + else post_pred_var ) h_line, h_patch = plot_xY( self.datapost.index, - post_mu_plot, + post_pred_plot, ax=ax[0], plot_hdi_kwargs={"color": "C1"}, ) @@ -661,7 +709,7 @@ def _bayesian_plot( else self.post_y[:, 0], "k.", ) - # Shaded causal effect + # Shaded causal effect - always use mu for the fill_between mean line post_pred_mu = az.extract( self.post_pred, group="posterior_predictive", var_names="mu" ) @@ -701,11 +749,25 @@ def _bayesian_plot( ax[0].set(title=title_str) # MIDDLE PLOT ----------------------------------------------- + # Calculate impact for plotting based on response_type + if response_type == "expectation": + # Use stored mu-based impact + pre_impact_for_plot = self.pre_impact + post_impact_for_plot = self.post_impact + else: + # Calculate y_hat-based impact on demand + pre_impact_for_plot = self.model.calculate_impact( + self.pre_y, self.pre_pred, response_type="prediction" + ) + post_impact_for_plot = self.model.calculate_impact( + self.post_y, self.post_pred, response_type="prediction" + ) + pre_impact_plot = ( - self.pre_impact.isel(treated_units=0) - if hasattr(self.pre_impact, "dims") - and "treated_units" in self.pre_impact.dims - else self.pre_impact + pre_impact_for_plot.isel(treated_units=0) + if hasattr(pre_impact_for_plot, "dims") + and "treated_units" in pre_impact_for_plot.dims + else pre_impact_for_plot ) plot_xY( self.datapre.index, @@ -714,10 +776,10 @@ def _bayesian_plot( plot_hdi_kwargs={"color": "C0"}, ) post_impact_plot = ( - self.post_impact.isel(treated_units=0) - if hasattr(self.post_impact, "dims") - and "treated_units" in self.post_impact.dims - else self.post_impact + post_impact_for_plot.isel(treated_units=0) + if hasattr(post_impact_for_plot, "dims") + and "treated_units" in post_impact_for_plot.dims + else post_impact_for_plot ) plot_xY( self.datapost.index, @@ -727,9 +789,9 @@ def _bayesian_plot( ) ax[1].axhline(y=0, c="k") post_impact_mean = ( - self.post_impact.mean(["chain", "draw"]) - if hasattr(self.post_impact, "mean") - else self.post_impact + post_impact_for_plot.mean(["chain", "draw"]) + if hasattr(post_impact_for_plot, "mean") + else post_impact_for_plot ) if ( hasattr(post_impact_mean, "dims") @@ -747,11 +809,19 @@ def _bayesian_plot( # BOTTOM PLOT ----------------------------------------------- ax[2].set(title="Cumulative Causal Impact") + # Calculate cumulative impact based on response_type + if response_type == "expectation": + post_impact_cumulative_for_plot = self.post_impact_cumulative + else: + post_impact_cumulative_for_plot = self.model.calculate_cumulative_impact( + post_impact_for_plot + ) + post_cum_plot = ( - self.post_impact_cumulative.isel(treated_units=0) - if hasattr(self.post_impact_cumulative, "dims") - and "treated_units" in self.post_impact_cumulative.dims - else self.post_impact_cumulative + post_impact_cumulative_for_plot.isel(treated_units=0) + if hasattr(post_impact_cumulative_for_plot, "dims") + and "treated_units" in post_impact_cumulative_for_plot.dims + else post_impact_cumulative_for_plot ) plot_xY( self.datapost.index, @@ -794,6 +864,10 @@ def _bayesian_plot( ) format_date_axes(ax, full_index) + # Add HDI type annotation to the top subplot's title + if show_hdi_annotation: + add_hdi_annotation(ax[0], response_type) + return fig, ax def _ols_plot( @@ -887,14 +961,31 @@ def _ols_plot( return (fig, ax) - def get_plot_data_bayesian(self, hdi_prob: float = 0.94) -> pd.DataFrame: + def get_plot_data_bayesian( + self, + hdi_prob: float = 0.94, + response_type: ResponseType = "expectation", + ) -> pd.DataFrame: """ Recover the data of the experiment along with the prediction and causal impact information. - :param hdi_prob: - Prob for which the highest density interval will be computed. The default value is defined as the default from the :func:`arviz.hdi` function. + Parameters + ---------- + hdi_prob : float, default=0.94 + Probability for which the highest density interval will be computed. + The default value is defined as the default from the :func:`arviz.hdi` function. + response_type : {"expectation", "prediction"}, default="expectation" + The response type to use for predictions and impact: + + - ``"expectation"``: Uses the model expectation (μ). Excludes observation + noise, focusing on the systematic causal effect. + - ``"prediction"``: Uses the full posterior predictive (ŷ). Includes + observation noise, showing the full predictive uncertainty. """ + if isinstance(self.model, PyMCModel): + # Map semantic response_type to internal variable name + var_name = "mu" if response_type == "expectation" else "y_hat" hdi_pct = int(round(hdi_prob * 100)) pred_lower_col = f"pred_hdi_lower_{hdi_pct}" @@ -906,10 +997,10 @@ def get_plot_data_bayesian(self, hdi_prob: float = 0.94) -> pd.DataFrame: post_data = self.datapost.copy() pre_mu = az.extract( - self.pre_pred, group="posterior_predictive", var_names="mu" + self.pre_pred, group="posterior_predictive", var_names=var_name ) post_mu = az.extract( - self.post_pred, group="posterior_predictive", var_names="mu" + self.post_pred, group="posterior_predictive", var_names=var_name ) if "treated_units" in pre_mu.dims: pre_mu = pre_mu.isel(treated_units=0) @@ -919,10 +1010,10 @@ def get_plot_data_bayesian(self, hdi_prob: float = 0.94) -> pd.DataFrame: post_data["prediction"] = post_mu.mean("sample").values hdi_pre_pred = get_hdi_to_df( - self.pre_pred["posterior_predictive"].mu, hdi_prob=hdi_prob + self.pre_pred["posterior_predictive"][var_name], hdi_prob=hdi_prob ) hdi_post_pred = get_hdi_to_df( - self.post_pred["posterior_predictive"].mu, hdi_prob=hdi_prob + self.post_pred["posterior_predictive"][var_name], hdi_prob=hdi_prob ) # If treated_units present, select unit_0; otherwise use directly if ( @@ -943,15 +1034,28 @@ def get_plot_data_bayesian(self, hdi_prob: float = 0.94) -> pd.DataFrame: post_data.index ) + # Select impact based on response_type + if response_type == "expectation": + pre_impact = self.pre_impact + post_impact = self.post_impact + else: + # Calculate y_hat-based impact on demand + pre_impact = self.model.calculate_impact( + self.pre_y, self.pre_pred, response_type="prediction" + ) + post_impact = self.model.calculate_impact( + self.post_y, self.post_pred, response_type="prediction" + ) + pre_impact_mean = ( - self.pre_impact.mean(dim=["chain", "draw"]) - if hasattr(self.pre_impact, "mean") - else self.pre_impact + pre_impact.mean(dim=["chain", "draw"]) + if hasattr(pre_impact, "mean") + else pre_impact ) post_impact_mean = ( - self.post_impact.mean(dim=["chain", "draw"]) - if hasattr(self.post_impact, "mean") - else self.post_impact + post_impact.mean(dim=["chain", "draw"]) + if hasattr(post_impact, "mean") + else post_impact ) if ( hasattr(pre_impact_mean, "dims") @@ -971,10 +1075,10 @@ def get_plot_data_bayesian(self, hdi_prob: float = 0.94) -> pd.DataFrame: lower_q = alpha / 2 upper_q = 1 - alpha / 2 - pre_lower_da = self.pre_impact.quantile(lower_q, dim=["chain", "draw"]) - pre_upper_da = self.pre_impact.quantile(upper_q, dim=["chain", "draw"]) - post_lower_da = self.post_impact.quantile(lower_q, dim=["chain", "draw"]) - post_upper_da = self.post_impact.quantile(upper_q, dim=["chain", "draw"]) + pre_lower_da = pre_impact.quantile(lower_q, dim=["chain", "draw"]) + pre_upper_da = pre_impact.quantile(upper_q, dim=["chain", "draw"]) + post_lower_da = post_impact.quantile(lower_q, dim=["chain", "draw"]) + post_upper_da = post_impact.quantile(upper_q, dim=["chain", "draw"]) # If a treated_units dim remains for some models, select unit_0 if hasattr(pre_lower_da, "dims") and "treated_units" in pre_lower_da.dims: @@ -1021,6 +1125,7 @@ def analyze_persistence( self, hdi_prob: float = 0.95, direction: Literal["increase", "decrease", "two-sided"] = "increase", + response_type: ResponseType = "expectation", ) -> dict[str, Any]: """Analyze effect persistence between intervention and post-intervention periods. @@ -1038,6 +1143,13 @@ def analyze_persistence( Probability for HDI interval (Bayesian models only) direction : {"increase", "decrease", "two-sided"}, default="increase" Direction for tail probability calculation (Bayesian models only) + response_type : {"expectation", "prediction"}, default="expectation" + The response type to use for effect analysis (Bayesian models only): + + - ``"expectation"``: Uses the model expectation (μ). Excludes observation + noise, focusing on the systematic causal effect. + - ``"prediction"``: Uses the full posterior predictive (ŷ). Includes + observation noise, showing the full predictive uncertainty. Returns ------- @@ -1090,8 +1202,38 @@ def analyze_persistence( # PyMC: Compute statistics using xarray operations from causalpy.reporting import _extract_hdi_bounds + # Select impact based on response_type + if response_type == "expectation": + intervention_impact = self.intervention_impact + post_intervention_impact = self.post_intervention_impact + intervention_impact_cumulative = self.intervention_impact_cumulative + post_intervention_impact_cumulative = ( + self.post_intervention_impact_cumulative + ) + else: + # Calculate y_hat-based impact on demand + # Get y values for intervention and post-intervention periods from post_y + intervention_y = self.post_y.sel(obs_ind=self.data_intervention.index) + post_intervention_y = self.post_y.sel( + obs_ind=self.data_post_intervention.index + ) + intervention_impact = self.model.calculate_impact( + intervention_y, self.intervention_pred, response_type="prediction" + ) + post_intervention_impact = self.model.calculate_impact( + post_intervention_y, + self.post_intervention_pred, + response_type="prediction", + ) + intervention_impact_cumulative = self.model.calculate_cumulative_impact( + intervention_impact + ) + post_intervention_impact_cumulative = ( + self.model.calculate_cumulative_impact(post_intervention_impact) + ) + # Intervention period - intervention_avg = self.intervention_impact.mean(dim=time_dim) + intervention_avg = intervention_impact.mean(dim=time_dim) intervention_mean = float( intervention_avg.mean(dim=["chain", "draw"]).values ) @@ -1101,18 +1243,18 @@ def analyze_persistence( ) # Post-intervention period - post_avg = self.post_intervention_impact.mean(dim=time_dim) + post_avg = post_intervention_impact.mean(dim=time_dim) post_mean = float(post_avg.mean(dim=["chain", "draw"]).values) post_hdi = az.hdi(post_avg, hdi_prob=hdi_prob) post_lower, post_upper = _extract_hdi_bounds(post_hdi, hdi_prob) # Cumulative (total) impacts - intervention_cum = self.intervention_impact_cumulative.isel({time_dim: -1}) + intervention_cum = intervention_impact_cumulative.isel({time_dim: -1}) intervention_cum_mean = float( intervention_cum.mean(dim=["chain", "draw"]).values ) - post_cum = self.post_intervention_impact_cumulative.isel({time_dim: -1}) + post_cum = post_intervention_impact_cumulative.isel({time_dim: -1}) post_cum_mean = float(post_cum.mean(dim=["chain", "draw"]).values) # Persistence ratio: post_mean / intervention_mean (as decimal, not percentage) @@ -1216,6 +1358,7 @@ def effect_summary( treated_unit: str | None = None, period: Literal["intervention", "post", "comparison"] | None = None, prefix: str = "Post-period", + response_type: ResponseType = "expectation", **kwargs: Any, ) -> EffectSummary: """ @@ -1246,6 +1389,13 @@ def effect_summary( prefix : str, optional Prefix for prose generation (e.g., "During intervention", "Post-intervention"). Defaults to "Post-period". + response_type : {"expectation", "prediction"}, default="expectation" + Response type to compute effect sizes: + + - ``"expectation"``: Effect size HDI based on model expectation (μ). + Excludes observation noise, focusing on the systematic causal effect. + - ``"prediction"``: Effect size HDI based on posterior predictive (ŷ). + Includes observation noise, showing full predictive uncertainty. Returns ------- @@ -1264,7 +1414,10 @@ def effect_summary( _generate_table_ols, ) + # Log HDI type info once per session (for PyMC models only) is_pymc = isinstance(self.model, PyMCModel) + if is_pymc: + _log_response_type_effect_summary_once() # Handle period parameter for three-period designs if period is not None: @@ -1311,20 +1464,26 @@ def effect_summary( # Extract windowed impact data using calculated window windowed_impact, window_coords = _extract_window( - self, window, treated_unit=treated_unit + self, window, treated_unit=treated_unit, response_type=response_type ) # Extract counterfactual for relative effects counterfactual = _extract_counterfactual( - self, window_coords, treated_unit=treated_unit + self, + window_coords, + treated_unit=treated_unit, + response_type=response_type, ) else: # No period specified, use standard flow windowed_impact, window_coords = _extract_window( - self, window, treated_unit=treated_unit + self, window, treated_unit=treated_unit, response_type=response_type ) counterfactual = _extract_counterfactual( - self, window_coords, treated_unit=treated_unit + self, + window_coords, + treated_unit=treated_unit, + response_type=response_type, ) if is_pymc: diff --git a/causalpy/experiments/prepostnegd.py b/causalpy/experiments/prepostnegd.py index 8122bd013..3e6e67c9f 100644 --- a/causalpy/experiments/prepostnegd.py +++ b/causalpy/experiments/prepostnegd.py @@ -29,7 +29,12 @@ from causalpy.custom_exceptions import ( DataException, ) -from causalpy.plot_utils import plot_xY +from causalpy.plot_utils import ( + ResponseType, + _log_response_type_info_once, + add_hdi_annotation, + plot_xY, +) from causalpy.pymc_models import LinearRegression, PyMCModel from causalpy.reporting import EffectSummary, _effect_summary_did from causalpy.utils import _is_variable_dummy_coded, round_num @@ -235,9 +240,47 @@ def summary(self, round_to: int | None = None) -> None: self.print_coefficients(round_to) def _bayesian_plot( - self, round_to: int | None = None, **kwargs: dict + self, + round_to: int | None = None, + response_type: ResponseType = "expectation", + show_hdi_annotation: bool = False, + **kwargs: dict, ) -> tuple[plt.Figure, list[plt.Axes]]: - """Generate plot for ANOVA-like experiments with non-equivalent group designs.""" + """Generate plot for ANOVA-like experiments with non-equivalent group designs. + + Parameters + ---------- + round_to : int, optional + Number of decimals used to round results. Defaults to 2. + Use None to return raw numbers. + response_type : {"expectation", "prediction"}, default="expectation" + The response type to display in the HDI band: + + - ``"expectation"``: HDI of the model expectation (μ). This shows + uncertainty from model parameters only, excluding observation noise. + Results in narrower intervals that represent the uncertainty in + the expected value of the outcome. + - ``"prediction"``: HDI of the posterior predictive (ŷ). This includes + observation noise (σ) in addition to parameter uncertainty, resulting + in wider intervals that represent the full predictive uncertainty + for new observations. + show_hdi_annotation : bool, default=False + Whether to display a text annotation at the bottom of the figure + explaining what the HDI represents. Set to False to hide the annotation. + **kwargs : dict + Additional keyword arguments. + + Returns + ------- + tuple[plt.Figure, list[plt.Axes]] + The matplotlib figure and axes. + """ + # Log HDI type info once per session + _log_response_type_info_once() + + # Select the variable name based on response_type + var_name = "mu" if response_type == "expectation" else "y_hat" + fig, ax = plt.subplots( 2, 1, figsize=(7, 9), gridspec_kw={"height_ratios": [3, 1]} ) @@ -257,7 +300,7 @@ def _bayesian_plot( # plot posterior predictive of untreated h_line, h_patch = plot_xY( self.pred_xi, - self.pred_untreated["posterior_predictive"].mu.isel(treated_units=0), + self.pred_untreated["posterior_predictive"][var_name].isel(treated_units=0), ax=ax[0], plot_hdi_kwargs={"color": "C0"}, label="Control group", @@ -268,7 +311,7 @@ def _bayesian_plot( # plot posterior predictive of treated h_line, h_patch = plot_xY( self.pred_xi, - self.pred_treated["posterior_predictive"].mu.isel(treated_units=0), + self.pred_treated["posterior_predictive"][var_name].isel(treated_units=0), ax=ax[0], plot_hdi_kwargs={"color": "C1"}, label="Treatment group", @@ -282,9 +325,14 @@ def _bayesian_plot( fontsize=LEGEND_FONT_SIZE, ) - # Plot estimated caual impact / treatment effect + # Plot estimated causal impact / treatment effect az.plot_posterior(self.causal_impact, ref_val=0, ax=ax[1], round_to=round_to) ax[1].set(title="Estimated treatment effect") + + # Add HDI type annotation to the top subplot's title + if show_hdi_annotation: + add_hdi_annotation(ax[0], response_type) + return fig, ax def effect_summary( diff --git a/causalpy/experiments/regression_discontinuity.py b/causalpy/experiments/regression_discontinuity.py index ba485eae6..fc8c656eb 100644 --- a/causalpy/experiments/regression_discontinuity.py +++ b/causalpy/experiments/regression_discontinuity.py @@ -29,7 +29,12 @@ DataException, FormulaException, ) -from causalpy.plot_utils import plot_xY +from causalpy.plot_utils import ( + ResponseType, + _log_response_type_info_once, + add_hdi_annotation, + plot_xY, +) from causalpy.pymc_models import LinearRegression, PyMCModel from causalpy.utils import _is_variable_dummy_coded, convert_to_string, round_num @@ -296,9 +301,47 @@ def summary(self, round_to: int | None = None) -> None: self.print_coefficients(round_to) def _bayesian_plot( - self, round_to: int | None = 2, **kwargs: dict + self, + round_to: int | None = 2, + response_type: ResponseType = "expectation", + show_hdi_annotation: bool = False, + **kwargs: dict, ) -> tuple[plt.Figure, plt.Axes]: - """Generate plot for regression discontinuity designs.""" + """Generate plot for regression discontinuity designs. + + Parameters + ---------- + round_to : int, optional + Number of decimals used to round results. Defaults to 2. + Use None to return raw numbers. + response_type : {"expectation", "prediction"}, default="expectation" + The response type to display in the HDI band: + + - ``"expectation"``: HDI of the model expectation (μ). This shows + uncertainty from model parameters only, excluding observation noise. + Results in narrower intervals that represent the uncertainty in + the expected value of the outcome. + - ``"prediction"``: HDI of the posterior predictive (ŷ). This includes + observation noise (σ) in addition to parameter uncertainty, resulting + in wider intervals that represent the full predictive uncertainty + for new observations. + show_hdi_annotation : bool, default=False + Whether to display a text annotation at the bottom of the figure + explaining what the HDI represents. Set to False to hide the annotation. + **kwargs : dict + Additional keyword arguments. + + Returns + ------- + tuple[plt.Figure, plt.Axes] + The matplotlib figure and axes. + """ + # Log HDI type info once per session + _log_response_type_info_once() + + # Select the variable name based on response_type + var_name = "mu" if response_type == "expectation" else "y_hat" + fig, ax = plt.subplots() # Plot data: use two layers only when there are excluded observations @@ -324,7 +367,7 @@ def _bayesian_plot( # Plot model fit to data plot_xY( self.x_pred[self.running_variable_name], - self.pred["posterior_predictive"].mu.isel(treated_units=0), + self.pred["posterior_predictive"][var_name].isel(treated_units=0), ax=ax, plot_hdi_kwargs={"color": "C1"}, label="Posterior mean", @@ -351,7 +394,6 @@ def _bayesian_plot( color="r", label="treatment threshold", ) - # Add donut hole boundary lines if donut_hole > 0 if self.donut_hole > 0: ax.axvline( @@ -369,6 +411,10 @@ def _bayesian_plot( ) ax.legend(fontsize=LEGEND_FONT_SIZE) + + # Add HDI type annotation to the title + if show_hdi_annotation: + add_hdi_annotation(ax, response_type) return (fig, ax) def _ols_plot( diff --git a/causalpy/experiments/regression_kink.py b/causalpy/experiments/regression_kink.py index 4b892bcf7..5e7b22f7e 100644 --- a/causalpy/experiments/regression_kink.py +++ b/causalpy/experiments/regression_kink.py @@ -25,7 +25,12 @@ import seaborn as sns from patsy import build_design_matrices, dmatrices import xarray as xr -from causalpy.plot_utils import plot_xY +from causalpy.plot_utils import ( + ResponseType, + _log_response_type_info_once, + add_hdi_annotation, + plot_xY, +) from causalpy.pymc_models import LinearRegression, PyMCModel from causalpy.reporting import EffectSummary, _effect_summary_rkink @@ -248,9 +253,47 @@ def summary(self, round_to: int | None = 2) -> None: self.print_coefficients(round_to) def _bayesian_plot( - self, round_to: int | None = 2, **kwargs: dict + self, + round_to: int | None = 2, + response_type: ResponseType = "expectation", + show_hdi_annotation: bool = False, + **kwargs: dict, ) -> tuple[plt.Figure, plt.Axes]: - """Generate plot for regression kink designs.""" + """Generate plot for regression kink designs. + + Parameters + ---------- + round_to : int, optional + Number of decimals used to round results. Defaults to 2. + Use None to return raw numbers. + response_type : {"expectation", "prediction"}, default="expectation" + The response type to display in the HDI band: + + - ``"expectation"``: HDI of the model expectation (μ). This shows + uncertainty from model parameters only, excluding observation noise. + Results in narrower intervals that represent the uncertainty in + the expected value of the outcome. + - ``"prediction"``: HDI of the posterior predictive (ŷ). This includes + observation noise (σ) in addition to parameter uncertainty, resulting + in wider intervals that represent the full predictive uncertainty + for new observations. + show_hdi_annotation : bool, default=False + Whether to display a text annotation at the bottom of the figure + explaining what the HDI represents. Set to False to hide the annotation. + **kwargs : dict + Additional keyword arguments. + + Returns + ------- + tuple[plt.Figure, plt.Axes] + The matplotlib figure and axes. + """ + # Log HDI type info once per session + _log_response_type_info_once() + + # Select the variable name based on response_type + var_name = "mu" if response_type == "expectation" else "y_hat" + fig, ax = plt.subplots() # Plot raw data sns.scatterplot( @@ -264,7 +307,7 @@ def _bayesian_plot( # Plot model fit to data h_line, h_patch = plot_xY( self.x_pred[self.running_variable_name], - self.pred["posterior_predictive"].mu.isel(treated_units=0), + self.pred["posterior_predictive"][var_name].isel(treated_units=0), ax=ax, plot_hdi_kwargs={"color": "C1"}, ) @@ -296,6 +339,11 @@ def _bayesian_plot( labels=labels, fontsize=LEGEND_FONT_SIZE, ) + + # Add HDI type annotation to the title + if show_hdi_annotation: + add_hdi_annotation(ax, response_type) + return fig, ax def effect_summary( diff --git a/causalpy/experiments/synthetic_control.py b/causalpy/experiments/synthetic_control.py index 380135fbf..7bde07923 100644 --- a/causalpy/experiments/synthetic_control.py +++ b/causalpy/experiments/synthetic_control.py @@ -27,7 +27,14 @@ from causalpy.custom_exceptions import BadIndexException from causalpy.date_utils import _combine_datetime_indices, format_date_axes -from causalpy.plot_utils import get_hdi_to_df, plot_xY +from causalpy.plot_utils import ( + ResponseType, + _log_response_type_effect_summary_once, + _log_response_type_info_once, + add_hdi_annotation, + get_hdi_to_df, + plot_xY, +) from causalpy.pymc_models import PyMCModel, WeightedSumFitter from causalpy.reporting import EffectSummary from causalpy.utils import check_convex_hull_violation, round_num @@ -295,8 +302,14 @@ def algorithm(self) -> None: # get the model predictions of the observed (pre-intervention) data self.pre_pred = self.model.predict(X=self.datapre_control) - # calculate the counterfactual + # Calculate the counterfactual self.post_pred = self.model.predict(X=self.datapost_control) + + # TODO: REFACTOR TARGET - Currently, stored impacts use the model expectation + # (mu) by default. When users request response_type="prediction" in plot(), the + # y_hat-based impact is calculated on-the-fly in _bayesian_plot(). This works + # but is not ideal: consider storing both mu and y_hat based impacts, or + # refactoring to always calculate on-demand. See calculate_impact() for details. self.pre_impact = self.model.calculate_impact( self.datapre_treated, self.pre_pred ) @@ -383,19 +396,51 @@ def _bayesian_plot( self, round_to: int | None = None, treated_unit: str | None = None, + response_type: ResponseType = "expectation", + show_hdi_annotation: bool = False, **kwargs: dict, ) -> tuple[plt.Figure, list[plt.Axes]]: """ Plot the results for a specific treated unit - :param round_to: - Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. - :param treated_unit: + Parameters + ---------- + round_to : int, optional + Number of decimals used to round results. Defaults to 2. + Use None to return raw numbers. + treated_unit : str, optional Which treated unit to plot. Must be a string name of the treated unit. If None, plots the first treated unit. + response_type : {"expectation", "prediction"}, default="expectation" + The response type to display in the HDI band: + + - ``"expectation"``: HDI of the model expectation (μ). This shows + uncertainty from model parameters only, excluding observation noise. + Results in narrower intervals that represent the uncertainty in + the expected value of the outcome. + - ``"prediction"``: HDI of the posterior predictive (ŷ). This includes + observation noise (σ) in addition to parameter uncertainty, resulting + in wider intervals that represent the full predictive uncertainty + for new observations. + show_hdi_annotation : bool, default=False + Whether to display a text annotation at the bottom of the figure + explaining what the HDI represents. Set to False to hide the annotation. + **kwargs : dict + Additional keyword arguments. + + Returns + ------- + tuple[plt.Figure, list[plt.Axes]] + The matplotlib figure and axes. """ + # Log HDI type info once per session + _log_response_type_info_once() + counterfactual_label = "Counterfactual" + # Select the variable name based on response_type + var_name = "mu" if response_type == "expectation" else "y_hat" + fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8)) # TOP PLOT -------------------------------------------------- # pre-intervention period @@ -410,10 +455,10 @@ def _bayesian_plot( f"treated_unit '{treated_unit}' not found. Available units: {self.treated_units}" ) - pre_pred = self.pre_pred["posterior_predictive"].mu.sel( + pre_pred = self.pre_pred["posterior_predictive"][var_name].sel( treated_units=treated_unit ) - post_pred = self.post_pred["posterior_predictive"].mu.sel( + post_pred = self.post_pred["posterior_predictive"][var_name].sel( treated_units=treated_unit ) @@ -466,22 +511,38 @@ def _bayesian_plot( ax[0].set(title=f"{self._get_score_title(treated_unit, round_to)}") # MIDDLE PLOT ----------------------------------------------- + # Calculate impact for plotting based on response_type + if response_type == "expectation": + # Use stored mu-based impact + pre_impact_for_plot = self.pre_impact + post_impact_for_plot = self.post_impact + else: + # Calculate y_hat-based impact on demand + pre_impact_for_plot = self.model.calculate_impact( + self.datapre_treated, self.pre_pred, response_type="prediction" + ) + post_impact_for_plot = self.model.calculate_impact( + self.datapost_treated, self.post_pred, response_type="prediction" + ) + plot_xY( self.datapre.index, - self.pre_impact.sel(treated_units=treated_unit), + pre_impact_for_plot.sel(treated_units=treated_unit), ax=ax[1], plot_hdi_kwargs={"color": "C0"}, ) plot_xY( self.datapost.index, - self.post_impact.sel(treated_units=treated_unit), + post_impact_for_plot.sel(treated_units=treated_unit), ax=ax[1], plot_hdi_kwargs={"color": "C1"}, ) ax[1].axhline(y=0, c="k") ax[1].fill_between( self.datapost.index, - y1=self.post_impact.mean(["chain", "draw"]).sel(treated_units=treated_unit), + y1=post_impact_for_plot.mean(["chain", "draw"]).sel( + treated_units=treated_unit + ), color="C0", alpha=0.25, label="Causal impact", @@ -490,9 +551,17 @@ def _bayesian_plot( # BOTTOM PLOT ----------------------------------------------- ax[2].set(title="Cumulative Causal Impact") + # Calculate cumulative impact based on response_type + if response_type == "expectation": + post_impact_cumulative_for_plot = self.post_impact_cumulative + else: + post_impact_cumulative_for_plot = self.model.calculate_cumulative_impact( + post_impact_for_plot + ) + plot_xY( self.datapost.index, - self.post_impact_cumulative.sel(treated_units=treated_unit), + post_impact_cumulative_for_plot.sel(treated_units=treated_unit), ax=ax[2], plot_hdi_kwargs={"color": "C1"}, ) @@ -543,6 +612,10 @@ def _bayesian_plot( ) format_date_axes(ax, full_index) + # Add HDI type annotation to the top subplot's title + if show_hdi_annotation: + add_hdi_annotation(ax[0], response_type) + return fig, ax def _ols_plot( @@ -670,20 +743,36 @@ def get_plot_data_ols(self) -> pd.DataFrame: return self.plot_data def get_plot_data_bayesian( - self, hdi_prob: float = 0.94, treated_unit: str | None = None + self, + hdi_prob: float = 0.94, + treated_unit: str | None = None, + response_type: ResponseType = "expectation", ) -> pd.DataFrame: """ Recover the data of the PrePostFit experiment along with the prediction and causal impact information. - :param hdi_prob: - Prob for which the highest density interval will be computed. The default value is defined as the default from the :func:`arviz.hdi` function. - :param treated_unit: + Parameters + ---------- + hdi_prob : float, default=0.94 + Probability for which the highest density interval will be computed. + The default value is defined as the default from the :func:`arviz.hdi` function. + treated_unit : str, optional Which treated unit to extract data for. Must be a string name of the treated unit. If None, uses the first treated unit. + response_type : {"expectation", "prediction"}, default="expectation" + The response type to use for predictions and impact: + + - ``"expectation"``: Uses the model expectation (μ). Excludes observation + noise, focusing on the systematic causal effect. + - ``"prediction"``: Uses the full posterior predictive (ŷ). Includes + observation noise, showing the full predictive uncertainty. """ if not isinstance(self.model, PyMCModel): raise ValueError("Unsupported model type") + # Map semantic response_type to internal variable name + var_name = "mu" if response_type == "expectation" else "y_hat" + hdi_pct = int(round(hdi_prob * 100)) pred_lower_col = f"pred_hdi_lower_{hdi_pct}" @@ -706,10 +795,10 @@ def get_plot_data_bayesian( # Extract predictions - handle multi-unit case pre_pred_vals = az.extract( - self.pre_pred, group="posterior_predictive", var_names="mu" + self.pre_pred, group="posterior_predictive", var_names=var_name ).mean("sample") post_pred_vals = az.extract( - self.post_pred, group="posterior_predictive", var_names="mu" + self.post_pred, group="posterior_predictive", var_names=var_name ).mean("sample") # Extract predictions for the specified treated unit (always has treated_units dimension) @@ -718,11 +807,15 @@ def get_plot_data_bayesian( # HDI intervals for predictions (always use treated_units dimension) pre_hdi = get_hdi_to_df( - self.pre_pred["posterior_predictive"].mu.sel(treated_units=treated_unit), + self.pre_pred["posterior_predictive"][var_name].sel( + treated_units=treated_unit + ), hdi_prob=hdi_prob, ) post_hdi = get_hdi_to_df( - self.post_pred["posterior_predictive"].mu.sel(treated_units=treated_unit), + self.post_pred["posterior_predictive"][var_name].sel( + treated_units=treated_unit + ), hdi_prob=hdi_prob, ) @@ -733,23 +826,35 @@ def get_plot_data_bayesian( pre_data[[pred_lower_col, pred_upper_col]] = pre_lower_upper post_data[[pred_lower_col, pred_upper_col]] = post_lower_upper - # Impact data - always use primary unit for main dataframe + # Impact data - select based on response_type + if response_type == "expectation": + pre_impact = self.pre_impact + post_impact = self.post_impact + else: + # Calculate y_hat-based impact on demand + pre_impact = self.model.calculate_impact( + self.datapre_treated, self.pre_pred, response_type="prediction" + ) + post_impact = self.model.calculate_impact( + self.datapost_treated, self.post_pred, response_type="prediction" + ) + pre_data["impact"] = ( - self.pre_impact.mean(dim=["chain", "draw"]) + pre_impact.mean(dim=["chain", "draw"]) .sel(treated_units=treated_unit) .values ) post_data["impact"] = ( - self.post_impact.mean(dim=["chain", "draw"]) + post_impact.mean(dim=["chain", "draw"]) .sel(treated_units=treated_unit) .values ) # Impact HDI intervals (always use treated_units dimension) pre_impact_hdi = get_hdi_to_df( - self.pre_impact.sel(treated_units=treated_unit), hdi_prob=hdi_prob + pre_impact.sel(treated_units=treated_unit), hdi_prob=hdi_prob ) post_impact_hdi = get_hdi_to_df( - self.post_impact.sel(treated_units=treated_unit), hdi_prob=hdi_prob + post_impact.sel(treated_units=treated_unit), hdi_prob=hdi_prob ) # Extract only the lower and upper columns for impact HDI @@ -793,6 +898,7 @@ def effect_summary( treated_unit: str | None = None, period: Literal["intervention", "post", "comparison"] | None = None, prefix: str = "Post-period", + response_type: ResponseType = "expectation", **kwargs: Any, ) -> EffectSummary: """ @@ -822,6 +928,13 @@ def effect_summary( Ignored for Synthetic Control (two-period design only). prefix : str, optional Prefix for prose generation. Defaults to "Post-period". + response_type : {"expectation", "prediction"}, default="expectation" + Response type to compute effect sizes: + + - ``"expectation"``: Effect size HDI based on model expectation (μ). + Excludes observation noise, focusing on the systematic causal effect. + - ``"prediction"``: Effect size HDI based on posterior predictive (ŷ). + Includes observation noise, showing full predictive uncertainty. Returns ------- @@ -852,14 +965,21 @@ def effect_summary( is_pymc = isinstance(self.model, PyMCModel) + # Log HDI type info once per session (for PyMC models only) + if is_pymc: + _log_response_type_effect_summary_once() + # Extract windowed impact data windowed_impact, window_coords = _extract_window( - self, window, treated_unit=treated_unit + self, window, treated_unit=treated_unit, response_type=response_type ) # Extract counterfactual for relative effects counterfactual = _extract_counterfactual( - self, window_coords, treated_unit=treated_unit + self, + window_coords, + treated_unit=treated_unit, + response_type=response_type, ) if is_pymc: diff --git a/causalpy/plot_utils.py b/causalpy/plot_utils.py index df1d151c0..45940ed3a 100644 --- a/causalpy/plot_utils.py +++ b/causalpy/plot_utils.py @@ -15,7 +15,9 @@ Plotting utility functions. """ -from typing import Any +import logging +from functools import lru_cache +from typing import Any, Literal import arviz as az import matplotlib.pyplot as plt @@ -26,6 +28,92 @@ from matplotlib.lines import Line2D from pandas.api.extensions import ExtensionArray +# Type alias for response type parameter +ResponseType = Literal["expectation", "prediction"] + +# Module-level logger +logger = logging.getLogger(__name__) + + +@lru_cache(maxsize=1) +def _log_response_type_info_once() -> None: + """Log response type information for plots once per session. + + This function uses lru_cache to ensure the message is only logged once, + regardless of how many times plot() is called. + """ + logger.info( + "Plot intervals use response_type='expectation' by default (model mean, excluding " + "observation noise). For full predictive uncertainty including observation " + "noise, use response_type='prediction'. To annotate plots with this information, " + "use show_hdi_annotation=True." + ) + + +@lru_cache(maxsize=1) +def _log_response_type_effect_summary_once() -> None: + """Log response type information for effect_summary once per session. + + This function uses lru_cache to ensure the message is only logged once, + regardless of how many times effect_summary() is called. + """ + logger.info( + "Effect size intervals use response_type='expectation' by default (model mean, excluding " + "observation noise). For full predictive uncertainty including observation " + "noise, use response_type='prediction'." + ) + + +def add_hdi_annotation( + ax: plt.Axes, + response_type: ResponseType, + hdi_prob: float = 0.94, +) -> None: + """Add HDI type information to an axes title. + + This function appends a line to the existing title of the given axes + to indicate whether the HDI (Highest Density Interval) represents: + - Model expectation (μ): excludes observation noise + - Posterior predictive (ŷ): includes observation noise + + Parameters + ---------- + ax : plt.Axes + The matplotlib axes whose title should be updated. + response_type : {"expectation", "prediction"} + The response type used for the HDI: + - "expectation": HDI of the model expectation (μ), which excludes + observation noise. Shows uncertainty from model parameters only. + - "prediction": HDI of the posterior predictive (ŷ), which includes + observation noise. Shows the full predictive uncertainty. + hdi_prob : float, optional + The probability mass of the HDI. Default is 0.94. + + Examples + -------- + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots() + >>> ax.set_title("My Plot") + Text(...) + >>> add_hdi_annotation(ax, "expectation") # doctest: +SKIP + """ + hdi_pct = int(hdi_prob * 100) + + if response_type == "expectation": + annotation = ( + f"Shaded: {hdi_pct}% HDI of model expectation (μ), excl. observation noise" + ) + else: + annotation = ( + f"Shaded: {hdi_pct}% HDI of posterior predictive (ŷ), " + "incl. observation noise" + ) + + # Get existing title and append annotation + current_title = ax.get_title() + new_title = f"{current_title}\n{annotation}" if current_title else annotation + ax.set_title(new_title) + def plot_xY( x: pd.DatetimeIndex | np.ndarray | pd.Index | pd.Series | ExtensionArray, diff --git a/causalpy/pymc_models.py b/causalpy/pymc_models.py index 65b95bfec..2a670e65a 100644 --- a/causalpy/pymc_models.py +++ b/causalpy/pymc_models.py @@ -27,6 +27,9 @@ from patsy import dmatrix from pymc_extras.prior import Prior +# Type alias for response type - reuse from plot_utils for consistency +# Using semantic names ("expectation", "prediction") rather than internal variable names ("mu", "y_hat") +from causalpy.plot_utils import ResponseType from causalpy.utils import round_num from causalpy.variable_selection_priors import VariableSelectionPrior @@ -392,46 +395,66 @@ def score(self, X, y, coords: dict[str, Any] | None = None, **kwargs) -> pd.Seri return pd.Series(scores) def calculate_impact( - self, y_true: xr.DataArray, y_pred: az.InferenceData + self, + y_true: xr.DataArray, + y_pred: az.InferenceData, + response_type: ResponseType = "expectation", ) -> xr.DataArray: """ Calculate the causal impact as the difference between observed and predicted values. - The impact is calculated using the posterior expectation (`mu`) rather than the - posterior predictive (`y_hat`). This means the causal impact represents the - difference from the expected value of the model, excluding observation noise. - This approach provides a cleaner measure of the causal effect by focusing on - the systematic difference rather than including sampling variability from the - observation noise term. + By default, the impact is calculated using the posterior expectation (`mu`) + rather than the posterior predictive (`y_hat`). This means the causal impact + represents the difference from the expected value of the model, excluding + observation noise. This approach provides a cleaner measure of the causal + effect by focusing on the systematic difference rather than including sampling + variability from the observation noise term. Parameters ---------- y_true : xr.DataArray The observed outcome values with dimensions ["obs_ind", "treated_units"]. y_pred : az.InferenceData - The posterior predictive samples containing the "mu" variable, which - represents the expected value (mean) of the outcome. - + The posterior predictive samples containing the predicted values. + response_type : {"expectation", "prediction"}, default="expectation" + The response type to use for impact calculation: + + - ``"expectation"``: Uses the model expectation (μ). Excludes observation + noise, focusing on the systematic causal effect. + - ``"prediction"``: Uses the full posterior predictive (ŷ). Includes + observation noise, showing the full predictive uncertainty. Returns ------- xr.DataArray - The causal impact with dimensions ending in "obs_ind". The impact includes - posterior uncertainty from the model parameters but excludes observation noise. + The causal impact with dimensions ending in "obs_ind". Notes ----- - By using `mu` (the posterior expectation) rather than `y_hat` (the posterior - predictive with observation noise), the uncertainty in the impact reflects: + When using ``response_type="expectation"`` (the posterior expectation), the + uncertainty in the impact reflects: + - Parameter uncertainty in the fitted model - Uncertainty in the counterfactual prediction But excludes: + - Observation-level noise (sigma) - This makes the impact plots focus on the systematic causal effect rather than - individual observation variability. + When using ``response_type="prediction"``, the uncertainty also includes + observation-level noise, resulting in wider intervals. + + .. note:: + + **REFACTOR TARGET**: Currently, experiment classes store impacts using + the default ``response_type="expectation"`` at fit time. When users call + ``plot(response_type="prediction")``, the prediction-based impact is calculated + on-the-fly in the ``_bayesian_plot()`` methods. This approach works but + duplicates the decision logic. A future refactor could unify this by + either storing both variants or always calculating on-demand. """ - y_hat = y_pred["posterior_predictive"]["mu"] + # Map semantic response_type to internal variable name + var_name = "mu" if response_type == "expectation" else "y_hat" + y_hat = y_pred["posterior_predictive"][var_name] # Ensure the coordinate type and values match along obs_ind so xarray can align if "obs_ind" in y_hat.dims and "obs_ind" in getattr(y_true, "coords", {}): try: diff --git a/causalpy/reporting.py b/causalpy/reporting.py index c81114c77..bb8c03f99 100644 --- a/causalpy/reporting.py +++ b/causalpy/reporting.py @@ -32,6 +32,8 @@ import xarray as xr from scipy.stats import t +from causalpy.plot_utils import ResponseType + @dataclass class EffectSummary: @@ -535,7 +537,12 @@ def _select_treated_unit_numpy( return data[:, 0] -def _extract_window(result, window, treated_unit=None): +def _extract_window( + result, + window, + treated_unit=None, + response_type: ResponseType = "expectation", +): """Extract windowed impact data based on window specification. Assumes result.post_impact is properly shaped xarray or numpy array. @@ -548,6 +555,11 @@ def _extract_window(result, window, treated_unit=None): Window specification: "post", (start, end) tuple, or slice object treated_unit : str, optional For multi-unit experiments, specify which treated unit to analyze + response_type : {"expectation", "prediction"}, default="expectation" + Response type to compute effect sizes: + + - ``"expectation"``: Uses stored mu-based impact (default) + - ``"prediction"``: Calculates y_hat-based impact on demand Returns ------- @@ -555,7 +567,31 @@ def _extract_window(result, window, treated_unit=None): (windowed_impact, window_coords) where windowed_impact is the data and window_coords is the corresponding index """ - post_impact = result.post_impact + # Select impact based on response_type + if response_type == "prediction": + # Calculate y_hat-based impact on demand + from causalpy.pymc_models import PyMCModel + + if isinstance(result.model, PyMCModel): + # ITS uses post_y, SyntheticControl uses datapost_treated + if hasattr(result, "post_y"): + y_true = result.post_y + elif hasattr(result, "datapost_treated"): + y_true = result.datapost_treated + else: + # Fall back to stored impact if we can't find y_true + post_impact = result.post_impact + y_true = None + + if y_true is not None: + post_impact = result.model.calculate_impact( + y_true, result.post_pred, response_type="prediction" + ) + else: + # OLS doesn't support y_hat, fall back to stored impact + post_impact = result.post_impact + else: + post_impact = result.post_impact # Check if PyMC (xarray with chain/draw dims) or OLS is_pymc = isinstance(post_impact, xr.DataArray) and ( @@ -651,7 +687,12 @@ def _extract_window(result, window, treated_unit=None): return windowed_impact, window_coords -def _extract_counterfactual(result, window_coords, treated_unit=None): +def _extract_counterfactual( + result, + window_coords, + treated_unit=None, + response_type: ResponseType = "expectation", +): """Extract counterfactual predictions for the window. Reuses logic from _extract_window for consistency. @@ -664,6 +705,11 @@ def _extract_counterfactual(result, window_coords, treated_unit=None): Window coordinates from _extract_window treated_unit : str, optional For multi-unit experiments, specify which treated unit to analyze + response_type : {"expectation", "prediction"}, default="expectation" + Response type to compute effect sizes: + + - ``"expectation"``: Uses mu (model expectation) for counterfactual + - ``"prediction"``: Uses y_hat (posterior predictive) for counterfactual Returns ------- @@ -672,10 +718,13 @@ def _extract_counterfactual(result, window_coords, treated_unit=None): """ post_pred = result.post_pred + # Select variable name based on response_type + var_name = "mu" if response_type == "expectation" else "y_hat" + # PyMC: Extract from InferenceData if hasattr(post_pred, "posterior_predictive"): # PyMC model - InferenceData object - counterfactual = post_pred.posterior_predictive["mu"] + counterfactual = post_pred.posterior_predictive[var_name] # Handle treated_unit selection using helper if "treated_units" in counterfactual.dims: @@ -687,7 +736,7 @@ def _extract_counterfactual(result, window_coords, treated_unit=None): elif isinstance(post_pred, dict) and "posterior_predictive" in post_pred: # PyMC model - dict format (fallback) - counterfactual = post_pred["posterior_predictive"]["mu"] + counterfactual = post_pred["posterior_predictive"][var_name] # Handle treated_unit selection using helper if "treated_units" in counterfactual.dims: diff --git a/causalpy/tests/test_experiment_refactor.py b/causalpy/tests/test_experiment_refactor.py new file mode 100644 index 000000000..00acdf214 --- /dev/null +++ b/causalpy/tests/test_experiment_refactor.py @@ -0,0 +1,100 @@ +# Copyright 2026 - 2026 The PyMC Labs Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +""" +Tests for experiment init refactors and new helpers. +""" + +import numpy as np +import pandas as pd +import pytest + +import causalpy as cp + +SAMPLE_KWARGS = { + "tune": 5, + "draws": 10, + "chains": 1, + "progressbar": False, + "random_seed": 42, +} + + +def test_interrupted_time_series_datapre_datapost_properties() -> None: + """datapre/datapost should reflect treatment time split.""" + df = cp.load_data("its") + df["date"] = pd.to_datetime(df["date"]) + df = df.set_index("date") + + treatment_time = pd.Timestamp("2017-01-01") + result = cp.InterruptedTimeSeries( + df, + treatment_time=treatment_time, + formula="y ~ 1 + t", + model=cp.pymc_models.LinearRegression(sample_kwargs=SAMPLE_KWARGS), + ) + + assert result.data.index.name == "obs_ind" + assert result.datapre.index.max() < treatment_time + assert result.datapost.index.min() >= treatment_time + + +def test_synthetic_control_datapre_datapost_properties() -> None: + """datapre/datapost should reflect treatment time split.""" + df = cp.load_data("sc") + treatment_time = 70 + result = cp.SyntheticControl( + df, + treatment_time=treatment_time, + control_units=["a", "b", "c", "d", "e", "f", "g"], + treated_units=["actual"], + model=cp.pymc_models.WeightedSumFitter(sample_kwargs=SAMPLE_KWARGS), + ) + + assert result.data.index.name == "obs_ind" + assert result.datapre.index.max() < treatment_time + assert result.datapost.index.min() >= treatment_time + + +@pytest.mark.integration +def test_regression_kink_gradient_change_uses_epsilon() -> None: + """Gradient change should respect the instance epsilon value.""" + rng = np.random.default_rng(42) + n_obs = 50 + kink = 0.5 + beta = [0, -1, 0, 2, 0] + sigma = 0.05 + x = rng.uniform(-1, 1, n_obs) + y = ( + beta[0] + + beta[1] * x + + beta[2] * x**2 + + beta[3] * (x - kink) * (x >= kink) + + beta[4] * ((x - kink) ** 2) * (x >= kink) + + rng.normal(0, sigma, n_obs) + ) + df = pd.DataFrame({"x": x, "y": y, "treated": x >= kink}) + + result = cp.RegressionKink( # type: ignore[abstract] + df, + formula=f"y ~ 1 + x + I((x-{kink})*treated)", + kink_point=kink, + epsilon=0.01, + model=cp.pymc_models.LinearRegression(sample_kwargs=SAMPLE_KWARGS), + ) + + mu_kink_left, mu_kink, mu_kink_right = result._probe_kink_point() + expected = result._eval_gradient_change( + mu_kink_left, mu_kink, mu_kink_right, result.epsilon + ) + np.testing.assert_allclose(result.gradient_change.values, expected.values) diff --git a/causalpy/tests/test_hdi_type.py b/causalpy/tests/test_hdi_type.py new file mode 100644 index 000000000..a3b2adc7e --- /dev/null +++ b/causalpy/tests/test_hdi_type.py @@ -0,0 +1,319 @@ +# Copyright 2025 - 2026 The PyMC Labs Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +""" +Tests for response_type parameter functionality across experiment classes. + +These tests specifically target the response_type="prediction" and show_hdi_annotation=True +branches to ensure code coverage. +""" + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import pytest + +import causalpy as cp + +# Minimal sample kwargs for fast testing +SAMPLE_KWARGS = { + "tune": 5, + "draws": 10, + "chains": 1, + "progressbar": False, + "random_seed": 42, +} + + +def _setup_regression_kink_data(kink): + """Set up data for regression kink design tests.""" + seed = 42 + rng = np.random.default_rng(seed) + N = 50 + beta = [0, -1, 0, 2, 0] + sigma = 0.05 + x = rng.uniform(-1, 1, N) + y = _reg_kink_function(x, beta, kink) + rng.normal(0, sigma, N) + return pd.DataFrame({"x": x, "y": y, "treated": x >= kink}) + + +def _reg_kink_function(x, beta, kink): + """Utility function for regression kink design.""" + return ( + beta[0] + + beta[1] * x + + beta[2] * x**2 + + beta[3] * (x - kink) * (x >= kink) + + beta[4] * ((x - kink) ** 2) * (x >= kink) + ) + + +@pytest.fixture +def its_result(): + """Create a minimal ITS result for testing.""" + df = cp.load_data("its") + df["date"] = pd.to_datetime(df["date"]) + df = df.set_index("date") + + result = cp.InterruptedTimeSeries( + df, + treatment_time=pd.Timestamp("2017-01-01"), + formula="y ~ 1 + t", + model=cp.pymc_models.LinearRegression(sample_kwargs=SAMPLE_KWARGS), + ) + return result + + +@pytest.fixture +def its_three_period_result(): + """Create a minimal three-period ITS result for testing.""" + df = cp.load_data("its") + df["date"] = pd.to_datetime(df["date"]) + df = df.set_index("date") + + result = cp.InterruptedTimeSeries( + df, + treatment_time=pd.Timestamp("2017-01-01"), + treatment_end_time=pd.Timestamp("2017-06-01"), + formula="y ~ 1 + t", + model=cp.pymc_models.LinearRegression(sample_kwargs=SAMPLE_KWARGS), + ) + return result + + +@pytest.fixture +def sc_result(): + """Create a minimal Synthetic Control result for testing.""" + df = cp.load_data("sc") + result = cp.SyntheticControl( + df, + treatment_time=70, + control_units=["a", "b", "c", "d", "e", "f", "g"], + treated_units=["actual"], + model=cp.pymc_models.WeightedSumFitter(sample_kwargs=SAMPLE_KWARGS), + ) + return result + + +@pytest.fixture +def did_result(): + """Create a minimal DiD result for testing.""" + df = cp.load_data("did") + result = cp.DifferenceInDifferences( + df, + formula="y ~ 1 + group*post_treatment", + time_variable_name="t", + group_variable_name="group", + model=cp.pymc_models.LinearRegression(sample_kwargs=SAMPLE_KWARGS), + ) + return result + + +@pytest.fixture +def rd_result(): + """Create a minimal RD result for testing.""" + df = cp.load_data("rd") + result = cp.RegressionDiscontinuity( + df, + formula="y ~ 1 + x + treated", + treatment_threshold=0.5, + model=cp.pymc_models.LinearRegression(sample_kwargs=SAMPLE_KWARGS), + running_variable_name="x", + ) + return result + + +@pytest.fixture +def rkink_result(): + """Create a minimal RegressionKink result for testing.""" + kink = 0.5 + df = _setup_regression_kink_data(kink) + result = cp.RegressionKink( + df, + formula=f"y ~ 1 + x + I((x-{kink})*treated)", + kink_point=kink, + model=cp.pymc_models.LinearRegression(sample_kwargs=SAMPLE_KWARGS), + ) + return result + + +@pytest.fixture +def prepostnegd_result(): + """Create a minimal PrePostNEGD result for testing.""" + df = cp.load_data("anova1") + result = cp.PrePostNEGD( + df, + formula="post ~ 1 + pre + group", + group_variable_name="group", + pretreatment_variable_name="pre", + model=cp.pymc_models.LinearRegression(sample_kwargs=SAMPLE_KWARGS), + ) + return result + + +class TestResponseTypePlotting: + """Test response_type parameter for plotting methods.""" + + @pytest.mark.integration + def test_its_plot_prediction_hdi(self, its_result): + """Test ITS plot with response_type='prediction'.""" + fig, ax = its_result.plot(response_type="prediction") + assert fig is not None + plt.close(fig) + + @pytest.mark.integration + def test_its_plot_show_annotation(self, its_result): + """Test ITS plot with show_hdi_annotation=True.""" + fig, ax = its_result.plot(show_hdi_annotation=True) + # Check that annotation was added to the title + title = ax[0].get_title() + assert "HDI" in title or "94%" in title + plt.close(fig) + + @pytest.mark.integration + def test_its_plot_prediction_with_annotation(self, its_result): + """Test ITS plot with both response_type='prediction' and annotation.""" + fig, ax = its_result.plot(response_type="prediction", show_hdi_annotation=True) + title = ax[0].get_title() + assert "posterior predictive" in title or "ŷ" in title + plt.close(fig) + + @pytest.mark.integration + def test_sc_plot_prediction_hdi(self, sc_result): + """Test Synthetic Control plot with response_type='prediction'.""" + fig, ax = sc_result.plot(response_type="prediction") + assert fig is not None + plt.close(fig) + + @pytest.mark.integration + def test_sc_plot_show_annotation(self, sc_result): + """Test Synthetic Control plot with show_hdi_annotation=True.""" + fig, ax = sc_result.plot(show_hdi_annotation=True) + title = ax[0].get_title() + assert "HDI" in title or "94%" in title + plt.close(fig) + + @pytest.mark.integration + def test_did_plot_prediction_hdi(self, did_result): + """Test DiD plot with response_type='prediction'.""" + fig, ax = did_result.plot(response_type="prediction") + assert fig is not None + plt.close(fig) + + @pytest.mark.integration + def test_did_plot_show_annotation(self, did_result): + """Test DiD plot with show_hdi_annotation=True.""" + fig, ax = did_result.plot(show_hdi_annotation=True) + title = ax.get_title() + assert "HDI" in title or "94%" in title + plt.close(fig) + + @pytest.mark.integration + def test_rd_plot_prediction_hdi(self, rd_result): + """Test RD plot with response_type='prediction'.""" + fig, ax = rd_result.plot(response_type="prediction") + assert fig is not None + plt.close(fig) + + @pytest.mark.integration + def test_rd_plot_show_annotation(self, rd_result): + """Test RD plot with show_hdi_annotation=True.""" + fig, ax = rd_result.plot(show_hdi_annotation=True) + title = ax.get_title() + assert "HDI" in title or "94%" in title + plt.close(fig) + + @pytest.mark.integration + def test_rkink_plot_prediction_hdi(self, rkink_result): + """Test RegressionKink plot with response_type='prediction'.""" + fig, ax = rkink_result.plot(response_type="prediction") + assert fig is not None + plt.close(fig) + + @pytest.mark.integration + def test_rkink_plot_show_annotation(self, rkink_result): + """Test RegressionKink plot with show_hdi_annotation=True.""" + fig, ax = rkink_result.plot(show_hdi_annotation=True) + title = ax.get_title() + assert "HDI" in title or "94%" in title + plt.close(fig) + + @pytest.mark.integration + def test_prepostnegd_plot_prediction_hdi(self, prepostnegd_result): + """Test PrePostNEGD plot with response_type='prediction'.""" + fig, ax = prepostnegd_result.plot(response_type="prediction") + assert fig is not None + plt.close(fig) + + @pytest.mark.integration + def test_prepostnegd_plot_show_annotation(self, prepostnegd_result): + """Test PrePostNEGD plot with show_hdi_annotation=True.""" + fig, ax = prepostnegd_result.plot(show_hdi_annotation=True) + # PrePostNEGD returns a list of axes + title = ax[0].get_title() + assert "HDI" in title or "94%" in title + plt.close(fig) + + +class TestResponseTypeEffectSummary: + """Test response_type parameter for effect_summary methods.""" + + @pytest.mark.integration + def test_its_effect_summary_prediction(self, its_result): + """Test ITS effect_summary with response_type='prediction'.""" + summary = its_result.effect_summary(response_type="prediction") + assert summary is not None + assert hasattr(summary, "table") + + @pytest.mark.integration + def test_sc_effect_summary_prediction(self, sc_result): + """Test Synthetic Control effect_summary with response_type='prediction'.""" + summary = sc_result.effect_summary(response_type="prediction") + assert summary is not None + assert hasattr(summary, "table") + + +class TestResponseTypeAnalyzePersistence: + """Test response_type parameter for analyze_persistence method.""" + + @pytest.mark.integration + def test_its_analyze_persistence_prediction(self, its_three_period_result): + """Test ITS analyze_persistence with response_type='prediction'.""" + result = its_three_period_result.analyze_persistence(response_type="prediction") + assert isinstance(result, dict) + for key in [ + "mean_effect_during", + "mean_effect_post", + "persistence_ratio", + "total_effect_during", + "total_effect_post", + ]: + assert key in result + + +class TestResponseTypeGetPlotData: + """Test response_type parameter for get_plot_data_bayesian methods.""" + + @pytest.mark.integration + def test_its_get_plot_data_prediction(self, its_result): + """Test ITS get_plot_data_bayesian with response_type='prediction'.""" + df = its_result.get_plot_data_bayesian(response_type="prediction") + assert isinstance(df, pd.DataFrame) + assert "prediction" in df.columns + + @pytest.mark.integration + def test_sc_get_plot_data_prediction(self, sc_result): + """Test Synthetic Control get_plot_data_bayesian with response_type='prediction'.""" + df = sc_result.get_plot_data_bayesian(response_type="prediction") + assert isinstance(df, pd.DataFrame) + assert "prediction" in df.columns diff --git a/causalpy/tests/test_plot_utils.py b/causalpy/tests/test_plot_utils.py index c44b2f67e..dace13d3d 100644 --- a/causalpy/tests/test_plot_utils.py +++ b/causalpy/tests/test_plot_utils.py @@ -15,12 +15,18 @@ Tests for plot utility functions """ +import matplotlib.pyplot as plt import numpy as np import pandas as pd import pytest import xarray as xr -from causalpy.plot_utils import get_hdi_to_df +from causalpy.plot_utils import ( + _log_response_type_effect_summary_once, + _log_response_type_info_once, + add_hdi_annotation, + get_hdi_to_df, +) @pytest.mark.integration @@ -95,3 +101,73 @@ def test_get_hdi_to_df_with_coordinate_dimensions(): # Verify reasonable HDI values (should be around the mean of 5.0) assert result["lower"].min() > 3.0, "HDI lower bounds should be reasonable" assert result["higher"].max() < 7.0, "HDI upper bounds should be reasonable" + + +class TestAddHdiAnnotation: + """Tests for the add_hdi_annotation function.""" + + def test_add_hdi_annotation_expectation(self): + """Test adding HDI annotation for expectation type.""" + fig, ax = plt.subplots() + ax.set_title("Original Title") + + add_hdi_annotation(ax, "expectation") + + title = ax.get_title() + assert "Original Title" in title + assert "94% HDI of model expectation (μ)" in title + assert "excl. observation noise" in title + plt.close(fig) + + def test_add_hdi_annotation_prediction(self): + """Test adding HDI annotation for prediction type.""" + fig, ax = plt.subplots() + ax.set_title("Original Title") + + add_hdi_annotation(ax, "prediction") + + title = ax.get_title() + assert "Original Title" in title + assert "94% HDI of posterior predictive (ŷ)" in title + assert "incl. observation noise" in title + plt.close(fig) + + def test_add_hdi_annotation_custom_prob(self): + """Test adding HDI annotation with custom probability.""" + fig, ax = plt.subplots() + ax.set_title("My Plot") + + add_hdi_annotation(ax, "expectation", hdi_prob=0.89) + + title = ax.get_title() + assert "89% HDI" in title + plt.close(fig) + + def test_add_hdi_annotation_empty_title(self): + """Test adding HDI annotation when there's no existing title.""" + fig, ax = plt.subplots() + # No title set + + add_hdi_annotation(ax, "expectation") + + title = ax.get_title() + assert "94% HDI of model expectation (μ)" in title + plt.close(fig) + + +class TestResponseTypeLogging: + """Tests for the response type logging functions.""" + + def test_log_response_type_info_once_callable(self): + """Test that _log_response_type_info_once is callable without error.""" + # Clear the cache to ensure fresh state + _log_response_type_info_once.cache_clear() + # Should not raise + _log_response_type_info_once() + + def test_log_response_type_effect_summary_once_callable(self): + """Test that _log_response_type_effect_summary_once is callable without error.""" + # Clear the cache to ensure fresh state + _log_response_type_effect_summary_once.cache_clear() + # Should not raise + _log_response_type_effect_summary_once()