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:
data (DataFrameLike)
formula (str)
time_variable_name (str)
group_variable_name (str)
post_treatment_variable_name (str)
model (PyMCModel | RegressorMixin | None)
- Return type:
None
- algorithm()[source]#
Run the experiment algorithm: fit model, predict, and calculate causal impact.
- Return type:
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:
- 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 toNone, 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 toHDI_PROB(currently 0.94).kind (
Literal['ribbon','histogram','spaghetti']) – How posterior uncertainty is rendered viaplot_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 whenkind="ribbon". Defaults to"hdi".num_samples (
int) – Number of posterior draws whenkind="spaghetti". Defaults to 50. Ignored for other kinds.figsize (
tuple[float,float] |None) – Width and height of the figure in inches, passed tomatplotlib.pyplot.subplots(). Defaults toNone(use matplotlib’s default).show (
bool) – Whether to automatically display the plot. Defaults toTrue. Set toFalseif 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_transformis accepted alongsidebbox_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:
- 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:
- 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 bymethod’s constructor, exceptdata, which the pipeline supplies. This is a deliberately narrow dynamic forwarder:methodmay be an integrator-providedBaseExperimentsubclass, so its accepted constructor keys cannot be enumerated here. Built-in experiment constructors declare every supported key explicitly; unsupported, misspelled, or incomplete arguments raiseTypeErrorduring 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:
method (type[BaseExperiment])
kwargs (Any)
- Return type:
None
- run(context)[source]#
Instantiate and fit the experiment.
The experiment constructor receives
context.dataas its first positional argument, followed by all captured keyword arguments.- Parameters:
context (
PipelineContext) – Pipeline context.context.datais forwarded to the experiment constructor as the first positional argument.- Returns:
Updated context with
experiment,experiment_config, and (if available)effect_summarypopulated.- Return type:
- 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
BaseExperimentor supplied constructor arguments are incompatible with an inspectable constructor, including omitted required arguments.ValueError – If
datais 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]#
- run(context)[source]#
Generate the HTML report and store it in the context.
- Parameters:
context (
PipelineContext) – Pipeline context providingexperiment,effect_summary, andsensitivity_results(any of which may beNone).- Returns:
The same context with
reportpopulated.- Return type:
- 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:
instruments_data (DataFrameLike)
data (DataFrameLike)
instruments_formula (str)
formula (str)
model (InstrumentalVariableRegression | None)
priors (dict | None)
- 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:
- 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:
- 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).
- 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_timeis 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 passtime_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 exactlytreatment_timeare 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, aPyMCForecastModelwrapping apymc_forecastforecasting model can serve as the counterfactual backend (requires the optionalpymc-forecastdependency); seecausalpy.pymc_forecast_modelsfor when to prefer it.treatment_end_time (
int|float|Timestamp|None) – The time when treatment ended, enabling three-period analysis. Must be greater thantreatment_timeand within the data range. If None (default), the analysis assumes a permanent intervention (two-period design). INCLUSIVE: Observations at exactlytreatment_end_timeare 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 ofdatais 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
murather than noisy posterior-predictive drawsy_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, andeffect_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:
data (DataFrameLike)
formula (str)
model (PyMCModel | RegressorMixin | PyMCForecastModel | None)
time_column (str | None)
- 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:
- 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:
- 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:
- 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
predictionandimpact).
- input_validation(data, treatment_time, treatment_end_time=None)[source]#
Validate the input data and model formula for correctness.
- 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. UseNoneto 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 toHDI_PROB(currently 0.94).kind (
Literal['ribbon','histogram','spaghetti']) – How posterior uncertainty is rendered viaplot_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 whenkind="ribbon". Defaults to"hdi".num_samples (
int) – Number of posterior draws whenkind="spaghetti". Defaults to 50. Ignored for other kinds.figsize (
tuple[float,float]) – Width and height of the figure in inches, passed tomatplotlib.pyplot.subplots(). Defaults to(7, 8).show (
bool) – Whether to automatically display the plot. Defaults toTrue. Set toFalseif 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_transformis accepted alongsidebbox_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:
- 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:
data (DataFrameLike)
formula (str)
outcome_variable (str)
weighting_scheme (str)
model (PropensityScore | None)
- Return type:
None
- algorithm()[source]#
Run the experiment algorithm by fitting the propensity score model.
Delegates to
self.model.fitwith the covariate matrixself.X, treatment vectorself.t, and coordinate metadataself.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:
- 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:
- 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:
- input_validation()[source]#
Validate the input data and model formula for correctness.
Checks that the
outcome_variableexists inself.dataand 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 twoNonevalues 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 allpd.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 - psfor 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_ntrtandn_trtare both equal to the total number of observations (the raw scheme normalises by the full sample size).- Return type:
- 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, andn_ntrt/n_trtare the corresponding group sizes used for normalisation.- Return type:
- plot(*, show=True, legend_kwargs=None)[source]#
Plot the results.
- Parameters:
- Raises:
NotImplementedError – Always; call
plot_ate()orplot_balance_ecdf()instead.- Return type:
None
Notes
Inverse propensity weighting does not expose a unified
plot()view; instead, use the dedicated diagnosticsplot_ate()(treatment-effect distribution) andplot_balance_ecdf()(covariate-balance ECDF). This stub exists so every experiment subclass offers an explicit, kwarg-onlyplot()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_drawsposterior 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. IfNone, uses the fitted model backend’s DataTree.method (
str|None) – Weighting scheme to apply. One of'robust','raw','overlap', or'doubly_robust'. IfNone, falls back toself.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:
- 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. IfNone, uses the fitted model backend’s DataTree.weighting_scheme (
str|None) – Weighting scheme to apply. One of'raw','robust', or'overlap'. IfNone, falls back toself.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:
- weighted_percentile(data, weights, perc)[source]#
Compute a weighted percentile of the data.
Sorts
dataandweightstogether, builds a weighted empirical CDF, and linearly interpolates to find the value at the requested percentile.- Parameters:
- Returns:
The interpolated data value at the given weighted percentile.
- Return type:
- Raises:
ValueError – If
percis 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, includeC(unit_var)(and optionallyC(time_var)) in the formula. For the demeaned transformation, do NOT include thoseC(...)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.
Notes
Estimate extraction
PanelRegressionfits 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:
data (DataFrameLike)
formula (str)
unit_fe_variable (str)
time_fe_variable (str | None)
fe_method (Literal['dummies', 'demeaned'])
model (PyMCModel | RegressorMixin | None)
- 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:
- get_plot_data()[source]#
Get plot data with fitted values.
Bayesian models additionally return
y_fitted_lower/y_fitted_upper95% credible-interval columns.- Returns:
DataFrame with fitted values (and credible intervals when the model carries posterior draws).
- Return type:
pd.DataFrame
- 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 toHDI_PROB(currently 0.94).show (
bool) – Whether to automatically display the plot. Defaults toTrue.legend_kwargs (
dict[str,Any] |None) – Keyword arguments to adjust legend placement and styling. Supported keys:loc,bbox_to_anchor,fontsize,frameon,title(bbox_transformis accepted alongsidebbox_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:
- 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"). IfNone, 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 toHDI_PROB(currently 0.94).
- Returns:
Figure and axes objects
- Return type:
tuple[plt.Figure, plt.Axes]
- Raises:
ValueError – If
var_namesis empty orhdi_probis outside(0, 1).
- 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 toHDI_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.
- 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
InterruptedTimeSerieswhich 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 changesramp(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 thestep()orramp()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 onestep()orramp()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.
- 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()andramp()design-matrix column to zero, and the pointwise effect is the fitted conditional expectation minus that counterfactual expectation. Bayesian backends contrast posteriormuvalues, 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:
data (DataFrameLike)
formula (str)
model (PyMCModel | RegressorMixin | None)
- 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 (seeBaseExperiment.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; passNone.prefix (
str) – Prefix for prose generation.
- Return type:
- 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.
- 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. UseNoneto 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 toHDI_PROB(currently 0.94).kind (
Literal['ribbon','histogram','spaghetti']) – How posterior uncertainty is rendered viaplot_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 whenkind="ribbon". Defaults to"hdi".num_samples (
int) – Number of posterior draws whenkind="spaghetti". Defaults to 50. Ignored for other kinds.figsize (
tuple[float,float]) – Width and height of the figure in inches, passed tomatplotlib.pyplot.subplots(). Defaults to(10, 10).show (
bool) – Whether to automatically display the plot. Defaults toTrue.legend_kwargs (
dict[str,Any] |None) – Keyword arguments to adjust legend placement and styling. Supported keys:loc,bbox_to_anchor,fontsize,frameon,title(bbox_transformis accepted alongsidebbox_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:
- 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:
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()
- 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
SensitivityAnalysiscan derive experiment factories.- Type:
dict or None
- effect_summary#
The effect summary from the primary experiment.
- Type:
EffectSummary or None
- __init__(data, experiment=None, experiment_config=None, effect_summary=None, sensitivity_results=<factory>, report=None)#
- Parameters:
data (DataFrame)
experiment (BaseExperiment | None)
effect_summary (EffectSummary | None)
report (Any)
- Return type:
None
- effect_summary: EffectSummary | None = None#
- experiment: BaseExperiment | None = None#
- 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
- __init__(experiment, effect_summary, sensitivity_results, report)#
- Parameters:
experiment (BaseExperiment | None)
effect_summary (EffectSummary | None)
report (Any)
- Return type:
None
- effect_summary: EffectSummary | None#
- experiment: BaseExperiment | None#
- classmethod from_context(context)[source]#
Build a
PipelineResultfrom a completedPipelineContext.- 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:
- 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 indatafor the group variable; should be either binary or boolean.pretreatment_variable_name (
str) – Name of the column indatafor the pretreatment variable.model (
PyMCModel|None) – A PyMC model. Defaults toLinearRegression.
Notes
Estimate extraction
The reported
causal_impactis 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:
data (DataFrameLike)
formula (str)
group_variable_name (str)
pretreatment_variable_name (str)
model (PyMCModel | None)
- Return type:
None
- algorithm()[source]#
Run the experiment algorithm: fit model, predict, and calculate causal impact.
- Return type:
None
- effect_summary(*, direction='increase', alpha=0.05, min_effect=None)[source]#
Generate a decision-ready summary of causal effects for PrePostNEGD.
- Parameters:
- Returns:
Object with .table (DataFrame) and .text (str) attributes
- Return type:
- 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 toNone, 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 toHDI_PROB(currently 0.94).kind (
Literal['ribbon','histogram','spaghetti']) – How posterior uncertainty is rendered viaplot_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 whenkind="ribbon". Defaults to"hdi".num_samples (
int) – Number of posterior draws whenkind="spaghetti". Defaults to 50. Ignored for other kinds.figsize (
tuple[float,float]) – Width and height of the figure in inches, passed tomatplotlib.pyplot.subplots(). Defaults to(7, 9).show (
bool) – Whether to automatically display the plot. Defaults toTrue.legend_kwargs (
dict[str,Any] |None) – Keyword arguments to adjust legend placement and styling. Supported keys:loc,bbox_to_anchor,fontsize,frameon,title(bbox_transformis accepted alongsidebbox_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:
- 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 toLinearRegression.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 thanbandwidthifbandwidthis finite.
Notes
Estimate extraction
After fitting the regression on the selected bandwidth, the class predicts the conditional expectation immediately below the threshold with
treated=0and immediately above it withtreated=1.discontinuity_at_thresholdis the upper prediction minus the lower prediction, evaluated atthreshold ± 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:
data (DataFrameLike)
formula (str)
treatment_threshold (float)
model (PyMCModel | RegressorMixin | None)
running_variable_name (str)
epsilon (float)
bandwidth (float)
donut_hole (float)
- 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:
- 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. UseNoneto 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 toHDI_PROB(currently 0.94).kind (
Literal['ribbon','histogram','spaghetti']) – How posterior uncertainty is rendered viaplot_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 whenkind="ribbon". Defaults to"hdi".num_samples (
int) – Number of posterior draws whenkind="spaghetti". Defaults to 50. Ignored for other kinds.figsize (
tuple[float,float] |None) – Width and height of the figure in inches, passed tomatplotlib.pyplot.subplots(). Defaults toNone(use matplotlib’s default).show (
bool) – Whether to automatically display the plot. Defaults toTrue.legend_kwargs (
dict[str,Any] |None) – Keyword arguments to adjust legend placement and styling. Supported keys:loc,bbox_to_anchor,fontsize,frameon,title(bbox_transformis accepted alongsidebbox_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:
- 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 toLinearRegression.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, andkink_point + epsilon. It forms finite-difference slopes on the left and right and stores their difference asgradient_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]#
- 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:
- 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. UseNoneto 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 toHDI_PROB(currently 0.94).kind (
Literal['ribbon','histogram','spaghetti']) – How posterior uncertainty is rendered viaplot_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 whenkind="ribbon". Defaults to"hdi".num_samples (
int) – Number of posterior draws whenkind="spaghetti". Defaults to 50. Ignored for other kinds.figsize (
tuple[float,float] |None) – Width and height of the figure in inches, passed tomatplotlib.pyplot.subplots(). Defaults toNone(use matplotlib’s default).show (
bool) – Whether to automatically display the plot. Defaults toTrue.legend_kwargs (
dict[str,Any] |None) – Keyword arguments to adjust legend placement and styling. Supported keys:loc,bbox_to_anchor,fontsize,frameon,title(bbox_transformis accepted alongsidebbox_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:
- class causalpy.SensitivityAnalysis[source]#
Pipeline step that runs a suite of sensitivity / diagnostic checks.
Examples
>>> import causalpy as cp >>> step = cp.SensitivityAnalysis( ... checks=[ ... cp.checks.PlaceboInTime(n_folds=4), ... cp.checks.PriorSensitivity(priors=[...]), ... ] ... )
- classmethod default_for(method)[source]#
Create a
SensitivityAnalysispre-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:
- run(context)[source]#
Run all checks against the fitted experiment.
- Parameters:
context (
PipelineContext) – Pipeline context containing the fitted experiment and anyexperiment_configrequired by the checks.- Returns:
The same context with
sensitivity_resultsandreportpopulated.- Return type:
- 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
checksdoes not satisfy theCheckprotocol.- Return type:
None
- class causalpy.SensitivitySummary[source]#
Aggregate result of all sensitivity checks.
- results#
Individual check results.
- Type:
- all_passed#
Trueif every check with a pass/fail criterion passed,Falseif any failed, orNoneif no check had a pass/fail criterion.- Type:
bool or None
- __init__(results=<factory>, all_passed=None, text='')#
- Parameters:
results (list[CheckResult])
all_passed (bool | None)
text (str)
- Return type:
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:
- results: list[CheckResult]#
- 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
identifiedcolumn; non-identified cells haveNaNestimates.- Type:
pd.DataFrame
- att_event_time_#
Event-time ATT estimates: ATT(e) for each event-time e = t - G. Includes an
identifiedcolumn; non-identified cells haveNaNestimates.- Type:
pd.DataFrame
- non_identified_cohorts_#
Treatment cohorts with at least one non-identified post-treatment ATT(g, t).
- Type:
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:
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.
Parallel trends: In the absence of treatment, treated and control units would have followed parallel outcome trajectories.
No anticipation: Units do not change their behavior in anticipation of future treatment.
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)andATT(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:
- 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 toHDI_PROB(currently 0.94).- Returns:
DataFrame with
event_timeandattcolumns plusatt_lower/att_upperHDI bounds (posterior draws) oratt_std/n_obsdispersion columns (point estimates). Includes both pre-treatment (placebo) and post-treatment effects.- Return type:
pd.DataFrame
- 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_probfor 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 cachedhdi_prob_; otherwise aValueErroris raised. PassNone(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 tomatplotlib.pyplot.subplots(). Defaults to(10, 6).show (
bool) – Whether to automatically display the plot. Defaults toTrue.legend_kwargs (
dict[str,Any] |None) – Keyword arguments to adjust legend placement and styling. Supported keys:loc,bbox_to_anchor,fontsize,frameon,title(bbox_transformis accepted alongsidebbox_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:
- 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 withplot(), BayesianATT(g, t)bounds are cached during effect aggregation. If supplied here, the value must match the cachedhdi_prob_; otherwise aValueErroris raised. PassNone(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 anATT(g, e)view derived fromATT(g, t)."calendar_time"plots each cohort against calendar timet. Defaults to"event_time".include_placebo (
bool) – Whether to include pre-treatment residual estimates for eventually-treated cohorts as placebo diagnostics. Defaults toTrue.figsize (
tuple[float,float] |None) – Width and height of the figure in inches, passed tomatplotlib.pyplot.subplots(). Defaults to a height scaled by the number of cohorts whenlayout="facet"and(10, 6)whenlayout="overlay".show (
bool) – Whether to automatically display the plot. Defaults toTrue.legend_kwargs (
dict[str,Any] |None) – Keyword arguments to adjust legend placement and styling. Supported keys:loc,bbox_to_anchor,fontsize,frameon,title(bbox_transformis accepted alongsidebbox_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 whenlayout="overlay".
- Return type:
- 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 sharedPipelineContext, 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:
- 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 passtime_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 toWeightedSumFitter.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 aUserWarning. Defaults to0.0(warn on negatively correlated donors).auto_scale_sigma (
bool) – IfTrue(default) and the model still carries the weighted-sum fitters’ stocky_hatprior, thatsigma ~ HalfNormal(1)default is replaced bysigma ~ 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 toFalseto keep the originalHalfNormal(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 explicity_hatprior 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 ofdatais 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
murather than noisy posterior-predictive drawsy_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]#
- 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:
- 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
predictionandimpact).- Parameters:
hdi_prob (
float) – Probability mass of the highest density interval. Defaults to the project-wideHDI_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. IfNone, uses the first treated unit.
- Return type:
- input_validation(data, treatment_time)[source]#
Validate the input data and model formula for correctness.
- 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 toNone, in which case 2 significant figures are used.treated_unit (
str|None) – Which treated unit to plot. Must be one of the names supplied viatreated_unitsat construction time. Defaults toNone, 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 toHDI_PROB(currently 0.94).kind (
Literal['ribbon','histogram','spaghetti']) – How posterior uncertainty is rendered viaplot_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 whenkind="ribbon". Defaults to"hdi".num_samples (
int) – Number of posterior draws whenkind="spaghetti". Defaults to 50. Ignored for other kinds.plot_predictors (
bool) – Whether to overlay the donor (control) unit trajectories on the top panel. Defaults toFalse.figsize (
tuple[float,float]) – Width and height of the figure in inches, passed tomatplotlib.pyplot.subplots(). Defaults to(7, 8).show (
bool) – Whether to automatically display the plot. Defaults toTrue. Set toFalseif 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_transformis accepted alongsidebbox_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:
- 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 passtime_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) – ASyntheticDifferenceInDifferencesWeightFitterinstance. Defaults toSyntheticDifferenceInDifferencesWeightFitter.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 ofdatais 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_posteriorATT; the effect is not read from a regression coefficient or obtained by population-standardized g-computation. The time-indexedpost_impactconsumed byeffect_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, ... } ... ), ... )
- 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:
_build_weight_fitter_inputs()prepares the dict-basedX,yandcoordsinputs for the weight fitter.PyMCModel.fit()fits both the omega and lambda modules via MCMC._extract_weight_posteriors()pulls the posterior weight arrays out of the fitted model._compute_synthetic_and_gaps()builds the synthetic control trajectory and the gap between treated and synthetic._compute_tau()evaluates the double-difference ATT._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.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:
- 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. UseNonefor 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 toHDI_PROB(currently 0.94).kind (
Literal['ribbon','histogram','spaghetti']) – How posterior uncertainty is rendered viaplot_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 whenkind="ribbon". Defaults to"hdi".num_samples (
int) – Number of posterior draws whenkind="spaghetti". Defaults to 50. Ignored for other kinds.show (
bool) – Whether to callmatplotlib.pyplot.show()after drawing. Defaults toTrue.legend_kwargs (
dict[str,Any] |None) – Keyword arguments applied to the top-axis legend in place after the figure is built. Supported keys includeloc,bbox_to_anchor,fontsize,frameon,title, and optionallybbox_transformalongsidebbox_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.Axesinstances.
- Return type:
- 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:
- 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 distributionsigma: 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. IfNone, all numeric columns are used.method (
Literal['pearson','kendall','spearman']) – Correlation method passed topandas.DataFrame.corr().figsize (
tuple[float,float] |None) – Width and height in inches for the figure. Only used whenaxis not provided. IfNone, matplotlib’s default is used.ax (
Axes|None) – Axes on which to draw the heatmap. IfNone, a new figure and axes are created (sized according tofigsize).**kwargs (
Any) – Keyword arguments forwarded toseaborn.heatmap():vmin,vmax,cmap,center,robust,annot,fmt,annot_kws,linewidths,linecolor,cbar,cbar_kws,cbar_ax,square,xticklabels,yticklabels,mask, and thematplotlib.axes.Axes.pcolormesh()keywords supported by the installed seaborn version.dataandaxare 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,
xandthresholdare supplied tomemorize_chunk()andtransform()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,
xandthresholdare supplied tomemorize_chunk()andtransform()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#
Shared constants for the CausalPy package. |
|
Code for loading datasets. |
|
Dataframe-agnostic input handling. |
|
Custom PyMC models for causal inference. |
|
Custom scikit-learn models for causal inference. |
|
Adapter that lets a |
|
CausalPy experiment module. |
|
Pipeline orchestration for composable causal inference workflows. |
|
Reporting utilities for causal inference experiments. |
|
Pipeline steps for causal inference workflows. |
|
Sensitivity and diagnostic checks for causal inference experiments. |