`statspai.datasets`¶

datasets ¶

Canonical econometrics datasets with documented expected estimates.

This subpackage provides deterministic, reproducible datasets used throughout the causal-inference literature, consolidated under a single import path sp.datasets:

import statspai as sp df = sp.datasets.nsw_lalonde() # real extract, no network df.shape (614, 11) sp.datasets.nsw_lalonde(simulated=True).attrs['expected_experimental_att'] 1794

One rule: if StatsPAI ships the real published data for a dataset, a bare call returns it. simulated=True asks for the calibrated replica instead. list_datasets()['source'] says which you get. The one exception is nsw_dw(), whose bare call is the replica (parity fixtures are built on it); nsw_dw(simulated=False) is the real sample.

Shipped datasets never touch the network. A few real datasets that are too large for the wheel, or whose licence does not clearly allow redistribution, are fetched: the first call downloads the file from its canonical public URL, checks a pinned SHA-256 and caches it under $STATSPAI_DATA_DIR (default ~/.cache/statspai/datasets); later calls work offline. list_datasets()['delivery'] says which is which. Nothing downloads at import time.

Each function returns a pd.DataFrame with:

A fully-deterministic DGP (fixed seed, BSD/MIT redistributable).
df.attrs containing:
'paper' — original paper citation
'expected_*' — theoretically-anchored estimates (from the published paper on the ORIGINAL data, not necessarily the simulated replica)
'notes' — what to be careful about when using this replica

The simulated replicas are designed to match the structure and summary statistics of the original datasets. For numerical R / Stata parity against the original data, see tests/external_parity/PUBLISHED_REFERENCE_VALUES.md.

Available datasets

DID / panel mpdta() — Callaway-Sant'Anna teen employment (replica; simulated=False fetches the original) teen_employment() — alias of mpdta() castle_doctrine() — Cheng-Hoekstra castle-doctrine panel (real) bacon_divorce() — Goodman-Bacon divorce-reform panel (real, fetched) jel_medicaid() — Baker et al. JEL county mortality (real, fetched) lalonde_panel() — NSW + CPS + PSID 1975/1978 panel of Sant'Anna-Zhao (real, fetched)

RD lee_2008_senate() — US Senate RD (Lee 2008 design; CFT 2015 data)

IV card_1995() — IV returns-to-schooling angrist_krueger_1991() — quarter-of-birth IV

Matching / SOO nsw_lalonde() — LaLonde NSW job training (real MatchIt extract n=614) nsw_dw() — simulated replica of the Dehejia-Wahba NSW + PSID comparison; the real sample with simulated=False

Synthetic control california_prop99() — ADH tobacco (re-exported from synth) basque_terrorism() — Abadie-Gardeazabal (re-exported) german_reunification() — ADH 2015 (re-exported)

Heterogeneous effects in panels (simulated) currency_union_panel() — dyadic trade panel with staggered currency-union adoption and known pair-level effects (layout after Aytug 2026)

Public health / epidemiology (REAL data) nhefs() — Hernán-Robins What If NHEFS (g-methods canon: quit-smoking → weight / mortality) load_nhefs() — alias of nhefs()

basque_terrorism ¶

basque_terrorism() -> DataFrame

Basque Country terrorism dataset (simulated).

Returns a balanced panel of GDP per capita (thousands of 1986 USD) for 17 Spanish regions, 1955--1997. The Basque Country is the treated unit; treatment begins in 1970 (onset of ETA terrorism).

The simulated data reproduce the gradual widening of an approximately 10 % GDP gap between the Basque Country and its synthetic counterfactual after 1970.

References

Abadie, A. & Gardeazabal, J. (2003). "The Economic Costs of Conflict: A Case Study of the Basque Country." American Economic Review, 93(1), 113--132. [@abadie2003economic]

Returns:

Type	Description
`DataFrame`	Columns: `region`, `year`, `gdppc`, `treated`.

Examples:

>>> import statspai as sp
>>> df = sp.basque_terrorism()
>>> result = sp.synth(df, outcome='gdppc', unit='region',
...                   time='year', treated_unit='Basque Country',
...                   treatment_time=1970)
>>> bool(result.estimate is not None)
True

german_reunification ¶

german_reunification() -> DataFrame

German reunification dataset (simulated).

Returns a balanced panel of GDP per capita for 17 OECD countries, 1960--2003. West Germany is the treated unit; treatment begins in 1990 (reunification).

The simulated trajectories reproduce the key stylised facts: Luxembourg has the highest GDP per capita (~40 000), Portugal the lowest (~10 000), and all countries share a common upward growth trend. Post-1990, West Germany exhibits an approximately 1 500 GDP-per-capita decline relative to its synthetic counterfactual.

References

Abadie, A., Diamond, A. & Hainmueller, J. (2015). "Comparative Politics and the Synthetic Control Method." American Journal of Political Science, 59(2), 495--510. [@abadie2015comparative]

Returns:

Type	Description
`DataFrame`	Columns: `country`, `year`, `gdppc`, `treated`.

Examples:

>>> import statspai as sp
>>> df = sp.german_reunification()
>>> result = sp.synth(df, outcome='gdppc', unit='country',
...                   time='year', treated_unit='West Germany',
...                   treatment_time=1990)
>>> bool(result.estimate is not None)
True

angrist_krueger_1991 ¶

angrist_krueger_1991(seed: int = 42) -> DataFrame

Simulated replica of Angrist-Krueger (1991) quarter-of-birth IV.

Classical weak-instrument case. Quarter of birth predicts years of schooling because compulsory-schooling laws tie entry age to calendar date (Q1 borns are slightly older at entry so can drop out with fewer years of school). First-stage F is a few dozen on several million observations; point estimates are unstable on subsets.

Our replica uses n=5,000 (the original is ~329k). Published IV returns-to-schooling on the original: 0.08-0.11 depending on controls and birth cohort.

Returns:

Type	Description
`pd.DataFrame with columns: lwage, educ, q1, q2, q3, q4, year_of_birth.`

References

Angrist, J. & Krueger, A. (1991). Does Compulsory School Attendance Affect Schooling and Earnings? QJE 106(4), 979-1014. [@angrist1991does]

card_1995 ¶

card_1995(seed: int = 42, simulated: bool = False) -> DataFrame

Card (1995) NLS Young Men data — simulated replica or real extract.

Card uses proximity to a 4-year college (nearc4) as an instrument for years of education in a wage equation. Published OLS and IV point estimates (Card 1995 Table 2):

OLS: β_educ ≈ 0.075 (col 2)
IV (nearc4): β_educ ≈ 0.132 (col 5)

IV exceeds OLS — the "Card puzzle". The LATE interpretation is for compliers on the margin of attending college because of proximity.

Parameters:

Name	Type	Description	Default
`seed`	`int`	RNG seed for the simulated DGP (ignored when `simulated=False`).	`42`
`simulated`	`bool`	If True, return a deterministic simulated replica calibrated so StatsPAI estimators recover OLS ≈ 0.11 and IV ≈ 0.142. If False, load the real NLSYM extract bundled in `statspai/datasets/data/card_1995.csv` (n=3010, identical to R's `wooldridge::card` complete-cases subset on Card's modelling variables). StatsPAI on this real data recovers OLS ≈ 0.0740 (paper 0.075) and IV ≈ 0.1323 (paper 0.132).	`False`

Returns:

Type	Description
`DataFrame`	Simulated columns: `lwage, educ, exper, expersq, black, south, smsa, nearc4` (n=3010). Real columns: same plus `nearc2` (proximity to 2-year college).

References

Card, D. (1995). Using Geographic Variation in College Proximity to Estimate the Return to Schooling. In Christofides et al. (eds.), Aspects of Labour Market Behaviour. [@card1995using]

castle_doctrine ¶

castle_doctrine(region_year_fe: bool = False, state_trends: bool = False, event_time: bool = False) -> DataFrame

Cheng & Hoekstra (2013) castle-doctrine panel — real data.

The canonical staggered difference-in-differences teaching dataset: 50 US states x 11 years (2000-2010). Between 2005 and 2009, 21 states expanded "castle doctrine" self-defence law (no duty to retreat); 29 states never did, giving a clean never-treated control group. Cheng & Hoekstra (2013) find these expansions raised homicide by roughly 8 log points rather than deterring crime.

This is the dataset behind Chapter 9 of Cunningham's Causal Inference: The Mixtape [@cunningham2021causal], where it motivates the Goodman-Bacon decomposition and the modern staggered-adoption estimators.

Parameters:

Name	Type	Description	Default
`region_year_fe`	`bool`	Append the 44 `r20YYQ` Census-region x year dummies used in the paper's saturated specification.	`False`
`state_trends`	`bool`	Append the 51 `trend_j` state-specific linear trends used in the paper's saturated specification.	`False`
`event_time`	`bool`	Append `time_til` (`year - effyear`, NaN for never-treated) and `gvar` (adoption cohort with never-treated coded `0`, the encoding `sp.callaway_santanna` and Stata's `csdid` expect).	`False`

Returns:

Type Description

DataFrame

550 rows x 29 base columns. Key modelling variables:

state, sid, year — panel identifiers (sid is the state code) l_homicide — log homicide rate per 100,000 (the outcome) homicide — homicide rate per 100,000 (levels) post — the paper's treatment dummy (see Notes) cdl — fractional castle-doctrine exposure within year effyear — year the law took effect (NaN = never treated) popwt — state population weight (Stata aweight) l_police, unemployrt, poverty, l_income, l_prisoner, l_lagprisoner, blackm_15_24, whitem_15_24, blackm_25_44, whitem_25_44, l_exp_subsidy, l_exp_pubwelfare — the paper's time-varying controls

Notes

post is not 1{year >= effyear}. Cheng & Hoekstra code post = 1{year > effyear}: the adoption year itself is coded untreated because the law was in force for only part of it. The fractional exposure in that year lives in cdl (e.g. Alabama 2006 = 0.5808). Reconstructing the treatment as year >= effyear silently changes 21 observations and moves the headline estimate.

This also makes the adoption cohort ambiguous for group-time estimators: gvar = effyear keeps a clean pre-treatment base period but counts the partially-exposed adoption year as treated, while gvar = effyear + 1 (consistent with post) pushes that partial year into the base period instead. On this panel the choice moves the Callaway-Sant'Anna simple ATT from 0.1104 to 0.0194 — see sp.replicate('castle_2013') for the full discussion.

Reference values verified against Stata 18 MP (see tests/reference_parity/test_castle_stata_parity.py):

==================================== ========== ========== Specification beta SE ==================================== ========== ========== TWFE, unweighted 0.0693984 0.0558596 TWFE, aweight=popwt 0.0755332 0.0331936 TWFE, weighted + controls 0.0796349 0.0308756 TWFE, full (region x year + trends) 0.0769490 0.0339377 Callaway-Sant'Anna, gvar=effyear 0.1103830 0.0387242 ==================================== ========== ==========

Examples:

>>> import statspai as sp
>>> df = sp.datasets.castle_doctrine()
>>> r = sp.feols('l_homicide ~ post | sid + year', data=df,
...              weights='popwt', vcov={'CRV1': 'sid'})

References

Cheng, C. & Hoekstra, M. (2013). Does Strengthening Self-Defense Law Deter Crime or Escalate Violence?: Evidence from Expansions to Castle Doctrine. Journal of Human Resources, 48(3), 821-853. [@cheng2013does]

Cunningham, S. (2021). Causal Inference: The Mixtape. Yale University Press. [@cunningham2021causal]

lee_2008_senate ¶

lee_2008_senate(seed: int = 42, simulated: bool = False) -> DataFrame

U.S. Senate close-election RD — real extract or Lee (2008)-style replica.

Parameters:

Name	Type	Description	Default
`seed`	`int`	RNG seed for the simulated DGP (ignored when `simulated=False`).	`42`
`simulated`	`bool`	Default False: this dataset ships the real extract and hands back the real extract, so `lee_2008_senate()` returns `x, y`. If True, return a deterministic simulated panel (n=6558, `voteshare_next, margin, win`) on a 0-1 vote-share scale, calibrated to a 0.08 jump at the cutoff. If False, load the real `rdrobust::rdrobust_RDsenate` extract (n=1390, `x, y`): `x` is the party's vote-share margin in a Senate election at time t and `y` its vote share (percent points, 0-100) in the election at t+2, following `rdrobust`'s own illustration. 93 rows have a missing `y`.	`False`

Notes

The real extract is the Senate data constructed by Cattaneo, Frandsen and Titiunik (2015) and distributed with R rdrobust. It is not Lee's (2008) U.S. House data -- the loader keeps its historical name because it implements Lee's close-election design. On this extract the default sp.rdrobust(df, y='y', x='x', c=0) returns conventional 7.414 and robust 7.507 (bandwidths 17.754 / 28.028), matching R and Stata rdrobust (Track A parity module 06_rd).

Returns:

Type	Description
`DataFrame`	Simulated columns: `voteshare_next, margin, win` (0-1 scale). Real columns: `x, y` (running variable; vote share 0-100).

References

Lee, D. (2008). Randomized experiments from non-random selection in U.S. House elections. Journal of Econometrics 142, 675-697. [@lee2008randomized] Cattaneo, M.D., Frandsen, B.R. & Titiunik, R. (2015). Randomization inference in the regression discontinuity design: An application to party advantages in the U.S. Senate. Journal of Causal Inference 3(1), 1-24. [@cattaneo2015randomization] Calonico, S., Cattaneo, M.D. & Titiunik, R. (2014). Robust nonparametric confidence intervals for regression-discontinuity designs. Econometrica 82(6), 2295-2326. [@calonico2014robust]

mpdta ¶

mpdta(seed: int = 42, simulated: bool = True, download: bool = True) -> DataFrame

Simulated replica of the mpdta dataset from R's did package.

The original mpdta is a county-year panel of log teen-employment (2003-2007) where some counties raise their minimum wage in 2004, 2006, or 2007 (staggered adoption).

simulated=False returns the original panel instead. It is not shipped (R did distributes it under GPL-3): the first call downloads Stata csdid's mpdta.dta (67 kB, the same values as R's copy), checks a pinned SHA-256 and caches it under $STATSPAI_DATA_DIR (default ~/.cache/statspai/datasets). On it R did and Stata csdid report simple ATT -0.0400, group -0.0310, calendar -0.0417 and dynamic -0.0772.

Parameters:

Name	Type	Description	Default
`seed`	`int`	Seed of the simulated replica (ignored with `simulated=False`).	`42`
`simulated`	`bool`	`False` fetches the original data.	`True`
`download`	`bool`	With `simulated=False`: download when not cached. `False` raises `DataInsufficient` instead of touching the network.	`True`

Returns:

Type	Description
`pd.DataFrame with columns: countyreal, year, lemp, first_treat, treat`	`lemp` — log teen employment (outcome) `first_treat` — period of first treatment (0 if never) `treat` — binary on/off indicator (post × treated cohort)

Notes

Our replica preserves:

500 counties × 5 years = 2500 rows
Three treatment cohorts: 2004, 2006, 2007 + never-treated
Negative homogeneous ATT ≈ -0.04 log points (matches the published R did::att_gt aggregated ATT of roughly -0.045 on the original)
County-level clustering in residuals

df.attrs['expected_simple_att'] = -0.040 (R did::att_gt + aggte(type='simple') on the original data returns -0.03995 with est_method='reg', never-treated controls, and analytic SEs -- see tests/orig_parity/02_mpdta_original; our replica's target is -0.04).

Because this is a simulated DGP, numerical values will not match R did::att_gt to high precision on the original data; but they match the sign, order of magnitude, and aggregation pattern.

References

Callaway, B. & Sant'Anna, P.H.C. (2021). Difference-in-Differences with Multiple Time Periods. Journal of Econometrics 225(2), 200-230. [@callaway2021difference]

nhefs ¶

nhefs(complete_case: bool = False) -> DataFrame

NHEFS — the canonical dataset of Hernán & Robins, Causal Inference: What If (2020), bundled as real, public-domain data for exact replication of the book's g-methods examples.

The National Health and Nutrition Examination Survey I (NHANES I) Epidemiologic Followup Study (NHEFS) follows US adults from a 1971-1975 baseline to a 1982 re-examination. The book uses it throughout Part II to estimate the average causal effect of quitting smoking (qsmk) on 10-year weight change (wt82_71, kg) and on 10-year mortality (death).

Parameters:

Name	Type	Description	Default
`complete_case`	`bool`	If False, return the full NHEFS extract (n=1629, 67 columns). If True, restrict to subjects with a non-missing 1982 weight (`wt82_71` not null, n=1566) — the analytic sample used for the weight-change examples in Chapters 12-15 of the book.	`False`

Returns:

Type Description

DataFrame

67 columns. Key modelling variables used in the book:

qsmk — quit smoking 1971-1982 (1 = yes; the "treatment") wt82_71 — weight change 1971→1982 in kg (continuous outcome) death — died by 1992 (1 = yes; the survival outcome) yrdth, modth — year / month of death (for survival timing) sex, race, age — demographics (sex: 0 male / 1 female; race 0/1) education — 5-level education (1-5) smokeintensity — cigarettes/day at baseline smokeyrs — years smoked at baseline exercise — 3-level exercise (0 much / 1 moderate / 2 little) active — 3-level daily activity (0 / 1 / 2) wt71 — baseline weight (kg)

df.attrs records the book citation and the published reference estimates (see Notes).

Notes

This is real data (df.attrs['data_source'] == 'real'), unlike the simulated econometrics replicas in this module. Because the data are the genuine book extract, StatsPAI reproduces the book's published numbers — not merely their neighbourhood:

Crude (unadjusted) mean weight-change difference, quitters vs non-quitters: 2.54 kg (book §12.2; StatsPAI 2.5406).
IP-weighted average treatment effect (stabilized weights, Chapter 12 MSM): 3.4 kg, 95% CI (2.4, 4.5) (book Program 12.4; StatsPAI sp.ipw 3.48, gold statsmodels MSM 3.44).
Parametric g-formula / standardization (Chapter 13): 3.5 kg.
G-estimation of a structural nested mean model (Chapter 14): psi ≈ 3.4.

Strict numerical reproductions of the full chapter programs live in tests/external_parity/test_whatif_nhefs.py and the public-health validation notebooks under examples/public_health/.

Provenance & licence

NHEFS is a US Federal public-use survey (NCHS / NIH) and is therefore in the public domain as a US Government work. The specific analytic extract bundled here (n=1629 × 67) is the one distributed by Hernán & Robins with the book and re-packaged in the MIT-licensed causaldata package (Huntington-Klein); it is byte-reproducible from causaldata.nhefs. Redistribution here is consistent with both the public-domain status of the underlying survey and StatsPAI's policy of only bundling freely redistributable datasets.

References

Hernán, M.A. & Robins, J.M. (2020). Causal Inference: What If. Boca Raton: Chapman & Hall/CRC. [@hernan2020causal]

nsw_dw ¶

nsw_dw(seed: int = 42, simulated: bool = True) -> DataFrame

Dehejia-Wahba NSW + PSID-1 comparison: simulated replica, or the data.

The default is a simulated replica, not the Dehejia-Wahba data. Every row of it is drawn from a DGP calibrated to the published summary statistics: 185 "NSW treated" rows and 2,490 "PSID-1" comparison rows with the right covariate imbalance, and a treatment effect of $1,794 on latent earnings (earnings are censored at zero, so the realised effect on the treated is somewhat smaller, and the outcome is not linear in the covariates where the treated are). Use it to see how an estimator copes with a large selection bias and almost no overlap; do not report its numbers as a replication of LaLonde (1986) or Dehejia and Wahba (1999). df.attrs['simulated'] says which one you have.

simulated=False returns the real sample: the 185 treated men of the Dehejia-Wahba NSW subsample and the 2,490 men of the PSID-1 comparison group. The difference in mean 1978 earnings is -$15,205 (Dehejia and Wahba 1999, Table 3); the experimental benchmark, which needs the 260 experimental controls that are not in this file, is $1,794.

The replica stays the default, unlike the other loaders, because committed cross-language parity fixtures are built on it.

Parameters:

Name	Type	Description	Default
`seed`	`int`	Seed of the replica (ignored when `simulated=False`).	`42`
`simulated`	`bool`	`False` for the real data, bundled as `statspai/datasets/data/nsw_psid_dw.csv` (no network).	`True`

Returns:

Type	Description
`pd.DataFrame with columns: treat, age, education, black, hispanic,`	married, nodegree, re74, re75, re78 (2,675 rows).

Examples:

>>> import statspai as sp
>>> real = sp.datasets.nsw_dw(simulated=False)
>>> real.shape, real.attrs["simulated"]
((2675, 10), False)
>>> sp.datasets.nsw_dw().attrs["simulated"]
True

References

Dehejia, R. & Wahba, S. (1999). Causal Effects in Nonexperimental Studies. JASA 94(448), 1053-1062. [@dehejia1999causal]

nsw_lalonde ¶

nsw_lalonde(seed: int = 42, simulated: bool = False) -> DataFrame

LaLonde NSW data — real MatchIt extract (default) or simulated replica.

Parameters:

Name	Type	Description	Default
`seed`	`int`	RNG seed for the simulated replica (ignored when `simulated=False`).	`42`
`simulated`	`bool`	If False (the default since 1.21.0), load the real `MatchIt::lalonde` extract bundled in `statspai/datasets/data/lalonde_matchit.csv` — the DW NSW treated cohort (185) plus a 429-unit PSID-1 subset for observational comparisons (n=614 total, with race factor already split into `black` and `hispanic` indicators). If True, return a deterministic simulated NSW experimental subset (185 + 260 = 445 rows) calibrated so naive OLS recovers the Dehejia-Wahba experimental ATT of about $1,794. The default flipped from True to False in 1.21.0 so the honest path is the default one — see MIGRATION.md.	`False`

Notes

The bundled real data is MatchIt::lalonde (n=614), NOT the larger DW (1999) NSW + PSID-1 sample (n=2,675). On this smaller subset, naive OLS gives ATT roughly -$635 (less negative than DW Table 3's headline -$8,498, which uses the full PSID-1). For the headline naive-bias demonstration, use the simulated nsw_dw() panel instead.

Simulated replica calibration

sasp_panel ¶

sasp_panel(analytic_sample: bool = False) -> DataFrame

SASP sex-worker session panel — real data.

The Survey of Adult Service Providers, collected by Cunningham & Kendall [@cunningham2011prostitution]: sex workers reporting on individual client sessions, with the log hourly wage (lnw) and whether the session involved unprotected sex (unsafe). Chapter 8 of Cunningham's Causal Inference: The Mixtape [@cunningham2021causal] uses it to teach the within (fixed-effects) transformation: pooled OLS conflates who a provider is with what she does, and only within-provider variation identifies the compensating differential for unsafe sex.

Parameters:

Name	Type	Description	Default
`analytic_sample`	`bool`	If False, return the full survey (1787 rows). If True, apply the book's recipe: drop rows missing any modelling variable, then keep only providers with exactly four sessions — 1028 rows, 257 providers, balanced. Use this to reproduce the published numbers.	`False`

Returns:

Type Description

DataFrame

31 columns. Key modelling variables:

id, session — provider and session identifiers lnw — log hourly wage (the outcome) unsafe — session involved unprotected sex llength — log session length age, bmi, schooling, black, hispanic, asian, other, married, cohab, divorced, separated — provider characteristics (time-invariant) age_cl, asq_cl, appearance_cl, asian_cl, black_cl, hispanic_cl, othrace_cl, reg, hot, massage_cl, provider_second — client / session characteristics

Notes

Twelve of the 25 controls are provider-level and drop out of the within estimator — see :data:SASP_TIME_INVARIANT. Stata's xtreg, fe omits them silently; StatsPAI raises NumericalInstability if you hand demeaned constants to sp.regress, which is the intended behaviour: a regressor with no within variation is not identified, and saying so is better than quietly dropping it.

On the analytic sample, all three routes agree with Stata 18 MP to ~5e-10 (see tests/reference_parity/test_sasp_within_parity.py):

============================== =========== =========== specification b(unsafe) SE ============================== =========== =========== pooled OLS, HC1 0.0134074 0.0283001 within (FE), cluster by id 0.0510339 0.0282831 manual demeaning, cluster id 0.0510339 0.0282831 ============================== =========== ===========

Pooled OLS puts the unsafe-sex premium at 1.3 log points; controlling for the provider raises it to 5.1 — the selection runs the other way from the naive reading, which is the chapter's point. The manual demeaning row reproduces xtreg, fe exactly, standard error included, which is the algebraic identity the chapter is teaching.

Examples:

>>> import statspai as sp
>>> df = sp.datasets.sasp_panel(analytic_sample=True)
>>> fe = sp.feols('lnw ~ unsafe + llength + reg | id', data=df,
...               vcov={'CRV1': 'id'})

References

Cunningham, S. & Kendall, T. D. (2011). Prostitution 2.0: The changing face of sex work. Journal of Urban Economics, 69(3), 273-287. [@cunningham2011prostitution]

Cunningham, S. (2021). Causal Inference: The Mixtape. Yale University Press. [@cunningham2021causal]

texas_prison ¶

texas_prison() -> DataFrame

Texas 1993 prison-capacity expansion panel — real data.

51 US states (50 + DC) x 16 years (1985-2000). Starting in 1993 Texas expanded operational prison capacity by roughly 35% per year for three years, approximately doubling it — a natural experiment used throughout Chapter 10 of Cunningham's Causal Inference: The Mixtape [@cunningham2021causal] to teach the synthetic control method. The outcome is the count of Black male prisoners (bmprison); Texas is statefip == 48.

Returns:

Type	Description
`DataFrame`	816 rows x 24 columns. Key modelling variables: `state, statefip, year` — panel identifiers `bmprison, wmprison` — Black / white male prisoner counts `bmpop, wmpop` — Black / white male population `bmprate, wmprate` — the corresponding rates `alcohol, aidscapita, income, ur, poverty, black, perc1519` — the predictors used by the book's recipe

Notes

Classic SCM does not reproduce across implementations here, and that is the point of shipping it. The book's recipe uses four lagged outcomes among its predictors, which leaves the predictor-weight matrix V weakly identified (Kaul et al. 2022). The resulting nested V-W problem is non-convex, and Stata's synth and StatsPAI settle on different local optima:

===================== ================================== ============= implementation donor weights mean gap 1994-2000 ===================== ================================== ============= Stata synth CA .408 IL .360 LA .122 FL .109 23073.70 StatsPAI sp.synth FL .436 NY .311 IL .253 23779.41 ===================== ================================== =============

StatsPAI attains the lower pre-treatment RMSE (865.3 vs 1227.0), so neither is "wrong" on its own objective. The estimated effect agrees to about 3% despite entirely different donor sets — the effect is identified far more robustly than the weights are. Do not read the donor weights as a finding. See sp.replicate('texas_1993').

Examples:

>>> import statspai as sp
>>> df = sp.datasets.texas_prison()
>>> sc = sp.synth(data=df, outcome='bmprison', unit='state', time='year',
...               treated_unit='Texas', treatment_time=1993)

References

Cunningham, S. (2021). Causal Inference: The Mixtape. Yale University Press. [@cunningham2021causal]

Cunningham's data readme attributes the Texas natural experiment to a "Cornwell and Cunningham (2016)" manuscript and to Perkinson (2010), Texas Tough: The Rise of America's Prison Empire. The Perkinson book is verifiable; no "Cornwell and Cunningham (2016)" record could be located in Crossref, so it is reported here as the upstream author's own attribution rather than asserted as a publication.

thornton_hiv ¶

thornton_hiv(complete_case: bool = False) -> DataFrame

Thornton (2008) HIV-incentive experiment — real data.

A randomized experiment in rural Malawi [@thornton2008demand]: respondents who had been tested for HIV were randomly offered a cash incentive to collect their results at a nearby centre. The randomization makes this the Mixtape's worked example for randomization inference — the Fisherian sharp-null test that needs no asymptotics and no distributional assumption, only the randomization that actually happened.

Parameters:

Name	Type	Description	Default
`complete_case`	`bool`	If False, return all 4820 survey rows. If True, restrict to rows with both `got` and `any` observed (n=2834) — the analysis sample for the incentive effect.	`False`

Returns:

Type Description

DataFrame

17 columns, a modelling subset of the 133 in the published file:

got — collected HIV results (the outcome) any — offered any cash incentive (the treatment) tinc — incentive amount offered under, over— incentive below / above the sample median distvct — distance to the results centre villnum — village identifier (the randomization cluster) hiv2004, age, male, mar, educ2004, simaverage, followupsurvey, usecondom04, rumphi, balaka — respondent characteristics and study arms

Notes

Reference values verified against Stata 18 MP on the complete-case sample (see tests/reference_parity/test_thornton_ri_parity.py):

========================== ============ ====== quantity value n ========================== ============ ====== mean got | any = 1 0.789235640 2211 mean got | any = 0 0.338683788 623 simple difference (SDO) 0.450551852 2834 OLS got ~ any, HC1 0.450551852 SE 0.020857971 ========================== ============ ======

The effect is enormous — a 45-point increase on a 34-point base — so randomization inference returns p = 0 at any practical permutation count: no reassignment of the treatment label comes close to the observed difference. That is the honest answer, not a rounding artefact, and it makes the dataset a poor place to study borderline p-values but an excellent one for seeing what RI is.

Examples:

>>> import statspai as sp
>>> df = sp.datasets.thornton_hiv(complete_case=True)
>>> ri = sp.ri_test(df, y='got', treat='any', n_perms=1000, seed=42)
>>> round(ri['observed'], 6)
0.450552

References

Thornton, R. L. (2008). The Demand for, and Impact of, Learning HIV Status. American Economic Review, 98(5), 1829-1863. [@thornton2008demand]

Cunningham, S. (2021). Causal Inference: The Mixtape. Yale University Press. [@cunningham2021causal]

currency_union_panel ¶

currency_union_panel(seed: int = 0, late_adopters: bool = False) -> DataFrame

Simulated country-pair panel with staggered currency-union adoption.

Parameters:

Name	Type	Description	Default
`seed`	`int`		`0`
`late_adopters`	`bool`	`False`: 15 countries (12 adopt, 11 in 1999 and one in 2001; three never adopt), 105 pairs, 1995-2015. `True`: adds six countries adopting in 2007, 2008, 2009, 2011, 2014 and 2015 and seven that never adopt (378 pairs, 1995-2019), with a common downturn in 2009-2012 that coincides with the late adoptions.	`False`

Returns:

Type Description

DataFrame

One row per pair-year: pair, country_i, country_j, year, log_trade (outcome), euro (1 when both members have adopted), ever_euro, adopt_year (NaN if never), pre_trade (mean log_trade of the pair before 1999), log_gdp_prod and log_gdppc (time-varying), and tau_true (the effect of adoption on log_trade for that pair-year, defined for every row). attrs records the design and the true average effect on treated rows.

Notes

log_trade = 20 + a_ij + g_t + 0.6 (log_gdp_prod - mean) + tau_ij,t euro + e with AR(1) errors (rho = 0.5) within pairs. The pair effect a_ij rises with the members' latent "core" index, which also raises the chance of adopting, and tau rises with the pair's latent trade intensity and current size: tau = 0.05 + 0.12 z_trade + 0.06 z_size, clipped to [-0.2, 0.6].

Examples:

>>> import statspai as sp
>>> df = sp.datasets.currency_union_panel(seed=0)
>>> df.shape
(2205, 13)
>>> round(float(df.loc[df["euro"] == 1, "tau_true"].mean()), 2) > 0
True

References

[aytug2026euro], [kattenberg2023causal]

bacon_divorce ¶

bacon_divorce(download: bool = True) -> DataFrame

Unilateral divorce reform and female suicide, 1964-1996 — real data.

The application of Goodman-Bacon (2021) [@goodmanbacon2021difference] and the example data of Stata's bacondecomp [@goodmanbacon2019bacondecomp]: 49 states x 33 years (1,617 rows). asmrs is the age-adjusted female suicide rate per million and post the reform indicator. Twelve reform cohorts (1969-1985), 5 never-reformed states and 8 states that reformed before 1964 (treated in every year of the panel). Fetched on first use (81 kB) and cached; no licence is stated at the source, so it is not shipped.

Parameters:

Name	Type	Description	Default
`download`	`bool`	Download the file when it is not cached.	`True`

Returns:

Type	Description
`DataFrame`	The published columns (`stfips`, `year`, `_nfd`, `post`, `asmrs`, `pcinc`, `asmrh`, `cases`, `weight`, `copop`) in float64 / int64, plus `first_treat`: the first year with `post == 1` (0 = never reformed; 1964 marks the 8 states already treated when the panel starts).

Raises:

Type	Description
`DataInsufficient`	Offline and not cached, or a SHA-256 mismatch.

Examples:

>>> import statspai as sp
>>> df = sp.datasets.bacon_divorce()
>>> bd = sp.bacon_decomposition(
...     df, y="asmrs", treat="post", time="year", id="stfips"
... )
>>> round(bd["beta_twfe"], 4)
-3.0799

jel_medicaid ¶

jel_medicaid(sample: Literal['raw', '2x2', 'staggered'] = 'raw', download: bool = True) -> DataFrame

County mortality and the ACA Medicaid expansion (Baker et al., JEL) — real data.

The running example of Baker, Callaway, Cunningham, Goodman-Bacon and Sant'Anna's practitioner's guide to difference-in-differences [@baker2026difference]: US counties 2009-2019, the outcome is the crude mortality rate of adults aged 20-64 per 100,000. The file is fetched from the authors' replication repository on first use (6.3 MB, MIT licence) and cached; see :mod:statspai.datasets._fetch.

Parameters:

Name	Type	Description	Default
`sample`	`('raw', '2x2', 'staggered')`	`'raw'` is the published file unchanged (31,843 rows). `'2x2'` and `'staggered'` apply the cleaning of the authors' scripts: drop the five states that expanded before 2014 (DC, DE, MA, NY, VT), build the covariates (`perc_female`, `perc_white`, `perc_hispanic` and `unemp_rate` in percent, `median_income` in thousands), drop counties with a missing value, keep counties observed in all 11 years, and add `set_wt` (2013 population aged 20-64, the paper's weight) and `treat_year` (expansion year, 0 = not expanded by 2019). `'2x2'` keeps 2014 expanders and counties not expanded by 2019 (2,200 counties, 24,200 rows: Tables 2-7 and the 2 x T event study); `'staggered'` keeps every cohort (2014, 2015, 2016, 2019 and never; 2,604 counties).	`'raw'`
`download`	`bool`	Download the file when it is not cached. With `False` a missing file raises instead of touching the network.	`True`

Returns:

Type	Description
`DataFrame`	`attrs` carries the source URL, the pinned SHA-256 and the licence.

Raises:

Type	Description
`DataInsufficient`	Offline and not cached, or the download does not match the pinned SHA-256. The message gives the URL and the path to save it at.

Examples:

>>> import statspai as sp
>>> df = sp.datasets.jel_medicaid(sample="2x2")
>>> df.county_code.nunique()
2200

References

[@baker2026difference]

lalonde_panel ¶

lalonde_panel(download: bool = True) -> DataFrame

NSW experiment with CPS and PSID comparison groups, 1975 and 1978 — real data.

The long panel behind Sant'Anna and Zhao (2020) [@santanna2020doubly], Table 3, and the example data of Stata's drdid: 19,204 people observed in 1975 and 1978 (38,408 rows). sample is 1 for the NSW experiment (LaLonde sample, 722 people), 2 for CPS (15,992) and 3 for PSID (2,490); treated is the NSW random assignment, experimental marks NSW rows, dwincl the Dehejia-Wahba subsample and early_ra the early-random-assignment subsample; re is real earnings in the row's year. Fetched on first use (0.9 MB) and cached; no licence is stated at the source, so it is not shipped.

Unlike nsw_lalonde() (a 614-row cross-section from MatchIt, with re74 / re75 / re78 columns), this is the two-period panel used by the DR-DiD papers.

Parameters:

Name	Type	Description	Default
`download`	`bool`	Download the file when it is not cached.	`True`

Returns:

Type	Description
`DataFrame`	Covariates are float64. The `.dta` stores `age` and `educ` as int8, where `age ** 3` silently wraps around; the cast removes that trap.

Raises:

Type	Description
`DataInsufficient`	Offline and not cached, or a SHA-256 mismatch.

Examples:

>>> import statspai as sp
>>> df = sp.datasets.lalonde_panel()
>>> nsw_control = (df["sample"] == 1) & (df["treated"] == 0)
>>> d = df[nsw_control | (df["sample"] == 2)].copy()
>>> d["post"] = (d["year"] == 1978).astype(int)
>>> x = ["age", "educ", "black", "married", "nodegree", "hisp", "re74"]
>>> r = sp.drdid(
...     d, y="re", treat="experimental", time="post", id="id", covariates=x
... )
>>> round(r.estimate, 2)
-871.33

from_fred ¶

from_fred(payload: Union[JSONLike, Mapping[str, JSONLike]], *, series_id: Optional[str] = None) -> DataFrame

Normalise FRED series observations to a tidy time series.

Accepts:

a single series' {"observations": [{"date", "value"}, ...]} dict,
just the observations list,
a mapping {series_id: observations} for several series, merged on date into a wide frame (one column per series).

FRED's missing-value sentinel "." becomes NaN; dates are parsed to datetime64.

Parameters:

Name	Type	Description	Default
`payload`	`dict or list or mapping of series_id -> observations`		required
`series_id`	`str`	Column name for the value series when a single series is passed (default `"value"`).	`None`

Returns:

Type	Description
`DataFrame`	Columns `date` + one value column per series, sorted by date.

Examples:

>>> obs = {"observations": [{"date": "2020-01-01", "value": "1.5"},
...                         {"date": "2020-02-01", "value": "."}]}
>>> import statspai as sp
>>> df = sp.from_fred(obs, series_id="cpi")
>>> list(df.columns)
['date', 'cpi']
>>> bool(df['cpi'].isna().iloc[1])
True

from_sdmx ¶

from_sdmx(payload: JSONLike, *, value_name: str = 'value') -> DataFrame

Normalise an SDMX-JSON payload (OECD / Eurostat / IMF) to long form.

SDMX-JSON encodes each observation by integer indices into per-dimension code lists. This expands those indices back into human-readable dimension columns plus a value column — the shape sp.detect_design can read.

Accepts the SDMX-JSON 1.0 structure ({"dataSets": [...], "structure": {"dimensions": {...}}}) or a pre-tidied list of record dicts (returned as a DataFrame unchanged).

Parameters:

Name	Type	Description	Default
`payload`	`dict or list or DataFrame`		required
`value_name`	`str`		`"value"`

Returns:

Type	Description
`DataFrame`	One column per SDMX dimension (e.g. `LOCATION`, `TIME_PERIOD`) plus `<value_name>`. `df.attrs['source'] == 'sdmx'`.

Examples:

>>> payload = {
...   "dataSets": [{"series": {"0:0": {"observations": {"0": [3.2]}}}}],
...   "structure": {"dimensions": {
...      "series": [
...        {"id": "LOCATION", "values": [{"id": "USA", "name": "USA"}]},
...        {"id": "SUBJECT", "values": [{"id": "UNR", "name": "Unemp"}]}],
...      "observation": [
...        {"id": "TIME_PERIOD", "values": [{"id": "2020", "name": "2020"}]}]
...   }}}
>>> import statspai as sp
>>> df = sp.from_sdmx(payload)
>>> df.loc[0, "LOCATION"], df.loc[0, "TIME_PERIOD"], float(df.loc[0, "value"])
('USA', '2020', 3.2)

from_worldbank ¶

from_worldbank(payload: JSONLike, *, wide: bool = False, value_name: str = 'value') -> DataFrame

Normalise a World Bank Indicators API payload to a tidy panel.

Accepts any of the shapes a World Bank MCP / the v2 REST API returns:

the raw [metadata, rows] two-element list,
just the rows list of observation dicts,
a DataFrame already close to tidy.

Each observation dict looks like::

{"indicator": {"id": "NY.GDP.PCAP.KD", "value": "GDP per capita"},
 "country":   {"id": "US", "value": "United States"},
 "countryiso3code": "USA", "date": "2020", "value": 63027.7}

Parameters:

Name	Type	Description	Default
`payload`	`list or dict or DataFrame`		required
`wide`	`bool`	If True, pivot indicators to columns indexed by (iso3, year) — handy when several indicators were fetched and you want one regression frame.	`False`
`value_name`	`str`	Name of the value column in long form.	`"value"`

Returns:

Type	Description
`DataFrame`	Long form columns: `country`, `iso3`, `indicator`, `indicator_id`, `year`, `<value_name>`. `df.attrs['source']` records the provenance. Wide form: one row per (iso3, year), one column per indicator.

Examples:

>>> rows = [{"indicator": {"id": "NY.GDP", "value": "GDP"},
...          "country": {"id": "US", "value": "United States"},
...          "countryiso3code": "USA", "date": "2020", "value": 1.0}]
>>> import statspai as sp
>>> df = sp.from_worldbank(rows)
>>> list(df.columns)
['country', 'iso3', 'indicator', 'indicator_id', 'year', 'value']
>>> int(df.loc[0, 'year'])
2020

california_prop99 ¶

california_prop99(simulated: bool = False) -> DataFrame

California Proposition 99 panel (Abadie-Diamond-Hainmueller 2010).

Parameters:

Name	Type	Description	Default
`simulated`	`bool`	If True, return the simulated covariate-rich replica from `synth.california_tobacco` (39 states × 31 years, 1970-2000, ADH-shaped DGP). If False, load the real ADH (2010) panel bundled in `statspai/datasets/data/california_prop99.csv` (39 states × 31 years, with covariates `cigsale, retprice, lnincome, age15to24, beer`; identical to tidysynth's smoking dataset). Use this for exact paper replication.	`False`

Returns:

Type	Description
`DataFrame`	Columns (both branches): `state, year, cigsale, retprice, lnincome, age15to24, beer`. The simulated branch additionally provides `treated`; on the real branch we derive it as `(state == 'California') & (year >= 1989)`.

References

Abadie, A., Diamond, A. & Hainmueller, J. (2010). Synthetic Control Methods for Comparative Case Studies. Journal of the American Statistical Association 105(490), 493-505. [@abadie2010synthetic]

list_datasets ¶

list_datasets() -> DataFrame

Return a DataFrame describing all available datasets.

Columns: name, design, n_obs, paper, paper_original, expected_main, source, simulated, delivery, available_offline, real_variant, columns.

source says which variant a bare name() call returns: "bundled CSV" is the real extract shipped in datasets/data/, "simulated" a deterministic replica calibrated near the published estimate, "fetched" real data downloaded from its public source on first use. delivery is "shipped" (in the wheel, works offline) or "fetched"; available_offline is True for shipped data and, for fetched data, whether a verified copy is already in the cache ($STATSPAI_DATA_DIR). real_variant names the call that returns the real data when a bare call returns a replica. Listing the catalogue never downloads anything.

paper_original is the headline number from the published paper on the ORIGINAL data (what readers expect to see).
expected_main is what the canonical estimator recovers from the variant a bare name() call returns (what users will actually observe), and it is prefixed REAL data: wherever that default is the bundled extract rather than a replica. For a "simulated" row the two columns differ because the replica is a deterministic DGP calibrated to the neighbourhood of the published value, not the original data. For a "bundled CSV" row any gap is an estimator or extract difference and is named as such.

For the strict numerical neighbourhood proofs see tests/external_parity/test_published_replications.py and tests/external_parity/PUBLISHED_REFERENCE_VALUES.md.

Examples:

>>> import statspai as sp
>>> table = sp.list_datasets()
>>> bool(table.set_index("name").loc["mpdta", "simulated"])
True
>>> table.set_index("name").loc["card_1995", "source"]
'bundled CSV'

load_dataset ¶

load_dataset(name: str, **kwargs: Any) -> DataFrame

Load a bundled example dataset by name.

Parameters:

Name	Type	Description	Default
`name`	`str`	One of the names in :func:`list_datasets` (`"mpdta"`, `"card_1995"`, `"nsw_dw"`, `"lee_2008_senate"`, ...).	required
`**kwargs`	`Any`	Passed to the dataset's loader (e.g. `simulated=False` for `nsw_dw` or `mpdta`, `seed=` for the simulated replicas, `download=False` to refuse a download for a fetched dataset).	`{}`

Returns:

Type	Description
`DataFrame`	`df.attrs['simulated']` says whether it is a simulated replica or the real extract; `df.attrs['paper']` names the source study.

Examples:

>>> import statspai as sp
>>> df = sp.load_dataset("card_1995")
>>> df.attrs["simulated"]
False

statspai.datasets¶

datasets ¶

basque_terrorism ¶

german_reunification ¶

angrist_krueger_1991 ¶

card_1995 ¶

castle_doctrine ¶

lee_2008_senate ¶

mpdta ¶

nhefs ¶

nsw_dw ¶

nsw_lalonde ¶

sasp_panel ¶

texas_prison ¶

thornton_hiv ¶

currency_union_panel ¶

bacon_divorce ¶

jel_medicaid ¶

lalonde_panel ¶

from_fred ¶

from_sdmx ¶

from_worldbank ¶

california_prop99 ¶

list_datasets ¶

load_dataset ¶

`statspai.datasets`¶