API#

Top-level API#

CausalPy: causal inference for quasi-experiments in Python.

class causalpy.DifferenceInDifferences[source]#

A class to analyse data from Difference in Difference settings.

Note

There is no pre/post intervention data distinction for DiD, we fit all the data available.

Parameters:
  • data (NativeDataFrame) – Any eager dataframe Narwhals supports, such as pandas, Polars, or PyArrow. Converted to pandas internally.

  • formula (str) – A statistical model formula.

  • time_variable_name (str) – Name of the data column for the time variable.

  • group_variable_name (str) – Name of the data column for the group variable.

  • post_treatment_variable_name (str) – Name of the data column indicating post-treatment period. Defaults to “post_treatment”.

  • model (PyMCModel | RegressorMixin | None) – A PyMC model for difference in differences. Defaults to LinearRegression.

Notes

Estimate extraction

Both Bayesian and OLS backends store the fitted group-by-post interaction coefficient as causal_impact. The class also constructs treated-post counterfactual predictions for visualization by setting that interaction term to zero, but those predictions do not determine the reported scalar effect. In an additive identity-link model, the coefficient and the corresponding predicted contrast are algebraically identical.

Examples

>>> import causalpy as cp
>>> df = cp.load_data("did")
>>> seed = 42
>>> 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={
...             "target_accept": 0.95,
...             "random_seed": seed,
...             "progressbar": False,
...         }
...     ),
... )
__init__(data, formula, time_variable_name, group_variable_name, post_treatment_variable_name='post_treatment', model=None)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the experiment algorithm: fit model, predict, and calculate causal impact.

Return type:

None

causal_impact: DataArray | float | None#
effect_summary(*, direction='increase', alpha=0.05, min_effect=None)[source]#

Generate a decision-ready summary of causal effects for Difference-in-Differences.

Parameters:
  • direction (Literal['increase', 'decrease', 'two-sided']) – Direction for tail probability calculation (ignored for point-estimate predictions).

  • alpha (float) – Significance level for HDI/CI intervals (1-alpha confidence level).

  • min_effect (float | None) – Region of Practical Equivalence (ROPE) threshold (ignored for point-estimate predictions).

Returns:

Object with .table (DataFrame) and .text (str) attributes

Return type:

EffectSummary

input_validation()[source]#

Validate the input data and model formula for correctness.

Return type:

None

plot(*, round_to=None, ci_prob=0.94, kind='ribbon', ci_kind='hdi', num_samples=50, figsize=None, show=True, legend_kwargs=None)[source]#

Plot the difference-in-differences results.

Parameters:
  • round_to (int | None) – Number of decimals used to round numerical results in the figure title. Defaults to None, in which case 2 significant figures are used.

  • ci_prob (float) – Probability mass of the highest density interval drawn around the posterior predictive bands for the control, treatment, and counterfactual trajectories. Must be in (0, 1]. Ignored for OLS models. Defaults to HDI_PROB (currently 0.94).

  • kind (Literal['ribbon', 'histogram', 'spaghetti']) – How posterior uncertainty is rendered via plot_posterior_over_x(). Defaults to "ribbon". For "spaghetti", legends use draw lines rather than a shaded band. For "histogram", uncertainty is shown as a 2D density heatmap with a mean line overlay (no ribbon patch for legends).

  • ci_kind (Literal['hdi', 'eti']) – Credible interval type when kind="ribbon". Defaults to "hdi".

  • num_samples (int) – Number of posterior draws when kind="spaghetti". Defaults to 50. Ignored for other kinds.

  • figsize (tuple[float, float] | None) – Width and height of the figure in inches, passed to matplotlib.pyplot.subplots(). Defaults to None (use matplotlib’s default).

  • show (bool) – Whether to automatically display the plot. Defaults to True. Set to False if you want to modify the figure before displaying it.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments to adjust legend placement and styling. Supported keys: loc, bbox_to_anchor, fontsize, frameon, title (bbox_transform is accepted alongside bbox_to_anchor). The existing legend is modified in place so that custom handles are preserved.

Returns:

  • fig (matplotlib.figure.Figure) – The figure that was created.

  • ax (matplotlib.axes.Axes) – The axes object containing the plot.

Return type:

tuple[Figure, Axes]

summary(round_to=2)[source]#

Print summary of main results and model coefficients.

Parameters:

round_to (int | None) – Number of decimals used to round results. Defaults to 2. Use None to return raw numbers.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = True#
class causalpy.EffectSummary[source]#

Container for effect summary statistics and prose report.

table#

DataFrame containing summary statistics (mean, median, HDI, tail probabilities)

Type:

pd.DataFrame

text#

Detailed multi-paragraph narrative report with observed vs counterfactual breakdown, statistical credibility assessment, and assumptions/guidance.

Type:

str

__init__(table, text)#
Parameters:
Return type:

None

table: DataFrame#
text: str#
class causalpy.EstimateEffect[source]#

Pipeline step that fits a causal experiment.

Captures the experiment class and its keyword arguments. When the pipeline runs, instantiates the experiment with the pipeline’s data (which triggers fitting) and stores the result in the context.

Parameters:
  • method (type[BaseExperiment]) – The experiment class to instantiate (e.g. cp.InterruptedTimeSeries).

  • **kwargs (Any) – Keyword arguments accepted by method’s constructor, except data, which the pipeline supplies. This is a deliberately narrow dynamic forwarder: method may be an integrator-provided BaseExperiment subclass, so its accepted constructor keys cannot be enumerated here. Built-in experiment constructors declare every supported key explicitly; unsupported, misspelled, or incomplete arguments raise TypeError during pipeline validation.

Examples

>>> import causalpy as cp
>>> step = cp.EstimateEffect(
...     method=cp.InterruptedTimeSeries,
...     treatment_time=pd.Timestamp("2020-01-01"),
...     formula="y ~ 1 + t",
...     model=cp.pymc_models.LinearRegression(),
... )
__init__(method, **kwargs)[source]#
Parameters:
Return type:

None

run(context)[source]#

Instantiate and fit the experiment.

The experiment constructor receives context.data as its first positional argument, followed by all captured keyword arguments.

Parameters:

context (PipelineContext) – Pipeline context. context.data is forwarded to the experiment constructor as the first positional argument.

Returns:

Updated context with experiment, experiment_config, and (if available) effect_summary populated.

Return type:

PipelineContext

validate(context)[source]#

Check that the step is properly configured.

Parameters:

context (PipelineContext) – Pipeline context. Its data is used to validate the selected experiment constructor’s keyword arguments before execution.

Raises:
  • TypeError – If method is not a subclass of BaseExperiment or supplied constructor arguments are incompatible with an inspectable constructor, including omitted required arguments.

  • ValueError – If data is passed in kwargs (it comes from the pipeline).

Return type:

None

class causalpy.GenerateReport[source]#

Pipeline step that generates an HTML report from pipeline results.

Parameters:
  • include_plots (bool) – Whether to include diagnostic plots in the report.

  • include_effect_summary (bool) – Whether to include the effect summary section.

  • include_sensitivity (bool) – Whether to include sensitivity analysis results.

  • output_file (str | Path | None) – If provided, write the HTML report to this file.

Examples

>>> import causalpy as cp
>>> step = cp.GenerateReport(
...     include_plots=True, output_file="report.html"
... )
__init__(include_plots=True, include_effect_summary=True, include_sensitivity=True, output_file=None)[source]#
Parameters:
  • include_plots (bool)

  • include_effect_summary (bool)

  • include_sensitivity (bool)

  • output_file (str | Path | None)

Return type:

None

run(context)[source]#

Generate the HTML report and store it in the context.

Parameters:

context (PipelineContext) – Pipeline context providing experiment, effect_summary, and sensitivity_results (any of which may be None).

Returns:

The same context with report populated.

Return type:

PipelineContext

validate(context)[source]#

GenerateReport has no strict prerequisites; it gracefully handles missing data.

Parameters:

context (PipelineContext) – Pipeline context (unused; required by the step interface).

Return type:

None

class causalpy.InstrumentalVariable[source]#

A class to analyse instrumental variable style experiments.

Parameters:
  • instruments_data (NativeDataFrame) – Instruments for our treatment variable, as any eager dataframe Narwhals supports, such as pandas, Polars, or PyArrow. Should contain instruments Z, and treatment t. Converted to pandas internally.

  • data (NativeDataFrame) – Covariates for fitting the focal regression of interest, as any eager dataframe Narwhals supports. Should contain covariates X including treatment t and outcome y. Converted to pandas internally.

  • instruments_formula (str) – A statistical model formula for the instrumental stage regression, e.g. t ~ 1 + z1 + z2 + z3.

  • formula (str) – A statistical model formula for the focal regression, e.g. y ~ 1 + t + x1 + x2 + x3.

  • model (InstrumentalVariableRegression | None) – A PyMC model. Defaults to InstrumentalVariableRegression.

  • priors (dict | None) – Dictionary of priors for the mus and sigmas of both regressions. If priors are not specified we will substitute MLE estimates for the beta coefficients. Example: priors = {"mus": [0, 0], "sigmas": [1, 1], "eta": 2, "lkj_sd": 2}.

  • vs_prior_type (str or None, default=None) – Type of variable selection prior: ‘spike_and_slab’, ‘horseshoe’, or None. If None, uses standard normal priors.

  • vs_hyperparams (dict, optional) – Hyperparameters for variable selection priors. Only used if vs_prior_type is not None.

  • binary_treatment (bool, default=False) – A indicator for whether the treatment to be modelled is binary or not. Determines which PyMC model we use to model the joint outcome and treatment.

Notes

Estimate extraction

The class computes naive OLS and two-stage least-squares reference fits, then fits a joint Bayesian model for the treatment and outcome equations. Under the instrumental-variable assumptions, the causal quantity is read from the outcome-stage coefficient associated with the instrumented treatment; no counterfactual prediction or population standardization is performed. For binary treatments, its LATE interpretation applies to the complier population induced by the instrument; continuous treatments require the corresponding structural IV interpretation.

Examples

>>> import pandas as pd
>>> import causalpy as cp
>>> from causalpy.pymc_models import InstrumentalVariableRegression
>>> import numpy as np
>>> N = 100
>>> e1 = np.random.normal(0, 3, N)
>>> e2 = np.random.normal(0, 1, N)
>>> Z = np.random.uniform(0, 1, N)
>>> ## Ensure the endogeneity of the the treatment variable
>>> X = -1 + 4 * Z + e2 + 2 * e1
>>> y = 2 + 3 * X + 3 * e1
>>> test_data = pd.DataFrame({"y": y, "X": X, "Z": Z})
>>> sample_kwargs = {
...     "tune": 1,
...     "draws": 5,
...     "chains": 2,
...     "cores": 1,
...     "target_accept": 0.95,
...     "progressbar": False,
... }
>>> instruments_formula = "X  ~ 1 + Z"
>>> formula = "y ~  1 + X"
>>> instruments_data = test_data[["X", "Z"]]
>>> data = test_data[["y", "X"]]
>>> iv = cp.InstrumentalVariable(
...     instruments_data=instruments_data,
...     data=data,
...     instruments_formula=instruments_formula,
...     formula=formula,
...     model=InstrumentalVariableRegression(sample_kwargs=sample_kwargs),
... )
>>> # With variable selection
>>> iv = cp.InstrumentalVariable(
...     instruments_data=instruments_data,
...     data=data,
...     instruments_formula=instruments_formula,
...     formula=formula,
...     model=InstrumentalVariableRegression(sample_kwargs=sample_kwargs),
...     vs_prior_type="spike_and_slab",
...     vs_hyperparams={"slab_sigma": 5.0},
... )
__init__(instruments_data, data, instruments_formula, formula, model=None, priors=None, vs_prior_type=None, vs_hyperparams=None, binary_treatment=False)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the experiment algorithm: fit OLS, 2SLS, and Bayesian IV model.

Return type:

None

effect_summary()[source]#

Raise because unified effect summaries are unavailable.

Raises:

NotImplementedError – Instrumental-variable experiments do not implement a unified decision-ready effect summary.

Return type:

NoReturn

get_2SLS_fit()[source]#

Two Stage Least Squares Fit.

This function is called by the experiment, results are used for priors if none are provided.

Return type:

None

get_naive_OLS_fit()[source]#

Naive Ordinary Least Squares.

This function is called by the experiment.

Return type:

None

input_validation()[source]#

Validate the input data and model formula for correctness.

Return type:

None

plot(*, show=True, legend_kwargs=None)[source]#

Plot the results.

Parameters:
  • show (bool) – Reserved; ignored. Defaults to True.

  • legend_kwargs (dict[str, Any] | None) – Reserved; ignored.

Raises:

NotImplementedError – Always.

Return type:

None

Notes

Plotting is not yet implemented for instrumental variable experiments. This stub exists so every experiment subclass offers an explicit, kwarg-only plot() signature (issue #886).

summary(round_to=2)[source]#

Print summary of main results and model coefficients.

Parameters:

round_to (int | None) – Number of decimals used to round results. Defaults to 2. Use None to return raw numbers.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = False#
class causalpy.InterruptedTimeSeries[source]#

The class for interrupted time series analysis.

Supports both two-period (permanent intervention) and three-period (temporary intervention) designs. When treatment_end_time is provided, the analysis splits the post-intervention period into an intervention period and a post-intervention period, enabling analysis of effect persistence and decay.

Parameters:
  • data (NativeDataFrame) – Time series data as any eager dataframe Narwhals supports. For a pandas dataframe the index carries the time axis, and it should be either a DatetimeIndex or numeric (integer/float), with unique values in monotonically increasing order. Dataframes from other libraries have no index, so those callers must pass time_column.

  • treatment_time (int | float | Timestamp) – The time when treatment occurred, should be in reference to the data index. Must match the index type (DatetimeIndex requires pd.Timestamp). INCLUSIVE: Observations at exactly treatment_time are included in the post-intervention period (uses >= comparison).

  • formula (str) – A statistical model formula using patsy syntax (e.g., “y ~ 1 + t + C(month)”).

  • model (PyMCModel | RegressorMixin | PyMCForecastModel | None) – A PyMC (Bayesian) or sklearn (OLS) model. If None, defaults to a PyMC LinearRegression model. Alternatively, a PyMCForecastModel wrapping a pymc_forecast forecasting model can serve as the counterfactual backend (requires the optional pymc-forecast dependency); see causalpy.pymc_forecast_models for when to prefer it.

  • treatment_end_time (int | float | Timestamp | None) – The time when treatment ended, enabling three-period analysis. Must be greater than treatment_time and within the data range. If None (default), the analysis assumes a permanent intervention (two-period design). INCLUSIVE: Observations at exactly treatment_end_time are included in the post-intervention period (uses >= comparison).

  • time_column (str | None) – Column holding the time axis. It becomes the index of the data. Required for non-pandas inputs, which carry no index. If None (default), the pandas index of data is used. Passing it for data that already has a meaningful index raises, since only one of the two can be the time axis.

Notes

Estimate extraction

The model is fitted to pre-intervention observations and predicts the untreated trajectory after the intervention. Pointwise impact is the observed post-intervention outcome minus that one-sided counterfactual prediction, and cumulative impact is its running sum. Bayesian backends subtract the posterior conditional expectation mu rather than noisy posterior-predictive draws y_hat; OLS subtracts its point prediction.

This fit-predict-subtract procedure is a reduced-form estimator. From a Bayesian structural perspective, the same impact can be viewed as the response to an intervention shock in a state-space model of the outcome series; see the knowledgebase page on structural causal models for the reduced-form versus structural distinction.

The three-period design is useful for analyzing temporary interventions such as:

  • Marketing campaigns with defined start and end dates

  • Policy trials or pilot programs

  • Clinical treatments with limited duration

  • Seasonal interventions

Use effect_summary(period="intervention") to analyze effects during the intervention, and effect_summary(period="post") to analyze effect persistence after the intervention ends.

Examples

Two-period design (permanent intervention):

>>> import causalpy as cp
>>> df = (
...     cp.load_data("its")
...     .assign(date=lambda x: pd.to_datetime(x["date"]))
...     .set_index("date")
... )
>>> treatment_time = pd.to_datetime("2017-01-01")
>>> result = cp.InterruptedTimeSeries(
...     df,
...     treatment_time,
...     formula="y ~ 1 + t + C(month)",
...     model=cp.pymc_models.LinearRegression(
...         sample_kwargs={"random_seed": 42, "progressbar": False}
...     ),
... )

Three-period design (temporary intervention):

>>> treatment_time = pd.to_datetime("2017-01-01")
>>> treatment_end_time = pd.to_datetime("2017-06-01")
>>> result = cp.InterruptedTimeSeries(
...     df,
...     treatment_time,
...     formula="y ~ 1 + t + C(month)",
...     model=cp.pymc_models.LinearRegression(
...         sample_kwargs={"random_seed": 42, "progressbar": False}
...     ),
...     treatment_end_time=treatment_end_time,
... )
>>> # Get period-specific effect summaries
>>> intervention_summary = result.effect_summary(period="intervention")
>>> post_summary = result.effect_summary(period="post")
__init__(data, treatment_time, formula, model=None, treatment_end_time=None, time_column=None)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the experiment algorithm: fit model, predict, and calculate causal impact.

Return type:

None

analyze_persistence(hdi_prob=0.94, direction='increase')[source]#

Analyze effect persistence between intervention and post-intervention periods.

Computes mean effects, persistence ratio, and total (cumulative) impacts for both periods. The persistence ratio is the post-intervention mean effect divided by the intervention mean effect (as a decimal, e.g., 0.30 means 30% persistence, 1.5 means 150%). Note: The ratio can exceed 1.0 if the post-intervention effect is larger than the intervention effect.

Automatically prints a summary of the results.

Parameters:
  • hdi_prob (float) – Probability for the HDI interval (Bayesian models only). Defaults to HDI_PROB (currently 0.94).

  • direction (Literal['increase', 'decrease', 'two-sided']) – Direction for tail probability calculation (Bayesian models only)

Returns:

Dictionary containing:

  • ”mean_effect_during”: Mean effect during intervention period

  • ”mean_effect_post”: Mean effect during post-intervention period

  • ”persistence_ratio”: Post-intervention mean effect divided by intervention mean (decimal, can exceed 1.0)

  • ”total_effect_during”: Total (cumulative) effect during intervention period

  • ”total_effect_post”: Total (cumulative) effect during post-intervention period

Return type:

dict[str, Any]

Raises:

ValueError – If treatment_end_time is not provided (two-period design)

Examples

>>> import causalpy as cp
>>> import pandas as pd
>>> df = (
...     cp.load_data("its")
...     .assign(date=lambda x: pd.to_datetime(x["date"]))
...     .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 + C(month)",
...     model=cp.pymc_models.LinearRegression(
...         sample_kwargs={"random_seed": 42, "progressbar": False}
...     ),
... )
>>> persistence = result.analyze_persistence()
... # Note: Results are automatically printed to console
>>> persistence["persistence_ratio"]
-1.224
property datapost: DataFrame#

Data from on or after the treatment time (inclusive).

Post-period: index >= treatment_time

property datapre: DataFrame#

Data from before the treatment time (exclusive).

Pre-period: index < treatment_time

effect_summary(*, window='post', direction='increase', alpha=0.05, cumulative=True, relative=True, min_effect=None, treated_unit=None, period=None, prefix='Post-period')[source]#

Generate a decision-ready summary of causal effects for Interrupted Time Series.

Parameters:
  • window (Union[Literal['post'], tuple, slice]) –

    Time window for analysis:

    • ”post”: All post-treatment time points (default)

    • (start, end): Tuple of start and end times (handles both datetime and integer indices)

    • slice: Python slice object for integer indices

  • direction (Literal['increase', 'decrease', 'two-sided']) – Direction for tail probability calculation (PyMC only, ignored for OLS).

  • alpha (float) – Significance level for HDI/CI intervals (1-alpha confidence level).

  • cumulative (bool) – Whether to include cumulative effect statistics.

  • relative (bool) – Whether to include relative effect statistics (% change vs counterfactual).

  • min_effect (float | None) – Region of Practical Equivalence (ROPE) threshold (PyMC only, ignored for OLS).

  • treated_unit (str | None) – Ignored for Interrupted Time Series (single unit).

  • period (Optional[Literal['intervention', 'post', 'comparison']]) – For three-period designs (with treatment_end_time), specify which period to summarize. Defaults to None for standard behavior.

  • prefix (str) – Prefix for prose generation (e.g., “During intervention”, “Post-intervention”). Defaults to “Post-period”.

Returns:

Object with .table (DataFrame) and .text (str) attributes. The .text attribute contains a detailed multi-paragraph narrative report.

Return type:

EffectSummary

get_plot_data(*, hdi_prob=0.94)[source]#

Recover the data of the experiment along with the prediction and causal impact information.

HDI columns are included only when the prediction container carries posterior draws (point-estimate backends return just prediction and impact).

Parameters:

hdi_prob (float) – Probability mass of the highest density interval. Defaults to the project-wide HDI_PROB (currently 0.94). Ignored when the prediction container has no posterior draws.

Return type:

DataFrame

input_validation(data, treatment_time, treatment_end_time=None)[source]#

Validate the input data and model formula for correctness.

Parameters:
  • data (DataFrame) – The experiment data.

  • treatment_time (int | float | Timestamp) – Start of the treatment period.

  • treatment_end_time (int | float | Timestamp | None) – Optional end of the treatment period for three-period designs.

Return type:

None

plot(*, round_to=2, ci_prob=0.94, kind='ribbon', ci_kind='hdi', num_samples=50, figsize=(7, 8), show=True, legend_kwargs=None)[source]#

Plot the interrupted time-series results.

Parameters:
  • round_to (int | None) – Number of decimals used to round numerical results in the figure title (e.g. the Bayesian \(R^2\)). Defaults to 2. Use None to render raw numbers.

  • ci_prob (float) – Probability mass of the credible interval drawn around the posterior predictive, causal impact, and cumulative impact bands. Must be in (0, 1]. Ignored for OLS models. Defaults to HDI_PROB (currently 0.94).

  • kind (Literal['ribbon', 'histogram', 'spaghetti']) – How posterior uncertainty is rendered via plot_posterior_over_x(). Defaults to "ribbon". For "spaghetti", legends use draw lines rather than a shaded band. For "histogram", uncertainty is shown as a 2D density heatmap with a mean line overlay (no ribbon patch for legends).

  • ci_kind (Literal['hdi', 'eti']) – Credible interval type when kind="ribbon". Defaults to "hdi".

  • num_samples (int) – Number of posterior draws when kind="spaghetti". Defaults to 50. Ignored for other kinds.

  • figsize (tuple[float, float]) – Width and height of the figure in inches, passed to matplotlib.pyplot.subplots(). Defaults to (7, 8).

  • show (bool) – Whether to automatically display the plot. Defaults to True. Set to False if you want to modify the figure before displaying it.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments to adjust legend placement and styling. Supported keys: loc, bbox_to_anchor, fontsize, frameon, title (bbox_transform is accepted alongside bbox_to_anchor). The existing legend is modified in place so that custom handles are preserved.

Returns:

  • fig (matplotlib.figure.Figure) – The figure that was created.

  • ax (list[matplotlib.axes.Axes]) – The three axes (top: predictions, middle: causal impact, bottom: cumulative impact).

Return type:

tuple[Figure, list[Axes]]

post_design: Dataset#
pre_design: Dataset#
summary(round_to=None)[source]#

Print summary of main results and model coefficients.

Parameters:

round_to (int | None) – Number of decimals used to round results. Defaults to 2. Use None to return raw numbers.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = True#
supports_pymc_forecast: bool = True#
class causalpy.InversePropensityWeighting[source]#

A class to analyse inverse propensity weighting experiments.

Parameters:
  • data (NativeDataFrame) – Any eager dataframe Narwhals supports, such as pandas, Polars, or PyArrow. Converted to pandas internally.

  • formula (str) – A statistical model formula for the propensity model.

  • outcome_variable (str) – A string denoting the outcome variable in data to be reweighted.

  • weighting_scheme (str) – A string denoting which weighting scheme to use among: ‘raw’, ‘robust’, ‘doubly_robust’ or ‘overlap’. See Aronow and Miller “Foundations of Agnostic Statistics” for discussion and computation of these weighting schemes.

  • model (PropensityScore | None) – A PyMC model. Defaults to PropensityScore.

Notes

Estimate extraction

Fitting produces posterior propensity-score draws. get_ate() post-processes one draw at a time: "raw" and "robust" contrast inverse-probability-weighted mean outcomes for the treated and control potential outcomes, "overlap" contrasts overlap-weighted means for the overlap population, and "doubly_robust" augments inverse-probability weighting with separate OLS outcome regressions before averaging over all observations.

Examples

>>> import causalpy as cp
>>> df = cp.load_data("nhefs")
>>> seed = 42
>>> result = cp.InversePropensityWeighting(
...     df,
...     formula="trt ~ 1 + age + race",
...     outcome_variable="outcome",
...     weighting_scheme="robust",
...     model=cp.pymc_models.PropensityScore(
...         sample_kwargs={
...             "draws": 100,
...             "target_accept": 0.95,
...             "random_seed": seed,
...             "progressbar": False,
...         },
...     ),
... )
__init__(data, formula, outcome_variable, weighting_scheme, model=None)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the experiment algorithm by fitting the propensity score model.

Delegates to self.model.fit with the covariate matrix self.X, treatment vector self.t, and coordinate metadata self.coords.

Return type:

None

effect_summary()[source]#

Raise because unified effect summaries are unavailable.

Raises:

NotImplementedError – Inverse-propensity-weighting experiments do not implement a unified decision-ready effect summary.

Return type:

NoReturn

get_ate(i, idata, method='doubly_robust')[source]#

Compute the Average Treatment Effect for a single posterior sample.

Post-processes the sample posterior distribution for propensity scores, one sample at a time, using the specified weighting method.

Parameters:
  • i (int) – Index of the posterior sample to process.

  • idata (DataTree) – DataTree containing the posterior samples.

  • method (str) – Weighting scheme to use. One of ‘robust’, ‘raw’, ‘overlap’, or ‘doubly_robust’. Defaults to ‘doubly_robust’.

Returns:

A list of [ate, trt, ntrt] where:

  • ate: Average Treatment Effect

  • trt: Weighted mean outcome for treated group

  • ntrt: Weighted mean outcome for non-treated group

Return type:

list[float]

input_validation()[source]#

Validate the input data and model formula for correctness.

Checks that the outcome_variable exists in self.data and that the treatment values produced by the formula are binary. Note that a constant (single-value) treatment passes this check without error but may cause downstream estimators to fail or produce undefined results.

Raises:

DataException – If the outcome variable is missing from the data or if the formula produces treatment values other than zero and one.

Return type:

None

make_doubly_robust_adjustment(ps)[source]#

Compute doubly-robust adjusted outcomes.

The doubly-robust weighting scheme is discussed in Aronow and Miller’s Foundations of Agnostic Statistics. This implementation fixes the outcome model to ordinary least squares (OLS), so the compromise between the outcome model and the propensity model is always performed with a linear regression.

Parameters:

ps (ndarray) – Propensity scores for each observation.

Returns:

A tuple of (weighted_outcome0, weighted_outcome1, None, None). The two None values are returned for interface consistency with the other adjustment methods; no explicit group sizes are needed because the doubly-robust estimator averages over all observations.

Return type:

tuple[pd.Series, pd.Series, None, None]

make_overlap_adjustments(ps)[source]#

Compute inverse-propensity-weighted outcomes using the overlap scheme.

This weighting scheme was adapted from Lucy D’Agostino McGowan’s blog on propensity score weights, referenced in the primary CausalPy explainer notebook. Overlap weights target the population for which there is clinical equipoise (i.e., where propensity scores are near 0.5), reducing sensitivity to extreme scores.

Parameters:

ps (ndarray) – Propensity scores for each observation.

Returns:

A tuple of (weighted_outcome0, weighted_outcome1, n_ntrt, n_trt) where the weighted outcomes and normalisation terms are all pd.Series (unlike the raw/robust schemes which return integer counts).

Return type:

tuple[pd.Series, pd.Series, pd.Series, pd.Series]

make_raw_adjustments(ps)[source]#

Compute inverse-propensity-weighted outcomes using the raw (basic) scheme.

This is the simplest form of inverse propensity weighting, as discussed in Aronow and Miller’s Foundations of Agnostic Statistics. Each observation is weighted by the reciprocal of its propensity score (or 1 - ps for the control group).

Parameters:

ps (ndarray) – Propensity scores for each observation.

Returns:

A tuple of (weighted_outcome0, weighted_outcome1, n_ntrt, n_trt) where the weighted outcomes are the IPW-adjusted outcome values for the control and treated groups. n_ntrt and n_trt are both equal to the total number of observations (the raw scheme normalises by the full sample size).

Return type:

tuple[pd.Series, pd.Series, int, int]

make_robust_adjustments(ps)[source]#

Compute inverse-propensity-weighted outcomes using the robust (Horvitz-Thompson) scheme.

This estimator is discussed in Aronow and Miller’s Foundations of Agnostic Statistics as being related to the Horvitz-Thompson method.

Parameters:

ps (ndarray) – Propensity scores for each observation.

Returns:

A tuple of (weighted_outcome0, weighted_outcome1, n_ntrt, n_trt) where the weighted outcomes are the IPW-adjusted outcome values for the control and treated groups, and n_ntrt / n_trt are the corresponding group sizes used for normalisation.

Return type:

tuple[pd.Series, pd.Series, int, int]

plot(*, show=True, legend_kwargs=None)[source]#

Plot the results.

Parameters:
  • show (bool) – Reserved; ignored. Defaults to True.

  • legend_kwargs (dict[str, Any] | None) – Reserved; ignored.

Raises:

NotImplementedError – Always; call plot_ate() or plot_balance_ecdf() instead.

Return type:

None

Notes

Inverse propensity weighting does not expose a unified plot() view; instead, use the dedicated diagnostics plot_ate() (treatment-effect distribution) and plot_balance_ecdf() (covariate-balance ECDF). This stub exists so every experiment subclass offers an explicit, kwarg-only plot() signature (issue #886).

plot_ate(idata=None, method=None, prop_draws=100, ate_draws=300)[source]#

Plot the Average Treatment Effect and propensity score distributions.

Produces a three-panel figure:

  • Top panel – Weighted and unweighted histograms of posterior propensity scores for treated and control groups, drawn from prop_draws posterior samples. The "raw" method shows unweighted counts; "overlap" uses overlap weights; all other methods ("robust", "doubly_robust") share a common IPW-weighted histogram branch.

  • Bottom-left panel – Histograms of the reweighted potential outcomes E[Y(1)] and E[Y(0)].

  • Bottom-right panel – Histogram of the ATE distribution with a vertical line at its posterior mean.

Parameters:
  • idata (DataTree | None) – DataTree with posterior propensity-score samples. If None, uses the fitted model backend’s DataTree.

  • method (str | None) – Weighting scheme to apply. One of 'robust', 'raw', 'overlap', or 'doubly_robust'. If None, falls back to self.weighting_scheme.

  • prop_draws (int) – Number of posterior draws used for the propensity score histogram. Defaults to 100.

  • ate_draws (int) – Number of posterior draws used to compute ATE samples. Defaults to 300.

Returns:

The matplotlib Figure and a list of three Axes objects.

Return type:

tuple[plt.Figure, list[plt.Axes]]

plot_balance_ecdf(covariate, idata=None, weighting_scheme=None)[source]#

Plot the empirical CDF of a covariate before and after IPW adjustment.

Produces a two-panel figure comparing the raw (unweighted) ECDFs of the treated and control groups with the reweighted ECDFs. This serves as a visual balance diagnostic: well-balanced covariates should show overlapping ECDFs in the right-hand panel.

Parameters:
  • covariate (str) – Name of the covariate column (must be one of the model’s design matrix labels) to check for balance.

  • idata (DataTree | None) – DataTree with posterior propensity-score samples. If None, uses the fitted model backend’s DataTree.

  • weighting_scheme (str | None) – Weighting scheme to apply. One of 'raw', 'robust', or 'overlap'. If None, falls back to self.weighting_scheme.

Returns:

The matplotlib Figure and a list of two Axes objects (raw ECDF on the left, weighted ECDF on the right).

Return type:

tuple[plt.Figure, list[plt.Axes]]

supports_bayes: bool = True#
supports_ols: bool = False#
weighted_percentile(data, weights, perc)[source]#

Compute a weighted percentile of the data.

Sorts data and weights together, builds a weighted empirical CDF, and linearly interpolates to find the value at the requested percentile.

Parameters:
  • data (ndarray) – One-dimensional array of data values.

  • weights (ndarray) – Non-negative weights corresponding to each element of data.

  • perc (float) – Desired percentile expressed as a fraction in [0, 1] (e.g., 0.5 for the median).

Returns:

The interpolated data value at the given weighted percentile.

Return type:

float

Raises:

ValueError – If perc is not between 0 and 1.

class causalpy.PanelRegression[source]#

Panel regression with fixed effects estimation.

Enables panel-aware visualization and diagnostics, with support for both unpooled dummy-variable and demeaned (de-meaned) fixed effects.

Parameters:
  • data (NativeDataFrame) – Panel data as any eager dataframe Narwhals supports, such as pandas, Polars, or PyArrow. Each row is an observation for a unit at a time period. Converted to pandas internally.

  • formula (str) – A statistical model formula using patsy syntax. For the unpooled dummy-variable fixed-effects approach, include C(unit_var) (and optionally C(time_var)) in the formula. For the demeaned transformation, do NOT include those C(...) terms; fixed effects are removed by transformation before fitting.

  • unit_fe_variable (str) – Column name for the unit identifier (e.g., “state”, “id”, “country”).

  • time_fe_variable (str | None) – Column name for the time identifier (e.g., “year”, “wave”, “period”). If provided, time fixed effects will be included. Default is None.

  • fe_method (Literal['dummies', 'demeaned']) –

    Method for handling fixed effects:

    • ”dummies”: Use unpooled dummy-variable fixed effects (C(unit)/C(time) in formula). Gets individual unit effect estimates but creates N-1 dummy columns. Best for small N.

    • ”demeaned”: Use demeaned (de-meaned) transformation. Scales to large N but doesn’t directly estimate individual unit effects.

  • model (PyMCModel | RegressorMixin | None) – A PyMC (Bayesian) or sklearn (OLS) model. If None, a model must be provided.

n_units#

Number of unique units in the panel.

Type:

int

n_periods#

Number of unique time periods (None if time_fe_variable not provided).

Type:

int or None

fe_method#

The fixed effects method used (“dummies” or “demeaned”).

Type:

str

_group_means#

Stored group means for recovering unit effects (demeaned method only).

Type:

dict

Notes

Estimate extraction

PanelRegression fits the formula after representing fixed effects with dummy variables or demeaning, but it does not select a treatment term or compute a single built-in causal effect. Any causal estimand is extracted from the fitted coefficient chosen by the analyst, and its interpretation depends on the formula and identification assumptions.

The demeaned transformation (de-meaning by group) removes time-invariant confounders but also drops time-invariant covariates from the model. For the "dummies" approach (unpooled FE), individual unit effects can be extracted from the coefficients. For the demeaned approach, unit effects can be recovered post-hoc using the stored group means (_group_means), which are always computed from the original (pre-demeaning) data.

This class does not yet implement hierarchical/partial-pooling fixed effects. Those semantics are intentionally kept out of scope here so fe_method="dummies" remains an accurate label for the current unpooled estimator.

Two-way fixed effects (unit + time) control for both unit-specific and time-specific unobserved heterogeneity. This is the standard approach in difference-in-differences estimation.

Balanced vs unbalanced panels: A panel is balanced when every unit is observed in every time period; otherwise it is unbalanced (e.g. unit entry/exit, missing waves). When both unit and time fixed effects are requested with fe_method="demeaned", the sequential demeaning (first by unit, then by time) is algebraically equivalent to the standard two-way demeaned transformation only for balanced panels. For unbalanced panels, iterative alternating demeaning would be needed for exact convergence; the single-pass approximation used here may introduce small biases. Unbalanced panels are common in practice (e.g. firm or worker panels with attrition); for heavily unbalanced data, consider checking sensitivity or using dedicated FE packages that implement iterative two-way demeaning (e.g. reghdfe, pyfixest).

Examples

Small panel with dummy variables:

>>> import causalpy as cp
>>> import pandas as pd
>>> # Create small panel: 10 units, 20 time periods
>>> np.random.seed(42)
>>> units = [f"unit_{i}" for i in range(10)]
>>> periods = range(20)
>>> data = pd.DataFrame(
...     [
...         {
...             "unit": u,
...             "time": t,
...             "treatment": int(t >= 10 and u in units[:5]),
...             "x1": np.random.randn(),
...             "y": np.random.randn(),
...         }
...         for u in units
...         for t in periods
...     ]
... )
>>> result = cp.PanelRegression(
...     data=data,
...     formula="y ~ C(unit) + C(time) + treatment + x1",
...     unit_fe_variable="unit",
...     time_fe_variable="time",
...     fe_method="dummies",
...     model=cp.pymc_models.LinearRegression(
...         sample_kwargs={"random_seed": 42, "progressbar": False, "cores": 1}
...     ),
... )

Large panel with demeaned transformation:

>>> # Create larger panel: 1000 units, 10 time periods
>>> np.random.seed(42)
>>> units = [f"unit_{i}" for i in range(1000)]
>>> periods = range(10)
>>> data = pd.DataFrame(
...     [
...         {
...             "unit": u,
...             "time": t,
...             "treatment": int(t >= 5),
...             "x1": np.random.randn(),
...             "y": np.random.randn(),
...         }
...         for u in units
...         for t in periods
...     ]
... )
>>> result = cp.PanelRegression(
...     data=data,
...     formula="y ~ treatment + x1",  # No C(unit) needed
...     unit_fe_variable="unit",
...     time_fe_variable="time",
...     fe_method="demeaned",
...     model=cp.pymc_models.LinearRegression(
...         sample_kwargs={"random_seed": 42, "progressbar": False, "cores": 1}
...     ),
... )
__init__(data, formula, unit_fe_variable, time_fe_variable=None, fe_method='dummies', model=None)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the experiment algorithm: fit the model.

Return type:

None

effect_summary()[source]#

Raise because panel regression has no unified effect summary.

Raises:

NotImplementedError – Panel fixed-effects models estimate coefficients rather than time-varying causal impacts. Use summary() for coefficient-level inference.

Return type:

NoReturn

get_plot_data()[source]#

Get plot data with fitted values.

Bayesian models additionally return y_fitted_lower / y_fitted_upper 95% credible-interval columns.

Returns:

DataFrame with fitted values (and credible intervals when the model carries posterior draws).

Return type:

pd.DataFrame

input_validation()[source]#

Validate input parameters.

Return type:

None

plot(*, hdi_prob=0.94, show=True, legend_kwargs=None)[source]#

Plot the panel regression coefficients.

Bayesian models render a forest plot with HDI intervals; OLS models render a bar plot of point estimates. To plot only a subset of coefficients (or to customise the figure size), call plot_coefficients() directly.

Parameters:
  • hdi_prob (float) – Probability mass of the highest density interval drawn around each posterior coefficient. Must be in (0, 1]. Ignored for OLS models. Defaults to HDI_PROB (currently 0.94).

  • show (bool) – Whether to automatically display the plot. Defaults to True.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments to adjust legend placement and styling. Supported keys: loc, bbox_to_anchor, fontsize, frameon, title (bbox_transform is accepted alongside bbox_to_anchor). The existing legend is modified in place so that custom handles are preserved.

Returns:

  • fig (matplotlib.figure.Figure) – The figure that was created.

  • ax (matplotlib.axes.Axes) – The axes object containing the coefficient plot.

Return type:

tuple[Figure, Axes]

plot_coefficients(var_names=None, hdi_prob=0.94)[source]#

Plot coefficient estimates with credible/confidence intervals.

Automatically filters out fixed effect dummy coefficients to show only the treatment and control covariates.

Parameters:
  • var_names (list[str] | None) – Specific coefficient names to plot. Names must match the patsy design-matrix labels (e.g. "treatment", "x1"). If None, plots all non-FE coefficients.

  • hdi_prob (float) – Probability mass for the HDI interval when plotting Bayesian coefficients. Must be in (0, 1). Ignored for OLS models. Defaults to HDI_PROB (currently 0.94).

Returns:

Figure and axes objects

Return type:

tuple[plt.Figure, plt.Axes]

Raises:

ValueError – If var_names is empty or hdi_prob is outside (0, 1).

plot_residuals(kind='scatter')[source]#

Plot residual diagnostics.

Parameters:

kind (Literal['scatter', 'histogram', 'qq']) –

Type of residual plot:

  • ”scatter”: Residuals vs fitted values

  • ”histogram”: Distribution of residuals

  • ”qq”: Q-Q plot for normality check

Returns:

Figure and axes objects

Return type:

tuple[plt.Figure, plt.Axes]

plot_trajectories(units=None, n_sample=10, select='random', show_mean=True, hdi_prob=0.94, interval_type='mean')[source]#

Plot unit-level time series trajectories.

Shows actual vs fitted values for selected units over time. Useful for visualizing within-unit model fit and identifying problematic units.

Parameters:
  • units (list[str] | None) – Specific unit IDs to plot. If provided, ignores n_sample and select.

  • n_sample (int) – Number of units to sample if units not specified.

  • select (Literal['random', 'extreme', 'high_variance']) –

    Method for selecting units:

    • ”random”: Random sample of units

    • ”extreme”: Units with largest positive and negative effects

    • ”high_variance”: Units with most within-unit variation

  • show_mean (bool) – Whether to show the overall mean trajectory.

  • hdi_prob (float) – Probability mass for the HDI credible interval (Bayesian models only). Defaults to HDI_PROB (currently 0.94). Common alternative values are 0.89 or 0.5.

  • interval_type (Literal['mean', 'predictive']) –

    Which uncertainty interval to show for Bayesian models:

    • ”mean”: HDI of posterior mu (uncertainty in expected value)

    • ”predictive”: HDI of posterior predictive y_hat (includes observation noise)

Returns:

Figure and array of axes objects

Return type:

tuple[plt.Figure, np.ndarray]

Raises:

ValueError – If time_fe_variable is not provided (cannot plot trajectories without time)

plot_unit_effects(highlight=None, label_extreme=0)[source]#

Plot distribution of unit fixed effects.

Only available with fe_method=”dummies”. Shows histogram of estimated unit-specific intercepts.

Parameters:
  • highlight (list[str] | None) – List of unit IDs to highlight on the distribution.

  • label_extreme (int) – Number of extreme units to label (top N + bottom N).

Returns:

Figure and axes objects

Return type:

tuple[plt.Figure, plt.Axes]

Raises:

ValueError – If fe_method is not “dummies”

summary(round_to=None)[source]#

Print a summary of the panel regression results.

Parameters:

round_to (int | None) – Number of significant figures to round to. Defaults to None, in which case 2 significant figures are used.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = True#
class causalpy.PiecewiseITS[source]#

Piecewise Interrupted Time Series (Segmented Regression) experiment.

This class implements segmented-regression / piecewise linear models for Interrupted Time Series analysis with known interruption dates. Unlike the standard InterruptedTimeSeries which fits a model to pre-intervention data and forecasts a counterfactual, PiecewiseITS fits one model to the full time series and estimates explicit level and/or slope changes at each interruption.

The model uses patsy formulas with custom step() and ramp() transforms:

  • step(time, threshold): Creates a binary indicator (1 if time >= threshold) for level changes

  • ramp(time, threshold): Creates a ramp function (max(0, time - threshold)) for slope changes

Parameters:
  • data (NativeDataFrame) – Time series data as any eager dataframe Narwhals supports, such as pandas, Polars, or PyArrow. The time axis comes from the step() or ramp() column in the formula, not from the index, so a dataframe without an index works here. Converted to pandas internally.

  • formula (str) – A patsy formula specifying the model. Must include at least one step() or ramp() term, and all such terms must use the same time variable. Example: "y ~ 1 + t + step(t, 50) + ramp(t, 50)"

  • model (PyMCModel | RegressorMixin | None) – A PyMC (Bayesian) or sklearn (OLS) model. If None, defaults to a PyMC LinearRegression model.

formula#

The patsy formula used for the model.

Type:

str

interruption_times#

Canonicalized interruption thresholds extracted from the formula.

Type:

list

labels#

Names of all coefficients in the design matrix.

Type:

list[str]

effect#

Pointwise causal effect (fitted expectation - counterfactual expectation).

Type:

xr.DataArray or np.ndarray

cumulative_effect#

Cumulative causal effect over time.

Type:

xr.DataArray or np.ndarray

Notes

Estimate extraction

One model is fitted to the full time series. The no-intervention counterfactual is predicted after setting every step() and ramp() design-matrix column to zero, and the pointwise effect is the fitted conditional expectation minus that counterfactual expectation. Bayesian backends contrast posterior mu values, OLS contrasts point predictions, and the cumulative effect is the running sum.

The step and ramp transforms are patsy stateful transforms that handle both numeric and datetime time columns. For datetime, thresholds can be specified as strings (e.g., ‘2020-06-01’) or pd.Timestamp objects.

Bare datetime predictors are represented as continuous elapsed days. Use C(date) when date fixed effects are intended instead.

References

  • Wagner AK, et al. (2002). Segmented regression analysis of interrupted time series studies in medication use research. Journal of Clinical Pharmacy and Therapeutics.

  • Lopez Bernal J, et al. (2017). Interrupted time series regression for the evaluation of public health interventions: a tutorial. Int J Epidemiol.

Examples

>>> import causalpy as cp
>>> import pandas as pd
>>> import numpy as np
>>> # Generate simple piecewise data
>>> np.random.seed(42)
>>> t = np.arange(100)
>>> y = (
...     10
...     + 0.1 * t
...     + 5 * (t >= 50)
...     + 0.2 * np.maximum(0, t - 50)
...     + np.random.normal(0, 1, 100)
... )
>>> df = pd.DataFrame({"t": t, "y": y})
>>> result = cp.PiecewiseITS(
...     df,
...     formula="y ~ 1 + t + step(t, 50) + ramp(t, 50)",
...     model=cp.pymc_models.LinearRegression(
...         sample_kwargs={"random_seed": 42, "progressbar": False}
...     ),
... )

Different effects per intervention:

>>> # Level change only at t=50, level + slope change at t=100
>>> result = cp.PiecewiseITS(
...     df,
...     formula="y ~ 1 + t + step(t, 50) + step(t, 100) + ramp(t, 100)",
...     model=...,
... )

With datetime thresholds:

>>> df["date"] = pd.date_range("2020-01-01", periods=100, freq="D")
>>> result = cp.PiecewiseITS(
...     df,
...     formula="y ~ 1 + date + step(date, '2020-02-20') + ramp(date, '2020-02-20')",
...     model=...,
... )
__init__(data, formula, model=None)[source]#
Parameters:
Return type:

None

effect_summary(*, window='post', direction='increase', alpha=0.05, cumulative=True, relative=True, min_effect=None, treated_unit=None, period=None, prefix='Post-period')[source]#

Generate a decision-ready summary of PiecewiseITS causal effects.

Parameters:
  • window (Union[Literal['post'], tuple, slice]) – Time window for analysis (see BaseExperiment.effect_summary()).

  • direction (Literal['increase', 'decrease', 'two-sided']) – Direction for tail probability calculation (PyMC only).

  • alpha (float) – Significance level for HDI/CI intervals (1-alpha confidence).

  • cumulative (bool) – Whether to include cumulative effect statistics.

  • relative (bool) – Whether to include relative effect statistics.

  • min_effect (float | None) – Region of Practical Equivalence (ROPE) threshold (PyMC only).

  • treated_unit (str | None) – Multi-unit experiments select which unit to analyse.

  • period (Optional[Literal['intervention', 'post', 'comparison']]) – Not supported by PiecewiseITS; pass None.

  • prefix (str) – Prefix for prose generation.

Return type:

EffectSummary

get_plot_data(*, hdi_prob=0.94)[source]#

Recover the data of the experiment along with prediction and effect information.

HDI columns are included only when the prediction container carries posterior draws.

Parameters:

hdi_prob (float) – Probability for the highest density interval. Defaults to HDI_PROB (currently 0.94). Ignored when the prediction container has no posterior draws.

Returns:

DataFrame containing observed data, predictions, and effects.

Return type:

pd.DataFrame

plot(*, round_to=2, ci_prob=0.94, kind='ribbon', ci_kind='hdi', num_samples=50, figsize=(10, 10), show=True, legend_kwargs=None)[source]#

Plot the piecewise interrupted time-series results.

Parameters:
  • round_to (int | None) – Number of decimals used to round numerical results in the figure title. Defaults to 2. Use None to render raw numbers.

  • ci_prob (float) – Probability mass of the highest density interval drawn around the fitted, counterfactual, causal effect, and cumulative effect bands. Must be in (0, 1]. Ignored for OLS models. Defaults to HDI_PROB (currently 0.94).

  • kind (Literal['ribbon', 'histogram', 'spaghetti']) – How posterior uncertainty is rendered via plot_posterior_over_x(). Defaults to "ribbon". For "spaghetti", legends use draw lines rather than a shaded band. For "histogram", uncertainty is shown as a 2D density heatmap with a mean line overlay (no ribbon patch for legends).

  • ci_kind (Literal['hdi', 'eti']) – Credible interval type when kind="ribbon". Defaults to "hdi".

  • num_samples (int) – Number of posterior draws when kind="spaghetti". Defaults to 50. Ignored for other kinds.

  • figsize (tuple[float, float]) – Width and height of the figure in inches, passed to matplotlib.pyplot.subplots(). Defaults to (10, 10).

  • show (bool) – Whether to automatically display the plot. Defaults to True.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments to adjust legend placement and styling. Supported keys: loc, bbox_to_anchor, fontsize, frameon, title (bbox_transform is accepted alongside bbox_to_anchor). The existing legend is modified in place so that custom handles are preserved.

Returns:

  • fig (matplotlib.figure.Figure) – The figure that was created.

  • ax (list[matplotlib.axes.Axes]) – The three axes (top: observed, fitted and counterfactual; middle: causal effect; bottom: cumulative effect).

Return type:

tuple[Figure, list[Axes]]

summary(round_to=None)[source]#

Print summary of main results and model coefficients.

Parameters:

round_to (int | None) – Number of decimals used to round results. Defaults to 2.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = True#
class causalpy.Pipeline[source]#

Orchestrate a sequence of causal-inference steps.

The pipeline validates all steps before executing any of them, ensuring configuration errors are caught before potentially expensive model fitting.

Parameters:
  • data (DataFrame) – The dataset to analyse.

  • steps (list[Step]) – Ordered sequence of pipeline steps.

Examples

>>> import causalpy as cp
>>> result = cp.Pipeline(
...     data=df,
...     steps=[
...         cp.EstimateEffect(
...             method=cp.InterruptedTimeSeries,
...             treatment_time=pd.Timestamp("2020-01-01"),
...             formula="y ~ 1 + t",
...             model=cp.pymc_models.LinearRegression(),
...         ),
...     ],
... ).run()
__init__(data, steps)[source]#
Parameters:
Return type:

None

run()[source]#

Validate all steps, then execute them sequentially.

Returns:

The accumulated results of the pipeline.

Return type:

PipelineResult

Raises:

Exception – Re-raises any exception from validation or step execution.

class causalpy.PipelineContext[source]#

Mutable container that accumulates results as pipeline steps execute.

Each step reads from and writes to this context, building up a complete record of the analysis.

data#

The input dataset.

Type:

pd.DataFrame

experiment#

The fitted experiment object, populated by EstimateEffect.

Type:

BaseExperiment or None

experiment_config#

The configuration used to create the experiment (method class + keyword arguments), so that downstream steps like SensitivityAnalysis can derive experiment factories.

Type:

dict or None

effect_summary#

The effect summary from the primary experiment.

Type:

EffectSummary or None

sensitivity_results#

Accumulated sensitivity / diagnostic check results.

Type:

list

report#

Generated report artifact, populated by GenerateReport.

Type:

object or None

__init__(data, experiment=None, experiment_config=None, effect_summary=None, sensitivity_results=<factory>, report=None)#
Parameters:
Return type:

None

data: DataFrame#
effect_summary: EffectSummary | None = None#
experiment: BaseExperiment | None = None#
experiment_config: dict[str, Any] | None = None#
report: Any = None#
sensitivity_results: list[Any]#
class causalpy.PipelineResult[source]#

Immutable result returned by Pipeline.run().

experiment#

The fitted experiment.

Type:

BaseExperiment or None

effect_summary#

The effect summary from the experiment.

Type:

EffectSummary or None

sensitivity_results#

Results of all sensitivity / diagnostic checks.

Type:

list

report#

Generated report artifact.

Type:

object or None

__init__(experiment, effect_summary, sensitivity_results, report)#
Parameters:
Return type:

None

effect_summary: EffectSummary | None#
experiment: BaseExperiment | None#
classmethod from_context(context)[source]#

Build a PipelineResult from a completed PipelineContext.

Parameters:

context (PipelineContext) – Completed pipeline context to extract user-facing results from.

Returns:

Snapshot containing the experiment, effect summary, sensitivity results, and report.

Return type:

PipelineResult

report: Any#
sensitivity_results: list[Any]#
class causalpy.PrePostNEGD[source]#

A class to analyse data from pretest/posttest designs.

Parameters:
  • data (NativeDataFrame) – Any eager dataframe Narwhals supports, such as pandas, Polars, or PyArrow. Converted to pandas internally.

  • formula (str) – A statistical model formula.

  • group_variable_name (str) – Name of the column in data for the group variable; should be either binary or boolean.

  • pretreatment_variable_name (str) – Name of the column in data for the pretreatment variable.

  • model (PyMCModel | None) – A PyMC model. Defaults to LinearRegression.

Notes

Estimate extraction

The reported causal_impact is the posterior coefficient on the treatment-group term, conditional on the pretreatment outcome and any other formula covariates. Treated and untreated prediction curves are also computed for visualization, but they do not determine the reported scalar effect. With the current additive identity-link model, the treatment coefficient equals the corresponding conditional prediction contrast.

Examples

>>> import causalpy as cp
>>> df = cp.load_data("anova1")
>>> seed = 42
>>> result = cp.PrePostNEGD(
...     df,
...     formula="post ~ 1 + C(group) + pre",
...     group_variable_name="group",
...     pretreatment_variable_name="pre",
...     model=cp.pymc_models.LinearRegression(
...         sample_kwargs={
...             "target_accept": 0.95,
...             "random_seed": seed,
...             "progressbar": False,
...         }
...     ),
... )
>>> result.summary(round_to=1)
==================Pretest/posttest Nonequivalent Group Design===================
Formula: post ~ 1 + C(group) + pre

Results:
Causal impact = 2, $CI_{94%}$[2, 2]
Model coefficients:
    Intercept      -0.5, 94% HDI [-1, 0.2]
    C(group)[T.1]  2, 94% HDI [2, 2]
    pre            1, 94% HDI [1, 1]
    y_hat_sigma    0.5, 94% HDI [0.5, 0.6]
__init__(data, formula, group_variable_name, pretreatment_variable_name, model=None)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the experiment algorithm: fit model, predict, and calculate causal impact.

Return type:

None

causal_impact: DataArray#
effect_summary(*, direction='increase', alpha=0.05, min_effect=None)[source]#

Generate a decision-ready summary of causal effects for PrePostNEGD.

Parameters:
  • direction (Literal['increase', 'decrease', 'two-sided']) – Direction for tail probability calculation (PyMC only).

  • alpha (float) – Significance level for HDI/CI intervals (1-alpha confidence level).

  • min_effect (float | None) – Region of Practical Equivalence (ROPE) threshold (PyMC only).

Returns:

Object with .table (DataFrame) and .text (str) attributes

Return type:

EffectSummary

input_validation()[source]#

Validate the input data and model formula for correctness.

Return type:

None

plot(*, round_to=None, ci_prob=0.94, kind='ribbon', ci_kind='hdi', num_samples=50, figsize=(7, 9), show=True, legend_kwargs=None)[source]#

Plot the pre-post non-equivalent group design results.

Parameters:
  • round_to (int | None) – Number of decimals used to round numerical results in the figure. Defaults to None, in which case 2 significant figures are used.

  • ci_prob (float) – Probability mass of the highest density interval drawn around the posterior predictive bands for the control and treatment groups, and around the posterior of the estimated treatment effect. Must be in (0, 1]. Defaults to HDI_PROB (currently 0.94).

  • kind (Literal['ribbon', 'histogram', 'spaghetti']) – How posterior uncertainty is rendered via plot_posterior_over_x(). Defaults to "ribbon". For "spaghetti", legends use draw lines rather than a shaded band. For "histogram", uncertainty is shown as a 2D density heatmap with a mean line overlay (no ribbon patch for legends).

  • ci_kind (Literal['hdi', 'eti']) – Credible interval type when kind="ribbon". Defaults to "hdi".

  • num_samples (int) – Number of posterior draws when kind="spaghetti". Defaults to 50. Ignored for other kinds.

  • figsize (tuple[float, float]) – Width and height of the figure in inches, passed to matplotlib.pyplot.subplots(). Defaults to (7, 9).

  • show (bool) – Whether to automatically display the plot. Defaults to True.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments to adjust legend placement and styling. Supported keys: loc, bbox_to_anchor, fontsize, frameon, title (bbox_transform is accepted alongside bbox_to_anchor). The existing legend is modified in place so that custom handles are preserved.

Returns:

  • fig (matplotlib.figure.Figure) – The figure that was created.

  • ax (list[matplotlib.axes.Axes]) – The two axes (top: scatter and posterior predictive bands, bottom: estimated treatment effect posterior).

Return type:

tuple[Figure, list[Axes]]

pred_treated: DataArray#
pred_untreated: DataArray#
pred_xi: ndarray#
summary(round_to=None)[source]#

Print summary of main results and model coefficients.

Parameters:

round_to (int | None) – Number of decimals used to round results. Defaults to 2. Use None to return raw numbers.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = False#
class causalpy.RegressionDiscontinuity[source]#

A class to analyse sharp regression discontinuity experiments.

Parameters:
  • data (NativeDataFrame) – Any eager dataframe Narwhals supports, such as pandas, Polars, or PyArrow. Converted to pandas internally.

  • formula (str) – A statistical model formula.

  • treatment_threshold (float) – A scalar threshold value at which the treatment is applied.

  • model (PyMCModel | RegressorMixin | None) – A PyMC or sklearn model. Defaults to LinearRegression.

  • running_variable_name (str) – The name of the predictor variable that the treatment threshold is based upon.

  • epsilon (float) – A small scalar value which determines how far above and below the treatment threshold to evaluate the causal impact.

  • bandwidth (float) – Data outside of the bandwidth (relative to the discontinuity) is not used to fit the model.

  • donut_hole (float) – Observations within this distance from the treatment threshold are excluded from model fitting. Used as a robustness check when observations closest to the threshold may be problematic (e.g., due to manipulation or heaping). Must be non-negative and less than bandwidth if bandwidth is finite.

Notes

Estimate extraction

After fitting the regression on the selected bandwidth, the class predicts the conditional expectation immediately below the threshold with treated=0 and immediately above it with treated=1. discontinuity_at_threshold is the upper prediction minus the lower prediction, evaluated at threshold ± epsilon. This is a local prediction contrast, not a population-standardized effect.

Examples

>>> import causalpy as cp
>>> df = cp.load_data("rd")
>>> seed = 42
>>> result = cp.RegressionDiscontinuity(
...     df,
...     formula="y ~ 1 + x + treated + x:treated",
...     model=cp.pymc_models.LinearRegression(
...         sample_kwargs={
...             "draws": 100,
...             "target_accept": 0.95,
...             "random_seed": seed,
...             "progressbar": False,
...         },
...     ),
...     treatment_threshold=0.5,
... )
__init__(data, formula, treatment_threshold, model=None, running_variable_name='x', epsilon=0.001, bandwidth=inf, donut_hole=0.0)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the experiment algorithm: fit model, predict, and calculate discontinuity.

Return type:

None

effect_summary(*, direction='increase', alpha=0.05, min_effect=None)[source]#

Generate a decision-ready summary of causal effects for Regression Discontinuity.

Parameters:
  • direction (Literal['increase', 'decrease', 'two-sided']) – Direction for tail probability calculation (PyMC only, ignored for OLS).

  • alpha (float) – Significance level for HDI/CI intervals (1-alpha confidence level).

  • min_effect (float | None) – Region of Practical Equivalence (ROPE) threshold (PyMC only, ignored for OLS).

Returns:

Object with .table (DataFrame) and .text (str) attributes

Return type:

EffectSummary

input_validation()[source]#

Validate the input data and model formula for correctness.

Return type:

None

plot(*, round_to=2, ci_prob=0.94, kind='ribbon', ci_kind='hdi', num_samples=50, figsize=None, show=True, legend_kwargs=None)[source]#

Plot the regression discontinuity results.

Parameters:
  • round_to (int | None) – Number of decimals used to round numerical results in the figure title (e.g. the Bayesian \(R^2\)). Defaults to 2. Use None to render raw numbers.

  • ci_prob (float) – Probability mass of the highest density interval drawn around the posterior predictive band, and the central credible interval reported in the figure title for the discontinuity at threshold. Must be in (0, 1]. Ignored for OLS models. Defaults to HDI_PROB (currently 0.94).

  • kind (Literal['ribbon', 'histogram', 'spaghetti']) – How posterior uncertainty is rendered via plot_posterior_over_x(). Defaults to "ribbon". For "spaghetti", legends use draw lines rather than a shaded band. For "histogram", uncertainty is shown as a 2D density heatmap with a mean line overlay (no ribbon patch for legends).

  • ci_kind (Literal['hdi', 'eti']) – Credible interval type when kind="ribbon". Defaults to "hdi".

  • num_samples (int) – Number of posterior draws when kind="spaghetti". Defaults to 50. Ignored for other kinds.

  • figsize (tuple[float, float] | None) – Width and height of the figure in inches, passed to matplotlib.pyplot.subplots(). Defaults to None (use matplotlib’s default).

  • show (bool) – Whether to automatically display the plot. Defaults to True.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments to adjust legend placement and styling. Supported keys: loc, bbox_to_anchor, fontsize, frameon, title (bbox_transform is accepted alongside bbox_to_anchor). The existing legend is modified in place so that custom handles are preserved.

Returns:

  • fig (matplotlib.figure.Figure) – The figure that was created.

  • ax (matplotlib.axes.Axes) – The axes object containing the plot.

Return type:

tuple[Figure, Axes]

summary(round_to=None)[source]#

Print summary of main results and model coefficients.

Parameters:

round_to (int | None) – Number of decimals used to round results. Defaults to 2. Use None to return raw numbers.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = True#
class causalpy.RegressionKink[source]#

A class to analyse regression kink designs.

Parameters:
  • data (NativeDataFrame) – Any eager dataframe Narwhals supports, such as pandas, Polars, or PyArrow. Converted to pandas internally.

  • formula (str) – A statistical model formula.

  • kink_point (float) – A scalar value at which the kink occurs.

  • model (PyMCModel | None) – A PyMC model. Defaults to LinearRegression.

  • running_variable_name (str) – The name of the running variable column.

  • epsilon (float) – A small scalar for evaluating the causal impact above/below the kink.

  • bandwidth (float) – Data outside of the bandwidth (relative to the kink) is not used to fit the model.

Notes

Estimate extraction

The class predicts the conditional expectation at kink_point - epsilon, kink_point, and kink_point + epsilon. It forms finite-difference slopes on the left and right and stores their difference as gradient_change. This is a local prediction contrast on derivatives, not a population-standardized effect.

__init__(data, formula, kink_point, model=None, running_variable_name='x', epsilon=0.001, bandwidth=inf)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the experiment algorithm: fit model, predict, and evaluate gradient change.

Return type:

None

effect_summary(*, direction='increase', alpha=0.05, min_effect=None)[source]#

Generate a decision-ready summary of causal effects for Regression Kink.

Parameters:
  • direction (Literal['increase', 'decrease', 'two-sided']) – Direction for tail probability calculation (PyMC only, ignored for OLS).

  • alpha (float) – Significance level for HDI/CI intervals (1-alpha confidence level).

  • min_effect (float | None) – Region of Practical Equivalence (ROPE) threshold (PyMC only, ignored for OLS).

Returns:

Object with .table (DataFrame) and .text (str) attributes

Return type:

EffectSummary

input_validation()[source]#

Validate the input data and model formula for correctness.

Return type:

None

plot(*, round_to=2, ci_prob=0.94, kind='ribbon', ci_kind='hdi', num_samples=50, figsize=None, show=True, legend_kwargs=None)[source]#

Plot the regression kink results.

Parameters:
  • round_to (int | None) – Number of decimals used to round numerical results in the figure title (e.g. the Bayesian \(R^2\)). Defaults to 2. Use None to render raw numbers.

  • ci_prob (float) – Probability mass of the highest density interval drawn around the posterior predictive band, and the central credible interval reported in the figure title for the change in gradient at the kink point. Must be in (0, 1]. Defaults to HDI_PROB (currently 0.94).

  • kind (Literal['ribbon', 'histogram', 'spaghetti']) – How posterior uncertainty is rendered via plot_posterior_over_x(). Defaults to "ribbon". For "spaghetti", legends use draw lines rather than a shaded band. For "histogram", uncertainty is shown as a 2D density heatmap with a mean line overlay (no ribbon patch for legends).

  • ci_kind (Literal['hdi', 'eti']) – Credible interval type when kind="ribbon". Defaults to "hdi".

  • num_samples (int) – Number of posterior draws when kind="spaghetti". Defaults to 50. Ignored for other kinds.

  • figsize (tuple[float, float] | None) – Width and height of the figure in inches, passed to matplotlib.pyplot.subplots(). Defaults to None (use matplotlib’s default).

  • show (bool) – Whether to automatically display the plot. Defaults to True.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments to adjust legend placement and styling. Supported keys: loc, bbox_to_anchor, fontsize, frameon, title (bbox_transform is accepted alongside bbox_to_anchor). The existing legend is modified in place so that custom handles are preserved.

Returns:

  • fig (matplotlib.figure.Figure) – The figure that was created.

  • ax (matplotlib.axes.Axes) – The axes object containing the plot.

Return type:

tuple[Figure, Axes]

summary(round_to=2)[source]#

Print summary of main results and model coefficients.

Parameters:

round_to (int | None) – Number of decimals used to round results. Defaults to 2. Use None to return raw numbers.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = False#
class causalpy.SensitivityAnalysis[source]#

Pipeline step that runs a suite of sensitivity / diagnostic checks.

Parameters:

checks (list[Any] | None) – The checks to run against the fitted experiment.

Examples

>>> import causalpy as cp
>>> step = cp.SensitivityAnalysis(
...     checks=[
...         cp.checks.PlaceboInTime(n_folds=4),
...         cp.checks.PriorSensitivity(priors=[...]),
...     ]
... )
__init__(checks=None)[source]#
Parameters:

checks (list[Any] | None)

Return type:

None

checks: list[Any]#
classmethod default_for(method)[source]#

Create a SensitivityAnalysis pre-loaded with all registered default checks for method.

Parameters:

method (type[BaseExperiment]) – The experiment class to look up defaults for.

Returns:

Instance with applicable default checks instantiated.

Return type:

SensitivityAnalysis

run(context)[source]#

Run all checks against the fitted experiment.

Parameters:

context (PipelineContext) – Pipeline context containing the fitted experiment and any experiment_config required by the checks.

Returns:

The same context with sensitivity_results and report populated.

Return type:

PipelineContext

Raises:
  • RuntimeError – If no experiment has been fitted (context.experiment is None).

  • TypeError – If a check is not applicable to the experiment type.

validate(context)[source]#

Validate that checks are well-formed.

At validation time the experiment may not yet be fitted, so we only check structural issues (e.g. that each object satisfies the Check protocol).

Parameters:

context (PipelineContext) – Pipeline context (unused at validation time but required by the pipeline step interface).

Raises:

TypeError – If any item in checks does not satisfy the Check protocol.

Return type:

None

class causalpy.SensitivitySummary[source]#

Aggregate result of all sensitivity checks.

results#

Individual check results.

Type:

list[CheckResult]

all_passed#

True if every check with a pass/fail criterion passed, False if any failed, or None if no check had a pass/fail criterion.

Type:

bool or None

text#

Combined prose summary.

Type:

str

__init__(results=<factory>, all_passed=None, text='')#
Parameters:
Return type:

None

all_passed: bool | None = None#
classmethod from_results(results)[source]#

Build a summary from a list of check results.

Parameters:

results (list[CheckResult]) – Individual results to aggregate.

Returns:

Aggregated summary covering all supplied results.

Return type:

SensitivitySummary

results: list[CheckResult]#
text: str = ''#
class causalpy.StaggeredDifferenceInDifferences[source]#

A class to analyse data from staggered adoption Difference-in-Differences settings.

This class implements the Borusyak, Jaravel, and Spiess (BJS, 2024) imputation estimator for staggered adoption settings. It fits a model on untreated observations only (pre-treatment periods for eventually-treated units plus all periods for never-treated units), then predicts counterfactual outcomes for all observations. Treatment effects are computed as the difference between observed and predicted outcomes for treated observations.

Parameters:
  • data (NativeDataFrame) – Panel data (unit x time observations) as any eager dataframe Narwhals supports, such as pandas, Polars, or PyArrow. Converted to pandas internally.

  • formula (str) – A statistical model formula. Recommended: “y ~ 1 + C(unit) + C(time)” for unit and time fixed effects.

  • unit_variable_name (str) – Name of the column identifying units.

  • time_variable_name (str) – Name of the column identifying time periods.

  • treated_variable_name (str) – Name of the column indicating treatment status (0/1). Defaults to “treated”.

  • treatment_time_variable_name (str | None) – Name of the column containing unit-level treatment time (G_i). If None, treatment time is inferred from the treated_variable_name column.

  • never_treated_value (Any) – Value indicating never-treated units in treatment_time column. Defaults to np.inf.

  • model (PyMCModel | RegressorMixin | None) – A model for the untreated outcome. Defaults to LinearRegression.

  • event_window (tuple[int, int] | None) – Tuple (min_event_time, max_event_time) to restrict event-time aggregation. If None, uses all available event-times.

  • reference_event_time (int) – Event-time index associated with plots (reserved for future use). Defaults to -1.

data_#

Augmented data with G (treatment time), event_time, y_hat0 (counterfactual), and tau_hat (treatment effect) columns.

Type:

pd.DataFrame

att_group_time_#

Group-time ATT estimates: ATT(g, t) for each cohort g and calendar time t. Includes an identified column; non-identified cells have NaN estimates.

Type:

pd.DataFrame

att_event_time_#

Event-time ATT estimates: ATT(e) for each event-time e = t - G. Includes an identified column; non-identified cells have NaN estimates.

Type:

pd.DataFrame

non_identified_periods_#

Calendar periods with no untreated observations.

Type:

set

non_identified_cohorts_#

Treatment cohorts with at least one non-identified post-treatment ATT(g, t).

Type:

set

Notes

Estimate extraction

The Borusyak-Jaravel-Spiess imputation estimator fits the untreated outcome model using only observations that are not yet treated or never treated. It predicts each treated observation’s untreated potential outcome, subtracts that prediction from the observed outcome, and averages the resulting one-sided contrasts into group-time and event-time ATTs. Bayesian aggregation retains posterior uncertainty in mu; OLS aggregation uses point predictions and standard-error approximations.

Like Interrupted Time Series, this fit-predict-subtract procedure is a reduced-form estimator. The corresponding structural contrast is a saturated regression as in Wooldridge’s extended two-way fixed effects (ETWFE) framework, which CausalPy does not currently implement.

This estimator requires the following identifying assumptions:

  1. Absorbing treatment: Once a unit receives treatment, it must remain treated in all subsequent periods. Treatment cannot be reversed or temporarily suspended. This is validated at runtime.

  2. Parallel trends: In the absence of treatment, treated and control units would have followed parallel outcome trajectories.

  3. No anticipation: Units do not change their behavior in anticipation of future treatment.

  4. Untreated support at each calendar period: The time fixed effect \(\gamma_t\) for calendar period \(t\) is identified only if at least one unit is untreated in that period. Without never-treated units, post-treatment effects for the last-treated cohort (and any calendar periods where every unit is already treated) are not identified. CausalPy warns when this condition fails and marks the affected ATT(g, t) and ATT(e) cells as non-identified in the output tables.

Panel Balance: This implementation supports both balanced and unbalanced panel data. While balanced panels (where each unit is observed in every time period) are common in staggered DiD applications, the imputation-based approach of Borusyak et al. (2024) can accommodate unbalanced panels. The key requirement is that treatment timing is well-defined for each unit, not that all units are observed in all periods. Unit and observation counts in the summary output are computed without assuming balanced panels.

References

Borusyak, K., Jaravel, X., & Spiess, J. (2024). Revisiting Event Study Designs: Robust and Efficient Estimation. Review of Economic Studies.

Examples

>>> import causalpy as cp
>>> from causalpy.data.simulate_data import generate_staggered_did_data
>>> df = generate_staggered_did_data(n_units=30, n_time_periods=15, seed=42)
>>> result = cp.StaggeredDifferenceInDifferences(
...     df,
...     formula="y ~ 1 + C(unit) + C(time)",
...     unit_variable_name="unit",
...     time_variable_name="time",
...     treated_variable_name="treated",
...     treatment_time_variable_name="treatment_time",
...     model=cp.pymc_models.LinearRegression(
...         sample_kwargs={
...             "tune": 100,
...             "draws": 200,
...             "chains": 2,
...             "progressbar": False,
...         }
...     ),
... )
__init__(data, formula, unit_variable_name, time_variable_name, treated_variable_name='treated', treatment_time_variable_name=None, never_treated_value=inf, model=None, event_window=None, reference_event_time=-1)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the experiment algorithm: fit model, predict counterfactuals, and aggregate effects.

Return type:

None

effect_summary(*, direction='increase', alpha=0.05, min_effect=None)[source]#

Generate a decision-ready summary of causal effects for Staggered Difference-in-Differences.

Parameters:
  • direction (Literal['increase', 'decrease', 'two-sided']) – Direction for tail probability calculation (PyMC only, ignored for OLS).

  • alpha (float) – Significance level for HDI/CI intervals (1-alpha confidence level).

  • min_effect (float | None) – Region of Practical Equivalence (ROPE) threshold (PyMC only, ignored for OLS).

Returns:

Object with .table (DataFrame) and .text (str) attributes

Return type:

EffectSummary

get_plot_data(*, hdi_prob=0.94)[source]#

Get event-time plotting data.

Parameters:

hdi_prob (float) – Probability for HDI interval. Only used by models carrying posterior draws; when it differs from the value cached at fit time, the intervals are recomputed. Defaults to HDI_PROB (currently 0.94).

Returns:

DataFrame with event_time and att columns plus att_lower / att_upper HDI bounds (posterior draws) or att_std / n_obs dispersion columns (point estimates). Includes both pre-treatment (placebo) and post-treatment effects.

Return type:

pd.DataFrame

input_validation()[source]#

Validate the input data and parameters.

Return type:

None

plot(*, hdi_prob=None, figsize=(10, 6), show=True, legend_kwargs=None)[source]#

Plot the staggered difference-in-differences event study.

Parameters:
  • hdi_prob (float | None) – Probability mass of the highest density interval shown by the error bars. Unlike most other CausalPy experiments, hdi_prob for staggered DiD is fixed at fit time during effect aggregation and the resulting bounds are cached on the instance. If supplied here, the value must match the cached hdi_prob_; otherwise a ValueError is raised. Pass None (the default) to plot using the cached value. Ignored for OLS models.

  • figsize (tuple[float, float]) – Width and height of the figure in inches, passed to matplotlib.pyplot.subplots(). Defaults to (10, 6).

  • show (bool) – Whether to automatically display the plot. Defaults to True.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments to adjust legend placement and styling. Supported keys: loc, bbox_to_anchor, fontsize, frameon, title (bbox_transform is accepted alongside bbox_to_anchor). The existing legend is modified in place so that custom handles are preserved.

Returns:

  • fig (matplotlib.figure.Figure) – The figure that was created.

  • ax (list[matplotlib.axes.Axes]) – A single-element list containing the event-study axes.

Return type:

tuple[Figure, list[Axes]]

plot_group_time(*, hdi_prob=None, layout='facet', x_axis='event_time', include_placebo=True, figsize=None, show=True, legend_kwargs=None)[source]#

Plot cohort-specific ATT(g, t) trajectories.

Parameters:
  • hdi_prob (float | None) – Probability mass of the highest density interval shown by the uncertainty bands. As with plot(), Bayesian ATT(g, t) bounds are cached during effect aggregation. If supplied here, the value must match the cached hdi_prob_; otherwise a ValueError is raised. Pass None (the default) to plot using the cached value. Ignored for OLS models.

  • layout (Literal['facet', 'overlay']) – Plot layout. "facet" draws one row per cohort and "overlay" draws all cohorts on a single axes. Defaults to "facet".

  • x_axis (Literal['event_time', 'calendar_time']) – Time scale for the cohort trajectories. "event_time" plots each cohort against periods since treatment, giving an ATT(g, e) view derived from ATT(g, t). "calendar_time" plots each cohort against calendar time t. Defaults to "event_time".

  • include_placebo (bool) – Whether to include pre-treatment residual estimates for eventually-treated cohorts as placebo diagnostics. Defaults to True.

  • figsize (tuple[float, float] | None) – Width and height of the figure in inches, passed to matplotlib.pyplot.subplots(). Defaults to a height scaled by the number of cohorts when layout="facet" and (10, 6) when layout="overlay".

  • show (bool) – Whether to automatically display the plot. Defaults to True.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments to adjust legend placement and styling. Supported keys: loc, bbox_to_anchor, fontsize, frameon, title (bbox_transform is accepted alongside bbox_to_anchor). The existing legend is modified in place so that custom handles are preserved.

Returns:

  • fig (matplotlib.figure.Figure) – The figure that was created.

  • ax (list[matplotlib.axes.Axes]) – Axes containing the cohort trajectories. The list has one axes per cohort when layout="facet" and one axes when layout="overlay".

Return type:

tuple[Figure, list[Axes]]

summary(round_to=2, include_group_time=False)[source]#

Print summary of main results.

Parameters:
  • round_to (int | None) – Number of decimals for rounding. Defaults to 2.

  • include_group_time (bool) – Whether to print the disaggregated cohort-by-calendar-time ATT(g, t) table after the event-time estimates. Defaults to False.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = True#
class causalpy.Step[source]#

Protocol that all pipeline steps must satisfy.

Implementations must provide two methods:

  • validate – called before any step runs. Should raise on configuration errors (wrong types, missing parameters, etc.).

  • run – called sequentially. Receives the shared PipelineContext, mutates it, and returns it.

__init__(*args, **kwargs)#
run(context)[source]#

Execute the step, mutating and returning the context.

Parameters:

context (PipelineContext) – Shared pipeline context, which the step is allowed to mutate.

Returns:

The same context, returned for chaining convenience.

Return type:

PipelineContext

validate(context)[source]#

Check configuration before execution.

Parameters:

context (PipelineContext) – Shared pipeline context.

Return type:

None

class causalpy.SyntheticControl[source]#

The class for the synthetic control experiment.

Parameters:
  • data (NativeDataFrame) – Any eager dataframe Narwhals supports. For a pandas dataframe the index carries the time axis. Dataframes from other libraries have no index, so those callers must pass time_column.

  • treatment_time (int | float | Timestamp) – The time when treatment occurred, in reference to the data index.

  • control_units (list[str]) – A list of control units to be used in the experiment.

  • treated_units (list[str]) – A list of treated units to be used in the experiment.

  • model (PyMCModel | RegressorMixin | None) – A PyMC or sklearn model. Defaults to WeightedSumFitter.

  • min_donor_correlation (float) – Minimum acceptable Pearson correlation between each control unit and treated unit in the pre-treatment period. Control units below this threshold trigger a UserWarning. Defaults to 0.0 (warn on negatively correlated donors).

  • auto_scale_sigma (bool) – If True (default) and the model still carries the weighted-sum fitters’ stock y_hat prior, that sigma ~ HalfNormal(1) default is replaced by sigma ~ Exponential(2/s). The scale is computed per treated unit, with s the standard deviation of that unit’s pre-treatment data, so units on different scales are each calibrated separately. Set to False to keep the original HalfNormal(1) default; the experiment then fits a copy of the model with that prior pinned explicitly, leaving the instance you passed in untouched. A model constructed with an explicit y_hat prior is never rescaled either way.

  • time_column (str | None) – Column holding the time axis. It becomes the index of the data. Required for non-pandas inputs, which carry no index. If None (default), the pandas index of data is used. Passing it for data that already has a meaningful index raises, since only one of the two can be the time axis.

Notes

Estimate extraction

The model learns control-unit weights from pre-intervention outcomes and applies them to post-intervention controls to construct a synthetic untreated trajectory. Pointwise impact is the observed treated outcome minus this synthetic counterfactual, and cumulative impact is its running sum. Bayesian backends subtract the posterior conditional expectation mu rather than noisy posterior-predictive draws y_hat; OLS subtracts its weighted point prediction.

Examples

>>> import causalpy as cp
>>> df = cp.load_data("sc")
>>> treatment_time = 70
>>> seed = 42
>>> result = cp.SyntheticControl(
...     df,
...     treatment_time,
...     control_units=["a", "b", "c", "d", "e", "f", "g"],
...     treated_units=["actual"],
...     model=cp.pymc_models.WeightedSumFitter(
...         sample_kwargs={
...             "target_accept": 0.95,
...             "random_seed": seed,
...             "progressbar": False,
...         }
...     ),
... )
__init__(data, treatment_time, control_units, treated_units, model=None, min_donor_correlation=0.0, auto_scale_sigma=True, time_column=None)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the experiment algorithm: fit model, predict, and calculate causal impact.

Return type:

None

property datapost: DataFrame#

Data from on or after the treatment time (inclusive).

Post-period: index >= treatment_time

property datapre: DataFrame#

Data from before the treatment time (exclusive).

Pre-period: index < treatment_time

effect_summary(*, window='post', direction='increase', alpha=0.05, cumulative=True, relative=True, min_effect=None, treated_unit=None, period=None, prefix='Post-period')[source]#

Generate a decision-ready summary of causal effects for Synthetic Control.

Parameters:
  • window (Union[Literal['post'], tuple, slice]) –

    Time window for analysis:

    • ”post”: All post-treatment time points (default)

    • (start, end): Tuple of start and end times (handles both datetime and integer indices)

    • slice: Python slice object for integer indices

  • direction (Literal['increase', 'decrease', 'two-sided']) – Direction for tail probability calculation (PyMC only, ignored for OLS).

  • alpha (float) – Significance level for HDI/CI intervals (1-alpha confidence level).

  • cumulative (bool) – Whether to include cumulative effect statistics.

  • relative (bool) – Whether to include relative effect statistics (% change vs counterfactual).

  • min_effect (float | None) – Region of Practical Equivalence (ROPE) threshold (PyMC only, ignored for OLS).

  • treated_unit (str | None) – For multi-unit experiments, specify which treated unit to analyze. If None and multiple units exist, uses first unit.

  • period (Optional[Literal['intervention', 'post', 'comparison']]) – Ignored for Synthetic Control (two-period design only).

  • prefix (str) – Prefix for prose generation. Defaults to “Post-period”.

Returns:

Object with .table (DataFrame) and .text (str) attributes. The .text attribute contains a detailed multi-paragraph narrative report.

Return type:

EffectSummary

get_plot_data(*, hdi_prob=0.94, treated_unit=None)[source]#

Recover the data of the experiment along with the prediction and causal impact information.

HDI columns are included only when the prediction container carries posterior draws (point-estimate backends return just prediction and impact).

Parameters:
  • hdi_prob (float) – Probability mass of the highest density interval. Defaults to the project-wide HDI_PROB. Ignored when the prediction container has no posterior draws.

  • treated_unit (str | None) – Which treated unit to extract data for. Must be a string name of the treated unit. If None, uses the first treated unit.

Return type:

DataFrame

input_validation(data, treatment_time)[source]#

Validate the input data and model formula for correctness.

Parameters:
  • data (DataFrame) – The experiment data.

  • treatment_time (int | float | Timestamp) – The treatment time, expected to be compatible with data.index.

Return type:

None

plot(*, round_to=None, treated_unit=None, ci_prob=0.94, kind='ribbon', ci_kind='hdi', num_samples=50, plot_predictors=False, figsize=(7, 8), show=True, legend_kwargs=None)[source]#

Plot the synthetic control results for a specific treated unit.

Parameters:
  • round_to (int | None) – Number of decimals used to round numerical results in the figure title (e.g. the Bayesian \(R^2\)). Defaults to None, in which case 2 significant figures are used.

  • treated_unit (str | None) – Which treated unit to plot. Must be one of the names supplied via treated_units at construction time. Defaults to None, which selects the first treated unit.

  • ci_prob (float) – Probability mass of the highest density interval drawn around the posterior predictive, causal impact, and cumulative impact bands. Must be in (0, 1]. Ignored for OLS models. Defaults to HDI_PROB (currently 0.94).

  • kind (Literal['ribbon', 'histogram', 'spaghetti']) – How posterior uncertainty is rendered via plot_posterior_over_x(). Defaults to "ribbon". For "spaghetti", legends use draw lines rather than a shaded band. For "histogram", uncertainty is shown as a 2D density heatmap with a mean line overlay (no ribbon patch for legends).

  • ci_kind (Literal['hdi', 'eti']) – Credible interval type when kind="ribbon". Defaults to "hdi".

  • num_samples (int) – Number of posterior draws when kind="spaghetti". Defaults to 50. Ignored for other kinds.

  • plot_predictors (bool) – Whether to overlay the donor (control) unit trajectories on the top panel. Defaults to False.

  • figsize (tuple[float, float]) – Width and height of the figure in inches, passed to matplotlib.pyplot.subplots(). Defaults to (7, 8).

  • show (bool) – Whether to automatically display the plot. Defaults to True. Set to False if you want to modify the figure before displaying it.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments to adjust legend placement and styling. Supported keys: loc, bbox_to_anchor, fontsize, frameon, title (bbox_transform is accepted alongside bbox_to_anchor). The existing legend is modified in place so that custom handles are preserved.

Returns:

  • fig (matplotlib.figure.Figure) – The figure that was created.

  • ax (list[matplotlib.axes.Axes]) – The three axes (top: predictions, middle: causal impact, bottom: cumulative impact).

Return type:

tuple[Figure, list[Axes]]

summary(round_to=None)[source]#

Print summary of main results and model coefficients.

Parameters:

round_to (int | None) – Number of decimals used to round results. Defaults to 2. Use None to return raw numbers.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = True#
class causalpy.SyntheticDifferenceInDifferences[source]#

Bayesian Synthetic Difference-in-Differences experiment.

Combines the synthetic control method’s unit weighting with difference-in-differences time weighting. The treatment effect (tau) is computed analytically from the posterior weight distributions via the double-difference formula, rather than being estimated inside the MCMC model (cut-posterior formulation).

Parameters:
  • data (NativeDataFrame) – Any eager dataframe Narwhals supports, in wide format (columns = units, rows = time periods). For a pandas dataframe the index carries the time axis. Dataframes from other libraries have no index, so those callers must pass time_column.

  • treatment_time (int | float | Timestamp) – The time when treatment occurred, should be in reference to the data index.

  • control_units (list[str]) – A list of control unit column names.

  • treated_units (list[str]) – A list of treated unit column names.

  • model (PyMCModel | RegressorMixin | None) – A SyntheticDifferenceInDifferencesWeightFitter instance. Defaults to SyntheticDifferenceInDifferencesWeightFitter.

  • time_column (str | None) – Column holding the time axis. It becomes the index of the data. Required for non-pandas inputs, which carry no index. If None (default), the pandas index of data is used. Passing it for data that already has a meaningful index raises, since only one of the two can be the time axis.

Notes

Estimate extraction

The Bayesian weight model produces posterior draws of synthetic-control unit weights and pre-period time weights. For each draw, the class constructs treated-minus-synthetic gaps and evaluates the weighted double-difference analytically to obtain the scalar tau_posterior ATT; the effect is not read from a regression coefficient or obtained by population-standardized g-computation. The time-indexed post_impact consumed by effect_summary() is the post-period treated-minus-synthetic trajectory rather than this time-weighted scalar.

This implements Bayesian SDiD method. The model fits two weight modules via MCMC:

  • Unit weights (omega): balance control units against treated units in the pre-treatment period, similar to synthetic control.

  • Time weights (lambda): balance pre-treatment periods against post-treatment periods for control units.

The treatment effect is then computed analytically via the double-difference:

\[\tau = \bar{\Delta}_{\text{post}} - \boldsymbol{\lambda}^\top \boldsymbol{\Delta}_{\text{pre}}\]

where \(\Delta_t = y_{\text{tr},t} - (\omega_0 + \boldsymbol{\omega}^\top \mathbf{Y}_{\text{co},t})\) is the gap between the observed treated outcome and the synthetic control at time t.

References

Examples

>>> import causalpy as cp
>>> df = cp.load_data("sc")
>>> treatment_time = 70
>>> result = cp.SyntheticDifferenceInDifferences(
...     df,
...     treatment_time,
...     control_units=["a", "b", "c", "d", "e", "f", "g"],
...     treated_units=["actual"],
...     model=cp.pymc_models.SyntheticDifferenceInDifferencesWeightFitter(
...         sample_kwargs={
...             "tune": 20,
...             "draws": 20,
...             "chains": 2,
...             "cores": 2,
...             "progressbar": False,
...         }
...     ),
... )
__init__(data, treatment_time, control_units, treated_units, model=None, time_column=None)[source]#
Parameters:
Return type:

None

algorithm()[source]#

Run the SDiD algorithm: fit weight modules, compute tau analytically.

The method is a thin orchestrator that delegates each step to a private helper so that the individual pieces can be unit tested in isolation:

  1. _build_weight_fitter_inputs() prepares the dict-based X, y and coords inputs for the weight fitter.

  2. PyMCModel.fit() fits both the omega and lambda modules via MCMC.

  3. _extract_weight_posteriors() pulls the posterior weight arrays out of the fitted model.

  4. _compute_synthetic_and_gaps() builds the synthetic control trajectory and the gap between treated and synthetic.

  5. _compute_tau() evaluates the double-difference ATT.

  6. _build_reporting_objects() constructs the xarray objects required by the reporting helpers.

Return type:

None

property datapost: DataFrame#

Data from on or after the treatment time (inclusive).

Post-period: index >= treatment_time

property datapre: DataFrame#

Data from before the treatment time (exclusive).

Pre-period: index < treatment_time

effect_summary(*, window='post', direction='increase', alpha=0.05, cumulative=True, relative=True, min_effect=None, treated_unit=None, period=None, prefix='Post-period')[source]#

Generate a decision-ready summary of causal effects for SDiD.

Parameters:
  • window (Union[Literal['post'], tuple, slice]) – Time window for analysis.

  • direction (Literal['increase', 'decrease', 'two-sided']) – Direction for tail probability calculation.

  • alpha (float) – Significance level for HDI intervals.

  • cumulative (bool) – Whether to include cumulative effect statistics.

  • relative (bool) – Whether to include relative effect statistics.

  • min_effect (float | None) – ROPE threshold.

  • treated_unit (str | None) – Which treated unit to analyze. If None, uses first unit.

  • period (Optional[Literal['intervention', 'post', 'comparison']]) – Ignored for SDiD (two-period design only).

  • prefix (str) – Prefix for prose generation. Defaults to “Post-period”.

Returns:

Object with .table (DataFrame) and .text (str) attributes.

Return type:

EffectSummary

input_validation(data, treatment_time)[source]#

Validate the input data for correctness.

Parameters:
  • data (DataFrame) – A dataframe in wide format (columns = units, rows = time periods).

  • treatment_time (int | float | Timestamp) – The time when treatment occurred, should be in reference to the data index.

Return type:

None

plot(*, round_to=None, ci_prob=0.94, kind='ribbon', ci_kind='hdi', num_samples=50, show=True, legend_kwargs=None)[source]#

Plot SDiD results: counterfactual, period impact, and cumulative impact.

Parameters:
  • round_to (int | None) – Number of decimals used to round the ATT in the title. Defaults to 2. Use None for raw values.

  • ci_prob (float) – Probability mass of the highest density interval drawn around the posterior predictive, causal impact, and cumulative impact bands. Must be in (0, 1]. Defaults to HDI_PROB (currently 0.94).

  • kind (Literal['ribbon', 'histogram', 'spaghetti']) – How posterior uncertainty is rendered via plot_posterior_over_x(). Defaults to "ribbon". For "spaghetti", legends use draw lines rather than a shaded band. For "histogram", uncertainty is shown as a 2D density heatmap with a mean line overlay (no ribbon patch for legends).

  • ci_kind (Literal['hdi', 'eti']) – Credible interval type when kind="ribbon". Defaults to "hdi".

  • num_samples (int) – Number of posterior draws when kind="spaghetti". Defaults to 50. Ignored for other kinds.

  • show (bool) – Whether to call matplotlib.pyplot.show() after drawing. Defaults to True.

  • legend_kwargs (dict[str, Any] | None) – Keyword arguments applied to the top-axis legend in place after the figure is built. Supported keys include loc, bbox_to_anchor, fontsize, frameon, title, and optionally bbox_transform alongside bbox_to_anchor. See _render_plot().

Returns:

  • fig (matplotlib.figure.Figure) – The figure containing the three stacked panels.

  • ax (numpy.ndarray) – Array of the three matplotlib.axes.Axes instances.

Return type:

tuple[Figure, ndarray]

summary(round_to=None)[source]#

Print summary of main results.

Parameters:

round_to (int | None) – Number of decimals used to round results. Defaults to 2. Use None to return raw numbers.

Return type:

None

supports_bayes: bool = True#
supports_ols: bool = True#
causalpy.create_causalpy_compatible_class(estimator)[source]#

This function takes a scikit-learn estimator and returns a new class that is compatible with CausalPy.

Parameters:

estimator (type[RegressorMixin]) – A scikit-learn estimator class to augment.

Return type:

type[RegressorMixin]

causalpy.extract_lift_for_mmm(sc_result, channel, x, delta_x, aggregate='mean')[source]#

Extract lift test results from a Synthetic Control analysis for MMM calibration.

This function extracts lift estimates from a fitted SyntheticControl model in a format compatible with PyMC-Marketing’s add_lift_test_measurements() method. This enables using geo-level lift test results to calibrate Media Mix Models.

Parameters:
  • sc_result (SyntheticControl) – A fitted SyntheticControl model with one or more treated units. The model must have been fit with a Bayesian (PyMC) model to provide posterior distributions for uncertainty quantification.

  • channel (str) – Name of the marketing channel being tested (e.g., “tv”, “radio”, “digital”). This should match the channel names used in your MMM.

  • x (float) – Baseline spend level for the channel before the test period. For channels with zero pre-test spend, use 0.0.

  • delta_x (float) – The change in spend during the test period (i.e., test spend minus baseline spend). For a new channel activation, this equals the total test spend.

  • aggregate (str) –

    How to aggregate the causal impact across post-intervention time periods:

    • ”mean”: Average lift per time period. Use this for rate-based outcomes (e.g., weekly sales rate) or when your MMM operates at the same time granularity as the experiment.

    • ”median”: Median lift per time period. More robust to outliers than the mean; useful when the impact distribution across time periods is skewed.

    • ”sum”: Total cumulative lift across all post-intervention periods. Use this for cumulative outcomes or when you want total campaign impact.

Returns:

DataFrame with one row per treated geo, containing columns:

  • channel: The marketing channel name (from input parameter)

  • geo: The treated geo identifier (from sc_result.treated_units)

  • x: Pre-test spend level (from input parameter)

  • delta_x: Spend change during test (from input parameter)

  • delta_y: Mean lift estimate from the posterior distribution

  • sigma: Standard deviation of the lift estimate from the posterior

Return type:

pd.DataFrame

Raises:

ValueError – If the model is not a Bayesian (PyMC) model, as uncertainty quantification requires posterior samples.

Notes

This function is designed for integration with PyMC-Marketing’s MMM calibration workflow. The output DataFrame can be passed directly to MMM.add_lift_test_measurements() to inform the model’s saturation curves with experimental evidence.

For more information on lift test calibration in MMMs, see the PyMC-Marketing documentation: pymc-labs/pymc-marketing. Reference workflow: https://www.pymc-marketing.io/en/stable/notebooks/mmm/mmm_lift_test.html

Examples

import causalpy as cp

# Fit a multi-geo synthetic control model
result = cp.SyntheticControl(
    df,
    treatment_time,
    control_units=["geo_a", "geo_b", "geo_c"],
    treated_units=["geo_x", "geo_y"],
    model=cp.pymc_models.WeightedSumFitter(
        sample_kwargs={"progressbar": False}
    ),
)

# Extract lift results for MMM calibration
df_lift = cp.extract_lift_for_mmm(
    result,
    channel="tv_campaign",
    x=0.0,  # No pre-test TV spend
    delta_x=50000,  # $50k test spend
    aggregate="mean",
)

# The resulting DataFrame can be used with PyMC-Marketing:
# mmm.add_lift_test_measurements(df_lift)
causalpy.load_data(dataset)[source]#

Load example datasets for causal inference analysis.

This function loads pre-packaged datasets that are used in CausalPy’s documentation and examples. These datasets demonstrate various causal inference methods including difference-in-differences, regression discontinuity, synthetic control, interrupted time series, and more.

Parameters:

dataset (str) –

Name of the dataset to load. Available datasets are:

  • "banks" - Historic banking closures data for difference-in-differences

  • "brexit" - UK GDP data for estimating causal impact of Brexit

  • "covid" - Deaths and temperature data for England and Wales

  • "did" - Difference-in-differences example dataset

  • "drinking" - Minimum legal drinking age data for regression discontinuity

  • "its" - Interrupted time series example dataset

  • "its simple" - Simplified interrupted time series dataset

  • "rd" - Regression discontinuity example dataset

  • "sc" - Synthetic control example dataset

  • "anova1" - ANCOVA example with pre/post treatment nonequivalent groups

  • "geolift1" - Single treatment geo-lift dataset for synthetic control

  • "geolift_multi_cell" - Multi-cell geo-lift dataset for synthetic control

  • "risk" - Acemoglu, Johnson & Robinson (2001) data for instrumental variables

  • "nhefs" - National Health and Nutrition Examination Survey data

  • "schoolReturns" - Schooling returns data for instrumental variable analysis

  • "pisa18" - PISA 2018 sample data

  • "nets" - National Supported Work Demonstration dataset

  • "lalonde" - LaLonde dataset for propensity score analysis

  • "zipcodes" - Geo-experimentation zipcode data for comparative interrupted time series analysis. Based on synthetic data from Juan Orduz’s blog post on time-based regression for geo-experiments.

  • "nevo" - Berry, Levinsohn, and Pakes (1995) cereal data for BLP estimation

  • "california_prop99" - California Proposition 99 cigarette sales data (Abadie, Diamond & Hainmueller, 2010). Wide-format panel of annual per-capita cigarette sales (packs) for 39 US states, 1970–2000. Treatment: California in 1989.

Returns:

The requested dataset as a pandas DataFrame.

Return type:

pd.DataFrame

Raises:

ValueError – If the requested dataset name is not found in the available datasets.

Examples

Load the difference-in-differences example dataset:

>>> import causalpy as cp
>>> df = cp.load_data("did")

Load the regression discontinuity dataset:

>>> df = cp.load_data("rd")
causalpy.plot_correlations(data, columns=None, method='pearson', figsize=None, ax=None, **kwargs)[source]#

Plot a pairwise correlation heatmap for panel data columns.

Computes the pairwise correlation matrix between the specified columns (typically geographic units or time series) and displays it as a lower-triangle heatmap. This is a pre-experiment diagnostic for synthetic control analyses: markets that are highly correlated in the pre-treatment period are more likely to produce reliable counterfactuals.

Parameters:
  • data (DataFrame) – Wide-format panel data with time as the index and locations/units as columns.

  • columns (list[str] | None) – Subset of columns to include. If None, all numeric columns are used.

  • method (Literal['pearson', 'kendall', 'spearman']) – Correlation method passed to pandas.DataFrame.corr().

  • figsize (tuple[float, float] | None) – Width and height in inches for the figure. Only used when ax is not provided. If None, matplotlib’s default is used.

  • ax (Axes | None) – Axes on which to draw the heatmap. If None, a new figure and axes are created (sized according to figsize).

  • **kwargs (Any) – Keyword arguments forwarded to seaborn.heatmap(): vmin, vmax, cmap, center, robust, annot, fmt, annot_kws, linewidths, linecolor, cbar, cbar_kws, cbar_ax, square, xticklabels, yticklabels, mask, and the matplotlib.axes.Axes.pcolormesh() keywords supported by the installed seaborn version. data and ax are supplied by CausalPy. This narrow third-party forwarder lets callers override CausalPy’s heatmap defaults without duplicating seaborn’s evolving forwarding surface; unknown keys are rejected by seaborn or matplotlib rather than ignored.

Returns:

The correlation matrix and the axes containing the heatmap.

Return type:

tuple[pd.DataFrame, matplotlib.axes.Axes]

Examples

import causalpy as cp

df = cp.load_data("geolift1")
corr, ax = cp.plot_correlations(df)

# Larger figure with smaller annotation text
corr, ax = cp.plot_correlations(df, figsize=(10, 8), annot_kws={"size": 7})
causalpy.ramp()#

Stateful transform for ramp function (slope change) at threshold.

Creates a ramp: max(0, time - threshold). For datetime, the ramp values are in days.

Works with both numeric and datetime time columns. For datetime, the threshold can be specified as a string (‘2020-01-01’) or pd.Timestamp.

Notes

Per the patsy stateful transform protocol, x and threshold are supplied to memorize_chunk() and transform() rather than to the constructor; see those methods for parameter details.

For datetime inputs, the ramp values represent days since the threshold. This means the slope coefficient will be interpreted as “change per day”.

Examples

>>> # Numeric time - ramp is in same units as t
>>> formula = "y ~ 1 + t + ramp(t, 50)"
>>> # Datetime time - ramp is in DAYS
>>> formula = "y ~ 1 + date + ramp(date, '2020-06-01')"
Return type:

None

causalpy.step()#

Stateful transform for step function (level change) at threshold.

Creates a binary indicator: 1 if time >= threshold, 0 otherwise.

Works with both numeric and datetime time columns. For datetime, the threshold can be specified as a string (‘2020-01-01’) or pd.Timestamp.

The transform is “stateful” because it remembers the datetime origin from the training data, ensuring consistent behavior when predicting on new data.

Notes

Per the patsy stateful transform protocol, x and threshold are supplied to memorize_chunk() and transform() rather than to the constructor; see those methods for parameter details.

Examples

>>> # Numeric time
>>> formula = "y ~ 1 + t + step(t, 50)"
>>> # Datetime time with string threshold
>>> formula = "y ~ 1 + date + step(date, '2020-06-01')"
>>> # Datetime time with Timestamp threshold
>>> formula = "y ~ 1 + date + step(date, pd.Timestamp('2020-06-01'))"
Return type:

None

Modules#

constants

Shared constants for the CausalPy package.

data

Code for loading datasets.

input_data

Dataframe-agnostic input handling.

pymc_models

Custom PyMC models for causal inference.

skl_models

Custom scikit-learn models for causal inference.

pymc_forecast_models

Adapter that lets a pymc_forecast forecasting model act as a model provider behind CausalPy's experiment API.

experiments

CausalPy experiment module.

pipeline

Pipeline orchestration for composable causal inference workflows.

reporting

Reporting utilities for causal inference experiments.

steps

Pipeline steps for causal inference workflows.

checks

Sensitivity and diagnostic checks for causal inference experiments.