`statspai.matching`¶

matching ¶

Matching module for StatsPAI.

Unified interface for matching estimators:

Nearest-neighbor matching (propensity score, Mahalanobis, Euclidean)
Exact matching
Coarsened Exact Matching (CEM)
Propensity score stratification / subclassification
Abadie-Imbens (2011) bias correction
Entropy balancing (Hainmueller 2012)
Covariate Balancing Propensity Score (Imai-Ratkovic 2014)
Genetic Matching (Diamond-Sekhon 2013)
Stable Balancing Weights (Zubizarreta 2015)
Optimal pair / full / cardinality matching (Rosenbaum 1989, 2012)
Overlap weights (Li-Morgan-Zaslavsky 2018)

The single entry point is :func:match — a method-aware dispatcher that routes method= to the correct estimator. Standalone functions (ebalance, cbps, genmatch, sbw, optimal_match, cardinality_match, overlap_weights) remain fully accessible for power users who need their estimator-specific parameters.

References

Rosenbaum, P.R. and Rubin, D.B. (1983). Biometrika, 70(1), 41-55. Abadie, A. and Imbens, G.W. (2006). Econometrica, 74(1), 235-267. Abadie, A. and Imbens, G.W. (2011). JBES, 29(1), 1-11. Iacus, S.M., King, G., and Porro, G. (2012). Political Analysis, 20(1), 1-24. Hainmueller, J. (2012). Political Analysis, 20(1), 25-46. Imai, K. and Ratkovic, M. (2014). JRSS-B, 76(1), 243-263. Diamond, A. and Sekhon, J.S. (2013). REStat, 95(3), 932-945. Zubizarreta, J.R. (2015). JASA, 110(511), 910-922. Li, F., Morgan, K.L., and Zaslavsky, A.M. (2018). JASA, 113(521), 390-400. Rosenbaum, P.R. (2012). JASA, 107(498), 691-700. Cunningham, S. (2021). Causal Inference: The Mixtape. Yale University Press. [@rosenbaum1983central]

BalTabResult ¶

Bases: ResultProtocolMixin

Balance table returned by :func:sp.bal_tab.

Attributes:

Name	Type	Description
`table`	`DataFrame`	One row per covariate term. Binary / multi-category treatments: `diff_un` / `diff_adj` (multi-category: the largest absolute difference across comparisons), `vr_` and `ks_`. Continuous treatments: `corr_` (Pearson) and `scorr_` (Spearman). `balanced` compares the adjusted statistic with `threshold`.
`pairs`	`dict`	Multi-category treatments: one table per comparison, keyed `'g1 vs. g2'`; differences are `g1 - g2`.
`observations`	`DataFrame`	Sample sizes (`unadjusted`) and effective sample sizes (`adjusted`) per treatment group.
`summary_stats`	`dict`	Largest adjusted statistic, number of terms over the threshold, the standardization used.

Examples:

>>> import numpy as np, pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> df = pd.DataFrame({'x': rng.normal(size=200)})
>>> df['g'] = np.where(df['x'] + rng.normal(size=200) > 0, 'b', 'a')
>>> bt = sp.bal_tab(df, 'g', ['x'])
>>> bool(abs(bt.table.loc['x', 'diff_un']) > 0.5)
True

BalanceRandomizationResult `dataclass` ¶

Bases: ResultProtocolMixin

Result of :func:balance_vs_randomization.

Attributes:

Name	Type	Description
`table`	`DataFrame`	One row per covariate and three summary rows (`min_p`, `truncated_product`, `n_below_alpha`). `actual` is the value in the matched sample, `sim_median` its median over the simulated randomized experiments, and `share_better` the share of those experiments that were better balanced on that row.
`sim`	`DataFrame`	The simulated values, one row per simulated experiment.
`n_sim`	`int`	Number of simulated experiments.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> df = pd.DataFrame({"z": np.tile([1, 0], 100),
...                    "age": rng.normal(50, 10, 200),
...                    "female": rng.integers(0, 2, 200)})
>>> res = sp.balance_vs_randomization(df, "z", ["age", "female"],
...                                   n_sim=200, random_state=1)
>>> res.table.index.tolist()
['age', 'female', 'min_p', 'truncated_product', 'n_below_alpha']

FullMatchResult `dataclass` ¶

Bases: ResultProtocolMixin

Result of :func:full_match.

Attributes:

Name	Type	Description
`estimate, se, ci, pvalue`	`float`	Effect on the outcome and its matched-set-clustered standard error; `nan` when no outcome was given.
`estimand`	`str`
`subclass`	`Series`	Matched-set label of every row of the input (`<NA>` for a row left out by the caliper or by missing values).
`weights`	`Series`	Matching weights for the estimand (0 for a row left out).
`balance`	`DataFrame`	Standardised mean differences before and after matching.
`total_distance`	`float`	Sum of the matched distances, the quantity that is minimised.
`n_sets, n_treated, n_control, n_unmatched`	`int`

Examples:

>>> import numpy as np, pandas as pd, statspai as sp
>>> rng = np.random.default_rng(1)
>>> n = 300
>>> x = rng.normal(size=n)
>>> d = rng.binomial(1, 1 / (1 + np.exp(-x)))
>>> df = pd.DataFrame({"x": x, "d": d, "y": 1 + x + 2 * d + rng.normal(size=n)})
>>> fit = sp.full_match(df, "y", "d", ["x"])
>>> isinstance(fit, sp.FullMatchResult)
True
>>> bool(fit.n_treated + fit.n_control == n)
True
>>> sorted(fit.matched_data(df).columns)
['d', 'subclass', 'weights', 'x', 'y']

matched_data ¶

matched_data(data: DataFrame) -> DataFrame

data with subclass and weights columns, matched rows only.

GenMatchResult `dataclass` ¶

Bases: ResultProtocolMixin

Output of :func:sp.genmatch (Diamond-Sekhon genetic matching).

Holds the ATT estimate and its standard error, the optimal covariate weight vector, the matched control indices, and a pre/post balance table. Call .summary() for a formatted report.

.. warning::

att_se is the matched-pair standard error sd(Y_t - Y_c) / sqrt(n_pairs), not a bootstrap — earlier versions of this docstring said "bootstrap SE", which was never what the code computed. It conditions on the realised match structure and on the fitted covariate weights, and genetic matching matches with replacement (on a typical run one control serves a dozen treated units), so it ignores exactly the dependence that makes this estimator's naive SE too small.

The same formula, measured on sp.match over 36 designs x 1000 replications (benchmarks/matching_se_coverage.py), runs 0.56-0.91x the true sampling SD and never reaches nominal coverage (0.71-0.92 against a nominal 0.95). Treat att_se as a lower bound; for inference that covers, bootstrap the whole pipeline or use sp.match(se_method='abadie_imbens').

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> p = 1.0 / (1.0 + np.exp(-(0.5 * x1 - 0.5 * x2 - 0.5)))
>>> d = rng.binomial(1, p)
>>> y = 1.0 + 2.0 * d + x1 + x2 + rng.normal(size=n)
>>> df = pd.DataFrame({'y': y, 'd': d, 'x1': x1, 'x2': x2})
>>> res = sp.genmatch(df, y='y', treat='d', covariates=['x1', 'x2'],
...                   population_size=10, generations=5)
>>> isinstance(res, sp.GenMatchResult)
True
>>> res.n_treated
111

MatchEstimator ¶

Unified matching estimator supporting multiple distance × method combinations.

This is the object-oriented backend behind :func:match. Most users should call :func:sp.match; construct MatchEstimator directly only when you want to hold the configured estimator and call .fit() yourself. .fit() returns a CausalResult.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Input data.	required
`y`	`str`	Outcome column.	required
`treat`	`str`	Binary (0/1) treatment column.	required
`covariates`	`list of str`	Variables to match on.	required
`distance`	`str`	`'propensity'`, `'mahalanobis'`, `'euclidean'` or `'exact'`.	`None`
`method`	`str`	`'nearest'`, `'stratify'` or `'cem'` (legacy `'psm'` / `'mahalanobis'` are also accepted).	`'nearest'`
`estimand`	`str`	`'ATT'` or `'ATE'`.	`'ATT'`

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> age = rng.normal(40, 8, n)
>>> edu = rng.normal(12, 2, n)
>>> ps = 1 / (1 + np.exp(-(0.05 * (age - 40) + 0.1 * (edu - 12))))
>>> training = rng.binomial(1, ps)
>>> wage = 20 + 0.3 * age + 0.5 * edu + 4.0 * training + rng.normal(0, 3, n)
>>> df = pd.DataFrame({"wage": wage, "training": training,
...                    "age": age, "edu": edu})
>>> est = sp.MatchEstimator(df, y="wage", treat="training",
...                         covariates=["age", "edu"], distance="propensity")
>>> result = est.fit()
>>> type(result).__name__
'CausalResult'

fit ¶

fit() -> CausalResult

Fit matching estimator and return results.

CardinalityMatchResult `dataclass` ¶

Bases: ResultProtocolMixin

Result of :func:cardinality_match (Zubizarreta cardinality matching).

Attributes:

Name	Type	Description
`treated_matched, control_matched`	`ndarray`	Row indices (into the cleaned data) of the matched treated and control units making up each pair.
`ate`	`float`	Matched-pair average treatment effect on the treated (ATT).
`se`	`float`	Analytic standard error of `ate` over the matched pairs.
`n_matched_pairs`	`int`	Number of matched pairs retained.
`balance`	`DataFrame`	Post-match balance table with columns `covariate`, `SMD`, `\|SMD\|` (standardised mean differences).

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> p = 1.0 / (1.0 + np.exp(-(0.5 * x1 - 0.5 * x2 - 0.5)))
>>> d = rng.binomial(1, p)
>>> y = 1.0 + 2.0 * d + x1 + x2 + rng.normal(size=n)
>>> df = pd.DataFrame({'y': y, 'd': d, 'x1': x1, 'x2': x2})
>>> res = sp.cardinality_match(df, treatment='d', outcome='y',
...                            covariates=['x1', 'x2'],
...                            smd_tolerance=0.1)
>>> isinstance(res, sp.CardinalityMatchResult)
True
>>> res.n_matched_pairs
106
>>> round(res.ate, 2)
1.84
>>> bool((res.balance['|SMD|'] <= 0.1).all())
True

OptimalMatchResult `dataclass` ¶

Bases: ResultProtocolMixin

Result of :func:optimal_match (optimal 1:1 Hungarian matching).

Attributes:

Name	Type	Description
`pairs`	`DataFrame`	One row per matched pair with columns `treated_idx`, `control_idx`, `distance`.
`distances`	`ndarray`	Matching distance for each matched pair.
`ate`	`float`	Matched-pair average treatment effect on the treated (ATT).
`se`	`float`	Analytic standard error of `ate` over the matched pairs.
`n_treated, n_matched`	`int`	Number of treated units and number of retained matched pairs.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> p = 1.0 / (1.0 + np.exp(-(0.5 * x1 - 0.5 * x2 - 0.5)))
>>> d = rng.binomial(1, p)
>>> y = 1.0 + 2.0 * d + x1 + x2 + rng.normal(size=n)
>>> df = pd.DataFrame({'y': y, 'd': d, 'x1': x1, 'x2': x2})
>>> res = sp.optimal_match(df, treatment='d', outcome='y',
...                        covariates=['x1', 'x2'])
>>> isinstance(res, sp.OptimalMatchResult)
True
>>> res.n_matched
111
>>> round(res.ate, 2)
1.88
>>> res.pairs.columns.tolist()
['treated_idx', 'control_idx', 'distance']
>>> bool(len(res.distances) == res.n_matched)
True

att `property` ¶

att: float

Matched-pair effect, named for the estimand it actually targets.

1:1 matching on the treated retains every treated unit and reweights controls to them, so the estimand is the ATT. ate is retained as an alias for backward compatibility.

estimate `property` ¶

estimate: float

Alias used by the shared result protocol.

BalanceDiagnosticsResult ¶

Bases: ResultProtocolMixin

Container for raw/weighted matching balance diagnostics.

Returned by :func:sp.balance_diagnostics. Holds a per-covariate table (raw vs. weighted SMDs, variance ratios, KS stats) and a summary_stats dict (max/mean SMDs, imbalance counts, effective sample size, propensity overlap). Call .summary() for a report.

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> bal = sp.balance_diagnostics(
...     df, treatment='union',
...     covariates=['education', 'experience', 'tenure'])
>>> isinstance(bal, sp.BalanceDiagnosticsResult)
True
>>> bal.summary_stats['n_obs']
3000

PSBalanceResult ¶

Bases: ResultProtocolMixin

Container for propensity score balance diagnostics.

Attributes:

Name	Type	Description
`table`	`DataFrame`	Balance statistics per covariate: mean_treat, mean_control, smd_raw, smd_weighted, variance_ratio, ks_stat.
`ps`	`Series`	Estimated propensity scores.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> p = 1.0 / (1.0 + np.exp(-(0.5 * x1 - 0.5 * x2 - 0.5)))
>>> d = rng.binomial(1, p)
>>> df = pd.DataFrame({'d': d, 'x1': x1, 'x2': x2})
>>> bal = sp.ps_balance(df, treatment='d', covariates=['x1', 'x2'])
>>> isinstance(bal, sp.PSBalanceResult)
True
>>> bal.table['smd_weighted'].round(2).tolist()
[0.02, -0.06]

summary ¶

summary() -> str

Formatted balance summary table.

love_plot ¶

love_plot(threshold: float = 0.1, **kwargs: Any) -> Tuple[Any, Any]

Convenience method: calls love_plot() from balance data.

PScoreResult ¶

Bases: ResultProtocolMixin

What :func:pscore returns.

Attributes:

Name	Type	Description
`pscore`	`Series`	The fitted score, on the index of the data (missing where a row was not in the estimation sample).
`block`	`Series`	The block number of each row, from 1; missing outside the estimation sample and outside the common support.
`support`	`Series`	Whether the row is in the region of common support (all `True` without `common_support`).
`support_range`	`tuple of float`	Smallest and largest score among the treated.
`n_blocks`	`int`	The final number of blocks.
`balanced`	`bool`	Whether no covariate differs between treated and controls in any block at `level`. `False` when no block holds both treated and control units, since nothing could then be tested (`n_tested_blocks` is zero).
`n_tested_blocks`	`int`	Blocks holding both treated and control units.
`separated`	`bool`	The covariates separate treated from controls, so the treatment model has no maximum-likelihood fit and the scores pile up at 0 and 1.
`unbalanced`	`DataFrame`	The rejections: `variable`, `block` and `pvalue`.
`blocks`	`DataFrame`	One row per non-empty block: its `lower` and `upper` bound and the numbers of controls and treated units in it.
`coefficients`	`DataFrame`	The treatment model: `coef`, `se`, `z` and `pvalue`, the constant last. A covariate that is collinear with earlier ones is listed with a zero coefficient and a missing standard error.
`loglik`	`float`
`n_obs`	`int`	Rows in the estimation sample.

Examples:

>>> import statspai as sp
>>> df = sp.datasets.nsw_lalonde()
>>> res = sp.pscore(df, "treat", ["age", "educ", "married", "re74"])
>>> type(res).__name__, res.n_blocks == int(res.block.max())
('PScoreResult', True)
>>> list(res.blocks.columns)
['block', 'lower', 'upper', 'n_control', 'n_treated']

assign ¶

assign(data: DataFrame, pscore: str = 'pscore', block: Optional[str] = 'block', support: Optional[str] = 'comsup') -> DataFrame

A copy of data with the score, block and support columns.

Pass None for a column that is not wanted. The names are the ones Stata's pscore takes in pscore() and blockid() and the comsup it creates.

PSMatch2Result ¶

Bases: ResultProtocolMixin

Container for a sp.psmatch2 run.

Attributes:

Name	Type	Description
`matched_data`	`DataFrame`	The input data plus the psmatch2 columns (`_pscore`, `_treated`, `_support`, `_weight`, `_y`; plus `_n1` …, `_nn`, `_pdif` for nearest-neighbour matching). Also available as `.data`.
`att, se, pvalue, ci`	`float / tuple`	Average treatment effect on the treated and its inference.
`estimand`	`str`	Always `'ATT'` for `psmatch2`.
`result`	`CausalResult`	The underlying :func:`sp.match` result.

Methods:

Name	Description
`matched_sample`	Rows that entered the matched sample (`_weight` not missing).
`balance`	Post-matching covariate balance on the weighted matched sample.
`psplot`	Propensity-score density before/after matching.
`psm_did`	Frequency-weighted PSM-DID regression.

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> m = sp.psmatch2(df, outcome='log_wage', treat='union',
...                 covariates=['education', 'experience', 'tenure'])
>>> '_weight' in m.matched_data.columns
True
>>> bal = m.balance()                      # post-matching balance
>>> fig, ax = m.psplot()

matched_sample ¶

matched_sample(*, on_support: bool = True, drop_unmatched: bool = True) -> DataFrame

Return the rows that make up the matched sample.

Parameters:

Name	Type	Description	Default
`on_support`	`bool`	Keep only rows with `_support == 1`. Has no effect when matching was run with `common_support='none'` (every row is on support).	`True`
`drop_unmatched`	`bool`	Drop rows with a missing `_weight` — i.e. controls never used as a match and treated units that found no match. This is the estimation sample Stata restricts to before a weighted post-matching regression.	`True`

Returns:

Type	Description
`DataFrame`

balance ¶

balance(covariates: Optional[Sequence[str]] = None, *, threshold: float = 0.1) -> BalanceDiagnosticsResult

Covariate balance before vs after matching (StatsPAI conventions).

Standardized mean differences are reported two ways:

smd_raw — before matching: unweighted SMD over the full treated vs control sample.
smd_weighted — after matching: SMD with the _weight frequency weights, so a control used twice counts twice and unmatched / off-support units drop out (weight 0).

.. warning::

This is not Stata's pstest table, despite the family resemblance. pstest keeps the unmatched pooled standard deviation in the denominator of the post-matching bias; this method uses the matched-sample standard deviation, the convention most non-Stata packages follow. The two disagree by several percent on real data. Use :meth:pstest when you need to reproduce a printed Stata table.

Parameters:

Name	Type	Description	Default
`covariates`	`list of str`	Variables to assess. Defaults to the matching covariates.	`None`
`threshold`	`float`	\|SMD\| balance threshold.	`0.1`

Returns:

Type	Description
`BalanceDiagnosticsResult`	`.table` (per-covariate before vs after SMD, variance ratio, KS) and `.summary_stats`.

pstest ¶

pstest(covariates: Optional[Sequence[str]] = None) -> 'PSTestResult'

Stata pstest <covariates>, both — the printed balance table.

:meth:balance reports StatsPAI's own diagnostics. This reports Stata's, digit for digit, so a ported result can be checked against a pstest table without translating conventions. The two differ deliberately: pstest keeps the unmatched pooled standard deviation in the denominator of the post-matching bias (so the two rows are comparable), computes matched moments with Stata importance weights (variance divides by Σw - 1), and forms Rubin's B and R on the linear index rather than the propensity score.

Parameters:

Name	Type	Description	Default
`covariates`	`sequence of str`	Variables to test. Defaults to the matching covariates.	`None`

Returns:

Type	Description
`PSTestResult`	`.table` — per-covariate means, `%bias` before/after, `%reduct \|bias\|`, t / p, and V(T)/V(C) before/after. `.summary_stats` — the `Ps R2 / LR chi2 / p>chi2 / MeanBias / MedBias / B / R` block for the unmatched and matched samples.

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> m = sp.psmatch2(df, outcome='log_wage', treat='union',
...                 covariates=['education', 'experience'])
>>> t = m.pstest()
>>> bool(t.summary_stats['matched']['rubin_b']
...      < t.summary_stats['unmatched']['rubin_b'])
True

psplot ¶

psplot(*, before: bool = True, n_grid: int = 300, ax: Any = None, figsize: tuple[float, float] = (8.0, 4.5), title: Optional[str] = None) -> tuple[Any, Any]

Propensity-score density by treatment group, after matching.

Controls are reweighted by _weight so the plotted control density reflects the matched sample, not the raw pool. With before=True the raw (unweighted) densities are overlaid as dashed lines so the user can see how matching tightened overlap.

Returns:

Type	Description
`(fig, ax)`

psm_did ¶

psm_did(panel: DataFrame, *, id: str, y: str, time: Optional[str] = None, post: Optional[str] = None, treat: Optional[str] = None, treat_time: Optional[Any] = None, covariates: Optional[Sequence[str]] = None, fixed_effects: Optional[Sequence[str]] = None, cluster: Optional[Union[str, List[str]]] = None, on_support: bool = True, weight: str = 'aweight', alpha: float = 0.05) -> CausalResult

Weighted PSM-DID on a panel.

Implements the Stata workflow

.. code-block:: stata

psmatch2 d x1 x2, out(y) ...        // produces _weight
// merge _weight back onto the panel by id, then
reg y i.treat##i.post [aweight=_weight] if _support==1

The matching _weight (and _support) are merged onto panel by id, the matched sample is selected, and the weighted difference-in-differences regression

y ~ treat + post + treat:post (+ covariates | fixed_effects)

is fitted with :func:sp.feols. The treat:post coefficient is the PSM-DID treatment effect.

.. note::

Which Stata weight regime? aweight and fweight give the same coefficient but different standard errors, because fweight treats a control reused w times as w independent observations and so uses df = sum(w) - k instead of df = n_rows - k. On the reference fixture the DiD SE is 0.250051 under aweight and 0.214797 under fweight — the fweight interval is ~14% narrower purely from the degrees of freedom.

Choosing between them is mostly a reproducibility question: match the Stata line you are reconciling against. 'aweight' is the default because it is the more conservative of the two and the only one defined when k > 1 makes _weight fractional. 'fweight' is implemented by physically replicating rows — which is what a frequency weight means — and reproduces Stata bit-for-bit (see tests/reference_parity/test_psmdid_weight_parity.py).

Neither regime gives correct inference under matching with replacement. Both condition on the realised match structure and treat the matched sample as independently drawn, while a control serving several treated units induces dependence across those comparisons that no degrees-of-freedom convention captures. 'aweight' is the less optimistic of the two, not the correct one. When the standard error is load-bearing, cluster on the matched control's identity or bootstrap the pipeline.

.. versionchanged:: 1.22 weight='fweight' now computes genuine Stata fweight degrees of freedom. Before 1.22 it silently computed aweight numbers. The default changed from 'fweight' to 'aweight', so results from the default call are unchanged.

Parameters:

Name	Type	Description	Default
`panel`	`DataFrame`	Long panel (one row per unit-period).	required
`id`	`str`	Unit identifier. Must also exist in the matching data so the per-unit `_weight` can be merged in.	required
`y`	`str`	Outcome in the panel.	required
`time`	`str`	Time variable. Used with `treat_time` to build `post` if `post` is not supplied directly.	`None`
`post`	`str`	Binary post-period indicator. Provide this or `time` + `treat_time`.	`None`
`treat`	`str`	Time-invariant treated-group indicator in the panel. Defaults to the matching treatment variable.	`None`
`treat_time`	`scalar`	First treated period; `post = time >= treat_time`.	`None`
`covariates`	`list of str`	Additional time-varying controls.	`None`
`fixed_effects`	`list of str`	Columns absorbed as fixed effects (e.g. `[id, time]` for TWFE).	`None`
`cluster`	`str or list`	Cluster variable(s) for the standard errors.	`None`
`on_support`	`bool`	Keep only matched units on common support.	`True`
`weight`	`('aweight', 'fweight', 'none')`	Weight regime for the DiD regression, following Stata: `'aweight'` — WLS on the matched rows, `df = n_rows - k` (Stata `[aweight=_weight]`). The default, and the right choice when `_weight` is fractional (k > 1 neighbours). `'fweight'` — each row stands for `_weight` identical rows, `df = sum(_weight) - k` (Stata `[fweight=_weight]`). Requires integer weights, exactly as Stata does; raises :class:`~statspai.exceptions.MethodIncompatibility` otherwise. `'none'` — unweighted regression on the matched sample.	`'aweight'`
`alpha`	`float`	Significance level for the returned CI.	`0.05`

Returns:

Type	Description
`CausalResult`	`.estimate` is the DiD (`treat:post`) coefficient; the full weighted regression is stored in `model_info['feols_result']`.

summary ¶

summary() -> SummaryText

Stata-style text summary of the matched ATT.

cite ¶

cite(format: str = 'bibtex') -> Any

Citation for the matching estimator (delegates to the result).

PSTestResult ¶

Bases: ResultProtocolMixin

Stata pstest-shaped balance table (see :meth:PSMatch2Result.pstest).

Attributes:

Name	Type	Description
`table`	`DataFrame`	One row per covariate: unmatched / matched arm means, `%bias` before and after, `%reduct \|bias\|`, t and p, and V(T)/V(C).
`summary_stats`	`dict`	`{'unmatched': {...}, 'matched': {...}}` with `ps_r2`, `lr_chi2`, `p_chi2`, `mean_bias`, `median_bias`, `rubin_b`, `rubin_r` and `rubin_balanced`.

summary ¶

summary() -> SummaryText

Render the two blocks Stata's pstest ..., both prints.

SBWResult ¶

Bases: CausalResult

Stable balancing weights with a diagnostic panel.

Thin subclass of :class:CausalResult that attaches the weight vector, effective sample size, and covariate balance table. Returned by :func:sbw.

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage().iloc[:400].copy()
>>> res = sp.sbw(df, treat="union",
...              covariates=["education", "experience", "tenure"],
...              y="log_wage", delta=0.05)
>>> isinstance(res, sp.SBWResult)
True
>>> res.estimand
'ATT'
>>> list(res.balance.columns)
['mean_treated', 'mean_control', 'SMD_before', 'SMD_after']

TwoCriteriaMatchResult `dataclass` ¶

Bases: ResultProtocolMixin

Result of :func:two_criteria_match and :func:tighten_blocks.

Attributes:

Name	Type	Description
`matched`	`DataFrame`	The rows of the input that were matched, treated unit first within each set, with a column `mset` numbering the sets (and `pscore` when a propensity score was fitted).
`pairs`	`DataFrame`	One row per treated-control pair: index labels of the two units, `mset` and the pairing cost.
`balance`	`DataFrame`	Covariate means of the treated, the matched controls and all controls, with standardized differences before and after.
`pair_cost, balance_cost, control_cost, total_cost`	`float`	The parts of the minimised objective and their sum.
`n_treated, n_sets, n_unmatched`	`int`	Treated units, matched sets formed, treated units left out.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> df = pd.DataFrame({"age": rng.normal(50, 10, n),
...                    "female": rng.integers(0, 2, n)})
>>> df["z"] = rng.binomial(1, 1 / (1 + np.exp(-(df.age - 62) / 6)))
>>> m = sp.two_criteria_match(
...     df, "z",
...     pair=[{"type": "mahalanobis", "on": ["age", "female"]}],
...     balance=[{"type": "near_exact", "on": "female"}])
>>> type(m).__name__
'TwoCriteriaMatchResult'
>>> m.n_sets == int(df.z.sum())
True

WeightItResult ¶

Bases: ResultProtocolMixin

Balancing weights returned by :func:sp.weightit.

Attributes:

Name	Type	Description
`weights`	`Series`	One weight per row of the input data (`NaN` for rows dropped for missing values).
`ps`	`Series or DataFrame or None`	Binary: `P(treat = treated level \| X)`. Multi-category: one column of generalized propensity scores per level. Continuous `'glm'`: the fitted conditional mean of the treatment. `None` for entropy balancing.
`treat_type, method, estimand, focal`	`str`	What was estimated.
`diagnostics`	`dict`	Weight range, coefficient of variation, effective sample sizes (overall and per group), number of rows dropped.

Examples:

>>> import numpy as np, pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> x = rng.normal(size=300)
>>> g = np.where(x + rng.normal(size=300) > 0.5, 'b', 'a')
>>> df = pd.DataFrame({'g': g, 'x': x})
>>> W = sp.weightit(df, treat='g', covariates=['x'], method='ebal')
>>> bool(abs(W.balance().table.loc['x', 'diff_adj']) < 1e-8)
True

balance ¶

balance(**kwargs: Any) -> Any

Balance table of these weights (:func:sp.bal_tab).

l1_imbalance ¶

l1_imbalance(data: DataFrame, treat: str, covariates: Sequence[str], weights: Optional[Any] = None, n_bins: Optional[int] = None) -> float

Multivariate L1 imbalance of Iacus, King and Porro (2011).

Each numeric covariate with more than n_bins distinct values is cut into n_bins equal-width bins between its minimum and maximum (bins closed on the right, the first on both sides, as sp.match(method= 'cem') coarsens); other covariates keep their values. L1 = 1/2 sum_cells |f_cell - g_cell| with f and g the relative (weighted) frequencies of the treated and the controls.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`		required
`treat`	`str`	0/1 treatment.	required
`covariates`	`sequence of str`		required
`weights`	`array - like or str`	Matching or balancing weights (`NaN` = unmatched, weight 0); the default is the unweighted sample. The bins are set on every row with complete covariates whatever its weight, so a weighted (after matching) and an unweighted (before matching) L1 use the same cells. Negative or infinite weights are refused.	`None`
`n_bins`	`int`	Bins per covariate; default Sturges' rule `ceil(log2(n) + 1)` (the default of R `cem`'s `imbalance` is a different, histogram-based coarsening, so the numbers are comparable only under the same coarsening).	`None`

Returns:

Type	Description
`float`	In `[0, 1]`.

Examples:

>>> import numpy as np, pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> df = pd.DataFrame({'d': rng.integers(0, 2, 500), 'x': rng.normal(size=500)})
>>> bool(0 <= sp.l1_imbalance(df, 'd', ['x']) < 0.3)
True

References

[@iacus2011multivariate]

balance_vs_randomization ¶

balance_vs_randomization(data: DataFrame, treat: str, covariates: Sequence[str], *, n_sim: int = 1000, test: str = 'auto', max_levels: int = 2, trunc: float = 0.2, alpha: float = 0.05, random_state: Optional[int] = None) -> BalanceRandomizationResult

Compare covariate balance in a matched sample with randomization.

For each covariate a two-sample test compares the treated and control groups of the matched sample. The treatment labels are then permuted n_sim times, as if the same individuals had been assigned at random with the same group sizes, and the tests are repeated. A match whose p-values look like those of the simulated experiments, or larger, has balanced the observed covariates about as well as randomization would have [@hansen2008covariate; @pimentel2015large].

The p-values serve as a yardstick here. They are not tests of a hypothesis anyone holds: a matched sample was not randomized, and matching commonly balances observed covariates better than randomization does.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	The matched sample, one row per individual.	required
`treat`	`str`	0/1 treatment indicator.	required
`covariates`	`list of str`	Numeric covariates to assess. A categorical or text column is refused: expand it into 0/1 indicators first, one test per level.	required
`n_sim`	`int`	Number of simulated randomized experiments.	`1000`
`test`	`('auto', 'wilcoxon', 't')`	`"wilcoxon"` is the two-sample rank-sum test, `"t"` the unequal-variance t-test. `"auto"` uses the chi-square test for a covariate with at most `max_levels` distinct values and the rank-sum test otherwise.	`"auto"`
`max_levels`	`int`	See `test`.	`2`
`trunc`	`float`	Truncation point for the product of the p-values, which summarises balance across covariates; `1` gives Fisher's product.	`0.2`
`alpha`	`float`	Threshold for counting small p-values.	`0.05`
`random_state`	`int`	Seed for the permutations.	`None`

Returns:

Type	Description
`BalanceRandomizationResult`

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> df = pd.DataFrame({"z": np.tile([1, 0], 100),
...                    "age": rng.normal(50, 10, 200),
...                    "female": rng.integers(0, 2, 200)})
>>> res = sp.balance_vs_randomization(df, "z", ["age", "female"],
...                                   n_sim=200, random_state=1)
>>> res.tests_used
{'age': 'wilcoxon', 'female': 'chi2'}
>>> bool(0 <= res.table.loc["min_p", "share_better"] <= 1)
True

References

[@hansen2008covariate], [@pimentel2015large], [@yu2021evaluating], [@zaykin2002truncated]

full_match ¶

full_match(data: DataFrame, y: Optional[str] = None, treat: Optional[str] = None, covariates: Optional[Sequence[str]] = None, *, distance: str = 'propensity', pscore: Optional[str] = None, estimand: str = 'ATT', caliper: Optional[float] = None, caliper_scale: str = 'sd', ps_model: str = 'logit', alpha: float = 0.05) -> FullMatchResult

Optimal full matching.

Places every treated and every comparison unit in a matched set with one treated and several comparison units, or one comparison and several treated units, so that the total within-set distance is minimal. Unlike pair matching it uses the whole sample, and unlike subclassification the sets are chosen by the data.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`		required
`y`	`str`	Outcome. Without it only the matching (sets, weights, balance) is returned.	`None`
`treat`	`str`	0/1 treatment indicator.	`None`
`covariates`	`sequence of str`	Variables to match on, and the variables of the balance table.	`None`
`distance`	`('propensity', 'logit', 'mahalanobis', 'euclidean')`	`'propensity'`: absolute difference of the fitted propensity score (`MatchIt`'s default, `distance = "glm"`). `'logit'`: difference on the linear-predictor scale. `'mahalanobis'` uses the pooled within-group covariance, as `MatchIt` does.	`'propensity'`
`pscore`	`str`	Column holding a score computed elsewhere (a boosted or forest propensity score, say). Matching is on the absolute difference of this column and `covariates` are used for the balance table only.	`None`
`estimand`	`('ATT', 'ATC', 'ATE')`	Sets the matching weights; the matched sets are the same.	`'ATT'`
`caliper`	`float`	Largest permitted distance. A unit with no partner inside the caliper is left unmatched and the estimand then refers to the matched units.	`None`
`caliper_scale`	`('sd', 'raw')`	`'sd'` reads `caliper` in standard deviations of the score (score distances only), `'raw'` in the units of the distance.	`'sd'`
`ps_model`	`('logit', 'probit')`		`'logit'`
`alpha`	`float`		`0.05`

Returns:

Type	Description
`FullMatchResult`

Notes

The matching is exact: the reported total_distance is the minimum over all full matchings. optmatch::fullmatch (and MatchIt with method = "full", which calls it) solves the same problem on distances rounded to a tolerance, so its total can be slightly larger and the matched sets, hence the estimate, can differ a little. When several matchings attain the minimum total distance, which of them is returned, and so the matched sets and the estimate, can depend on the order of the rows.

When the covariates separate treated from comparison units the propensity model has no maximum-likelihood fit; a ConvergenceWarning is issued and the match on such a score balances nothing. Match on the covariates themselves (distance='mahalanobis') or drop the separating covariate.

Matched sets can be very unequal in size when overlap is poor; a few comparison units then carry most of the weight. Inspect model_info['set_sizes'] and the effective sample size model_info['ess_control'].

Examples:

>>> import numpy as np, pandas as pd, statspai as sp
>>> rng = np.random.default_rng(1)
>>> n = 300
>>> x = rng.normal(size=n)
>>> d = rng.binomial(1, 1 / (1 + np.exp(-x)))
>>> df = pd.DataFrame({"x": x, "d": d, "y": 1 + x + 2 * d + rng.normal(size=n)})
>>> fit = sp.full_match(df, "y", "d", ["x"])
>>> int(fit.n_unmatched)
0
>>> bool(abs(fit.balance.loc["x", "smd_matched"]) < 0.1)
True

References

[@rosenbaum1991characterization], [@hansen2004full]

balanceplot ¶

balanceplot(result: CausalResult, threshold: float = 0.1, ax: Any = None, figsize: tuple = (8, None), title: Optional[str] = None) -> Tuple[Any, Any]

Love plot: covariate balance visualization (SMD dot plot).

Displays standardized mean differences (SMD) for each covariate. The standard threshold for good balance is |SMD| < 0.1.

Parameters:

Name	Type	Description	Default
`result`	`CausalResult`	Result from `match()` or `ebalance()`.	required
`threshold`	`float`	SMD threshold lines.	`0.1`
`ax`	`matplotlib Axes`		`None`
`figsize`	`tuple`	Height auto-scales with number of covariates if None.	`(8, None)`
`title`	`str`		`None`

Returns:

Type	Description
`(fig, ax)`

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> age = rng.normal(40, 8, n)
>>> edu = rng.normal(12, 2, n)
>>> ps = 1 / (1 + np.exp(-(0.05 * (age - 40) + 0.1 * (edu - 12))))
>>> training = rng.binomial(1, ps)
>>> wage = 20 + 0.3 * age + 0.5 * edu + 4.0 * training + rng.normal(0, 3, n)
>>> df = pd.DataFrame({"wage": wage, "training": training,
...                    "age": age, "edu": edu})
>>> result = sp.match(df, y="wage", treat="training",
...                   covariates=["age", "edu"])
>>> fig, ax = sp.balanceplot(result)
>>> fig.savefig("balance.png")
>>> type(ax).__name__
'Axes'

psplot ¶

psplot(data: DataFrame, treat: str, covariates: List[str], *, n_bins: int = 40, ax: Any = None, figsize: tuple = (8, 5), title: Optional[str] = None, labels: tuple = ('Control', 'Treated'), colors: tuple = ('#3498DB', '#E74C3C'), trim: Optional[float] = None) -> Tuple[Any, Any]

Propensity score distribution plot (common support diagnostic).

Overlays histograms of the estimated propensity score for treated and control groups, so the user can visually assess whether the common support (overlap) assumption holds.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`		required
`treat`	`str`	Binary treatment column.	required
`covariates`	`list of str`	Covariates used to estimate the propensity score.	required
`n_bins`	`int`	Number of histogram bins.	`40`
`ax`	`matplotlib Axes`		`None`
`figsize`	`tuple`		`(8, 5)`
`title`	`str`		`None`
`labels`	`tuple of str`	Labels for (control, treated).	`('Control', 'Treated')`
`colors`	`tuple of str`	Colors for (control, treated).	`('#3498DB', '#E74C3C')`
`trim`	`float`	If set, draw vertical lines at (trim, 1-trim) to show the recommended trimming region.	`None`

Returns:

Type	Description
`(fig, ax)`

Examples:

>>> import statspai as sp, numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> x1, x2 = rng.normal(size=n), rng.normal(size=n)
>>> D = rng.binomial(1, 1 / (1 + np.exp(-(x1 + 0.5 * x2))))
>>> df = pd.DataFrame({"D": D, "x1": x1, "x2": x2})
>>> fig, ax = sp.psplot(df, treat="D", covariates=["x1", "x2"])

cardinality_match ¶

cardinality_match(data: DataFrame, treatment: str, outcome: str, covariates: List[str], smd_tolerance: float = 0.1, time_limit: float = 30.0) -> CardinalityMatchResult

Cardinality matching — maximise the number of matched pairs subject to a standardised-mean-difference tolerance on every covariate.

Formulation (Zubizarreta 2014):

maximise   sum_j z_j
s.t.       |mean(X_k | T=1) - sum_j z_j X_{jk} / sum_j z_j|
            <= smd_tolerance * SD(X_k)   ∀ k
           z_j ∈ {0, 1}  for each control j

Solved exactly as a binary integer program (scipy.optimize.milp, HiGHS), so the returned set always satisfies the tolerance. Matched pairs are the matched controls each assigned to a treated unit by optimal (Hungarian) assignment on the Mahalanobis distance.

.. versionchanged:: 1.22 Previously relaxed to a continuous LP and rounded the weights by a threshold. Rounding does not preserve the balance constraints, so the returned sample could — and usually did — violate smd_tolerance: on a 12-cell seed x tolerance grid, 9 solutions were infeasible, by up to 26% of the requested tolerance. Matched sets and effect estimates therefore change.

Examples:

Simulated observational data with two confounders (true ATT = 2):

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> p = 1.0 / (1.0 + np.exp(-(0.5 * x1 - 0.5 * x2 - 0.5)))
>>> d = rng.binomial(1, p)
>>> y = 1.0 + 2.0 * d + x1 + x2 + rng.normal(size=n)
>>> df = pd.DataFrame({'y': y, 'd': d, 'x1': x1, 'x2': x2})
>>> res = sp.cardinality_match(df, treatment='d', outcome='y',
...                            covariates=['x1', 'x2'],
...                            smd_tolerance=0.1)
>>> res.n_matched_pairs
106
>>> round(res.ate, 2)
1.84
>>> bool((res.balance['|SMD|'] <= 0.1).all())
True

optimal_match ¶

optimal_match(data: DataFrame, treatment: str, outcome: str, covariates: List[str], metric: str = 'mahalanobis', caliper: Optional[float] = None) -> OptimalMatchResult

Optimal 1:1 matching via the Hungarian algorithm.

Each treated unit is matched to exactly one control; the total sum of matched distances is globally minimised.

When there are fewer controls than treated units, every control is matched to a distinct treated unit and the remaining treated units go unmatched (the behaviour of R MatchIt::matchit(method = "optimal"), which warns "Fewer control units than treated units"). The estimate is then the effect on the matched treated units, not the ATT; a warning says how many treated units were dropped and n_unmatched_treated records it. sp.full_match keeps every unit.

Parameters:

Name	Type	Description	Default
`caliper`	`float`	Drop any pair with distance greater than `caliper`.	`None`

balance_diagnostics ¶

balance_diagnostics(data: DataFrame, treatment: str, covariates: List[str], weights: Optional[Union[ndarray, Series, str]] = None, ps: Optional[Union[ndarray, Series, str]] = None, method: str = 'logit', threshold: float = 0.1, sd_denom: str = 'weighted') -> BalanceDiagnosticsResult

Unified balance diagnostics for matching and weighting estimators.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Analysis frame.	required
`treatment`	`str`	Binary treatment indicator.	required
`covariates`	`list of str`	Covariates to audit.	required
`weights`	`array - like or str`	Observation weights after matching/weighting. If omitted, ATE inverse-propensity weights are computed from `ps`.	`None`
`ps`	`array - like or str`	Propensity scores. If omitted, estimated with `method`.	`None`
`method`	`(logit, probit, gbm)`	Propensity-score model when `ps` is not supplied.	`'logit'`
`threshold`	`float`	Balance threshold for absolute standardized mean differences.	`0.1`
`sd_denom`	`(weighted, unweighted)`	Standard deviation under the weighted mean difference. `'weighted'` uses the weighted variances of the two groups (Austin and Stuart 2015 [@austin2015moving]). `'unweighted'` keeps the variances of the unweighted sample, the default of R `cobalt`: the scale is then the same before and after weighting and weights cannot improve the statistic by inflating the spread.	`'weighted'`

Returns:

Type	Description
`BalanceDiagnosticsResult`	`.table` has one row per covariate; `.summary_stats` records max/mean SMDs, imbalance counts, the effective sample size (overall and in each group), the energy distance between the two groups before and after weighting (a joint measure that also reacts to interactions and higher moments; see :func:`sp.energy_distance`), and propensity-score overlap.

Notes

Weighted variances use sum(w (x - m)^2) / (sum(w) - sum(w^2) / sum(w)), which does not depend on the scale of the weights and reduces to the n - 1 variance for equal weights, so weights=1 reproduces the unweighted columns exactly.

Examples:

With no weights, ATE inverse-propensity weights are computed from the estimated propensity scores:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> bal = sp.balance_diagnostics(
...     df, treatment='union',
...     covariates=['education', 'experience', 'tenure'])
>>> bal.summary_stats['n_obs']
3000
>>> bool(bal.summary_stats['n_imbalanced_weighted']
...      <= bal.summary_stats['n_imbalanced_raw'])
True

Typical post-estimation flow — audit your own weights and scores:

>>> import numpy as np
>>> ps = sp.propensity_score(
...     df, 'union', ['education', 'experience', 'tenure'])
>>> w = np.where(df['union'] == 1, 1 / ps, 1 / (1 - ps))
>>> bal = sp.balance_diagnostics(
...     df, treatment='union',
...     covariates=['education', 'experience', 'tenure'],
...     weights=w, ps=ps)
>>> bool(bal.summary_stats['effective_sample_size'] > 0)
True

love_plot ¶

love_plot(data: Any, treatment: Optional[str] = None, covariates: Optional[List[str]] = None, weights: Optional[Union[ndarray, Series]] = None, threshold: float = 0.1, ps_method: str = 'logit', ax: Any = None, figsize: Tuple[float, Optional[float]] = (7, None), title: str = 'Covariate Balance (Love Plot)') -> Tuple[Any, Any]

Love plot: dot plot of standardized mean differences before/after.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Input data.	required
`treatment`	`str`	Binary treatment column.	`None`
`covariates`	`list of str`	Covariate columns.	`None`
`weights`	`array - like`	IPW or matching weights. If None, inverse-PS weights are computed. A missing weight counts as zero (an unmatched row).	`None`
`threshold`	`float`	SMD threshold for the vertical dashed line (default 0.1).	`0.1`
`ps_method`	`str`	PS estimation method for balance computation.	`'logit'`
`ax`	`matplotlib Axes`		`None`
`figsize`	`tuple`	(width, height). Height defaults to 0.4 * n_covariates + 1.	`(7, None)`
`title`	`str`	Plot title.	`'Covariate Balance (Love Plot)'`

Returns:

Type	Description
`(fig, ax) : tuple`

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> fig, ax = sp.love_plot(df, treatment='union',
...                        covariates=['education', 'experience',
...                                    'tenure'])
>>> fig.savefig('love_plot.png')

>>> # Tighter balance threshold and custom title
>>> fig, ax = sp.love_plot(df, treatment='union',
...                        covariates=['education', 'experience',
...                                    'tenure'],
...                        threshold=0.05,
...                        title='Balance: union vs non-union')

overlap_plot ¶

overlap_plot(data: DataFrame, treatment: str, covariates: List[str], ps: Optional[Series] = None, method: str = 'logit', ax: Any = None, figsize: Tuple[float, float] = (8, 4), title: str = 'Propensity Score Overlap') -> Tuple[Any, Any]

Mirrored density plot of propensity scores by treatment group.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Input data.	required
`treatment`	`str`	Binary treatment column.	required
`covariates`	`list of str`	Covariates for PS estimation (ignored if ps supplied).	required
`ps`	`Series`	Pre-estimated propensity scores.	`None`
`method`	`str`	PS estimation method if ps is None.	`'logit'`
`ax`	`matplotlib Axes`	Axes to plot on. If None, a new figure is created.	`None`
`figsize`	`tuple`	Figure size (width, height).	`(8, 4)`
`title`	`str`	Plot title.	`'Propensity Score Overlap'`

Returns:

Type	Description
`(fig, ax) : tuple`	Matplotlib figure and axes.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> p = 1.0 / (1.0 + np.exp(-(0.5 * x1 - 0.5 * x2 - 0.5)))
>>> d = rng.binomial(1, p)
>>> df = pd.DataFrame({'d': d, 'x1': x1, 'x2': x2})
>>> fig, ax = sp.overlap_plot(df, treatment='d',
...                           covariates=['x1', 'x2'])
>>> fig.savefig('overlap.png')

Reuse pre-estimated propensity scores and set a custom title:

>>> ps = sp.propensity_score(df, treatment='d',
...                          covariates=['x1', 'x2'])
>>> fig, ax = sp.overlap_plot(df, treatment='d',
...                           covariates=['x1', 'x2'], ps=ps,
...                           title='PS overlap')

propensity_score ¶

propensity_score(data: DataFrame, treatment: str, covariates: List[str], method: str = 'logit', trimming: Optional[str] = None) -> Series

Estimate propensity scores P(D=1|X).

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Input data.	required
`treatment`	`str`	Name of binary treatment column (0/1).	required
`covariates`	`list of str`	Covariate column names.	required
`method`	`(logit, probit, gbm)`	Estimation method. `'logit'` uses IRLS (no sklearn needed). `'probit'` uses scipy.optimize. `'gbm'` tries sklearn GradientBoostingClassifier, falling back to logit with interactions.	`'logit'`
`trimming`	`(None, crump)`	If `'crump'`, apply Crump et al. (2009) trimming after estimation. Trimmed observations receive `NaN` scores.	`None`

Returns:

Type	Description
`Series`	Propensity scores indexed like data. A row with a missing treatment or covariate value is left out of the fit and has a `NaN` score (with a warning that says how many).

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> ps = sp.propensity_score(df, treatment='union',
...                          covariates=['education', 'experience',
...                                      'tenure'])
>>> round(float(ps.mean()), 3)  # matches the union share of 0.177
0.177

GBM scores with Crump trimming — poorly overlapping observations receive NaN:

>>> ps_trim = sp.propensity_score(df, treatment='union',
...                               covariates=['education', 'experience',
...                                           'tenure'],
...                               method='gbm', trimming='crump')

Typical diagnostics flow afterwards: pass the scores (or derived IPW weights) to :func:sp.overlap_plot, :func:sp.ps_balance, or :func:sp.balance_diagnostics.

ps_balance ¶

ps_balance(data: DataFrame, treatment: str, covariates: List[str], weights: Optional[Union[ndarray, Series]] = None, method: str = 'logit', sd_denom: str = 'weighted') -> PSBalanceResult

Compute comprehensive propensity score balance table.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Input data.	required
`treatment`	`str`	Binary treatment column.	required
`covariates`	`list of str`	Covariate columns to assess balance for.	required
`weights`	`array - like`	IPW or matching weights. If None, inverse-PS weights are computed automatically from estimated propensity scores. A missing weight means the row is outside the weighted sample (unmatched rows of a matched frame) and counts as zero.	`None`
`method`	`str`	PS estimation method ('logit', 'probit', 'gbm').	`'logit'`
`sd_denom`	`(weighted, unweighted)`	Standard deviation under the weighted mean difference. `'weighted'` uses the weighted variances of the two groups (Austin and Stuart 2015 [@austin2015moving]). `'unweighted'` keeps the variances of the unweighted sample, the default of R `cobalt`: the scale is then the same before and after weighting and weights cannot improve the statistic by inflating the spread.	`'weighted'`

Returns:

Type	Description
`PSBalanceResult`	Object with `.table`, `.ps`, `.summary()`, `.love_plot()`.

Examples:

Simulated data with confounded treatment assignment:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> p = 1.0 / (1.0 + np.exp(-(0.5 * x1 - 0.5 * x2 - 0.5)))
>>> d = rng.binomial(1, p)
>>> df = pd.DataFrame({'d': d, 'x1': x1, 'x2': x2})

Without weights, ATE inverse-propensity weights are computed from the estimated propensity scores:

>>> bal = sp.ps_balance(df, treatment='d', covariates=['x1', 'x2'])
>>> bal.table['smd_raw'].round(2).tolist()
[0.68, -0.45]
>>> bal.table['smd_weighted'].round(2).tolist()
[0.02, -0.06]
>>> fig, ax = bal.love_plot()

trimming ¶

trimming(data: DataFrame, treatment: str, covariates: List[str], method: str = 'crump', ps: Optional[Series] = None, ps_method: str = 'logit') -> DataFrame

Trim sample to optimal overlap region.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Input data.	required
`treatment`	`str`	Binary treatment column.	required
`covariates`	`list of str`	Covariates for PS estimation (if ps not supplied).	required
`method`	`(crump, sturmer)`	`'crump'` uses Crump et al. (2009) optimal rule. `'sturmer'` trims at the fixed [0.1, 0.9] interval.	`'crump'`
`ps`	`Series`	Pre-estimated propensity scores. If None, estimated via ps_method.	`None`
`ps_method`	`str`	Method for PS estimation if ps is None.	`'logit'`

Returns:

Type	Description
`DataFrame`	Trimmed data (rows with PS in the overlap region).

Examples:

Strong selection on covariates creates limited overlap:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> p = 1.0 / (1.0 + np.exp(-(2.0 * x1 - 2.0 * x2)))
>>> d = rng.binomial(1, p)
>>> df = pd.DataFrame({'d': d, 'x1': x1, 'x2': x2})

Crump et al. (2009) optimal rule drops poor-overlap rows:

>>> trimmed = sp.trimming(df, treatment='d',
...                       covariates=['x1', 'x2'])
>>> (len(df), len(trimmed))
(300, 192)

Fixed [0.1, 0.9] trimming keeps a narrower sample:

>>> trimmed_s = sp.trimming(df, treatment='d',
...                         covariates=['x1', 'x2'],
...                         method='sturmer')
>>> len(trimmed_s)
180

energy_distance ¶

energy_distance(data: DataFrame, treat: str, covariates: Sequence[str], weights: Any = None) -> float

Energy distance between the covariate distributions of two groups.

2 E|X1 - X0| - E|X1 - X1'| - E|X0 - X0'| with Euclidean distance on standardised covariates and expectations taken over the (weighted) samples. It is zero only when the two weighted joint distributions coincide, so unlike one standardized difference per covariate it picks up imbalance in interactions and higher moments.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`		required
`treat`	`str`	Binary 0/1 column.	required
`covariates`	`sequence of str`	Numeric columns. A categorical or text column is refused: expand it into 0/1 indicators first (the codes of a category are not distances).	required
`weights`	`str or array - like`	Weights; omitted, the unweighted distance. A Series is matched to the rows of `data` on its index and is an error if the labels differ; an array is taken to be in row order.	`None`

Returns:

Type	Description
`float`

Notes

Each covariate is centred and divided by its unweighted standard deviation (sqrt(p (1 - p)) for a 0/1 column), so the scaling is the same before and after weighting and two sets of weights can be compared on one scale. This is the between-group statistic that R halfmoon::bal_energy reports.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> x = rng.normal(size=400)
>>> d = rng.binomial(1, 1 / (1 + np.exp(-x)))
>>> df = pd.DataFrame({'d': d, 'x': x})
>>> raw = sp.energy_distance(df, 'd', ['x'])
>>> ps = sp.propensity_score(df, 'd', ['x'])
>>> w = sp.ps_weights(ps, df['d'])
>>> bool(sp.energy_distance(df, 'd', ['x'], weights=w) < raw)
True

References

[@szekely2013energy], [@huling2024energy]

ess ¶

ess(weights: Any, by: Any = None) -> Union[float, Series]

Effective sample size of a set of weights.

ESS = (sum w)^2 / sum(w^2): the number of equally weighted observations that would give a weighted mean of the same variance. It equals n when all weights are equal and falls as a few rows come to dominate.

Parameters:

Name	Type	Description	Default
`weights`	`array - like`	Weights. Negative entries are allowed, since the weights a regression implies can be negative (:func:`sp.implied_weights`); the formula is applied as it stands. When every weight (of a group, with `by`) is zero the value is NaN and a warning is issued.	required
`by`	`array - like`	Group labels (typically the exposure). The effective sample size is then returned for each group, which is what matters for a contrast: the precision of each group mean depends on its own ESS. Missing labels are an error.	`None`

Returns:

Type	Description
`float or Series`	A float, or with `by` a Series indexed by group.

Examples:

>>> import numpy as np
>>> import statspai as sp
>>> sp.ess(np.ones(50))
50.0
>>> round(sp.ess([1, 1, 1, 9]), 2)
1.71
>>> sp.ess([1, 1, 1, 9], by=[0, 0, 1, 1]).round(2).tolist()
[2.0, 1.22]

implied_weights ¶

implied_weights(data: DataFrame, treat: str, covariates: Sequence[str], *, interactions: bool = False, estimand: str = 'ATE') -> Series

Weights that a linear regression adjustment implicitly puts on each row.

The coefficient on a binary treatment in y ~ treat + covariates is a difference of two weighted means of y, with weights that depend only on the treatment and the covariates. Looking at those weights, without the outcome, shows what the regression is doing: which population it represents, how many observations effectively carry the estimate, and whether some rows enter with a negative weight, that is, are extrapolated.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`		required
`treat`	`str`	Binary 0/1 column.	required
`covariates`	`sequence of str`	Regressors other than the treatment. Categorical columns (`category` dtype, strings, `C(col)`) are expanded into indicators, the first level as the base, as the regression would enter them; integer columns are taken as numeric.	required
`interactions`	`bool`	`False`: the regression `y ~ treat + covariates` with one set of covariate slopes ("uniform regression imputation"). `True`: the treatment interacted with every centred covariate, i.e. a separate regression in each arm ("multi regression imputation").	`False`
`estimand`	`('ATE', 'ATT', 'ATC')`	With `interactions=True`, the population the two arm regressions are averaged over. The regression without interactions has no choice here: its weights target a population of its own, between the two groups, and `estimand` must stay `'ATE'`.	`'ATE'`

Returns:

Type	Description
`Series`	One weight per row, scaled to sum to the group size within each treatment group (so that 1 means "counts as one observation"). The regression coefficient equals the weighted mean of `y` among the treated minus the weighted mean among the controls.

Notes

Without interactions the weights are n_g * r_i / sum(r^2) in absolute value, with r the residual of the treatment regressed on the covariates. With interactions the weights of arm g are n_g * xbar' (X_g' X_g)^{-1} x_i with xbar the covariate mean of the target population. Both match R lmw::lmw.

The weights do not exist when the covariates predict the treatment perfectly (no interactions) or are collinear within a group (interactions): the regression then has no treatment coefficient, and an error is raised.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> x = rng.normal(size=300)
>>> d = rng.binomial(1, 1 / (1 + np.exp(-x)))
>>> y = 1 + 2 * d + x + rng.normal(size=300)
>>> df = pd.DataFrame({'y': y, 'd': d, 'x': x})
>>> w = sp.implied_weights(df, 'd', ['x'])
>>> ols = sp.regress('y ~ d + x', data=df).params['d']
>>> t, c = df['d'] == 1, df['d'] == 0
>>> diff = (np.average(df['y'][t], weights=w[t])
...         - np.average(df['y'][c], weights=w[c]))
>>> bool(abs(diff - ols) < 1e-10)
True

References

[@chattopadhyay2023implied]

tighten_blocks ¶

tighten_blocks(data: DataFrame, treat: str, block: str, *, covariates: Optional[Sequence[str]] = None, fine_balance: Optional[Sequence[str]] = None, ratio: int = 1, subset_cost: Optional[float] = None, penalty_scale: float = 10.0) -> TwoCriteriaMatchResult

Tighten a block design into smaller or fewer blocks.

Starts from blocks that each hold one treated individual and several controls, already matched for some covariates, and keeps ratio of the controls in each block (and, with subset_cost, only some of the blocks) so that the retained sample is also comparable on covariates the original blocks did not control. Controls are always taken from the treated individual's own block. Among such choices the function first balances the fine_balance covariates between the retained treated and control groups, and then minimises the within-block rank-based Mahalanobis distance on covariates.

Typical uses are a secondary analysis that adjusts for a covariate possibly affected by the treatment, which the primary design deliberately left out, and a differential comparison within the blocks that have a particular pattern of a second exposure.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	One row per individual of the block design.	required
`treat`	`str`	0/1 treatment indicator and block identifier. Each block must hold exactly one treated individual and at least one control.	required
`block`	`str`	0/1 treatment indicator and block identifier. Each block must hold exactly one treated individual and at least one control.	required
`covariates`	`list of str`	Covariates for the within-block Mahalanobis distance.	`None`
`fine_balance`	`list of str`	Nominal covariates whose distributions should agree between the retained treated and control groups. Each contributes a penalty for every treated-control mismatch that cannot be avoided.	`None`
`ratio`	`int`	Controls kept per block.	`1`
`subset_cost`	`float`	Price of dropping a block. Smaller values drop more blocks in exchange for better balance; `None` keeps every block. Requires `ratio=1`.	`None`
`penalty_scale`	`float`	Factor separating the priorities: staying within the block outweighs fine balance, which outweighs covariate distance.	`10.0`

Returns:

Type	Description
`TwoCriteriaMatchResult`	`.matched` holds the tightened design; `mset` numbers its blocks.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> blocks = np.repeat(np.arange(80), 4)
>>> df = pd.DataFrame({"block": blocks,
...                    "z": np.tile([1, 0, 0, 0], 80),
...                    "bmi": rng.normal(27, 4, 320),
...                    "smoker": rng.integers(0, 2, 320)})
>>> tight = sp.tighten_blocks(df, "z", "block", covariates=["bmi"],
...                           fine_balance=["smoker"], ratio=2)
>>> tight.matched.groupby("mset").size().unique().tolist()
[3]
>>> bool(tight.matched.groupby("mset")["block"].nunique().eq(1).all())
True

References

[@rosenbaum2012optimal], [@zhang2023matching], [@rosenbaum2010design]

two_criteria_match ¶

two_criteria_match(data: DataFrame, treat: str, *, pair: Union[Sequence[Term], ndarray, None] = None, balance: Union[Sequence[Term], ndarray, None] = None, ratio: int = 1, ps: Union[str, Sequence[str], None] = None, control_cost: Union[str, Sequence[float], None] = None, subset_cost: Optional[float] = None) -> TwoCriteriaMatchResult

Optimal matching that pairs on one distance and balances on another.

Chooses ratio distinct controls for each treated unit so as to minimise the sum of two costs [@zhang2023matching]. The pair cost is the total distance between each treated unit and its own controls. The balance cost is the smallest total distance at which the chosen controls can be assigned to the treated units when that assignment is free to ignore the pairing, so it depends only on which controls were chosen. A covariate placed in pair is matched closely within sets; a covariate placed in balance has its distribution equalised between the groups without constraining who is paired with whom. A large near_exact or integer penalty in balance is fine balance when it can be met and near-fine balance when it cannot.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	One row per unit. Rows with a missing `treat` are ignored.	required
`treat`	`str`	0/1 treatment indicator.	required
`pair`	`list of dict, or ndarray`	The two distances, each the sum of its terms. A term is a dict with a `"type"`, the column(s) it is `"on"`, and for all but the first type a `"penalty"` (default 1000): `{"type": "mahalanobis", "on": [...]}`: squared rank-based Mahalanobis distance on the listed covariates. `{"type": "near_exact", "on": col, "penalty": p}`: `p` when the two units differ on a nominal covariate. `{"type": "integer", "on": col, "penalty": p}`: `p` times the absolute difference of an integer-coded covariate. `{"type": "caliper", "on": col, "width": w, "penalty": p, "two_step": True}`: `p` when treated minus control falls outside the caliper, and `p` more when it falls outside twice the caliper. `w` is a half-width, a pair `(low, high)` for an asymmetric caliper that tolerates mismatches in one direction, or omitted for 0.2 standard deviations. `{"type": "quantile", "on": col, "probs": [...], "penalty": p}`: an `integer` term on the covariate cut at its quantiles. The column name `"pscore"` refers to the propensity score requested through `ps`. A treated-by-control matrix (rows and columns in data order) may be passed instead of a list.	`None`
`balance`	`list of dict, or ndarray`	The two distances, each the sum of its terms. A term is a dict with a `"type"`, the column(s) it is `"on"`, and for all but the first type a `"penalty"` (default 1000): `{"type": "mahalanobis", "on": [...]}`: squared rank-based Mahalanobis distance on the listed covariates. `{"type": "near_exact", "on": col, "penalty": p}`: `p` when the two units differ on a nominal covariate. `{"type": "integer", "on": col, "penalty": p}`: `p` times the absolute difference of an integer-coded covariate. `{"type": "caliper", "on": col, "width": w, "penalty": p, "two_step": True}`: `p` when treated minus control falls outside the caliper, and `p` more when it falls outside twice the caliper. `w` is a half-width, a pair `(low, high)` for an asymmetric caliper that tolerates mismatches in one direction, or omitted for 0.2 standard deviations. `{"type": "quantile", "on": col, "probs": [...], "penalty": p}`: an `integer` term on the covariate cut at its quantiles. The column name `"pscore"` refers to the propensity score requested through `ps`. A treated-by-control matrix (rows and columns in data order) may be passed instead of a list.	`None`
`ratio`	`int`	Controls per treated unit.	`1`
`ps`	`str or list of str`	A column holding a propensity score, or the covariates of a logistic regression that is fitted to produce one.	`None`
`control_cost`	`str or array - like`	Non-negative cost of using each control at all (a column name, or one value per control in data order): a way to steer the match away from controls that are unlike any treated unit.	`None`
`subset_cost`	`float`	Allow a treated unit to be left out at this cost, so that the match keeps the subset of treated units that can be matched well [@rosenbaum2012optimal]. Requires `ratio=1`. By default every treated unit is matched.	`None`

Returns:

Type	Description
`TwoCriteriaMatchResult`	`.matched` is the matched sample in long format, ready for an outcome analysis by matched set.

Notes

The costs are used as given. R's iTOS::makematch passes them through rcbalance::callrelax, which truncates every cost to an integer before solving, so a Mahalanobis term enters that solution only through its integer part.

Examples:

Pair on age and sex, keep the sexes balanced, and stay within a propensity-score caliper:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> df = pd.DataFrame({"age": rng.normal(50, 10, n),
...                    "female": rng.integers(0, 2, n)})
>>> df["z"] = rng.binomial(1, 1 / (1 + np.exp(-(df.age - 62) / 6)))
>>> m = sp.two_criteria_match(
...     df, "z", ps=["age", "female"],
...     pair=[{"type": "mahalanobis", "on": ["age", "female"]},
...           {"type": "caliper", "on": "pscore", "penalty": 100}],
...     balance=[{"type": "near_exact", "on": "female"}])
>>> m.matched["mset"].value_counts().unique().tolist()
[2]
>>> bool(abs(m.balance.loc["age", "smd_after"])
...      < abs(m.balance.loc["age", "smd_before"]))
True

References

[@zhang2023matching], [@rosenbaum1989optimal], [@rosenbaum2007minimum], [@yang2012optimal], [@yu2019directional], [@rosenbaum2012optimal]

statspai.matching¶

matching ¶

BalTabResult ¶

BalanceRandomizationResult dataclass ¶

FullMatchResult dataclass ¶

matched_data ¶

GenMatchResult dataclass ¶

MatchEstimator ¶

fit ¶

CardinalityMatchResult dataclass ¶

OptimalMatchResult dataclass ¶

att property ¶

estimate property ¶

BalanceDiagnosticsResult ¶

PSBalanceResult ¶

summary ¶

love_plot ¶

PScoreResult ¶

assign ¶

PSMatch2Result ¶

matched_sample ¶

balance ¶

pstest ¶

psplot ¶

psm_did ¶

summary ¶

cite ¶

PSTestResult ¶

summary ¶

SBWResult ¶

TwoCriteriaMatchResult dataclass ¶

WeightItResult ¶

balance ¶

l1_imbalance ¶

balance_vs_randomization ¶

full_match ¶

balanceplot ¶

psplot ¶

cardinality_match ¶

optimal_match ¶

balance_diagnostics ¶

love_plot ¶

overlap_plot ¶

propensity_score ¶

ps_balance ¶

trimming ¶

energy_distance ¶

ess ¶

implied_weights ¶

tighten_blocks ¶

two_criteria_match ¶

`statspai.matching`¶

BalanceRandomizationResult `dataclass` ¶

FullMatchResult `dataclass` ¶

GenMatchResult `dataclass` ¶

CardinalityMatchResult `dataclass` ¶

OptimalMatchResult `dataclass` ¶

att `property` ¶

estimate `property` ¶

TwoCriteriaMatchResult `dataclass` ¶