statspai.did

did

Difference-in-Differences (DID) module for StatsPAI.

Provides estimators for: - Classic 2×2 DID (two groups, two periods) - Triple Differences / DDD (two groups, two periods, within-unit subgroup) - Callaway & Sant'Anna (2021) — staggered DID with DR/IPW/REG - Sun & Abraham (2021) — interaction-weighted event study - Synthetic DID (Arkhangelsky et al. 2021) - Goodman-Bacon (2021) — TWFE decomposition diagnostic - Honest DID (Rambachan & Roth 2023) — parallel trends sensitivity - de Chaisemartin & D'Haultfoeuille (2020) — DID with treatment switching - Borusyak, Jaravel & Spiess (2024) — imputation DID estimator - Stacked DID (Cengiz, Dube, Lindner & Zipperer, 2019) - did_analysis() — one-call DID workflow - Wooldridge (2021) — extended TWFE with cohort × time interactions - Sant'Anna & Zhao (2020) — doubly robust DID - TWFE decomposition — Bacon (2021) + de Chaisemartin–D'Haultfoeuille (2020) weights

AbsorbingCheck

Bases: NamedTuple

Whether treatment is absorbing, and where it is not.

summary

summary() -> str

One-line reading suitable for a warning or a report row.

EquivalenceResult

Bases: NamedTuple

Joint pre-trend test plus its two equivalence counterparts.

verdict

verdict(alpha: float = 0.05) -> str

One-line reading of the detection and equivalence tests together.

Equivalence is judged on the TOST when a threshold was supplied, because that bound is in interpretable outcome units. The non-central-F test uses a dimensionless bound (0.6 by default) that rejects easily, so it is only used as the fallback.

StaggeredRolloutResult

Bases: NamedTuple

Point estimate, both standard errors, and the fitted control weights.

DIDAnalysis dataclass

Bundled results from a full DID analysis workflow.

Returned by :func:sp.did_analysis; bundles the detected design, the main ATT estimate, optional event study / Bacon decomposition / honest-DID sensitivity, and a step-by-step log. Use .summary() for a human-readable report and .plot() for the event study.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> post = np.tile([0, 1], n)
>>> policy = np.repeat(rng.integers(0, 2, size=n), 2)
>>> wage = (2.0 + 0.5 * post + 1.0 * (policy * post)
...         + rng.normal(size=2 * n))
>>> df = pd.DataFrame({'post': post, 'policy': policy, 'wage': wage})
>>> report = sp.did_analysis(df, y='wage', treat='policy',
...                          time='post', run_sensitivity=False)
>>> isinstance(report, sp.DIDAnalysis)
True
>>> report.design
'2x2'

summary

summary() -> str

Print comprehensive analysis summary.

plot

plot(**kwargs: Any) -> Any

Plot event study if available, else main result.

DiDBalanceResult dataclass

Balance table for a DiD design.

Attributes:

Name	Type	Description
table	DataFrame	Long form, one row per (panel, covariate). Columns: panel ('levels' / 'changes'), covariate, mean_treated, mean_comparison, norm_diff, abs_norm_diff, flagged, and — when ω is supplied — the same three statistics again with a w_ prefix.
levels, changes	DataFrame	The two panels, split out for convenience.
threshold	float	Absolute normalized-difference cutoff used for flagged.
n_treated, n_comparison	int
base_period, comparison_period	Any
weighted	bool
diagnostics	dict

flagged property

flagged: List[str]

Covariates breaching the threshold in either panel.

summary

summary() -> str

Fixed-width report in the shape of the paper's Table 4.

to_latex

to_latex(caption: Optional[str] = None) -> str

booktabs fragment mirroring the paper's Table 4 layout.

DiDClusterDiagnostics dataclass

Cluster counts for a staggered design, graded against known evidence.

DiDDesignContract dataclass

What a fitted DiD result pins down, against Baker et al.'s recipe.

EventStudyConventionResult dataclass

Paths and convention diagnostics for one panel.

Attributes:

Name	Type	Description
paths	DataFrame	Long format: key, relative_time, att.
table	DataFrame	One row per estimator with the registry fields plus shift_pre, shift_post, asymmetry, shape_gap (split into shape_gap_pre / shape_gap_post) and matches_twfe.
reference	str	Registry key of the benchmark path (dynamic TWFE).
tolerance	float	Scale-aware threshold used for matches_twfe.

summary

summary() -> str

Human-readable report.

DistributionalDiDResult dataclass

Bases: ResultProtocolMixin

Result of :func:distributional_did.

Attributes:

Name	Type	Description
table	DataFrame	One row per outcome bin: bin_lower, bin_upper, level, estimate (the effect on the probability of landing in that bin), se, and used.
n_units	int	Units contributing to the influence functions.
aggregation	str	Which aggte aggregation produced the per-bin estimates.
alpha	float
diagnostics	dict

plot

plot(ax: Any = None, figsize: tuple = (8, 5), **kwargs: Any) -> Any

Bar chart of the per-bin effect with pointwise confidence intervals.

Parameters:

Name	Type	Description	Default
ax	matplotlib Axes		None
figsize	tuple		(8, 5)
**kwargs	Any	Passed to ax.bar.	{}

Returns:

Type	Description
Axes

FunctionalFormResult dataclass

Bases: ResultProtocolMixin

Result of :func:functional_form_test.

Attributes:

Name	Type	Description
pvalue	float	P-value for H0: the implied counterfactual density is non-negative everywhere. Large = no evidence against functional-form insensitivity.
table	DataFrame	One row per outcome bin: bin_lower, bin_upper, implied_density and se.
statistic	float	The max-t statistic max_b (-density_b / se_b).
n_bins	int	Number of bins actually used (bins with a degenerate influence function are dropped, as in the reference implementation).
n_units	int	Units contributing to the influence functions.
aggregation	str	Which aggte aggregation produced the per-bin estimates.
n_sims	int	Simulation draws behind the critical value.
alpha	float
diagnostics	dict

negative_bins property

negative_bins: DataFrame

The bins whose implied density came out negative.

plot

plot(ax: Any = None, figsize: tuple = (8, 5), lb: Optional[float] = None, ub: Optional[float] = None, **kwargs: Any) -> Any

Bar chart of the implied counterfactual density, negatives flagged.

This is the figure the test is really about: the p-value compresses into one number what the bars show directly — whether the distribution the design implies for the treated group's untreated outcome is a distribution at all. Bars below zero are the violation.

Parameters:

Name	Type	Description	Default
ax	matplotlib Axes		None
figsize	tuple		(8, 5)
lb	float	Restrict the plotted outcome range, as didFF's lb_graph / ub_graph do. Bins outside are dropped from the picture only — the test already used all of them.	None
ub	float	Restrict the plotted outcome range, as didFF's lb_graph / ub_graph do. Bins outside are dropped from the picture only — the test already used all of them.	None
**kwargs	Any	Passed to ax.bar.	{}

Returns:

Type	Description
Axes

Examples:

>>> import statspai as sp
>>> res = sp.functional_form_test(df, y="lemp", g="first_treat",
...                               t="year", i="countyreal")
>>> ax = res.plot()

HarvestDIDResult dataclass

Bases: ResultProtocolMixin

Full diagnostic output of :func:harvest_did.

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> import pandas as pd
>>> es = pd.DataFrame({
...     "relative_time": [0, 1], "att": [0.4, 0.5],
...     "se": [0.1, 0.1], "pvalue": [0.01, 0.02], "n_comparisons": [3, 3],
... })
>>> res = sp.HarvestDIDResult(
...     estimate=0.45, se=0.07, ci=(0.31, 0.59), alpha=0.05,
...     n_comparisons=6,
...     comparisons=pd.DataFrame({"att": [0.4]}),
...     event_study=es,
...     pretrend_test={"pvalue": 0.6},
... )
>>> float(res.estimate)
0.45

SensitivityResult dataclass

Bases: ResultProtocolMixin

Result of Rambachan & Roth (2023) sensitivity analysis.

Attributes:

Name	Type	Description
mbar_grid	ndarray	Grid of M-bar values tested.
ci_lower	ndarray	Lower bound of the honest CI at each M-bar.
ci_upper	ndarray	Upper bound of the honest CI at each M-bar.
breakdown_mbar	float	Smallest M-bar for which the CI includes zero (sign reversal).
att	float	Point estimate of the ATT.
att_se	float	Standard error of the ATT.
method	str	Extrapolation method used ('C-LF').
alpha	float	Significance level.

Methods:

Name	Description
summary	Print a formatted summary table.
plot	Matplotlib sensitivity plot (M-bar vs CI).

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(60):
...     g = 5 if i < 30 else 0
...     for t in range(1, 9):
...         post = 1 if (g and t >= 5) else 0
...         y = 1.0 + 0.2 * t + i / 120 + 2.0 * post + rng.normal(0, 0.5)
...         rows.append({"unit": i, "time": t, "y": y, "g": g})
>>> df = pd.DataFrame(rows)
>>> result = sp.event_study(df, y="y", treat_time="g", time="time",
...                         unit="unit", window=(-3, 3))
>>> sens = sp.sensitivity_rr(result, Mbar=[0.0, 0.5, 1.0])
>>> type(sens).__name__
'SensitivityResult'
>>> bool(isinstance(sens.summary(), str))
True
>>> import matplotlib.pyplot as plt
>>> fig, ax = plt.subplots()
>>> ax = sens.plot(ax=ax)
>>> fig.savefig("sensitivity.png")

References

Rambachan, A. & Roth, J. (2023). [@rambachan2023more]

summary

summary() -> str

Return a formatted summary string.

plot

plot(ax: Any = None, figsize: tuple[float, float] = (8, 5), **kwargs: Any) -> Any

Sensitivity plot: M-bar on x-axis, honest CI band on y-axis.

Parameters:

Name	Type	Description	Default
ax	matplotlib Axes		None
figsize	tuple		(8, 5)
**kwargs	passed to ``ax.fill_between``.		{}

Returns:

Type	Description
Axes

CSReport dataclass

Structured output of :func:cs_report.

Attributes are plain pandas objects so downstream users can export to LaTeX, Markdown, or Excel without any custom converters.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for u in range(60):
...     g = rng.choice([3, 4, 0])  # cohort: treated at t=3, t=4, or never
...     ui = rng.normal(0, 1)
...     for t in range(1, 7):
...         post = 1 if (g != 0 and t >= g) else 0
...         y = ui + 0.3 * t + 2.0 * post + rng.normal(0, 0.5)
...         rows.append({"id": u, "t": t, "g": g, "y": y})
>>> df = pd.DataFrame(rows)
>>> rpt = sp.cs_report(
...     df, y="y", g="g", t="t", i="id",
...     n_boot=50, random_state=42, verbose=False)
>>> isinstance(rpt, sp.CSReport)
True
>>> isinstance(rpt.dynamic, pd.DataFrame)
True
>>> isinstance(rpt.breakdown, pd.DataFrame)
True

forward_engineering_checklist

forward_engineering_checklist() -> DataFrame

The eight steps of Baker et al. (2026, §6), with what was run.

The paper closes with an eight-step recipe for any DiD study. Steps 1-6 are things this report can execute or record; steps 7-8 are judgement calls it can only prompt. Returning them as a frame makes the checklist machine-readable for agent consumers, which is the point of running the workflow rather than eyeballing it.

to_text

to_text() -> str

Return the human-readable report as a single string.

plot

plot(figsize: Any = (14, 10), suptitle: Optional[str] = None) -> Any

Render a 2×2 summary figure of the report.

The four quadrants show:

• Top-left: event study (dynamic) with uniform confidence band
• Top-right: θ(g) per-cohort aggregation with uniform band
• Bottom-left: θ(t) per-calendar-time aggregation
• Bottom-right: Rambachan–Roth breakdown M* across post event times

Requires matplotlib. Returns (fig, axes).

to_markdown

to_markdown(float_format: str = '%.4f') -> str

Render the report as GitHub-Flavoured Markdown.

Suitable for pasting directly into a pull request, blog post, or Jupyter notebook Markdown cell.

to_excel

to_excel(path: Any, float_format: Optional[str] = '%.6f', engine: Optional[str] = None) -> str

Dump the report to a multi-sheet Excel workbook.

Creates one sheet per block — Summary, Dynamic, Group, Calendar, Breakdown, Meta — so downstream Excel consumers (policy briefs, regulatory reports) can link to or copy from the individual tables directly.

Parameters:

Name	Type	Description	Default
path	str | Path	Destination .xlsx path.	required
float_format	str	Passed through to :meth:pandas.DataFrame.to_excel. Pass None to preserve full precision.	'%.6f'
engine	str	Excel writer engine ('openpyxl' or 'xlsxwriter'). If None pandas picks an installed one; raises a clear ImportError here if none is available.	None

Returns:

Type	Description
str	The path written.

to_latex

to_latex(float_format: str = '%.4f', caption: Optional[str] = None, label: Optional[str] = None) -> str

Render the report as a LaTeX fragment.

Uses the booktabs package for each sub-table and wraps the result in a single table float. Requires \usepackage{booktabs} in the preamble of the consuming document.

ParallelTrendsRobustnessResult dataclass

Bases: ResultProtocolMixin

Bundled pre-trends power + honest-DiD sensitivity for one DiD result.

Attributes:

Name	Type	Description
power_table	DataFrame	One row per quantity from the Roth (2022) pre-test analysis (joint statistic, p-value, power, non-centrality).
ci_grid	DataFrame	Long-format robust-CI grid: family, M, ci_lower, ci_upper, rejects_zero.
breakdown	dict	family -> Mbar*, the largest violation magnitude at which the effect stays significant.
verdict	str	One-line plain-language reading of the table.
att, att_se	float	Point estimate and standard error at relative time e.
e	int	Relative time the sensitivity analysis targets.
alpha	float	Significance level.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=80, n_periods=8, staggered=False, seed=0)
>>> df["first_treat"] = df["first_treat"].fillna(0)
>>> es = sp.event_study(df, y="y", treat_time="first_treat",
...                     time="time", unit="unit", window=(-3, 3))
>>> rob = sp.parallel_trends_robustness(es)
>>> bool(isinstance(rob.summary(), str))
True

References

Rambachan & Roth (2023) [@rambachan2023more]; Roth (2022) [@roth2022pretest].

summary

summary() -> str

Return a formatted multi-section summary string.

to_latex

to_latex(*, caption: Optional[str] = None, label: Optional[str] = None) -> str

Booktabs LaTeX table of the robust-CI grid and breakdown values.

plot

plot(ax: Any = None, figsize: Tuple[float, float] = (8, 5), **kwargs: Any) -> Any

Robust-CI bands against M, one band per restriction family.

Parameters:

Name	Type	Description	Default
ax	matplotlib Axes		None
figsize	tuple		(8, 5)
**kwargs	passed to ``ax.fill_between``.		{}

Returns:

Type	Description
Axes

DIDInputTypeError

Bases: MethodIncompatibility, TypeError

DID input type error that preserves the historical TypeError catch.

check_absorbing

check_absorbing(data: DataFrame, unit: str, time: str, treat: str, strict: bool = False) -> AbsorbingCheck

Check whether a time-varying treatment indicator is absorbing.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long panel.	required
unit	str	Column names. treat must be the time-varying 0/1 indicator, not a cohort / first-treatment column — a cohort column cannot express reversal, so checking one is meaningless.	required
time	str	Column names. treat must be the time-varying 0/1 indicator, not a cohort / first-treatment column — a cohort column cannot express reversal, so checking one is meaningless.	required
treat	str	Column names. treat must be the time-varying 0/1 indicator, not a cohort / first-treatment column — a cohort column cannot express reversal, so checking one is meaningless.	required
strict	bool	Raise :class:MethodIncompatibility instead of returning when treatment reverts. Useful as a guard in front of a cohort-based estimator.	False

Returns:

Type	Description
AbsorbingCheck	is_absorbing plus the location and extent of any reversals.

Examples:

>>> import numpy as np, pandas as pd, statspai as sp
>>> rows = []
>>> for u in range(6):
...     for t in range(1, 7):
...         on = 1 if (u < 3 and 3 <= t < 5) else 0   # units 0-2 revert
...         rows.append({"i": u, "t": t, "d": on})
>>> chk = sp.check_absorbing(pd.DataFrame(rows), "i", "t", "d")
>>> chk.is_absorbing
False
>>> chk.n_reverting_units
3

References

dechaisemartin2024difference

pretrend_equivalence

pretrend_equivalence(pre_estimates: ndarray, pre_cov: ndarray, n_bar: int, pre_se: Optional[ndarray] = None, f_threshold: float = 0.6, tost_threshold: Optional[float] = None, event_times: Optional[ndarray] = None) -> EquivalenceResult

Joint and equivalence tests on a vector of pre-period effects.

Parameters:

Name	Type	Description	Default
pre_estimates	ndarray	Pre-period event-study coefficients, excluding the normalised reference period.	required
pre_cov	ndarray	Their joint covariance matrix. The joint covariance is what makes the F test valid; a diagonal built from the standard errors would ignore the (substantial) cross-period correlation.	required
n_bar	int	Number of treated units contributing to the pre-periods. This enters the finite-sample scaling and the non-centrality parameter.	required
pre_se	Optional[ndarray]	Per-period standard errors used by the TOST. Defaults to sqrt(diag(pre_cov)).	None
f_threshold	float	Dimensionless effect-size bound for the F equivalence test (fect's default is 0.6).	0.6
tost_threshold	Optional[float]	Equivalence bound in outcome units. Required for the TOST; when omitted the TOST is reported as None.	None
event_times	Optional[ndarray]	Optional labels, carried through for reporting.	None

pretrends_equivalence

pretrends_equivalence(result, f_threshold: float = 0.6, tost_threshold: Optional[float] = None, alpha: float = 0.05) -> EquivalenceResult

Pre-trend equivalence tests for a fitted DiD result.

Answers the question a plain pre-trend test cannot: is the pre-trend demonstrably small, as opposed to merely not detected? A failure to reject "no pre-trend" is frequently just low power (Roth 2022), so reporting it alone overstates the evidence for parallel trends.

Needs the joint covariance of the pre-period event-study coefficients, which is recovered from a Callaway-Sant'Anna fit's influence functions. Results without influence functions cannot support the joint F test and raise rather than silently substituting a diagonal.

Parameters:

Name	Type	Description	Default
result	CausalResult	A DiD fit carrying an event study and influence functions (e.g. :func:sp.callaway_santanna).	required
f_threshold	float	Dimensionless effect-size bound for the F equivalence test (fect's default).	0.6
tost_threshold	float	Equivalence bound in outcome units. Without it the TOST is skipped — there is no defensible universal default for "how big a pre-trend is too big", and inventing one would be a substantive judgement made on the user's behalf. fect uses 0.36 * residual SD.	None
alpha	float	Level used by :meth:EquivalenceResult.verdict.	0.05

Returns:

Type	Description
EquivalenceResult	f_stat / f_pvalue (conventional joint test), f_equivalence_pvalue and tost_pvalue (equivalence tests, where small is the reassuring outcome), plus .verdict().

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=200, n_periods=8, staggered=True, seed=42)
>>> df['first_treat'] = df['first_treat'].fillna(0)
>>> cs = sp.callaway_santanna(df, y='y', g='first_treat', t='time', i='unit')
>>> eq = sp.pretrends_equivalence(cs)
>>> 0.0 <= eq.f_pvalue <= 1.0
True
>>> isinstance(eq.verdict(), str)
True

References

liu2024practical, roth2022pretest

staggered_cs

staggered_cs(data: DataFrame, y: str, i: str, t: str, g: str, estimand: str = 'simple', event_time: Union[float, Sequence[float]] = 0.0, se_type: str = 'neyman', fisher: bool = False, n_fisher: int = 500, random_state: Optional[int] = None, alpha: float = 0.05) -> 'CausalResult'

Callaway-Sant'Anna's estimand with design-based inference.

Same weights as sp.callaway_santanna — every not-yet-treated cohort serves as control — but the standard error comes from random adoption timing rather than parallel trends. Use it when timing was randomised and you want the familiar CS estimand; use sp.callaway_santanna when it was not.

This is R staggered::staggered_cs: the plug-in weights (beta = 1) with every not-yet-treated cohort as control, after dropping units already treated in the first period.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	DataFrame		required
i	DataFrame		required
t	DataFrame		required
g	DataFrame		required
estimand	DataFrame		required
event_time	DataFrame		required
se_type	DataFrame		required
fisher	DataFrame		required
n_fisher	DataFrame		required
random_state	Optional[int]	As in :func:staggered_rollout.	None
alpha	Optional[int]	As in :func:staggered_rollout.	None

Returns:

Type	Description
CausalResult

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=200, n_periods=6, staggered=True, seed=7)
>>> res = sp.staggered_cs(df, y='y', i='unit', t='time', g='first_treat')
>>> res.se > 0
True

References

roth2023efficient callaway2021difference

staggered_rollout

staggered_rollout(data: DataFrame, y: str, i: str, t: str, g: str, estimand: str = 'simple', efficient: bool = True, event_time: Union[float, Sequence[float]] = 0.0, use_last_treated_only: bool = False, use_did_a0: bool = True, se_type: str = 'neyman', fisher: bool = False, n_fisher: int = 500, random_state: Optional[int] = None, alpha: float = 0.05) -> 'CausalResult'

Efficient DiD for a randomised staggered rollout (Roth-Sant'Anna 2023).

Use this when treatment timing was randomly assigned — a policy lottery, a phased platform launch, an RCT rolled out in waves. Every other DiD estimator in StatsPAI identifies off parallel trends, which under random timing is both unnecessary and wasteful: this estimator uses the randomisation directly and is efficient in the class of linear combinations of cohort-period means.

.. warning:: This is a different estimand and identifying assumption from sp.callaway_santanna. On canonical did::mpdta (where timing is not randomised) this returns −0.0471 against CS's −0.0400; the gap is not a discrepancy, it is what happens when you apply a design-based estimator to an observational rollout. If timing was not randomised, use a parallel-trends estimator instead.

Parameters:

Name	Type	Description	Default
data	DataFrame	Balanced panel in long form.	required
y	str	Outcome, unit id, period, and first-treatment period. Never-treated units may be coded 0, NaN or inf; all three are accepted.	required
i	str	Outcome, unit id, period, and first-treatment period. Never-treated units may be coded 0, NaN or inf; all three are accepted.	required
t	str	Outcome, unit id, period, and first-treatment period. Never-treated units may be coded 0, NaN or inf; all three are accepted.	required
g	str	Outcome, unit id, period, and first-treatment period. Never-treated units may be coded 0, NaN or inf; all three are accepted.	required
estimand	('simple', 'cohort', 'calendar', 'eventstudy')	'simple' weights each identified cohort-period cell by cohort size; 'cohort' averages within a cohort first, then weights cohorts by size; 'calendar' averages within a calendar period; 'eventstudy' reports the ATT event_time periods after adoption.	'simple'
efficient	bool	True uses the optimal pre-period control weights; False gives the plug-in estimator (beta = 1), which is the simple difference-in-means analogue and matches R's beta = 1.	True
event_time	float or sequence of float	Only read when estimand='eventstudy'. A sequence returns one row per event time in .detail, with the joint covariance in model_info['vcov']; .estimate then averages the requested non-negative event times (all of them if none is non-negative), with the SE taken from that covariance.	0
use_last_treated_only	bool	Restrict the comparison group to the last-treated cohort, which is what Sun & Abraham's estimator does. The default uses every not-yet-treated cohort.	False
use_did_a0	bool	Which controls the efficient weights are chosen over. True uses the single DiD contrast at each cohort's last pre-period g - 1. False uses every pre-period as a separate control, so beta becomes a vector — the general form of the estimator, weakly more efficient at the cost of estimating more nuisance weights. Requires efficient=True: the plug-in fixes beta = 1, which is a single contrast and has no meaning against a vector control set.	True
se_type	('neyman', 'adjusted')	Which standard error lands in .se. 'neyman' is the conservative bound that treats the fitted control weights as fixed; 'adjusted' subtracts the part of the variance the randomisation identifies and is what R staggered prints. Both are always available in model_info. The adjusted SE is never larger, so the default is the cautious one.	'neyman'
fisher	bool	Run a Fisher randomisation test: permute adoption dates across units and compare studentised statistics. The p-value lands in model_info['fisher_pvalue'].	False
n_fisher	int	Permutation draws for the randomisation test.	500
random_state	int	Seed for the permutation draws.	None
alpha	float	Level for the reported confidence interval.	0.05

Returns:

Type	Description
CausalResult	.estimate / .se carry the ATT and its standard error; model_info holds beta (the fitted control weights), both standard errors, and any randomisation p-values.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=200, n_periods=6, staggered=True, seed=7)
>>> res = sp.staggered_rollout(df, y='y', i='unit', t='time',
...                            g='first_treat')
>>> res.se > 0
True

See Also

staggered_cs : the Callaway-Sant'Anna estimand with design-based inference. staggered_sa : the Sun-Abraham estimand with design-based inference.

References

roth2023efficient

staggered_sa

staggered_sa(data: DataFrame, y: str, i: str, t: str, g: str, estimand: str = 'simple', event_time: Union[float, Sequence[float]] = 0.0, se_type: str = 'neyman', fisher: bool = False, n_fisher: int = 500, random_state: Optional[int] = None, alpha: float = 0.05) -> 'CausalResult'

Sun-Abraham's estimand with design-based inference.

Identical to :func:staggered_cs except that only the last-treated cohort serves as control, which is what Sun & Abraham's interaction- weighted estimator does.

This is R staggered::staggered_sa.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	DataFrame		required
i	DataFrame		required
t	DataFrame		required
g	DataFrame		required
estimand	DataFrame		required
event_time	DataFrame		required
se_type	DataFrame		required
fisher	DataFrame		required
n_fisher	DataFrame		required
random_state	Optional[int]	As in :func:staggered_rollout.	None
alpha	Optional[int]	As in :func:staggered_rollout.	None

Returns:

Type	Description
CausalResult

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=200, n_periods=6, staggered=True, seed=7)
>>> res = sp.staggered_sa(df, y='y', i='unit', t='time', g='first_treat')
>>> res.se > 0
True

References

roth2023efficient sun2021estimating

did_analysis

did_analysis(data: DataFrame, y: str, treat: str, time: str, id: Optional[str] = None, covariates: Optional[List[str]] = None, method: str = 'auto', estimator: str = 'dr', control_group: str = 'nevertreated', run_bacon: bool = True, run_event_study: bool = True, run_sensitivity: bool = True, event_window: Optional[tuple[int, int]] = None, cluster: Optional[str] = None, robust: bool = True, alpha: float = 0.05, **kwargs: Any) -> DIDAnalysis

Comprehensive DID analysis workflow.

Runs the full DID analysis pipeline in one call: design detection, Bacon decomposition (staggered), estimation, event study, and honest_did sensitivity analysis.

Parameters:

Name	Type	Description	Default
data	DataFrame	Input dataset.	required
y	str	Outcome variable name.	required
treat	str	Treatment variable. For 2×2: binary (0/1). For staggered: first treatment period (0 = never treated).	required
time	str	Time period variable.	required
id	str	Unit identifier. Required for staggered designs.	None
covariates	list of str	Control variables.	None
method	str	Estimation method: 'auto', '2x2', 'cs', 'sa', 'bjs', 'sdid'.	'auto'
estimator	str	For CS: 'dr', 'ipw', or 'reg'.	'dr'
control_group	str	For CS/SA: 'nevertreated' or 'notyettreated'.	'nevertreated'
run_bacon	bool	Run Bacon decomposition for staggered designs.	True
run_event_study	bool	Run event study for dynamic effects + pre-trend test.	True
run_sensitivity	bool	Run honest_did sensitivity analysis.	True
event_window	tuple of (int, int)	Event study window, e.g. (-5, 5). Auto-detected if None.	None
cluster	str	Cluster variable for standard errors.	None
robust	bool	HC1 robust standard errors.	True
alpha	float	Significance level.	0.05

Returns:

Type	Description
DIDAnalysis	Bundled results with .summary(), .plot() methods.

Examples:

Classic 2×2:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> post = np.tile([0, 1], n)
>>> policy = np.repeat(rng.integers(0, 2, size=n), 2)
>>> wage = 2.0 + 0.5 * post + 1.0 * (policy * post) + rng.normal(size=2 * n)
>>> df = pd.DataFrame({'post': post, 'policy': policy, 'wage': wage})
>>> report = sp.did_analysis(df, y='wage', treat='policy', time='post')
>>> type(report).__name__
'DIDAnalysis'
>>> print(report.summary())

Staggered — full pipeline (design auto-detected from id):

>>> rows = []
>>> cohorts = rng.choice([0, 3, 4], size=60)
>>> for i in range(60):
...     g = cohorts[i]
...     for t in range(1, 7):
...         treated = 1 if (g > 0 and t >= g) else 0
...         earn = 1.0 + 0.3 * t + 1.5 * treated + rng.normal()
...         rows.append({'worker': i, 'year': t,
...                      'first_treat': g, 'earnings': earn})
>>> panel = pd.DataFrame(rows)
>>> report = sp.did_analysis(panel, y='earnings', treat='first_treat',
...                          time='year', id='worker')
>>> report.design
'staggered'
>>> report.plot()

Quick estimate only (skip diagnostics):

>>> report = sp.did_analysis(panel, y='earnings', treat='first_treat',
...                          time='year', id='worker',
...                          run_bacon=False, run_sensitivity=False)

bacon_decomposition

bacon_decomposition(data: DataFrame, y: str, treat: str, time: str, id: str, alpha: float = 0.05) -> Dict[str, Any]

Goodman-Bacon (2021) decomposition of the TWFE DID estimator.

Decomposes the overall TWFE coefficient into a weighted sum of 2×2 DID comparisons between different treatment timing groups.

Parameters:

Name	Type	Description	Default
data	DataFrame	Balanced panel data.	required
y	str	Outcome variable.	required
treat	str	Binary treatment indicator (0 before treatment, 1 after).	required
time	str	Time period variable.	required
id	str	Unit identifier.	required
alpha	float	Significance level.	0.05

Returns:

Type

Description

dict

Keys: - beta_twfe: overall TWFE estimate - decomposition: pd.DataFrame with columns [type, treated, control, estimate, weight] - weighted_sum: Σ(weight × estimate) — equals beta_twfe under the same dyad conventions as R bacondecomp and Stata bacondecomp - n_comparisons: number of 2×2 sub-comparisons - negative_weight_share: fraction of signed weight mass on truly negative Bacon weights. This is usually zero; already-treated control comparisons are reported separately as already_treated_control_weight_share.

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> import pandas as pd
>>> rng = np.random.default_rng(0)
>>> cohorts = {1: 4, 2: 4, 3: 7, 4: 7, 5: 99, 6: 99}  # 99 = never treated
>>> rows = []
>>> for unit, g in cohorts.items():
...     for year in range(1, 11):
...         treated = int(year >= g)
...         y = unit + 0.3 * year + 2.0 * treated + rng.normal(0, 0.5)
...         rows.append({'unit': unit, 'year': year,
...                      'outcome': y, 'treated': treated})
>>> df = pd.DataFrame(rows)
>>> result = sp.bacon_decomposition(df, y='outcome', treat='treated',
...                                 time='year', id='unit')
>>> sorted(result['decomposition']['type'].unique())
['Earlier vs Later Treated', 'Later vs Earlier Treated', 'Treated vs Untreated']
>>> bool(abs(result['weighted_sum'] - result['beta_twfe']) < 1e-8)
True

Notes

The decomposition identifies three types of comparisons:

1. Earlier vs Later treated: Units treated at time g₁ vs units treated later at g₂ (g₁ < g₂). These are "good" comparisons.
2. Later vs Earlier treated: Units treated at g₂ vs already-treated units at g₁. These are "forbidden" — they use treated units as controls and can introduce negative weighting bias.
3. Treated vs Never treated: Always valid comparisons.

A large already_treated_control_weight_share signals that TWFE is relying heavily on comparisons with already-treated controls, so a heterogeneity-robust estimator (C&S, Sun-Abraham) should be used.

See Goodman-Bacon (2021, JEcon), Theorem 1.

did_balance

did_balance(data: DataFrame, covariates: Sequence[str], *, g: str, t: str, i: str, weights: Optional[str] = None, base_period: Optional[Any] = None, comparison_period: Optional[Any] = None, cohort: Optional[Any] = None, control_group: str = 'nevertreated', threshold: float = _IR_THRESHOLD) -> DiDBalanceResult

Covariate balance in levels and changes for a DiD design.

Reproduces the shape of Table 4 in Baker et al. (2026): for each covariate, the treated and comparison group means plus the Imbens-Rubin normalized difference, computed once on baseline levels and once on changes across the treatment date, optionally both weighted and unweighted.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long-format panel.	required
covariates	sequence of str	Columns to audit. May be time-varying; the changes panel is only meaningful for those that are.	required
g	str	First-treatment period per unit; 0 (or inf) marks never-treated.	required
t	str	Time period column.	required
i	str	Unit identifier.	required
weights	str	Unit weights ω. When given, the weighted statistics are reported alongside the unweighted ones, exactly as the paper does — the two answer different questions and neither is a check on the other. .. note:: The weighted statistic uses weighted variances in the denominator, following Baker et al. (2026, §4.1): "S²_ω,T and S²_ω,C are the sample weighted or unweighted variances". This differs from cobalt::col_w_smd, which deliberately holds the denominator at the unweighted pooled SD so that balance can be compared before and after reweighting. Both are defensible for their own purpose and they disagree by a few percent; unweighted, the two conventions coincide exactly. Pinned in tests/reference_parity/test_did_balance_parity.py.	None
base_period	optional	Pre-treatment period for the levels panel. Defaults to g-1 for the cohort under audit.	None
comparison_period	optional	Second period for the changes panel. Defaults to the cohort's treatment date g.	None
cohort	optional	Which treated cohort to audit. Defaults to the largest one. Balance is a two-group statistic, so a staggered design has one table per cohort rather than one overall.	None
control_group	('nevertreated', 'notyettreated')	'nevertreated' Which units form the comparison group. Should match the comparison group of the estimator you intend to run — balance against a group you will not use is not evidence about your design.	'nevertreated'
threshold	float	Absolute normalized difference above which a covariate is flagged (Imbens and Rubin 2015, p. 277).	0.25

Returns:

Type	Description
DiDBalanceResult

Examples:

>>> import statspai as sp
>>> mpdta = sp.datasets.mpdta()
>>> bal = sp.did_balance(
...     mpdta, ["lpop"], g="first_treat", t="year", i="countyreal"
... )
>>> print(bal.summary())

References

baker2026difference, imbens2015causal

bjs_pretrend_joint

bjs_pretrend_joint(result: CausalResult, data: DataFrame, y: str, group: str, time: str, first_treat: str, controls: Optional[List[str]] = None, cluster: Optional[str] = None, horizon: Optional[List[int]] = None, n_boot: int = 300, seed: Optional[int] = None) -> Dict[str, Any]

Cluster-bootstrap joint Wald test for BJS pre-treatment coefficients.

Parameters:

Name	Type	Description	Default
result	CausalResult	Output of :func:did_imputation on data with a non-trivial horizon that covers negative values. Only its model_info['event_study'] frame is consulted, to look up the observed pre-period point estimates that we re-test with a covariance-aware statistic.	required
data	DataFrame	Same arguments you passed to the original :func:did_imputation call. Needed to re-run BJS on each cluster-bootstrap resample.	required
y	DataFrame	Same arguments you passed to the original :func:did_imputation call. Needed to re-run BJS on each cluster-bootstrap resample.	required
group	DataFrame	Same arguments you passed to the original :func:did_imputation call. Needed to re-run BJS on each cluster-bootstrap resample.	required
time	DataFrame	Same arguments you passed to the original :func:did_imputation call. Needed to re-run BJS on each cluster-bootstrap resample.	required
first_treat	DataFrame	Same arguments you passed to the original :func:did_imputation call. Needed to re-run BJS on each cluster-bootstrap resample.	required
controls	DataFrame	Same arguments you passed to the original :func:did_imputation call. Needed to re-run BJS on each cluster-bootstrap resample.	required
cluster	DataFrame	Same arguments you passed to the original :func:did_imputation call. Needed to re-run BJS on each cluster-bootstrap resample.	required
horizon	list of int	If omitted, inferred from result.model_info['event_study'].	None
n_boot	int	Cluster-bootstrap replications. Clusters are sampled with replacement; unit ids are reassigned in the resampled frame so BJS refits cleanly.	300
seed	int	RNG seed for reproducibility.	None

Returns:

Type	Description
dict	{'statistic', 'df', 'pvalue', 'method', 'n_boot', 'pre_cov'} where pre_cov is the bootstrap covariance matrix used for the Wald quadratic form.

Notes

Cost: n_boot full BJS re-fits. On a 10 000-row balanced panel with |horizon|=10, expect roughly n_boot × 0.3 s = 90 s for the default n_boot=300 — the function is therefore opt-in, not run by default inside :func:did_imputation.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True,
...                 seed=0)
>>> df['first_treat'] = df['first_treat'].fillna(0).astype(int)
>>> imp = sp.did_imputation(
...     df, y='y', group='unit', time='time',
...     first_treat='first_treat', horizon=[-3, -2, -1, 0, 1, 2])
>>> jt = sp.bjs_pretrend_joint(
...     imp, df, y='y', group='unit', time='time',
...     first_treat='first_treat', n_boot=50, seed=0)
>>> jt['df']
3
>>> jt['method']
'cluster-bootstrap'

cgs_continuous_did

cgs_continuous_did(data: DataFrame, y: str, *, dose: str, time: str, unit: str, cohort: str, degree: int = 3, num_knots: int = 0, knots: Optional[Sequence[float]] = None, dose_grid: Optional[Sequence[float]] = None, control_group: str = 'nevertreated', curve_basis: str = 'fitted', alpha: float = 0.05) -> ContinuousDoseResult

ATT(d) and ACRT(d) for a continuous treatment.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long-format panel.	required
y	str	Outcome column.	required
dose	str	Continuous treatment intensity. Zero for untreated units.	required
time	str	Period column.	required
unit	str	Unit identifier.	required
cohort	str	First-treatment period; 0 marks never-treated units.	required
degree	int	B-spline degree for the dose. degree=1 with no interior knots gives a constant ACRT -- the "effect per unit of dose" reading -- and is the right starting point when the sample is small.	3
num_knots	int	Interior knots, placed at equally spaced quantiles of the positive doses. Ignored when knots is given. There is no automatic selector: more knots buy flexibility at the cost of variance, and the choice is recorded on the result.	0
knots	sequence of float	Explicit interior knots.	None
dose_grid	sequence of float	Doses at which to report the curves. Defaults to the 10th-99th percentiles of the positive doses, which is where there are data to support them.	None
control_group	('nevertreated', 'notyettreated')		"nevertreated"
curve_basis	('fitted', 'reference')	Which basis the reported curves are evaluated on. "fitted" uses the same basis the regression was fitted on, so acrt_d is the derivative of att_d and both are the fitted function. "reference" rebuilds the basis with its boundary at the ends of the dose grid instead, which is what contdid 0.1.1 reports. Because the fitted coefficients are then applied to a differently scaled basis, the reported curves are a rescaled version of the fitted dose response -- on the parity fixture the reported ACRT sits 10% above the overall_acrt the same call returns, and at degree 1 the gap is exactly the ratio of the two ranges. Use it only to reproduce output from that package.	"fitted"
alpha	float		0.05

Returns:

Type	Description
ContinuousDoseResult

Notes

The overall ACRT standard error comes from the influence function of the per-cell regression. contdid routes its standard errors through the pte package's aggregation layer, which is not replicated here, so the two differ by a few percent even though the point estimates agree exactly.

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(400):
...     g = 2 if i < 200 else 0
...     d = float(rng.uniform(0.1, 1.0)) if g else 0.0
...     fe = rng.normal()
...     for t in (1, 2):
...         eff = 2.0 * d if (g and t >= g) else 0.0
...         rows.append((i, t, g, d, fe + 0.5 * t + eff + rng.normal(0, 0.3)))
>>> df = pd.DataFrame(rows, columns=["id", "t", "g", "dose", "y"])
>>> res = sp.cgs_continuous_did(
...     df, y="y", dose="dose", time="t", unit="id", cohort="g", degree=1
... )
>>> bool(abs(res.overall_acrt - 2.0) < 0.4)
True

References

callaway2024difference

cohort_anchored_event_study

cohort_anchored_event_study(data: DataFrame, y: str, treat: str, time: str, id: str, leads: int = 4, lags: int = 4, cluster: Optional[str] = None, alpha: float = 0.05) -> CausalResult

Cohort-anchored event-study estimator.

Reports per-cohort event-time effects averaged with cohort weights. Standard errors are cluster-robust and carry no protection against violations of parallel trends; the robust-inference half of Liu (2025) is not implemented here. See the module docstring.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long-format panel.	required
y	str		required
treat	str	First-treatment-period column (0 = never-treated).	required
time	str		required
id	str		required
leads	int	Number of pre/post event-time periods to estimate.	4
lags	int	Number of pre/post event-time periods to estimate.	4
cluster	str	Cluster column for SE; defaults to id.	None
alpha	float		0.05

Returns:

Type	Description
CausalResult	estimate is the average post-treatment effect across event times 0..lags. Per-event-time coefficients in model_info['event_study'] (DataFrame with columns rel_time, att, se, ci_low, ci_high).

References

Liu, Z. (2025). "Cohort-Anchored Robust Inference for Event-Study with Staggered Adoption." arXiv:2509.01829. [@liu2025cohort] This function implements the cohort-anchored estimator, not that paper's block-bias robust-inference procedure.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True,
...                 seed=0)
>>> df['first_treat'] = df['first_treat'].fillna(0).astype(int)
>>> res = sp.cohort_anchored_event_study(
...     df, y='y', treat='first_treat', time='time', id='unit',
...     leads=2, lags=2)
>>> round(float(res.estimate), 4)
0.3668
>>> list(res.model_info['event_study'].columns)
['rel_time', 'att', 'se', 'ci_low', 'ci_high']

did_cluster_diagnostics

did_cluster_diagnostics(data: DataFrame, unit: str, first_treat: str, cluster: Optional[str] = None, *, warn: bool = True) -> DiDClusterDiagnostics

Count the clusters treatment is assigned at, and grade the count.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long panel.	required
unit	str	Unit identifier.	required
first_treat	str	First treated period; 0, NaN or inf for never-treated.	required
cluster	str	The level treatment is assigned at -- a state, a provider organisation, a school district. Defaults to unit with a warning, because the two coincide only when each unit is independently assigned, and assuming so is the optimistic error.	None
warn	bool	Emit a warning when the design sits in or below the weakest cell of the reference grid.	True

Returns:

Type	Description
DiDClusterDiagnostics

Notes

The grading is a statement about the evidence that exists, not a power calculation for this design. A "largest-cell-or-above" verdict means the closest published simulation found coverage near nominal at that cluster count for the estimators it evaluated; it is not a guarantee about the estimator or the data at hand.

Examples:

>>> import statspai as sp
>>> mpdta = sp.datasets.mpdta()
>>> diag = sp.did_cluster_diagnostics(
...     mpdta, unit="countyreal", first_treat="first_treat", warn=False
... )
>>> diag.n_cohorts
3

did_design_contract

did_design_contract(result: Any) -> DiDDesignContract

Report which forward-engineering steps a fitted DiD result pins down.

Baker et al. (2026, section 6) argue that a DiD study should be forward engineered -- target parameter, identifying assumption, estimator, inference frame, in that order -- rather than reverse engineered from whichever regression is familiar. This reads a fitted result and reports which of those the call actually determined.

The function never invents a status. A slot that cannot be read off the result is undetermined, which is the informative answer: it means the choice was made by a default rather than by the analyst, and that the write-up still owes the reader a statement of it.

Parameters:

Name	Type	Description	Default
result	CausalResult	Output of a DiD estimator.	required

Returns:

Type	Description
DiDDesignContract

Examples:

>>> import statspai as sp
>>> mpdta = sp.datasets.mpdta()
>>> fit = sp.callaway_santanna(
...     mpdta, y="lemp", g="first_treat", t="year", i="countyreal"
... )
>>> contract = sp.did_design_contract(fit)
>>> contract.steps[0]["status"] in {"determined", "undetermined"}
True

design_robust_event_study

design_robust_event_study(data: DataFrame, y: str, treat: str, time: str, id: str, leads: int = 4, lags: int = 4, cluster: Optional[str] = None, alpha: float = 0.05) -> CausalResult

Design-robust event-study with negative-weight diagnostics.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str	Same conventions as :func:callaway_santanna.	required
treat	str	Same conventions as :func:callaway_santanna.	required
time	str	Same conventions as :func:callaway_santanna.	required
id	str	Same conventions as :func:callaway_santanna.	required
leads	int	Event-time window.	4
lags	int	Event-time window.	4
cluster	str		None
alpha	float		0.05

Returns:

Type	Description
CausalResult	Headline = average post-treatment effect. model_info['weights'] reports the implicit TWFE weights per (cohort, time); negative entries flag contamination.

References

Wright, C. S. (2026). arXiv 2601.18801. See design_robust_es2026 bibkey at the bottom of this module for the full citation.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True,
...                 seed=0)
>>> df['first_treat'] = df['first_treat'].fillna(0).astype(int)
>>> res = sp.design_robust_event_study(
...     df, y='y', treat='first_treat', time='time', id='unit',
...     leads=2, lags=2)
>>> round(float(res.estimate), 4)
0.2886
>>> diag = res.model_info['diagnostics']
>>> diag['n_negative_weight_periods']
1

compare_event_study_conventions

compare_event_study_conventions(data: DataFrame, y: str, unit: str, time: str, first_treat: str, *, estimators: Optional[Sequence[str]] = None, window: Optional[Tuple[int, int]] = None, cluster: Optional[str] = None, tolerance: Optional[float] = None) -> EventStudyConventionResult

Measure how each estimator's event-study path departs from TWFE.

Runs the requested estimators on one panel, aligns their event-study paths on shared relative times, and decomposes the difference from the dynamic TWFE benchmark into a common vertical shift within each half of the path and a residual. A symmetric estimator differs from TWFE by one shift and nothing else; an asymmetric one shows a different shift before and after the treatment date, which is the jump or kink Roth (2026) documents.

The diagnostic is only interpretable in a non-staggered design, because that is where every estimator here targets the same object and Roth's analysis applies, so a staggered panel raises rather than returning a number that cannot be read.

Parameters:

Name	Type	Description	Default
data	DataFrame	Panel with unit, time, outcome and cohort columns.	required
y	str	Column names. first_treat holds the first treated period, with 0 (or a non-positive value) for never-treated units.	required
unit	str	Column names. first_treat holds the first treated period, with 0 (or a non-positive value) for never-treated units.	required
time	str	Column names. first_treat holds the first treated period, with 0 (or a non-positive value) for never-treated units.	required
first_treat	str	Column names. first_treat holds the first treated period, with 0 (or a non-positive value) for never-treated units.	required
estimators	sequence of str	Registry keys to run. Defaults to every estimator with a runner.	None
window	(int, int)	Relative-time window. Defaults to the widest the panel supports.	None
cluster	str	Cluster column; defaults to unit.	None
tolerance	float	Threshold for matches_twfe. Defaults to 1e-6 * max(1, max|beta_twfe|), which keeps the verdict scale-free.	None

Returns:

Type	Description
EventStudyConventionResult

Raises:

Type	Description
MethodIncompatibility	If the panel is staggered, or has no never-treated units.

Examples:

>>> import statspai as sp
>>> panel = sp.datasets
>>> res = sp.compare_event_study_conventions(
...     df, y="y", unit="unit", time="time", first_treat="g"
... )
>>> print(res.summary())

event_study_convention

event_study_convention(estimator: Optional[str] = None) -> 'pd.DataFrame | Dict[str, Any]'

Report how an estimator builds its event-study reference periods.

Parameters:

Name	Type	Description	Default
estimator	str	Either a registry key such as 'callaway_santanna[base_period=varying]' or a bare estimator name such as 'did_imputation'. A bare name returns every option-specific row for that estimator. None returns the whole registry.	None

Returns:

Type	Description
DataFrame or dict	A dict for an exact key, otherwise a DataFrame with one row per convention.

Examples:

>>> import statspai as sp
>>> table = sp.event_study_convention()
>>> bool((~table["twfe_comparable"]).any())
True
>>> sp.event_study_convention(
...     "callaway_santanna[base_period=varying]"
... )["symmetric"]
False

distributional_did

distributional_did(data: DataFrame, y: str, *, g: str, t: str, i: str, n_bins: Union[int, str, None] = 'auto', binpoints: Optional[Sequence[float]] = None, aggregation: str = 'group', estimator: str = 'dr', control_group: str = 'nevertreated', x: Optional[List[str]] = None, weights: Optional[str] = None, anticipation: int = 0, panel: bool = True, allow_unbalanced_panel: bool = False, balance_e: Optional[int] = None, min_e: float = -inf, max_e: float = inf, alpha: float = 0.05) -> DistributionalDiDResult

Treatment effect on the distribution of the outcome, bin by bin.

Where :func:functional_form_test asks whether the design's implied counterfactual density is a density at all, this asks a different question: how did treatment move probability mass around? Bin the outcome, run Callaway-Sant'Anna on each bin indicator, and read off the effect on P(Y in bin).

The per-bin effects sum to approximately zero by construction — treatment redistributes mass, it does not create it — so the informative content is the shape: which parts of the outcome distribution gained and which lost. A mean ATT of zero is perfectly consistent with large offsetting movements in the tails, and this is what shows them.

This is R didFF::distDD. Unlike :func:functional_form_test it runs no test and returns no p-value: the reference reports point estimates and standard errors only.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long-format panel.	required
y	str	Outcome column.	required
g	str	First-treatment period; 0 marks never-treated units.	required
t	str	Time period column.	required
i	str	Unit identifier.	required
n_bins	Union[int, str, None]	Binning, exactly as in :func:functional_form_test. Note the bins here are built over the whole panel, not the untreated rows only: the object of interest is the treated group's own outcome distribution, so post-treatment values must be inside the support.	'auto'
binpoints	Union[int, str, None]	Binning, exactly as in :func:functional_form_test. Note the bins here are built over the whole panel, not the untreated rows only: the object of interest is the treated group's own outcome distribution, so post-treatment values must be inside the support.	'auto'
aggregation	str		'group'
estimator	str		'group'
control_group	str		'group'
x	str		'group'
weights	str		'group'
anticipation	str		'group'
panel	str		'group'
allow_unbalanced_panel	bool	Passed through to :func:callaway_santanna and :func:aggte.	False
balance_e	bool	Passed through to :func:callaway_santanna and :func:aggte.	False
min_e	bool	Passed through to :func:callaway_santanna and :func:aggte.	False
max_e	bool	Passed through to :func:callaway_santanna and :func:aggte.	False
alpha	float	Level for the confidence intervals drawn by .plot().	0.05

Returns:

Type	Description
DistributionalDiDResult

Examples:

>>> import statspai as sp
>>> df = sp.datasets.mpdta()
>>> res = sp.distributional_did(
...     df, y="lemp", g="first_treat", t="year", i="countyreal", n_bins=6
... )
>>> len(res.table)
6
>>> bool(abs(res.diagnostics["effect_sum"]) < 1e-8)  # mass is conserved
True

See Also

functional_form_test : is parallel trends sensitive to the outcome scale?

References

roth2023when

functional_form_test

functional_form_test(data: DataFrame, y: str, *, g: str, t: str, i: str, n_bins: Union[int, str, None] = 'auto', binpoints: Optional[Sequence[float]] = None, aggregation: str = 'group', estimator: str = 'dr', control_group: str = 'nevertreated', x: Optional[List[str]] = None, weights: Optional[str] = None, anticipation: int = 0, panel: bool = True, allow_unbalanced_panel: bool = False, balance_e: Optional[int] = None, min_e: float = -inf, max_e: float = inf, n_sims: int = 100000, alpha: float = 0.05, random_state: Optional[int] = 0) -> FunctionalFormResult

Test whether parallel trends can hold for every monotonic transform.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long-format panel.	required
y	str	Outcome column.	required
g	str	First-treatment period; 0 marks never-treated units.	required
t	str	Time period column.	required
i	str	Unit identifier.	required
n_bins	int or 'auto'	Number of equal-width outcome bins. "auto" follows the reference implementation: an outcome with fewer than 20 distinct untreated values is treated as discrete (one bin per value, with a warning), otherwise it is cut into min(20, n_distinct) bins. An explicit integer always cuts. The test has no power against violations that are invisible at the chosen resolution, so a too-coarse binning buys a large p-value for nothing.	"auto"
binpoints	sequence of float	Explicit bin edges, for when the outcome has natural cut points. Padded to cover the outcome range (with a warning) if they fall short. Cannot be combined with an explicit n_bins.	None
aggregation	('group', 'simple', 'dynamic', 'calendar')	Which :func:aggte aggregation defines "the" implied density. "group" is the reference implementation's default.	"group"
estimator	str	Passed through to :func:callaway_santanna. weights names a sampling-weight column; leaving it unset weights every unit equally.	'dr'
control_group	str	Passed through to :func:callaway_santanna. weights names a sampling-weight column; leaving it unset weights every unit equally.	'dr'
x	str	Passed through to :func:callaway_santanna. weights names a sampling-weight column; leaving it unset weights every unit equally.	'dr'
weights	str	Passed through to :func:callaway_santanna. weights names a sampling-weight column; leaving it unset weights every unit equally.	'dr'
anticipation	str	Passed through to :func:callaway_santanna. weights names a sampling-weight column; leaving it unset weights every unit equally.	'dr'
panel	str	Passed through to :func:callaway_santanna. weights names a sampling-weight column; leaving it unset weights every unit equally.	'dr'
allow_unbalanced_panel	str	Passed through to :func:callaway_santanna. weights names a sampling-weight column; leaving it unset weights every unit equally.	'dr'
balance_e	Optional[int]	Passed through to :func:aggte; they only bite when aggregation="dynamic", where they restrict which event times enter the aggregate.	None
min_e	Optional[int]	Passed through to :func:aggte; they only bite when aggregation="dynamic", where they restrict which event times enter the aggregate.	None
max_e	Optional[int]	Passed through to :func:aggte; they only bite when aggregation="dynamic", where they restrict which event times enter the aggregate.	None
n_sims	int	Draws used for the least-favourable critical value.	100000
alpha	float		0.05
random_state	int	Seed for the simulation. Fixed by default so the p-value is reproducible; pass None for a fresh draw.	0

Returns:

Type	Description
FunctionalFormResult

Notes

A failure to reject is not evidence for functional-form insensitivity. The test compares an estimated density against zero, so with few units, coarse bins, or a short panel it will fail to reject almost regardless of the truth.

Examples:

>>> import statspai as sp
>>> df = sp.datasets.mpdta()
>>> res = sp.functional_form_test(
...     df, y="lemp", g="first_treat", t="year", i="countyreal", n_bins=6
... )
>>> res.pvalue > 0.05
True

References

roth2023when

gardner_did

gardner_did(data: DataFrame, y: str, group: str, time: str, first_treat: str, controls: Optional[List[str]] = None, event_study: bool = False, horizon: Optional[List[int]] = None, cluster: Optional[str] = None, alpha: float = 0.05, vce: str = 'analytic', n_boot: int = 199, boot_seed: int = 0) -> CausalResult

Gardner (2021) two-stage DID estimator.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long-format panel.	required
y	str	Outcome column name.	required
group	str	Unit (panel-id) column.	required
time	str	Time column.	required
first_treat	str	First-treatment-period column. Never-treated units should be encoded as 0, NaN, or +inf.	required
controls	list of str	Additional covariates included in both stages.	None
event_study	bool	If True, Stage 2 reports coefficients by relative time k = t - first_treat_i.	False
horizon	list of int	Relative-time leads/lags to report when event_study=True; defaults to range(-5, 6) intersected with available support.	None
cluster	str	Cluster variable for Stage-2 SEs. Defaults to group.	None
alpha	float	Two-sided CI level.	0.05
vce	('analytic', 'bootstrap')	Standard-error mode. 'analytic' clusters the Stage-2 residuals (fast) but ignores the variance from estimating the Stage-1 fixed effects and is anti-conservative (empirically ~0.78 coverage at a nominal 95% level); a UserWarning recommends 'bootstrap'. 'bootstrap' resamples whole clusters and re-runs the full two-step procedure (Gardner 2021 / did2s), substantially improving coverage (≈0.90 vs ≈0.78 in simulations; it approaches nominal as the number of clusters grows). Point estimates are identical either way.	'analytic'
n_boot	int	Number of cluster-bootstrap replications when vce='bootstrap'.	199
boot_seed	int	Seed for the cluster bootstrap (deterministic results).	0

Returns:

Type	Description
CausalResult	.estimate is the overall ATT; .model_info['event_study'] carries the event-study dict when requested. Supplies .summary(), .cite(), and is compatible with sp.outreg2().

Notes

Identification requires the usual staggered-DID conditions (parallel trends, no anticipation) plus a linear two-way FE + additive covariate structure for the untreated potential outcome. Stage-2 standard errors cluster by unit — bootstrapping the whole two-step procedure gives a conservative covariance when covariate models are heavy.

References

Gardner, J. (2022). Two-stage differences in differences. arXiv:2207.05943. [@gardner2022twostage]

Examples:

Staggered panel with a never-treated group (first_treat = 0). sp.did_2stage is an alias of this function.

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n_units, n_periods = 40, 8
>>> unit = np.repeat(np.arange(n_units), n_periods)
>>> time = np.tile(np.arange(1, n_periods + 1), n_units)
>>> first = np.where(unit < 20, 5, 0)  # 0 = never treated
>>> d = ((first > 0) & (time >= first)).astype(float)
>>> y = (0.5 * unit + 0.3 * time + 2.0 * d
...      + rng.normal(0, 0.5, unit.size))
>>> df = pd.DataFrame(
...     {"y": y, "unit": unit, "time": time, "g": first}
... )
>>> res = sp.gardner_did(
...     df, y="y", group="unit", time="time", first_treat="g"
... )
>>> round(res.estimate, 2)  # true ATT = 2.0
2.03

harvest_did

harvest_did(data: DataFrame, *, unit: str, time: str, outcome: str, treat: Optional[str] = None, cohort: Optional[str] = None, never_value: Any = 0, horizons: Optional[Sequence[int]] = None, reference: int = -1, alpha: float = 0.05, weighting: str = 'precision') -> CausalResult

Harvest every valid 2×2 DID comparison and aggregate them.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long-format panel.	required
unit	str	Column names.	required
time	str	Column names.	required
outcome	str	Column names.	required
treat	str	Binary treatment indicator. If provided, the cohort (first treatment time) is inferred per unit.	None
cohort	str	Alternative to treat: a column containing the already-computed cohort (first treatment time) per unit.	None
never_value	any	Value that marks "never treated" in the cohort column. If you use treat, units without any treated observation are mapped to this value automatically.	0
horizons	sequence of int	Event-time horizons to evaluate. Defaults to [-3, -2, -1, 0, 1, 2, 3, 4]. Positive values are post-treatment, 0 is the first treated period, negative values are placebo/pre-trend.	None
reference	int	Pre-treatment reference horizon relative to each cohort's treatment time. -1 = period immediately before treatment (standard event-study convention).	-1
alpha	float		0.05
weighting	('precision', 'equal', 'n_treated')	How to aggregate the harvested 2×2 estimates. precision uses inverse-variance weights (minimum-variance aggregate under independence). equal averages without weights. n_treated weights each comparison by its treated-unit count.	'precision'

Returns:

Type	Description
CausalResult	estimand is the aggregated post-treatment ATT (average over non-negative horizons). detail exposes the full 2×2 table; model_info['event_study'] is the per-horizon aggregation.

Notes

Inference assumes independence across units within each cohort (unit-level cluster-robust SEs), but the cross-horizon covariance induced by shared units is ignored when aggregating the event study into a single ATT. For strict inference, wrap this call in :func:sp.inference.bootstrap at the unit level, or use the per-comparison table to feed :func:sp.inference.multiway_cluster_vcov.

Examples:

>>> import statspai as sp
>>> df = sp.utils.dgp_did(n_units=80, n_periods=12, seed=0)
>>> out = sp.harvest_did(
...     df, unit='unit', time='time', outcome='y',
...     treat='treated', horizons=range(-3, 5),
... )
>>> out.estimate

breakdown_m

breakdown_m(result: CausalResult, e: int = 0, method: str = 'smoothness', alpha: float = 0.05) -> float

Compute the breakdown value of M.

The breakdown M is the largest violation magnitude under which the treatment effect at relative time e remains statistically significant. Larger M = more robust.

Parameters:

Name	Type	Description	Default
result	CausalResult	DID result with event study.	required
e	int	Relative time period.	0
method	str		'smoothness'
alpha	float		0.05

Returns:

Type	Description
float	Breakdown value M. The effect is significant for all M < M.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> rows = []  # staggered panel: half treated at t=5, half never
>>> for i in range(30):
...     g = 5 if i < 15 else 0
...     ui = rng.normal(0, 1)
...     for t in range(1, 9):
...         post = 1 if (g != 0 and t >= g) else 0
...         y = ui + 0.3 * t + 2.0 * post + rng.normal(0, 0.5)
...         rows.append({"unit": i, "time": t, "y": y, "g": g})
>>> df = pd.DataFrame(rows)
>>> r = sp.sun_abraham(df, y='y', g='g', t='time', i='unit')
>>> m_star = sp.breakdown_m(r, e=0)
>>> bool(m_star >= 0)  # parallel trends can deviate by up to M* per period
True

Notes

Formally, M* = sup{M : 0 ∉ CI(M)}.

For the smoothness restriction with n_drift periods: M* = (|θ̂| - z_{α/2} × SE) / n_drift

See Rambachan & Roth (2023, ReStud), Definition 2.

aggte_from_influence

aggte_from_influence(source: Union[DataFrame, str, Path], type: str = 'simple', **aggte_kwargs: Any) -> CausalResult

Aggregate group-time ATTs from exported influence functions.

The post-hoc half of the Stata csdid saverif() workflow: rebuild the ATT(g, t) grid and influence-function matrix from a frame written by :func:influence_functions and run :func:statspai.aggte on it — no refit, no original data needed.

Parameters:

Name	Type	Description	Default
source	DataFrame, str, or Path	A frame produced by :func:influence_functions, or a path to one (.parquet or CSV).	required
type	('simple', 'dynamic', 'group', 'calendar')	Aggregation scheme, forwarded to :func:statspai.aggte.	'simple'
**aggte_kwargs	Any	Any other :func:statspai.aggte options — min_e / max_e, balance_e, bstrap, n_boot, cband, alpha, random_state.	{}

Returns:

Type	Description
CausalResult	Same shape as sp.aggte output.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=60, n_periods=6, staggered=True, seed=1)
>>> df['first_treat'] = df['first_treat'].fillna(0)
>>> cs = sp.callaway_santanna(df, y='y', g='first_treat', t='time',
...                           i='unit')
>>> rif = sp.influence_functions(cs)
>>> es = sp.aggte_from_influence(rif, type='dynamic', n_boot=200,
...                              random_state=0)
>>> es.estimand
'ATT'

References

Callaway, B. and Sant'Anna, P. H. C. (2021). Difference-in-differences with multiple time periods. Journal of Econometrics, 225(2), 200-230. [@callaway2021difference]

influence_functions

influence_functions(result: CausalResult, path: Optional[Union[str, Path]] = None) -> DataFrame

Export per-unit influence functions of a Callaway–Sant'Anna fit.

Equivalent to Stata csdid, saverif(): the returned frame is self-contained — it carries everything :func:aggte_from_influence needs to recompute any aggregation without refitting.

Parameters:

Name	Type	Description	Default
result	CausalResult	Output of :func:statspai.callaway_santanna.	required
path	str or Path	If given, also write the frame to disk — .parquet via to_parquet, anything else via to_csv(index=False).	None

Returns:

Type

Description

DataFrame

Long format, one row per unit × (g, t) pair, columns: unit — unit identifier (observation index for RCS fits) unit_cohort — the unit's own first-treatment cohort (0 = never treated); used to rebuild aggregation weights group, time, relative_time — the (g, t) cell att — the cell's point estimate (repeated across units) influence — the unit's influence-function value ψᵢ(g, t) cluster — only when the fit used clustervars

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=60, n_periods=6, staggered=True, seed=1)
>>> df['first_treat'] = df['first_treat'].fillna(0)
>>> cs = sp.callaway_santanna(df, y='y', g='first_treat', t='time',
...                           i='unit')
>>> rif = sp.influence_functions(cs)
>>> set(rif.columns) >= {'unit', 'group', 'time', 'influence'}
True

References

Callaway, B. and Sant'Anna, P. H. C. (2021). Difference-in-differences with multiple time periods. Journal of Econometrics, 225(2), 200-230. [@callaway2021difference]

did_misclassified

did_misclassified(data: DataFrame, y: str, treat: str, time: str, id: str, pi_misclass: float = 0.0, anticipation_periods: int = 0, cluster: Optional[str] = None, alpha: float = 0.05) -> CausalResult

Staggered DiD robust to timing misclassification + anticipation.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str		required
treat	str		required
time	str		required
id	str		required
pi_misclass	float in [0, 0.5]	Probability that the recorded first-treatment period g is off by ±1 (symmetric). Pass 0 to skip this correction.	0.0
anticipation_periods	int	Number of leads to absorb as anticipation (subtracts the average of pre-event coefficients k = -1..-anticipation_periods from the post ATT estimate).	0
cluster	str		None
alpha	float		0.05

Returns:

Type	Description
CausalResult	estimate is the corrected ATT; model_info reports both the naive and the corrected estimate, plus the anticipation offset and misclassification adjustment factor.

References

arXiv 2507.20415, Staggered Adoption DiD Designs with Misclassification and Anticipation (2025).

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True,
...                 seed=0)
>>> df['first_treat'] = df['first_treat'].fillna(0).astype(int)
>>> res = sp.did_misclassified(df, y='y', treat='first_treat',
...                            time='time', id='unit')
>>> round(float(res.estimate), 4)
0.4354

dl_propensity_score

dl_propensity_score(data: DataFrame, *, treatment: str, covariates: Sequence[str], hidden_sizes: Sequence[int] = (64, 32), max_iter: int = 300, random_state: int = 0) -> ndarray

Neural-net propensity score with balance-targeted loss.

Fits a small multi-layer perceptron e(X) = P(T=1 | X); if torch is available uses a proper MLP, otherwise falls back to scikit-learn's :class:MLPClassifier (lbfgs optimiser, ReLU).

Parameters:

Name	Type	Description	Default
data	DataFrame		required
treatment	str		required
covariates	sequence of str		required
hidden_sizes	sequence of int		(64, 32)
max_iter	int		300
random_state	int		0

Returns:

Type	Description
ndarray of shape (n,)	Estimated propensity scores clipped to (0.02, 0.98).

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> treat = rng.binomial(1, 1 / (1 + np.exp(-(0.5 * x1 + 0.3 * x2))))
>>> df = pd.DataFrame({"treat": treat, "x1": x1, "x2": x2})
>>> e = sp.dl_propensity_score(df, treatment="treat",
...                            covariates=["x1", "x2"], random_state=0)
>>> e.shape == (n,)
True
>>> bool(((e >= 0.02) & (e <= 0.98)).all())
True

References

Peng, Li, Wu & Li (arXiv:2404.04794, 2024). [@peng2024local]

overlap_weighted_did

overlap_weighted_did(data: DataFrame, *, y: str, treat: str, time: str, covariates: Optional[Sequence[str]] = None, ps_model: Any = 'logit', alpha: float = 0.05) -> CausalResult

Overlap-weighted 2x2 DID.

Parameters:

Name	Type	Description	Default
data	DataFrame	Two-period panel with a binary treat indicator and a binary post/pre time indicator.	required
y	str		required
treat	str		required
time	str		required
covariates	sequence of str	Pre-treatment covariates for the propensity score. If omitted, reduces to standard (unweighted) 2x2 DID.	None
ps_model	('logit', 'gbm', 'dl')	How to estimate e(X) = P(treat=1 | X). 'dl' uses :func:dl_propensity_score.	'logit'
alpha	float		0.05

Returns:

Type	Description
CausalResult	estimand = 'ATT (overlap)'. Uses a sandwich-style bootstrap-ready SE derived from weighted residuals.

References

Li, Morgan & Zaslavsky (JASA 2018). "Overlap-weighted difference-in-differences" (Economics Letters 2025).

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 150
>>> x = rng.normal(0, 1, n)
>>> treat = rng.binomial(1, 1 / (1 + np.exp(-x)))
>>> base = 1.0 + 0.5 * x + rng.normal(0, 1, n)
>>> post = base + 0.4 + 1.5 * treat + rng.normal(0, 1, n)
>>> df = pd.DataFrame({
...     "y": np.concatenate([base, post]),
...     "treat": np.tile(treat, 2),
...     "time": np.repeat([0, 1], n),
...     "x": np.tile(x, 2),
... })
>>> res = sp.overlap_weighted_did(
...     df, y="y", treat="treat", time="time",
...     covariates=["x"],
... )
>>> round(res.estimate, 2)  # true effect = 1.5
1.8
>>> res.estimand
'ATT (overlap)'

bacon_plot

bacon_plot(bacon_result: Dict[str, Any], ax: Any = None, figsize: Tuple[float, float] = (10, 6), title: Optional[str] = None, colors: Optional[Dict[str, str]] = None, **kwargs: Any) -> Tuple[Any, Any]

Scatter plot of Goodman-Bacon decomposition.

Each point is a 2×2 sub-comparison: x = weight, y = DD estimate. Color distinguishes comparison types (Treated vs Never-treated, Earlier vs Later, Later vs Already-treated).

Parameters:

Name	Type	Description	Default
bacon_result	dict	Output from bacon_decomposition(). Must contain 'decomposition' DataFrame and 'beta_twfe'.	required
ax	matplotlib Axes		None
figsize	tuple		(10, 6)
title	str		None
colors	dict	Map comparison type → color. Defaults provided.	None

Returns:

Type	Description
(fig, ax)

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True, seed=0)
>>> bacon = sp.bacon_decomposition(df, y='y', treat='treated',
...                                time='time', id='unit')
>>> fig, ax = sp.bacon_plot(bacon)
>>> type(fig).__name__
'Figure'

References

goodmanbacon2021difference

cohort_event_study_plot

cohort_event_study_plot(result: Any, ax: Any = None, figsize: Tuple[float, float] = (12, 7), title: Optional[str] = None, palette: Optional[List[str]] = None, show_aggregate: bool = True, aggregate_color: str = '#2C3E50', ci: bool = True, ci_alpha: float = 0.08, **kwargs: Any) -> Tuple[Any, Any]

Per-cohort event study plot (overlay).

Plots a separate event study line for each treatment cohort, showing heterogeneity in treatment effects across cohorts. Optionally overlays the aggregate event study.

Parameters:

Name	Type	Description	Default
result	CausalResult	Result from callaway_santanna() or did(method='cs'). Must have detail with 'group', 'relative_time', 'att' columns, and model_info['event_study'] for aggregate.	required
ax	matplotlib Axes		None
figsize	tuple		(12, 7)
title	str		None
palette	list of str	Colors for each cohort. Auto-generated if None.	None
show_aggregate	bool	Overlay the aggregate event study line.	True
aggregate_color	str	Color for aggregate line.	'#2C3E50'
ci	bool	Show confidence intervals for each cohort.	True
ci_alpha	float	CI band transparency.	0.08

Returns:

Type	Description
(fig, ax)

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True, seed=0)
>>> result = sp.did(df, y='y', treat='first_treat', time='time',
...                 id='unit', method='cs')
>>> fig, ax = sp.cohort_event_study_plot(result)
>>> type(fig).__name__
'Figure'

References

callaway2021difference

did_plot

did_plot(data: DataFrame, y: str, time: str, treat: str, treat_time: Optional[Union[int, float]] = None, show_counterfactual: bool = True, labels: Optional[Dict[str, str]] = None, colors: Optional[Tuple[str, str, str]] = None, ax: Any = None, figsize: Tuple[float, float] = (10, 6), title: Optional[str] = None, annotate_effect: bool = True, **kwargs: Any) -> Tuple[Any, Any]

Classic DID diagram showing treatment effect with counterfactual.

Plots group means over time and adds a dashed counterfactual line for the treatment group (extrapolated from pre-treatment trend parallel to the control group).

Parameters:

Name	Type	Description	Default
data	DataFrame	Input dataset.	required
y	str	Outcome variable.	required
time	str	Time period variable.	required
treat	str	Binary treatment group indicator (0/1).	required
treat_time	int or float	Treatment onset time. If None, inferred as the midpoint.	None
show_counterfactual	bool	Draw the dashed counterfactual line.	True
labels	dict	Custom labels: {'treat': ..., 'control': ..., 'counterfactual': ...}.	None
colors	tuple	(treatment, control, counterfactual) colors.	None
ax	matplotlib Axes		None
figsize	tuple		(10, 6)
title	str		None
annotate_effect	bool	Annotate the treatment effect arrow on the plot.	True

Returns:

Type	Description
(fig, ax)

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=False, seed=1)
>>> fig, ax = sp.did_plot(df, y='y', time='time', treat='treated',
...                       treat_time=5)
>>> type(fig).__name__
'Figure'

did_summary_plot

did_summary_plot(result: Any, ax: Any = None, figsize: Tuple[float, float] = (9, 5), color: str = '#2C3E50', highlight_color: str = '#C0392B', reference: Optional[float] = None, title: Optional[str] = None, sort_by: Optional[str] = None) -> Tuple[Any, Any]

Forest plot of DID method-robustness summary.

Plots each method's point estimate with its confidence interval as a horizontal errorbar. Designed to consume the CausalResult returned by :func:statspai.did.did_summary.

Parameters:

Name	Type	Description	Default
result	CausalResult	Output of :func:did_summary. Must have a detail DataFrame with columns estimate, ci_low, ci_high, and either method or estimator.	required
ax	matplotlib Axes	Existing axes to draw on. If None a new figure is created.	None
figsize	tuple	Figure size when creating a new figure.	``(9, 5)``
color	str	Color for point estimates and CIs.	``"#2C3E50"``
highlight_color	str	Color for the cross-method mean line.	``"#C0392B"``
reference	float	Horizontal reference value (e.g. 0 for 'no effect'). Defaults to 0.	None
title	str	Plot title. Defaults to "DID Method-Robustness Summary".	None
sort_by	('estimate', None)	If 'estimate', sort methods by point estimate ascending. Otherwise keep the order in result.detail.	'estimate'

Returns:

Type	Description
(fig, ax) : matplotlib figure and axes.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True, seed=0)
>>> out = sp.did_summary(df, y='y', time='time',
...                      first_treat='first_treat', group='unit')
>>> fig, ax = sp.did_summary_plot(out)
>>> type(fig).__name__
'Figure'

enhanced_event_study_plot

enhanced_event_study_plot(result: Any, ax: Any = None, figsize: Tuple[float, float] = (10, 6), title: Optional[str] = None, color: str = '#2C3E50', sig_color: Optional[str] = '#E74C3C', ci_alpha: float = 0.15, shade_pre: bool = True, shade_post: bool = True, pre_color: str = '#EBF5FB', post_color: str = '#FDEDEC', show_zero: bool = True, marker: str = 'o', markersize: int = 6, alpha_level: float = 0.05, **kwargs: Any) -> Tuple[Any, Any]

Enhanced event study plot with pre/post shading and significance coloring.

Improvement over the basic CausalResult.event_study_plot() — adds optional background shading for pre/post periods and colors significant coefficients differently.

Parameters:

Name	Type	Description	Default
result	CausalResult	DID result with event study in model_info['event_study'].	required
ax	matplotlib Axes		None
figsize	tuple		(10, 6)
title	str		None
color	str	Default color for estimates.	'#2C3E50'
sig_color	str or None	Color for significant estimates. None disables coloring.	'#E74C3C'
ci_alpha	float	Confidence band transparency.	0.15
shade_pre	bool	Shade pre-treatment region.	True
shade_post	bool	Shade post-treatment region.	True
pre_color	str	Pre-treatment shading color.	'#EBF5FB'
post_color	str	Post-treatment shading color.	'#FDEDEC'
show_zero	bool	Show horizontal zero line.	True
marker	str		'o'
markersize	int		6
alpha_level	float		0.05

Returns:

Type	Description
(fig, ax)

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True, seed=0)
>>> result = sp.did(df, y='y', treat='first_treat', time='time',
...                 id='unit', method='cs')
>>> fig, ax = sp.enhanced_event_study_plot(result, shade_pre=True)
>>> type(fig).__name__
'Figure'

ggdid

ggdid(result: Any, ax: Any = None, figsize: Tuple[float, float] = (10, 6), title: Optional[str] = None, point_color: str = '#2E86AB', band_color: str = '#F18F01', show_pointwise: bool = True, show_uniform: bool = True) -> Tuple[Any, Any]

Plot an aggte() result, mirroring R :func:did::ggdid.

Automatically dispatches on result.model_info['aggregation']:

• simple : a single point with pointwise CI
• dynamic : event-study line with pointwise CI and uniform band
• group : horizontal bars of θ̂(g) per cohort
• calendar : time-series of θ̂(t) per calendar period

Uniform bands (sup-t simultaneous confidence bands) are drawn from the cband_lower / cband_upper columns created by :func:aggte.

Parameters:

Name	Type	Description	Default
result	CausalResult	Output of :func:aggte.	required
ax	matplotlib Axes		None
figsize	tuple		(10, 6)
title	str		None
point_color	str	Colours for the pointwise estimate and the uniform band.	'#2E86AB'
band_color	str	Colours for the pointwise estimate and the uniform band.	'#2E86AB'
show_pointwise	bool	Draw pointwise CI lines.	True
show_uniform	bool	Draw uniform band (shaded region).	True

Returns:

Type	Description
(fig, ax)

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True,
...                 seed=0)
>>> cs = sp.callaway_santanna(df, y='y', g='first_treat',
...                           t='time', i='unit')
>>> agg = sp.aggte(cs, type='dynamic')
>>> fig, ax = sp.ggdid(agg)
>>> type(fig).__name__
'Figure'

group_time_plot

group_time_plot(result: Any, plot_type: str = 'dot', ax: Any = None, figsize: Tuple[float, float] = (12, 7), title: Optional[str] = None, color: str = '#2C3E50', sig_color: str = '#E74C3C', insig_color: str = '#BDC3C7', alpha_level: float = 0.05, **kwargs: Any) -> Tuple[Any, Any]

Plot group-time ATT estimates from Callaway-Sant'Anna.

Two modes: - 'dot' — dot plot with CI error bars, colored by significance - 'heatmap' — (group × time) heatmap of ATT magnitudes

Parameters:

Name	Type	Description	Default
result	CausalResult	Result from callaway_santanna() or did(method='cs'). Must have detail DataFrame with 'group', 'time', 'att' columns.	required
plot_type	str	'dot' or 'heatmap'.	'dot'
ax	matplotlib Axes		None
figsize	tuple		(12, 7)
title	str		None
color	str	Default color for dot plot.	'#2C3E50'
sig_color	str	Color for significant estimates.	'#E74C3C'
insig_color	str	Color for insignificant estimates.	'#BDC3C7'
alpha_level	float	Significance threshold.	0.05

Returns:

Type	Description
(fig, ax)

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True, seed=0)
>>> result = sp.did(df, y='y', treat='first_treat', time='time',
...                 id='unit', method='cs')
>>> fig, ax = sp.group_time_plot(result)
>>> fig2, ax2 = sp.group_time_plot(result, plot_type='heatmap')
>>> type(fig).__name__
'Figure'

References

callaway2021difference

panel_view

panel_view(data: DataFrame, unit: str, time: str, treat: str, y: Optional[str] = None, type: str = 'treat', ax: Any = None, figsize: Tuple[float, float] = (10, 6), **kwargs: Any) -> Tuple[Any, Any, Dict[str, Any]]

panelView-style display of a panel's treatment status and outcomes.

A thin, dependency-light counterpart of Mou, Liu and Xu's panelView (R and Stata) for the three questions it answers before any panel causal estimator is fitted: which units are treated when (and whether treatment ever switches off), how outcomes move over time by treatment status, and how many cells are missing.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long panel.	required
unit	str	Unit and time identifiers.	required
time	str	Unit and time identifiers.	required
treat	str	0/1 treatment status in each unit-period.	required
y	str	Outcome column, required for type="outcome".	None
type	('treat', 'outcome')	'treat' draws the unit-by-period treatment-status tiles (sp.treatment_rollout_plot); 'outcome' draws every unit's outcome path, treated periods highlighted, with the mean path of the never-treated units.	'treat'
ax	matplotlib axis		None
figsize	tuple		(10, 6)

Returns:

Type	Description
(fig, ax, summary)	summary reports n_units, n_periods, n_treated_units, n_never_treated, first_treat (per unit), staggered, has_reversals, n_missing_cells and share_treated_cells, the facts that decide between a two-way FE, a staggered DiD, and a counterfactual (sp.fect) design.

Examples:

>>> import statspai as sp
>>> fig, ax, info = sp.panel_view(df, unit="id", time="t", treat="d")
>>> info["staggered"], info["has_reversals"]

References

[@mou2023panel]

parallel_trends_plot

parallel_trends_plot(data: DataFrame, y: str, time: str, treat: str, id: Optional[str] = None, treat_time: Optional[Union[int, float]] = None, agg: str = 'mean', labels: Optional[Dict[str, str]] = None, colors: Optional[Tuple[str, str]] = None, ci: bool = True, ax: Any = None, figsize: Tuple[float, float] = (10, 6), title: Optional[str] = None, **kwargs: Any) -> Tuple[Any, Any]

Plot raw outcome means over time for treatment and control groups.

The workhorse pre-analysis plot: shows whether parallel trends is plausible before running DID.

Parameters:

Name	Type	Description	Default
data	DataFrame	Input dataset.	required
y	str	Outcome variable.	required
time	str	Time period variable.	required
treat	str	Treatment group indicator. Binary (0/1) for 2×2, or first-treatment-period for staggered (0 = never treated).	required
id	str	Unit identifier (for panel data).	None
treat_time	int or float	Treatment onset time. Draws a vertical line if provided.	None
agg	str	Aggregation function: 'mean' or 'median'.	'mean'
labels	dict	Custom labels, e.g. {'treat': 'New Jersey', 'control': 'Pennsylvania'}.	None
colors	tuple of str	Colors for (treatment, control). Default: ('#E74C3C', '#2C3E50').	None
ci	bool	Show 95% confidence intervals (±1.96 SE of mean).	True
ax	matplotlib Axes	Existing axes to plot on.	None
figsize	tuple	Figure size.	(10, 6)
title	str	Plot title.	None

Returns:

Type	Description
(fig, ax)

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True, seed=0)
>>> fig, ax = sp.parallel_trends_plot(df, y='y', time='time',
...                                   treat='treated', treat_time=3)
>>> type(fig).__name__
'Figure'

sensitivity_plot

sensitivity_plot(sensitivity: DataFrame, original_ci: Optional[Tuple[float, float]] = None, original_estimate: Optional[float] = None, ax: Any = None, figsize: Tuple[float, float] = (10, 6), title: Optional[str] = None, color: str = '#2C3E50', breakdown_color: str = '#E74C3C', original_color: str = '#27AE60', **kwargs: Any) -> Tuple[Any, Any]

Plot Rambachan & Roth (2023) sensitivity analysis.

Shows how the robust confidence interval changes as the maximum allowed parallel trends violation (M) increases.

Parameters:

Name	Type	Description	Default
sensitivity	DataFrame	Output from honest_did(). Columns: M, ci_lower, ci_upper, rejects_zero.	required
original_ci	tuple of (float, float)	Original CI (at M=0) for comparison.	None
original_estimate	float	Original point estimate.	None
ax	matplotlib Axes		None
figsize	tuple		(10, 6)
title	str		None
color	str	CI band color.	'#2C3E50'
breakdown_color	str	Color for the breakdown point marker.	'#E74C3C'
original_color	str	Color for original estimate marker.	'#27AE60'

Returns:

Type	Description
(fig, ax)

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True, seed=0)
>>> result = sp.did(df, y='y', treat='first_treat', time='time',
...                 id='unit', method='cs')
>>> sens = sp.honest_did(result, e=0)
>>> fig, ax = sp.sensitivity_plot(sens, original_estimate=result.estimate,
...                               original_ci=result.ci)
>>> type(fig).__name__
'Figure'

References

rambachan2023more

treatment_rollout_plot

treatment_rollout_plot(data: DataFrame, time: str, treat: str, id: str, ax: Any = None, figsize: Tuple[float, float] = (12, 7), title: Optional[str] = None, treated_color: str = '#E74C3C', untreated_color: str = '#ECF0F1', never_color: str = '#BDC3C7', sort_by: str = 'treat_time', show_cohort_labels: bool = True, **kwargs: Any) -> Tuple[Any, Any]

Visualise staggered treatment adoption timing.

Draws a tile/heatmap where each row is a unit and each column is a time period. Treated periods are shaded, making the staggered rollout pattern immediately visible.

Parameters:

Name	Type	Description	Default
data	DataFrame	Panel data with unit, time, and treatment columns.	required
time	str	Time period variable.	required
treat	str	First-treatment-period column (0 = never treated), or binary treatment indicator.	required
id	str	Unit identifier.	required
ax	matplotlib Axes		None
figsize	tuple		(12, 7)
title	str		None
treated_color	str	Color for treated unit-periods.	'#E74C3C'
untreated_color	str	Color for untreated unit-periods.	'#ECF0F1'
never_color	str	Color for never-treated units.	'#BDC3C7'
sort_by	str	Sort units by: 'treat_time' (earliest first), 'id', or 'random'.	'treat_time'
show_cohort_labels	bool	Annotate cohort boundaries on the y-axis.	True

Returns:

Type	Description
(fig, ax)

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True, seed=0)
>>> fig, ax = sp.treatment_rollout_plot(df, time='time',
...                                     treat='first_treat', id='unit')
>>> type(fig).__name__
'Figure'

pretrends_power

pretrends_power(result: Any, delta: Optional[ndarray] = None, alpha: float = 0.05, test: str = 'individual') -> Dict[str, Any]

Power of the pre-trend test against a hypothesised violation.

Implements the power calculation from Roth (2022, AER: Insights). A non-significant pre-trend test is uninformative when the test has low power against economically meaningful violations of parallel trends.

Parameters:

Name	Type	Description	Default
result	CausalResult	Event-study result with pre-treatment estimates and SEs.	required
delta	array - like	Hypothesised trend violation in the pre-period (length = number of pre-periods). Default: linear trend delta[k] = (k+1) * min(|SE|) -- a violation equal to one SE at the furthest lag, declining linearly to near-zero.	None
alpha	float	Significance level of the pre-trend test.	0.05
test	('individual', 'joint')	Which pre-test the power refers to. "individual" is the practice Roth (2022) analyses and the one his pretrends R package implements: the analyst eyeballs the event-study plot and calls the pre-trends into question if any pre-period coefficient is individually significant. Power is then one minus the probability that every pre-period coefficient falls inside its own +/- z_{1-alpha/2} * SE band, integrated over the joint normal with mean delta -- a multivariate-normal rectangle probability. "joint" is the power of the joint Wald test that all pre-period coefficients are zero, chi2(K) with non-centrality delta' Sigma^-1 delta. Reported by :func:pretrends_test, and a strictly different quantity -- not a tighter or looser version of the same one. The two are not even comparable at face value: the joint test has size exactly alpha, while the coefficient-by-coefficient test rejects with probability above alpha under the null because each of the K coefficients gets its own alpha-level look. Which comes out more powerful against a given trend depends on the design. power_joint is always reported alongside, so both are available from one call. .. versionchanged:: 1.21.0 The default moved from "joint" to "individual" so the number matches Roth's pretrends package. This changes the returned power for existing calls -- see MIGRATION.md. Pass test="joint" to recover the previous behaviour.	"individual"

Returns:

Type	Description
dict	Always: power, power_under_null, bayes_factor, df, delta, alpha, test, warning. likelihood_ratio is present when the pre-period point estimates are available. test="joint" additionally reports noncentrality and critical_value; test="individual" reports threshold_tstat.

Notes

bayes_factor and likelihood_ratio follow the pretrends package: the Bayes factor is (1 - power) / (1 - power_under_null), the odds that a passed pre-test moves in favour of the hypothesised trend, and the likelihood ratio compares the observed pre-period coefficients under delta versus under no violation.

References

Roth, J. (2022). Pretest with Caution: Event-Study Estimates after Testing for Parallel Trends. AER: Insights, 4(3), 305--322. [@roth2022pretest]

Examples:

>>> import statspai as sp, numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(80):
...     cohort = 4 if i < 40 else 0          # 0 = never treated
...     for t in range(8):
...         post = cohort > 0 and t >= cohort
...         y = 0.3 * t + (2.0 if post else 0.0) + (i % 5) + rng.normal()
...         rows.append((i, t, cohort, y))
>>> df = pd.DataFrame(rows, columns=["id", "t", "cohort", "y"])
>>> es = sp.event_study(df, y="y", treat_time="cohort", time="t", unit="id")
>>> sp.pretrends_power(es)

pretrends_slope_for_power

pretrends_slope_for_power(result: Any, target_power: float = 0.5, alpha: float = 0.05, test: str = 'individual') -> Dict[str, Any]

Slope of a linear pre-trend the pre-test would detect target_power of the time.

The mirror image of :func:pretrends_power: instead of asking how much power the pre-test has against a chosen violation, it asks how large a violation has to be before the pre-test is even a coin flip. Roth's pretrends package exposes the same quantity as slope_for_power, and it is the number to quote when a reader asks what a passed pre-test actually rules out.

The hypothesised violation is linear in event time, delta_t = slope * (t - t_ref), with the reference period taken to be t = -1.

Parameters:

Name	Type	Description	Default
result	CausalResult	Event-study result with pre-treatment estimates and SEs. As with :func:pretrends_power, supply the full pre-period covariance via model_info['vcv_pre'] -- the diagonal fallback overstates the detectable slope.	required
target_power	float	Power the returned slope achieves. 0.5 is the pretrends default: the trend a pre-test catches half the time.	0.5
alpha	float	Significance level of the pre-test.	0.05
test	('individual', 'joint')	Which pre-test to solve against; see :func:pretrends_power.	"individual"

Returns:

Type	Description
dict	Keys: slope, target_power, achieved_power, delta, times, test, alpha.

Examples:

>>> import statspai as sp, numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(80):
...     cohort = 4 if i < 40 else 0          # 0 = never treated
...     for t in range(8):
...         post = cohort > 0 and t >= cohort
...         y = 0.3 * t + (2.0 if post else 0.0) + (i % 5) + rng.normal()
...         rows.append((i, t, cohort, y))
>>> df = pd.DataFrame(rows, columns=["id", "t", "cohort", "y"])
>>> es = sp.event_study(df, y="y", treat_time="cohort", time="t", unit="id")
>>> out = sp.pretrends_slope_for_power(es)
>>> round(out["slope"], 3)
0.245
>>> out["target_power"]
0.5

References

Roth, J. (2022). Pretest with Caution: Event-Study Estimates after Testing for Parallel Trends. AER: Insights, 4(3), 305--322. [@roth2022pretest]

pretrends_summary

pretrends_summary(result: Any, delta: Optional[ndarray] = None, alpha: float = 0.05) -> str

Print a combined pre-trends diagnostic report.

Runs pretrends_test and pretrends_power and formats the output in a single table.

Parameters:

Name	Type	Description	Default
result	CausalResult	Event-study result.	required
delta	array - like	Passed to pretrends_power.	None
alpha	float	Significance level.	0.05

Returns:

Type	Description
str	Formatted report.

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(60):
...     g = 5 if i < 30 else 0
...     for t in range(1, 9):
...         post = 1 if (g and t >= 5) else 0
...         y = 1.0 + 0.2 * t + i / 120 + 2.0 * post + rng.normal(0, 0.5)
...         rows.append({"unit": i, "time": t, "y": y, "g": g})
>>> df = pd.DataFrame(rows)
>>> result = sp.event_study(df, y="y", treat_time="g", time="time",
...                         unit="unit", window=(-3, 3))
>>> report = sp.pretrends_summary(result)  # also prints the report
>>> bool(isinstance(report, str))
True

pretrends_test

pretrends_test(result: Any, type: str = 'wald', alpha: float = 0.05) -> Dict[str, Any]

Joint test of pre-treatment coefficients.

Tests H0: beta_pre = 0 (all pre-treatment event-study coefficients are jointly zero).

Parameters:

Name	Type	Description	Default
result	CausalResult	Event-study result containing pre-treatment estimates and SEs.	required
type	``'wald'`` or ``'f'``	'wald': chi-squared test statistic. 'f': scaled F-statistic (requires df_resid in model_info).	'wald'
alpha	float	Significance level.	0.05

Returns:

Type	Description
dict	Keys: statistic, pvalue, df, type, reject, interpretation.

References

Standard Wald test; see Roth (2022) for caveats on interpretation.

Examples:

>>> import statspai as sp, numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(80):
...     cohort = 4 if i < 40 else 0          # 0 = never treated
...     for t in range(8):
...         post = cohort > 0 and t >= cohort
...         y = 0.3 * t + (2.0 if post else 0.0) + (i % 5) + rng.normal()
...         rows.append((i, t, cohort, y))
>>> df = pd.DataFrame(rows, columns=["id", "t", "cohort", "y"])
>>> es = sp.event_study(df, y="y", treat_time="cohort", time="t", unit="id")
>>> sp.pretrends_test(es)

sensitivity_rr

sensitivity_rr(result: Any, Mbar: Optional[Union[ndarray, List[float]]] = None, method: str = 'C-LF', alpha: float = 0.05, n_grid: int = 20) -> SensitivityResult

Rambachan & Roth (2023) honest confidence intervals.

Computes confidence intervals for the ATT that are valid under bounded departures from parallel trends. The conditional linear-in-relative-time (C-LF) restriction assumes the post-treatment violation is bounded by a linear extrapolation of the pre-trend plus an additional M-bar of slack.

Parameters:

Name	Type	Description	Default
result	CausalResult	Event-study result with pre- and post-treatment estimates.	required
Mbar	array - like	Grid of M-bar values. Default: np.linspace(0, 3 * max_pre_slope, n_grid).	None
method	``'C-LF'``	Extrapolation method. Currently only C-LF is implemented.	'C-LF'
alpha	float	Significance level.	0.05
n_grid	int	Number of grid points when Mbar is not supplied.	20

Returns:

Type	Description
SensitivityResult

Notes

.. versionchanged:: next The pre-period trend is now fitted by generalised least squares using the full pre-period covariance from result.model_info['vcv_pre'] when it is available (sp.event_study supplies it). Previously the fit always used diagonal 1/se**2 weights, i.e. it assumed the pre-treatment event-study coefficients were mutually independent -- they are not, since they share the omitted reference period and the unit/time fixed effects. Breakdown Mbar values therefore move slightly relative to earlier releases. When no covariance is available the diagonal fallback is still used, but it now warns loudly. Object with .summary(), .plot(), .mbar_grid, .ci_lower, .ci_upper, .breakdown_mbar.

References

Rambachan, A. & Roth, J. (2023). A More Credible Approach to Parallel Trends. Review of Economic Studies, 90(5), 2555--2591. [@rambachan2023more]

Examples:

>>> import statspai as sp, numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(80):
...     cohort = 4 if i < 40 else 0          # 0 = never treated
...     for t in range(8):
...         post = cohort > 0 and t >= cohort
...         y = 0.3 * t + (2.0 if post else 0.0) + (i % 5) + rng.normal()
...         rows.append((i, t, cohort, y))
>>> df = pd.DataFrame(rows, columns=["id", "t", "cohort", "y"])
>>> es = sp.event_study(df, y="y", treat_time="cohort", time="t", unit="id")
>>> sens = sp.sensitivity_rr(es, Mbar=[0, 0.01, 0.02, 0.05])
>>> sens.summary()

cs_report

cs_report(data_or_result: Union[DataFrame, CausalResult], y: Optional[str] = None, g: Optional[str] = None, t: Optional[str] = None, i: Optional[str] = None, x: Optional[List[str]] = None, weights: Optional[str] = None, estimator: str = 'dr', control_group: str = 'nevertreated', anticipation: int = 0, alpha: float = 0.05, n_boot: int = 1000, random_state: Optional[int] = 0, min_e: float = -inf, max_e: float = inf, rr_method: str = 'smoothness', balance: bool = True, triangulate: bool = True, functional_form: bool = True, verbose: bool = True, save_to: Optional[str] = None) -> CSReport

One-call staggered-DID workflow following Baker et al. (2026).

Runs the eight-step "forward-engineering" recipe the practitioner's guide closes with (§6): define the target parameter, state the identification assumption and generate evidence about it, choose an estimation method, declare the inference frame, estimate, and probe sensitivity. report.forward_engineering_checklist() returns the steps and what was actually run for each.

Parameters:

Name	Type	Description	Default
data_or_result	DataFrame | CausalResult	Either a long-format panel (then y, g, t, i are required and :func:callaway_santanna is run first), or an already-fitted :func:callaway_santanna result.	required
y	str	Outcome / cohort / time / unit id columns (required when data_or_result is a DataFrame).	None
g	str	Outcome / cohort / time / unit id columns (required when data_or_result is a DataFrame).	None
t	str	Outcome / cohort / time / unit id columns (required when data_or_result is a DataFrame).	None
i	str	Outcome / cohort / time / unit id columns (required when data_or_result is a DataFrame).	None
x	list of str	Covariates for conditional parallel trends.	None
weights	str	Unit weights ω. Changes the target parameter, not just precision — see :func:statspai.did.callaway_santanna.	None
balance	bool	Run :func:did_balance on x (step 2 evidence: covariate balance in levels and changes). Ignored when x is empty or when a pre-fitted result is passed without the source panel.	True
triangulate	bool	Also fit the other covariate strategies (regression adjustment, IPW, doubly robust) and report them side by side, as in the paper's Table 7 / Figure 4. Only meaningful with covariates.	True
functional_form	bool	Run :func:functional_form_test (Roth and Sant'Anna 2023b): is parallel trends insensitive to how the outcome is scaled?	True
estimator	('dr', 'ipw', 'reg')		'dr'
control_group	('nevertreated', 'notyettreated')		'nevertreated'
anticipation	int		0
alpha	float		0.05
n_boot	int	Multiplier-bootstrap replications for :func:aggte.	1000
random_state	int	Seed for the bootstrap (set to None for non-reproducibility).	0
min_e	float	Event-time window passed to the dynamic aggregation.	(-inf, inf)
max_e	float	Event-time window passed to the dynamic aggregation.	(-inf, inf)
rr_method	('smoothness', 'relative_magnitude')	Sensitivity restriction handed to :func:breakdown_m.	'smoothness'
verbose	bool	If True, print the report before returning.	True
save_to	str	When set, treats the value as a path prefix and writes the report in every supported format in one call: <prefix>.txt — fixed-width plain-text report <prefix>.md — GitHub-flavoured Markdown <prefix>.tex — booktabs LaTeX fragment <prefix>.xlsx — multi-sheet workbook <prefix>.png — 2×2 summary figure (only if matplotlib is installed; silently skipped otherwise) Missing parent directories are created on the fly.	None

Returns:

Type	Description
CSReport	Structured container; call .to_text() to re-render.

Examples:

>>> import statspai as sp, numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for unit in range(80):
...     g = 4 if unit < 40 else 0             # cohort; 0 = never treated
...     for t in range(8):
...         post = g > 0 and t >= g
...         y = 0.3 * t + (2.0 if post else 0.0) + (unit % 5) + rng.normal()
...         rows.append((unit, t, g, y))
>>> df = pd.DataFrame(rows, columns=["id", "t", "g", "y"])
>>> rpt = sp.cs_report(df, y="y", g="g", t="t", i="id", random_state=42)
>>> rpt.dynamic           # event-study DataFrame w/ uniform bands

parallel_trends_robustness

parallel_trends_robustness(result: Any, m_grid: Optional[Sequence[float]] = None, families: Sequence[str] = ('SD', 'RM'), alpha: float = 0.05, e: int = 0, delta: Optional[ndarray] = None) -> ParallelTrendsRobustnessResult

Run the full parallel-trends robustness pipeline on a DiD result.

Chains the joint pre-trend test, the Roth (2022) power calculation for that test, and Rambachan & Roth (2023) honest confidence intervals (plus their breakdown value Mbar*) for each requested restriction family, and reduces the whole thing to a one-line verdict.

Parameters:

Name	Type	Description	Default
result	CausalResult	A fitted DiD/event-study result carrying event-study estimates in result.model_info['event_study'] (e.g. from sp.event_study, sp.callaway_santanna, sp.sun_abraham). .. note:: If the result does not carry a pre-period covariance matrix in model_info['vcv_pre'], the pre-trend test and power fall back to assuming the pre-period coefficients are mutually independent and warn loudly. sp.event_study computes the full cluster-robust covariance; pass expose_pre_vcov=True to it to have this pipeline use the correct covariance instead of the diagonal fallback.	required
m_grid	sequence of float	Grid of violation magnitudes. Default: the honest_did default, multiples of the standard error at e.	None
families	sequence of str	Restriction families. "SD" maps to honest_did's method='smoothness' (bounded second differences); "RM" maps to method='relative_magnitude'.	``("SD", "RM")``
alpha	float	Significance level.	0.05
e	int	Relative time whose effect the sensitivity analysis targets.	0
delta	array - like	Hypothesised pre-trend violation passed to pretrends_power.	None

Returns:

Type	Description
ParallelTrendsRobustnessResult	With .summary(), .plot(), .to_latex(), and the power_table / ci_grid / breakdown / verdict fields.

References

Rambachan, A. and Roth, J. (2023). A More Credible Approach to Parallel Trends. Review of Economic Studies, 90(5), 2555-2591. [@rambachan2023more]

Roth, J. (2022). Pretest with Caution: Event-Study Estimates after Testing for Parallel Trends. AER: Insights, 4(3), 305-322. [@roth2022pretest]

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=80, n_periods=8, staggered=False, seed=0)
>>> df["first_treat"] = df["first_treat"].fillna(0)
>>> es = sp.event_study(df, y="y", treat_time="first_treat",
...                     time="time", unit="unit", window=(-3, 3))
>>> rob = sp.parallel_trends_robustness(es, families=("SD", "RM"))
>>> sorted(rob.breakdown)
['RM', 'SD']
>>> bool("Mbar" in rob.verdict or "not robust" in rob.verdict)
True

spillover_did

spillover_did(data: DataFrame, y: str, *, unit: str, time: str, cohort: str, coords: Optional[Sequence[str]] = None, distances: Optional[ndarray] = None, ring_edges: Sequence[float] = (0.0, 1.0), never_value: Any = 0, alpha: float = 0.05) -> SpilloverRingResult

Direct and spillover effects with distance-banded control groups.

Parameters:

Name	Type	Description	Default
data	DataFrame	Long-format panel.	required
y	str	Outcome column.	required
unit	str	Unit identifier, period, and first-treatment period (never_value marks never-treated units).	required
time	str	Unit identifier, period, and first-treatment period (never_value marks never-treated units).	required
cohort	str	Unit identifier, period, and first-treatment period (never_value marks never-treated units).	required
coords	sequence of str	Two columns giving each unit's position. Distances are Euclidean in whatever units these are. Ignored when distances is given.	None
distances	ndarray	Pre-computed (n_units, n_units) distance matrix, ordered by sorted unit id. Use this for great-circle or network distances.	None
ring_edges	sequence of float	Ring boundaries. (0, 1, 2) makes two rings, (0, 1] and (1, 2]; untreated units beyond the last edge are the clean controls. Rings are not estimated where no untreated unit falls in them.	(0.0, 1.0)
never_value	any	Value in cohort marking never-treated units.	0
alpha	float		0.05

Returns:

Type	Description
SpilloverRingResult

Notes

Standard errors are the influence-function form for a difference of group means, aggregated across (cohort, period) cells so the shared control units are accounted for. There is no reference implementation to pin them against; see the module docstring on how correctness is established instead.

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x = rng.uniform(0, 10, n); yc = rng.uniform(0, 10, n)
>>> treated = (x < 3) & (yc < 3)
>>> rows = []
>>> for i in range(n):
...     for t in (1, 2):
...         rows.append((i, t, 2 if treated[i] else 0, x[i], yc[i],
...                      rng.normal() + (1.0 if (treated[i] and t == 2) else 0)))
>>> df = pd.DataFrame(rows, columns=["i", "t", "g", "x", "y2", "y"])
>>> res = sp.spillover_did(df, y="y", unit="i", time="t", cohort="g",
...                        coords=["x", "y2"], ring_edges=(0.0, 2.0))
>>> bool(np.isfinite(res.direct))
True

References

butts2021difference

did_report

did_report(data: DataFrame, y: str, time: str, first_treat: str, group: str, save_to: str, methods: Union[str, List[str]] = 'auto', controls: Optional[List[str]] = None, cluster: Optional[str] = None, alpha: float = 0.05, include_sensitivity: bool = True, plot_sort_by: Optional[str] = 'estimate', verbose: bool = False) -> CausalResult

DID report bundle: fits selected methods and writes report artifacts.

Writes the following files to save_to:

• did_summary.txt : text dump of result.summary().
• did_summary.md : GitHub-Flavoured Markdown table.
• did_summary.tex : LaTeX booktabs fragment.
• did_summary.png : forest plot (requires matplotlib).
• did_summary.json : detail table + model_info in JSON.

Parameters:

Name	Type	Description	Default
data	DataFrame	Same as :func:did_summary.	required
y	DataFrame	Same as :func:did_summary.	required
time	DataFrame	Same as :func:did_summary.	required
first_treat	DataFrame	Same as :func:did_summary.	required
group	DataFrame	Same as :func:did_summary.	required
methods	DataFrame	Same as :func:did_summary.	required
controls	DataFrame	Same as :func:did_summary.	required
cluster	DataFrame	Same as :func:did_summary.	required
alpha	DataFrame	Same as :func:did_summary.	required
save_to	str	Directory path. Created if it does not exist.	required
include_sensitivity	bool	Whether to run Rambachan-Roth breakdown M*. Defaults to True in did_report (vs False in did_summary) because a report is expected to be comprehensive.	True
plot_sort_by	(estimate, None)	Sort the forest plot by estimate ascending.	'estimate'
verbose	bool	Print progress lines.	False

Returns:

Type	Description
CausalResult	The underlying :func:did_summary output. All side-effect files are written to save_to as a bundle.

Examples:

>>> import tempfile
>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True,
...                 seed=0)
>>> out_dir = tempfile.mkdtemp()
>>> res = sp.did_report(df, y='y', time='time',
...                     first_treat='first_treat', group='unit',
...                     save_to=out_dir, methods=['cs', 'bjs'],
...                     include_sensitivity=False)
>>> import os
>>> 'did_summary.md' in os.listdir(out_dir)
True

did_summary

did_summary(data: DataFrame, y: str, time: str, first_treat: str, group: str, methods: Union[str, List[str]] = 'auto', controls: Optional[List[str]] = None, cluster: Optional[str] = None, alpha: float = 0.05, include_sensitivity: bool = False, verbose: bool = False) -> CausalResult

One-call method-robustness comparison for staggered DID.

Fits every requested estimator to the same data and returns a single :class:CausalResult whose detail attribute is a tidy comparison table — one row per method, columns (method, estimator, estimate, se, pvalue, ci_low, ci_high, n_obs, note).

Parameters:

Name	Type	Description	Default
data	DataFrame	Panel dataset (long format).	required
y	str	Outcome variable.	required
time	str	Time / period variable (integer-valued).	required
first_treat	str	First-treatment period per unit; NaN (or 0) for never-treated.	required
group	str	Unit identifier.	required
methods	str or list of str	Methods to run. Valid keys: 'cs', 'sa', 'bjs', 'etwfe', 'stacked', or 'all' / 'auto' for all.	``'auto'``
controls	list of str	Time-varying covariates passed to methods that support them.	None
cluster	str	Cluster variable for SE (defaults to group in each sub-method).	None
alpha	float	Significance level for confidence intervals.	0.05
include_sensitivity	bool	If True and 'cs' is among the methods fit, compute the Rambachan–Roth (2023) breakdown M* — the largest relative violation of parallel trends under which the treatment effect is still significantly different from zero. The value is added to model_info['breakdown_m'] and to the breakdown_m column of detail (CS row only; other methods leave NaN).	False
verbose	bool	Print progress for each method.	False

Returns:

Type	Description
CausalResult	estimate : mean of successfully-fit overall ATTs. se : standard deviation across methods (not a standard error — a crude dispersion measure). detail : comparison DataFrame described above. model_info : {'methods_requested': [...], 'methods_fit': [...], 'methods_failed': {name: error_msg, ...}}.

Notes

Each method's overall ATT has slightly different interpretation:

• CS aggte(type='simple') averages ATT(g, t) for post-treatment :math:t \geq g, weighted by cohort size × exposure length.
• ETWFE reports the R/Stata treated-observation-weighted simple ATT.
• SA / BJS / Stacked report their estimator-specific overall ATT aggregations.

Differences across methods are informative about heterogeneity, model specification, and the sensitivity of conclusions to the estimator choice. Large disagreement is a red flag that deserves further investigation (e.g., via sp.bacon_decomposition or sp.honest_did).

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=200, n_periods=10, staggered=True, seed=0)
>>> out = sp.did_summary(df, y='y', time='time',
...                      first_treat='first_treat', group='unit')
>>> out.summary()
>>> print(out.detail[['method', 'estimate', 'se', 'pvalue']])

did_summary_to_latex

did_summary_to_latex(result: CausalResult, digits: int = 4, include_ci: bool = True, include_breakdown: bool = True, label: str = 'tab:did_summary', caption: str = 'DID method-robustness summary.') -> str

Render a :func:did_summary result as a LaTeX booktabs table.

Parameters:

Name	Type	Description	Default
result	CausalResult	Output of :func:did_summary.	required
digits	int	Decimal precision.	4
include_ci	bool	Include the 95 % CI column.	True
include_breakdown	bool	Include the Rambachan-Roth breakdown M* column when sensitivity was requested.	True
label	str	LaTeX label for the table.	``'tab:did_summary'``
caption	str	LaTeX caption.	``'DID method-robustness summary.'``

Returns:

Type	Description
str	Full \begin{table} ... \end{table} block using the booktabs package (\toprule, \midrule, \bottomrule).

Notes

Requires \usepackage{booktabs} in the LaTeX preamble.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True,
...                 seed=0)
>>> out = sp.did_summary(df, y='y', time='time',
...                      first_treat='first_treat', group='unit',
...                      methods=['cs', 'bjs'])
>>> tex = sp.did_summary_to_latex(out)
>>> isinstance(tex, str)
True
>>> tex.splitlines()[0].startswith('\\begin{table}')
True

did_summary_to_markdown

did_summary_to_markdown(result: CausalResult, digits: int = 4, include_ci: bool = True, include_breakdown: bool = True) -> str

Render a :func:did_summary result as a GitHub-Flavoured Markdown table.

Columns shown (in order): Method, Estimate, SE, 95 % CI, p-value, and optionally Breakdown M* (when sensitivity was requested).

Parameters:

Name	Type	Description	Default
result	CausalResult	Output of :func:did_summary.	required
digits	int	Decimal precision for numeric columns.	4
include_ci	bool	Include the 95 % CI column.	True
include_breakdown	bool	Include the Rambachan-Roth breakdown M* column (CS row only, blank for others). Ignored if sensitivity was not requested.	True

Returns:

Type	Description
str	Multi-line Markdown table, ready to paste into notebooks or PRs.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True,
...                 seed=0)
>>> out = sp.did_summary(df, y='y', time='time',
...                      first_treat='first_treat', group='unit',
...                      methods=['cs', 'bjs'])
>>> md = sp.did_summary_to_markdown(out)
>>> isinstance(md, str)
True
>>> md.splitlines()[0].startswith('| Method')
True

drdid

drdid(data: DataFrame, y: str, group: str, time: str, covariates: Optional[List[str]] = None, method: str = 'imp', alpha: float = 0.05, n_boot: int = 500, random_state: Optional[int] = None, seed: Optional[int] = None, id: Optional[str] = None, *, est_method: str = 'dr', normalized: bool = True, locally_efficient: bool = True, weights: Optional[str] = None, trim_level: float = 0.995) -> CausalResult

Doubly Robust Difference-in-Differences (Sant'Anna & Zhao 2020).

Combines outcome regression with inverse probability weighting for 2×2 DID with covariates. Consistent if either the outcome model or the propensity score model is correctly specified.

Parameters:

Name	Type	Description	Default
data	DataFrame	Dataset with one row per unit-period in 2x2 design.	required
y	str	Outcome variable.	required
group	str	Binary treatment-group indicator (1 = treated, 0 = control).	required
time	str	Binary time indicator (1 = post, 0 = pre).	required
covariates	list of str	Covariate names. If None, runs a simple (un-adjusted) DID.	None
method	str	Which nuisance estimators the DR variants use, following R DRDID::drdid(estMethod=): 'imp' — "improved": inverse probability tilting for the propensity score and propensity-odds-weighted least squares for the outcome model, so the DR moment is Neyman-orthogonal by construction and the influence function carries no nuisance-estimation terms. 'trad' — plain logit + OLS, with the estimation effects propagated explicitly. Only consulted when est_method='dr'. .. versionchanged:: 1.23.0 On repeated cross-sections this argument used to do nothing: both settings returned DRDID::drdid_rc1 regardless. It now selects drdid_imp_rc / drdid_rc as documented.	``'imp'``
est_method	(dr, ipw, reg, twfe)	Which estimator family to use. With method and the two flags below this reaches all 14 estimators of R DRDID 1.2.3: ============ ======================= ========================== est_method panel (id= given) repeated cross-sections ============ ======================= ========================== 'dr' drdid_[imp_]panel drdid_[imp_]rc[1] 'ipw' [std_]ipw_did_panel [std_]ipw_did_rc 'reg' reg_did_panel reg_did_rc 'twfe' twfe_did_panel twfe_did_rc ============ ======================= ========================== 'twfe' is included for comparison, not as a recommendation: with covariates it is exactly the specification Sant'Anna & Zhao (2020) and Caetano & Callaway (2024) warn about. .. versionadded:: 1.23.0	'dr'
normalized	bool	Only for est_method='ipw'. True gives the Hájek-normalised estimator (std_ipw_did_*), where the control arm is divided by its own realised weight mass. False gives Abadie (2005) (ipw_did_*), where both arms share the denominator E[D]. .. versionadded:: 1.23.0	True
locally_efficient	bool	Only for est_method='dr' on repeated cross-sections. False drops the terms that attain the semiparametric efficiency bound, giving drdid_rc1 / drdid_imp_rc1. Both are consistent; the rc1 variants avoid fitting outcome regressions on the treated cells. .. versionadded:: 1.23.0	True
weights	str	Column of observation weights (R DRDID's i.weights), renormalised to mean one. .. versionadded:: 1.23.0	None
trim_level	float	Drop control units whose estimated propensity score reaches this cutoff, matching DRDID's trim.level. Pass 1.0 to disable. .. versionadded:: 1.23.0	0.995
alpha	float	Significance level.	0.05
n_boot	int	Number of bootstrap replications for inference.	500
random_state	int	Seed for bootstrap reproducibility.	None
id	str	Unit identifier for a true two-period panel. When supplied, the improved estimator uses the Sant'Anna-Zhao panel formula with calibrated propensity scores and influence-function standard errors, matching DRDID::drdid_imp_panel and Stata drdid, drimp.	None

Returns:

Type	Description
CausalResult	estimate is the DR-DID ATT. detail contains influence-function diagnostics.

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(42)
>>> n = 500
>>> G = rng.integers(0, 2, n)
>>> T = rng.integers(0, 2, n)
>>> x = rng.normal(0, 1, n)
>>> y_val = 1 + 0.5*x + 2*G + 3*T + 4*G*T + rng.normal(0, 1, n)
>>> df = pd.DataFrame({'y': y_val, 'treated': G, 'post': T, 'x': x})
>>> result = sp.drdid(df, y='y', group='treated', time='post',
...                   covariates=['x'])
>>> abs(result.estimate - 4.0) < 1.0
True

etwfe

etwfe(data: DataFrame, y: str, group: str, time: str, first_treat: str, controls: Optional[List[str]] = None, cluster: Optional[str] = None, alpha: float = 0.05, xvar: Optional[Any] = None, panel: bool = True, cgroup: str = 'notyet', family: Optional[str] = None) -> CausalResult

Public sp.etwfe entry point — see _dispatch_etwfe_impl for the full docstring on options and behaviour.

family selects the outcome model. None/'gaussian' (default) is the historical linear ETWFE and is unchanged. 'poisson' and 'logit' fit Wooldridge (2023) nonlinear ETWFE by maximum likelihood and report the average marginal effect on the response scale, matching R etwfe::emfx(type='simple') — so a Poisson fit returns an effect in counts, not log points. The nonlinear branch uses not-yet-treated identification and does not currently accept xvar, panel=False, or cgroup='nevertreated'.

Thin wrapper around the 4-branch dispatcher (panel-with-xvar / panel-never-only / panel-notyet / repeated-cross-section) that attaches a :class:Provenance record to the returned result so downstream replication_pack / Quarto appendix / table footers can pick up the call without each branch having to opt in.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=120, n_periods=8, staggered=True, seed=42)
>>> res = sp.etwfe(df, y='y', group='unit', time='time',
...                first_treat='first_treat')
>>> res.estimate > 0  # R/Stata simple ATT (true effect 0.5)
True
>>> res.detail is not None  # cohort-specific ATTs
True

See Also

wooldridge_did : The same saturated cohort x period regression reported under a different headline aggregation — the cohort-size-weighted average of ATT(g) under a never-treated comparison group, i.e. R etwfe::etwfe(cgroup='never') + emfx(type='group'). sp.etwfe instead reports the treated-observation-weighted simple ATT that R emfx(type='simple') and Stata jwdid, estat simple print, under the not-yet-treated comparison group by default. The two are separate documented aggregations of one estimator, not two estimators, and they differ materially on the 17_etwfe parity bytes: the sp.etwfe default is 15.9% from the sp.wooldridge_did headline, and sp.etwfe(cgroup='nevertreated') — same comparison group, different weights — is still 10.5% from it. Pick the one your write-up claims.

etwfe_emfx

etwfe_emfx(result: CausalResult, type: str = 'simple', alpha: float = 0.05, include_leads: bool = False, weighting: str = 'treated') -> CausalResult

R etwfe::emfx-style aggregated marginal effects for an ETWFE fit.

Takes the result of :func:etwfe / :func:wooldridge_did and returns one of four aggregations used in applied work:

================ ======================================================== type Aggregation ================ ======================================================== 'simple' Overall treated-observation-weighted ATT (same as result.estimate for current sp.etwfe results). 'group' ATT per treatment cohort g. 'event' ATT per event time e = t - g, averaged across cohorts. 'calendar' ATT per calendar time t, averaged across cohorts for which t >= g. ================ ========================================================

Parameters:

Name	Type	Description	Default
result	CausalResult	Output of :func:etwfe or :func:wooldridge_did.	required
type	(simple, group, event, calendar)	Aggregation type.	'simple'
alpha	float	Significance level for confidence intervals.	0.05
include_leads	bool	For type='event' and type='calendar', whether to include pre-treatment relative times (rel_time < 0) in the output. These coefficients identify pre-trends and are informative for parallel-trends inspection. Default False for backward compatibility with earlier versions; set True for full event-study output matching the R etwfe::emfx(type='event') default. rel_time = -1 is always the reference category and is excluded.	False
weighting	(cohort, treated)	Aggregation weights for cohort-level marginal effects. 'treated' uses the number of treated post-period observations, matching R etwfe::emfx(type='simple') and Stata jwdid, estat simple. 'cohort' preserves the historical StatsPAI cohort-share weighting.	'cohort'

Returns:

Type	Description
CausalResult	estimate is the overall ATT (for type='simple') or the mean of the sub-aggregation (for the other types). detail contains one row per group/event-time/calendar-time with (estimate, se, pvalue, ci_low, ci_high).

Notes

For 'event' and 'calendar', the reported SE treats the per-cohort coefficients as independent — a standard approximation that matches R etwfe's default under classical vcov. Cluster-robust or fully-general SEs require the full regression vcov, which can be requested via sp.wooldridge_did + the model_info matrix in a future release.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=200, n_periods=10, staggered=True)
>>> fit = sp.etwfe(df, y='y', time='time',
...                first_treat='first_treat', group='unit')
>>> evt = sp.etwfe_emfx(fit, type='event')
>>> print(evt.detail)   # ATT by event time
>>> grp = sp.etwfe_emfx(fit, type='group')
>>> cal = sp.etwfe_emfx(fit, type='calendar')

twfe_decomposition

twfe_decomposition(data: DataFrame, y: str, group: str, time: str, first_treat: str, alpha: float = 0.05) -> CausalResult

TWFE decomposition: Goodman-Bacon (2021) + de Chaisemartin–D'Haultfoeuille weights.

Decomposes the standard two-way fixed effects estimator into all pairwise 2×2 DID comparisons, showing the weight and estimate for each. Also computes de Chaisemartin–D'Haultfoeuille (2020) weights to diagnose whether negative weights are present.

Parameters:

Name	Type	Description	Default
data	DataFrame	Panel dataset in long format.	required
y	str	Outcome variable.	required
group	str	Unit identifier.	required
time	str	Time period variable.	required
first_treat	str	Treatment timing column (NaN or 0 for never-treated).	required
alpha	float	Significance level.	0.05

Returns:

Type	Description
CausalResult	detail DataFrame has columns: type, treated_cohort, control_cohort, estimate, weight, weighted_est. model_info includes summary statistics and dCDH weights.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=200, n_periods=8, staggered=True, seed=0)
>>> result = sp.twfe_decomposition(df, y='y', group='unit',
...                                time='time',
...                                first_treat='first_treat')
>>> bool('weight' in result.detail.columns)  # 2x2 decomposition weights
True

did

did(data: DataFrame, y: str, treat: str, time: str, id: Optional[str] = None, covariates: Optional[Any] = None, method: str = 'auto', estimator: str = 'dr', control_group: str = 'nevertreated', base_period: str = 'universal', cluster: Optional[str] = None, robust: bool = True, alpha: float = 0.05, weights: Optional[str] = None, subgroup: Optional[str] = None, treat_unit: Any = None, treat_time: Any = None, se_method: str = 'placebo', aggregation: Optional[str] = None, n_boot: int = 1000, random_state: Optional[int] = None, panel: bool = True, allow_unbalanced_panel: bool = False, anticipation: int = 0, vce: Optional[str] = None, wild_reps: int = 999, wild_weight_type: str = 'rademacher', seed: Optional[int] = None, **kwargs: Any) -> CausalResult

Difference-in-Differences estimation.

Unified entry point that auto-detects design type and dispatches to the appropriate estimator.

Parameters:

Name	Type	Description	Default
data	DataFrame	Input dataset.	required
y	str	Outcome variable name.	required
treat	str	Treatment variable. The column semantics depend on the design — one of the most common pitfalls in DID: 2×2 DID (no id, two periods): a 0/1 group indicator (treated vs. control), identical to a standard DID dummy. Staggered DID (Callaway–Sant'Anna, Sun–Abraham, SDID, Borusyak–Jaravel–Spiess, de Chaisemartin–D'Haultfœuille, Wooldridge etwfe): the first treatment period for each unit (aka the cohort / g-variable in R's did package). Never-treated units must have 0 (or NaN), not 1. A plain 0/1 indicator will silently be interpreted as "everyone was first treated in period 1," producing nonsense estimates. If you only have a 0/1 treated column, construct first_treat per unit and broadcast it to all of that unit's rows:: # First treated year per unit (NaN for never-treated units) first = (df.loc[df['treated'] == 1] .groupby('id')['year'].min()) df['first_treat'] = df['id'].map(first).fillna(0).astype(int) sp.did(df, y='y', treat='first_treat', time='year', id='id')	required
time	str	Time period variable.	required
id	str	Unit identifier. Required for staggered DID and SDID.	None
covariates	list of str	Covariate names for conditional parallel trends / controls.	None
method	str	'auto' — 2×2 if id is None and treatment is binary, else Callaway-Sant'Anna. '2x2' — classic two-period, two-group DID. 'ddd' — triple differences (requires subgroup). 'callaway_santanna' or 'cs' — staggered DID. 'sun_abraham', 'sa', or 'sunab' — IW event study. 'bjs' or 'did_imputation' — Borusyak-Jaravel-Spiess imputation DID. 'sdid' — synthetic DID (Arkhangelsky et al. 2021).	'auto'
estimator	str	For staggered DID: 'dr' (doubly robust), 'ipw', 'reg'.	'dr'
control_group	str	For staggered DID: 'nevertreated' or 'notyettreated'.	'nevertreated'
base_period	str	For staggered DID: 'universal' or 'varying'.	'universal'
cluster	str	Cluster variable for standard errors.	None
robust	bool	HC1 robust standard errors (2×2 / DDD only).	True
alpha	float	Significance level for confidence intervals.	0.05
weights	str	Column name for analytical weights (e.g. population weights). Supported for '2x2', 'ddd', and event study methods. Equivalent to Stata's [aweight=...].	None
subgroup	str	For DDD: binary affected-subgroup indicator.	None
treat_unit	optional	For SDID: treated unit(s).	None
treat_time	optional	For SDID: treatment time.	None
se_method	str	For SDID: 'placebo', 'bootstrap', or 'jackknife'.	'placebo'
aggregation	str	When set and method is Callaway–Sant'Anna, the raw ATT(g,t) result is passed through :func:aggte with type=aggregation ('simple', 'dynamic', 'group', or 'calendar'), delivering the aggregated ATT with Mammen multiplier-bootstrap uniform confidence bands in a single call.	None
n_boot	int	Bootstrap replications for the multiplier bootstrap when aggregation is set.	1000
random_state	int	Seed for the multiplier bootstrap.	None
panel	bool	Forwarded to :func:callaway_santanna; set panel=False for repeated cross-sections.	True
allow_unbalanced_panel	bool	Forwarded to :func:callaway_santanna. True keeps units that are missing periods by switching to the repeated-cross-section estimators (R did::att_gt(allow_unbalanced_panel = TRUE)) instead of letting them drop out cell by cell. Only consulted on the staggered (method='cs') path. .. versionadded:: 1.23.0	False
anticipation	int	Forwarded to :func:callaway_santanna.	0

Returns:

Type	Description
CausalResult	Estimation results with .summary(), .plot(), .to_latex(), .cite() methods.

References

Callaway, B. and Sant'Anna, P. H. C. (2021). Difference-in-differences with multiple time periods. Journal of Econometrics, 225(2), 200-230. [@callaway2021difference]

Sant'Anna, P. H. C. and Zhao, J. (2020). Doubly Robust Difference-in-Differences Estimators. Journal of Econometrics, 219(1), 101-122. [@santanna2020doubly]

Goodman-Bacon, A. (2021). Difference-in-differences with variation in treatment timing. Journal of Econometrics, 225(2), 254-277. [@goodmanbacon2021difference]

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)

Classic 2x2 DID (one binary treatment, two periods):

>>> n = 200
>>> df = pd.DataFrame({
...     'treated': np.repeat(rng.integers(0, 2, n), 2),
...     'post': np.tile([0, 1], n),
... })
>>> df['wage'] = (1.0 + 2.0 * df['treated'] + 1.5 * df['post']
...               + 3.0 * df['treated'] * df['post']
...               + rng.normal(0, 1, len(df)))
>>> result = sp.did(df, y='wage', treat='treated', time='post')
>>> bool(result.estimate > 0)
True

Triple Differences (a third dimension via subgroup):

>>> df['low_wage'] = np.repeat(rng.integers(0, 2, n), 2)
>>> ddd = sp.did(df, y='wage', treat='treated', time='post',
...              method='ddd', subgroup='low_wage')

Staggered DID (cohort column gives first-treatment time; 0 = never):

>>> rows = []
>>> for u in range(30):
...     first = int(rng.choice([2003, 2005, 0]))
...     for yr in range(2000, 2008):
...         on = 1 if (first != 0 and yr >= first) else 0
...         rows.append({'unit': u, 'year': yr, 'first_treat': first,
...                      'y': 5 + 2.0 * on + rng.normal(0, 1)})
>>> panel = pd.DataFrame(rows)
>>> staggered = sp.did(panel, y='y', treat='first_treat',
...                    time='year', id='unit')
>>> bool(staggered.estimate == staggered.estimate)  # finite
True