Matching and balancing

statspai.matching covers classical matching, balancing weights, diagnostics, and Love plots behind a unified sp.match(...) dispatcher plus standalone estimator functions for power users.

See also the decision guide: Choosing a matching estimator, and the exhaustive auto-generated listing under Full API reference -> matching.

Choosing an entry point

import statspai as sp

# Default nearest-neighbour propensity-score matching.
r = sp.match(
    df,
    y="earnings",
    treat="training",
    covariates=["age", "education", "earnings_pre"],
)

# Balancing-weight estimators are available through method=...
r_ebal = sp.match(df, y="earnings", treat="training",
                  covariates=["age", "education"], method="ebalance")

# ...or as standalone functions with estimator-specific options.
w = sp.overlap_weights(df, treat="training",
                       covariates=["age", "education", "earnings_pre"])
diag = sp.balance_diagnostics(df, treat="training",
                              covariates=["age", "education"],
                              weights=w.weights)

Estimator families

Family	Functions	Typical use
Classical matching	sp.match(method="nearest" | "psm" | "mahalanobis" | "cem" | "stratify")	Matched samples or subclassification with transparent design choices.
Stata psmatch2	sp.psmatch2	Stata-faithful supported PSM paths with matched-sample variables (_weight/_support; plus _n1/_pdif for nearest-neighbour), post-matching balance, common-support plot, and frequency-weighted PSM-DID. See the PSM-DID guide.
Entropy / CBPS / SBW	sp.ebalance, sp.cbps, sp.sbw	Direct covariate balance by reweighting.
Genetic matching	sp.genmatch	Automated balance search over covariate weights.
Overlap weights	sp.overlap_weights	ATE-style overlap-population estimands with stable weights.
Diagnostics	sp.balance_diagnostics, sp.love_plot	Standardised mean differences, variance ratios, and Love plots.

Method-level API

sp.match(...)

match

Matching estimators for observational causal inference.

Unified interface supporting orthogonal design choices:

• distance: how to measure unit similarity
• 'propensity' — logit propensity score (Rosenbaum & Rubin 1983)
• 'mahalanobis' — Mahalanobis distance (Rubin 1980)
• 'euclidean' — normalized Euclidean distance

• 'exact' — exact covariate values (no approximation)

• method: how to use those distances

• 'nearest' — k-nearest-neighbor matching
• 'stratify' — subclassification / stratification

• 'cem' — coarsened exact matching (Iacus, King & Porro 2012)

• bias_correction: Abadie-Imbens (2011) regression adjustment for matching discrepancies in nearest-neighbor matching.

Backward compatible: method='psm', method='mahalanobis', and method='cem' still work and map to the new parameter space.

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. King, G. and Nielsen, R. (2019). Political Analysis, 27(4), 435-454. Cunningham, S. (2021). Causal Inference: The Mixtape. Yale University Press. Ch. 5: Matching and Subclassification. https://mixtape.scunning.com/ [@rosenbaum1983central]

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.

match

match(data: DataFrame, y: str, treat: str, covariates: List[str], *, distance: Optional[str] = None, method: str = 'nearest', estimand: str = 'ATT', n_matches: int = 1, caliper: Optional[float] = None, replace: bool = True, bias_correction: bool = False, ps_poly: int = 1, common_support: str = 'none', kernel: str = 'epan', bwidth: float = 0.06, se_method: str = 'auto', ai_matches: int = 1, n_strata: int = 5, n_bins: Optional[int] = None, alpha: float = 0.05) -> CausalResult

Estimate treatment effect using matching.

Parameters:

Name	Type	Description	Default
data	DataFrame	Input data.	required
y	str	Outcome variable.	required
treat	str	Binary treatment variable (0/1).	required
covariates	list of str	Variables to match on.	required
distance	str	Distance metric: 'propensity', 'mahalanobis', 'euclidean', 'exact'. Default is 'propensity' for method='nearest'/'stratify'.	None
method	str	Matching algorithm: 'nearest', 'stratify', 'cem'. Legacy values 'psm', 'mahalanobis' also accepted.	'nearest'
estimand	str	Target estimand: 'ATT' or 'ATE'.	'ATT'
n_matches	int	Number of nearest-neighbor matches per unit.	1
caliper	float	Maximum distance for a valid match.	None
replace	bool	Match with replacement (nearest-neighbor only).	True
bias_correction	bool	Apply Abadie-Imbens (2011) bias correction via regression adjustment on the matching discrepancy.	False
ps_poly	int	Polynomial degree for the propensity score logit model. ps_poly=1 uses linear terms only. ps_poly=2 adds all squared terms and pairwise interactions. ps_poly=3 adds cubic terms as well. Higher-order specifications are standard practice; see Cunningham (2021, Ch. 5) for worked examples with age + age^2 + age^3 + educ + educ^2 + educ*re74.	1
common_support	(none, minmax)	Common-support trimming for nearest-neighbor matching. 'none' (default) matches every treated unit and leaves the point estimate unchanged. 'minmax' mirrors Stata psmatch2 , common: treated units whose propensity score falls outside the [min, max] range of the control scores are dropped before matching and the ATT is taken over the on-support treated. The matched-sample frame (result.matched_data) records the common-support flag in _support either way.	'none'
kernel	str	Kernel type for method='kernel' — one of 'epan', 'normal', 'biweight', 'uniform', 'tricube' (matches Stata psmatch2 , kerneltype()). Ignored for other methods.	'epan'
bwidth	float	Kernel bandwidth on the propensity score for method='kernel' (Stata's bwidth() default). For method='radius' the bandwidth is taken from caliper instead.	0.06
se_method	(auto, ai, psmatch2, abadie_imbens)	Standard-error estimator. 'ai' is the simple matched-pair SE (the historical default for nearest-neighbour matching). It is anti-conservative: it treats matched pairs as independent and ignores the extra variance from reusing controls under matching with replacement (empirically ~0.68x the true sampling SD; ~81% coverage at a nominal 95% level), so for valid inference prefer 'abadie_imbens'. 'psmatch2' is Stata psmatch2's homoskedastic analytic ATT SE sqrt(var1/N1 + var0*Σw²/N1²). 'abadie_imbens' is the Abadie-Imbens (2006) heteroskedasticity-robust SE (Stata psmatch2 , ai(J)), with J = ai_matches within-arm matches, and is the recommended choice for nearest-neighbour inference. 'auto' keeps 'ai' for nearest-neighbour matching -- a deliberate JOSS-review-stability default that emits a UserWarning steering you to 'abadie_imbens' -- and uses 'psmatch2' for kernel / radius matching.	'auto'
ai_matches	int	Number of within-arm matches J used by the se_method='abadie_imbens' conditional-variance estimate (Stata's ai(J)).	1
n_strata	int	Number of strata for method='stratify'.	5
n_bins	int	Number of bins per covariate for method='cem'. Default uses Sturges' rule.	None
alpha	float	Significance level.	0.05

Returns:

Type	Description
CausalResult

Examples:

>>> # Propensity score matching (default)
>>> result = sp.match(df, y='wage', treat='training',
...                   covariates=['age', 'edu', 'exp'])

>>> # Mahalanobis distance + bias correction
>>> result = sp.match(df, y='wage', treat='training',
...                   covariates=['age', 'edu', 'exp'],
...                   distance='mahalanobis', bias_correction=True)

>>> # Exact matching
>>> result = sp.match(df, y='wage', treat='training',
...                   covariates=['age', 'edu'],
...                   distance='exact')

>>> # Propensity score stratification (5 strata)
>>> result = sp.match(df, y='wage', treat='training',
...                   covariates=['age', 'edu', 'exp'],
...                   method='stratify', n_strata=5)

>>> # CEM
>>> result = sp.match(df, y='wage', treat='training',
...                   covariates=['age', 'edu'],
...                   method='cem')

>>> # Quadratic PS model (Cunningham 2021, Ch. 5 style)
>>> result = sp.match(df, y='wage', treat='training',
...                   covariates=['age', 'edu', 'exp'],
...                   ps_poly=2)

>>> # Without-replacement matching
>>> result = sp.match(df, y='wage', treat='training',
...                   covariates=['age', 'edu'],
...                   replace=False)

>>> # Legacy API still works
>>> result = sp.match(df, y='wage', treat='training',
...                   covariates=['age', 'edu'], method='psm')

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"])

sp.psmatch2(...)

psmatch2

Stata psmatch2-faithful propensity-score matching with a full post-matching toolkit (balance, common-support plot, PSM-DID).

sp.psmatch2 is the migration-friendly front door for analysts coming from Stata's psmatch2 (Leuven & Sianesi 2003). It wraps the supported nearest-neighbour, kernel, and radius propensity-score matching paths in a :class:PSMatch2Result that exposes the matched-sample variables (_pscore, _treated, _support, _weight, _y; plus _n1 … _nk, _nn, _pdif for nearest-neighbour matching) and the three operations that sp.match alone could not previously support:

1. .balance() — covariate balance on the matched, weighted sample (the post-matching analogue of Stata pstest).
2. .psplot() — propensity-score density before/after matching, with the controls reweighted by _weight (the common-support diagnostic).
3. .psm_did() — frequency-weighted PSM-DID: merge _weight into a panel, keep the matched sample, and run the weighted y ~ treat * post regression (Stata's reg y i.treat##i.post [fweight=_weight]).

For the pinned Stata 18 psmatch2 paths (nearest-neighbour, Epanechnikov kernel, and radius matching), the point estimate, analytic SE, and emitted matched-frame columns match the reference fixtures — see tests/reference_parity.

References

Leuven, E. and Sianesi, B. (2003). PSMATCH2: Stata module to perform full Mahalanobis and propensity score matching, common support graphing, and covariate imbalance testing. Statistical Software Components S432001. Rosenbaum, P.R. and Rubin, D.B. (1983). Biometrika, 70(1), 41-55. Heckman, J.J., Ichimura, H. and Todd, P.E. (1997). Review of Economic Studies, 64(4), 605-654.

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 sample Stata uses for [fweight=_weight] regressions.	True

Returns:

Type	Description
DataFrame

balance

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

Covariate balance before vs after matching (Stata pstest).

Standardized mean differences are reported two ways, exactly like pstest:

• 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).

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.

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 = 'fweight', alpha: float = 0.05) -> CausalResult

Frequency-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 [fweight=_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.

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	('fweight', 'none')	'fweight' weights the regression by _weight; 'none' runs the matched-sample DiD unweighted.	'fweight'
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).

psmatch2

psmatch2(data: DataFrame, *, treat: Optional[str] = None, covariates: Optional[Union[Sequence[str], str]] = None, outcome: Optional[str] = None, y: Optional[str] = None, neighbor: int = 1, n_matches: Optional[int] = None, caliper: Optional[float] = None, common_support: str = 'none', method: str = 'neighbor', kernel: str = 'epan', bwidth: float = 0.06, se: str = 'psmatch2', ai: int = 0, replace: bool = True, ps_poly: int = 1, distance: str = 'propensity', alpha: float = 0.05) -> PSMatch2Result

Stata psmatch2-faithful supported propensity-score matching.

Runs nearest-neighbour propensity-score matching and returns a :class:PSMatch2Result carrying the psmatch2 matched-sample variables (_pscore _treated _support _weight _y; plus _n1 … _nn _pdif for nearest-neighbour matching) plus post-matching balance, common-support plotting, and PSM-DID helpers.

This is the Stata-migration front door over :func:sp.match; the pinned Stata 18 nearest-neighbour, kernel, and radius paths are covered by reference fixtures for the point estimate, analytic SE, and emitted matched-frame columns (Leuven & Sianesi 2003).

Parameters:

Name	Type	Description	Default
data	DataFrame	Cross-section with one row per unit (the matching sample).	required
treat	str	Binary treatment indicator (0/1). Stata's treated.	None
covariates	list of str	Pre-treatment covariates entering the propensity-score model.	None
outcome	str	Outcome variable (outcome mirrors Stata's outcome(); y is accepted as an alias). Optional, exactly like Stata: when omitted, the matched frame (_weight etc.) is still produced for downstream PSM-DID, but the cross-sectional ATT is left NaN.	None
y	str	Outcome variable (outcome mirrors Stata's outcome(); y is accepted as an alias). Optional, exactly like Stata: when omitted, the matched frame (_weight etc.) is still produced for downstream PSM-DID, but the cross-sectional ATT is left NaN.	None
neighbor	int	Number of nearest neighbours k (Stata neighbor(k)). n_matches is accepted as an alias.	1
caliper	float	Maximum propensity-score distance for a valid match (Stata caliper()).	None
common_support	('none', 'minmax')	'none' matches every treated unit (raw psmatch2). 'minmax' drops treated units outside the control PS range before matching (Stata common) and the ATT is taken over the on-support treated.	'none'
method	('neighbor', 'kernel', 'radius')	Matching algorithm. 'neighbor' is k-nearest-neighbour matching (Stata default; uses neighbor / caliper). 'kernel' is kernel matching (uses kernel + bwidth). 'radius' is radius matching (all controls within caliper; a uniform kernel).	'neighbor'
kernel	('epan', 'normal', 'biweight', 'uniform', 'tricube')	Kernel type for method='kernel' (Stata kerneltype()).	'epan'
bwidth	float	Kernel bandwidth on the propensity score for method='kernel' (Stata bwidth() default).	0.06
se	('psmatch2', 'ai', 'abadie_imbens')	Standard-error estimator. 'psmatch2' reproduces Stata's homoskedastic analytic ATT SE digit for digit; 'ai' is the simple matched-pair SE; 'abadie_imbens' is the Abadie-Imbens (2006) heteroskedasticity-robust SE (Stata psmatch2 , ai(J)).	'psmatch2'
ai	int	Shorthand for the Abadie-Imbens (2006) robust SE with J = ai within-arm matches (Stata's ai(J)). Any ai > 0 selects the robust SE and overrides se. Reproduces Stata's r(seatt) digit for digit.	0
replace	bool	Match with replacement (psmatch2 default).	True
ps_poly	int	Polynomial degree of the logit propensity-score model.	1
distance	str	Matching metric; 'propensity' reproduces psmatch2.	'propensity'
alpha	float	Significance level for the ATT confidence interval.	0.05

Returns:

Type	Description
PSMatch2Result

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> m = sp.psmatch2(df, outcome='log_wage', treat='union',
...                 covariates=['education', 'experience', 'tenure'])
>>> round(float(m.att), 4) == round(float(m.result.estimate), 4)
True
>>> sorted(c for c in m.matched_data.columns if c.startswith('_'))[:4]
['_id', '_n1', '_nn', '_pdif']

Post-matching balance (the analogue of Stata pstest — smd_raw is before matching, smd_weighted is the matched, _weight-weighted sample):

>>> bal = m.balance()
>>> {'smd_raw', 'smd_weighted'} <= set(bal.table.columns)
True

sp.ebalance(...)

ebalance

Entropy Balancing (Hainmueller 2012).

Reweights the control group so that weighted covariate moments (mean, variance, skewness) exactly match the treated group, without dropping observations or relying on propensity score models.

More robust than PSM because it directly targets balance rather than modeling the selection process.

References

Hainmueller, J. (2012). "Entropy Balancing for Causal Effects: A Multivariate Reweighting Method to Produce Balanced Samples in Observational Studies." Political Analysis, 20(1), 25-46. [@hainmueller2012entropy]

ebalance

ebalance(data: DataFrame, y: str, treat: str, covariates: List[str], moments: int = 1, alpha: float = 0.05) -> CausalResult

Entropy Balancing treatment effect estimator.

Reweights control units to exactly match treated covariate moments, then estimates ATT via weighted difference in means.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str	Outcome variable.	required
treat	str	Binary treatment indicator (0/1).	required
covariates	list of str	Covariates to balance on.	required
moments	int	Number of moments to balance: - 1: means only - 2: means and variances - 3: means, variances, and skewness	1
alpha	float		0.05

Returns:

Type	Description
CausalResult	ATT estimate with entropy-balanced weights and balance table.

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> age = rng.normal(40, 10, n)
>>> income = rng.normal(50, 15, n)
>>> education = rng.integers(8, 20, n).astype(float)
>>> ps = 1 / (1 + np.exp(-(0.03 * (age - 40) + 0.02 * (income - 50))))
>>> treated = (rng.uniform(size=n) < ps).astype(int)
>>> outcome = (2.0 * treated + 0.1 * age + 0.05 * income
...            + 0.2 * education + rng.normal(0, 1, n))
>>> df = pd.DataFrame({'outcome': outcome, 'treated': treated,
...                    'age': age, 'income': income,
...                    'education': education})
>>> result = sp.ebalance(df, y='outcome', treat='treated',
...                      covariates=['age', 'income', 'education'])
>>> bool(np.isfinite(result.estimate))
True
>>> 'balance' in result.model_info  # post-weighting balance table
True

Notes

Entropy balancing solves:

.. math:: \min_w \sum_i w_i \log(w_i / q_i)

subject to balance constraints (weighted moments match) and normalization (weights sum to 1).

Unlike PSM, this guarantees exact balance on specified moments without iteration or caliper tuning.

If the dual optimizer raises, uniform control weights are used as a fallback; a ConvergenceWarning is emitted and model_info['weights_fallback'] is set to True.

See Hainmueller (2012, Political Analysis).

sp.cbps(...)

cbps

Covariate-Balancing Propensity Score (Imai & Ratkovic 2014).

CBPS estimates the propensity score by solving a moment condition that jointly enforces:

(a) the logit score equation (standard MLE first-order condition);
(b) exact mean-balance of covariates under the implied IPW weights.

The "just-identified" (exact) variant uses K moment conditions where K equals the covariate dimension (drops the score equation). The "over-identified" variant stacks both sets and solves via GMM. This module implements both.

Mathematically, denote π(X; β) = 1 / (1 + exp(-X'β)). The over-identified moment vector for ATE is

g_i(β) = [ (T_i - π_i) * X_i ,                (MLE)
           (T_i - π_i) / (π_i (1 - π_i)) X_i ] (Balance)

CBPS minimises ḡ' W ḡ with W = identity for the exact case (K equations, K unknowns → method of moments) or with the efficient GMM weighting matrix for the over-identified case.

Treatment-effect point estimate uses the resulting weights in the standard (normalised Hajek) IPW formula; SEs come from a paired bootstrap re-estimation by default.

References

Imai, K., Ratkovic, M. (2014). "Covariate Balancing Propensity Score." JRSS-B, 76(1), 243-263. [@imai2014covariate]

Fong, C., Ratkovic, M., Imai, K. (2022). CBPS R package documentation.

cbps

cbps(data: DataFrame, y: str, treat: str, covariates: List[str], estimand: Literal['ATE', 'ATT'] = 'ATE', variant: Literal['exact', 'over'] = 'over', n_bootstrap: int = 500, alpha: float = 0.05, seed: Optional[int] = None, add_intercept: bool = True, trim: float = 0.0) -> CausalResult

Covariate-Balancing Propensity Score estimator (Imai-Ratkovic 2014).

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str	Outcome column.	required
treat	str	Binary 0/1 treatment column.	required
covariates	list of str	Covariates entering the logit score.	required
estimand	('ATE', 'ATT')		'ATE'
variant	('exact', 'over')	'exact': just-identified CBPS (only balance moments). 'over': over-identified CBPS (MLE + balance, solved via two-step GMM).	'exact'
n_bootstrap	int		500
alpha	float		0.05
seed	int		None
add_intercept	bool	Prepend a constant to the covariate matrix.	True
trim	float	Optional pscore clip for stability.	0.0

Returns:

Type	Description
CausalResult	estimate is the CBPS weighted treatment effect; model_info contains the estimated coefficients, balance diagnostics and effective sample size.

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> result = sp.cbps(df, y='log_wage', treat='union',
...                  covariates=['education', 'experience', 'tenure'],
...                  estimand='ATT', n_bootstrap=50, seed=42)
>>> result.summary()
>>> result.model_info['std_mean_diff_after']  # balance after weighting

>>> # Just-identified CBPS (balance moments only)
>>> result = sp.cbps(df, y='log_wage', treat='union',
...                  covariates=['education', 'experience', 'tenure'],
...                  variant='exact', n_bootstrap=50, seed=42)

sp.genmatch(...)

genmatch

Genetic Matching (Diamond & Sekhon 2013).

The user-supplied generalised Mahalanobis distance is

.. math::

d_W(x_i, x_j) = (x_i - x_j)^\top S^{-1/2}\, W\, S^{-1/2} (x_i - x_j),

where :math:S is the sample covariance of covariates and :math:W is a diagonal weight matrix found by a genetic (evolutionary) search that maximises the minimum across-covariate balance p-value (Kolmogorov-Smirnov + t-tests, following the Matching R package).

Outputs

• the optimal weight vector,
• matched treated-control pair indices,
• a balance table of standardised mean differences pre/post match,
• the ATT estimate + bootstrap SE.

References

Diamond, A. & Sekhon, J. S. (2013). "Genetic matching for estimating causal effects." Review of Economics and Statistics, 95(3), 932-945. [@diamond2013genetic]

GenMatchResult dataclass

Bases: ResultProtocolMixin

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

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

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

genmatch

genmatch(data: DataFrame, y: str, treat: str, covariates: Sequence[str], k: int = 1, population_size: int = 40, generations: int = 20, mutation_rate: float = 0.2, alpha: float = 0.05, random_state: int = 42) -> GenMatchResult

Genetic Matching for ATT estimation.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str		required
treat	str	Binary treatment indicator.	required
covariates	sequence of str		required
k	int	Number of matches per treated unit.	1
population_size	int		40
generations	int		20
mutation_rate	float		0.2
alpha	float		0.05
random_state	int		42

Returns:

Type	Description
GenMatchResult

Notes

Degenerate Kolmogorov-Smirnov balance tests fall back to p=1.0 with a StatsPAIWarning; failures in the final balance table are counted in result.detail['ks_test_failures'].

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})

Small genetic-search settings keep the example fast; prefer the defaults (population_size=40, generations=20) in practice:

>>> res = sp.genmatch(df, y='y', treat='d',
...                   covariates=['x1', 'x2'],
...                   population_size=10, generations=5)
>>> res.n_treated
111
>>> round(res.att, 2)
1.72
>>> res.balance.columns.tolist()
['variable', 'smd_pre', 'smd_post', 'ks_p_pre', 'ks_p_post']

sp.sbw(...)

sbw

Stable Balancing Weights (Zubizarreta 2015, JASA).

Finds weights that minimise dispersion (e.g. variance, or KL divergence from the uniform distribution) while imposing user-specified covariate balance tolerances. Unlike entropy balancing, SBW allows approximate balance via per-covariate tolerance δ_j, which is essential when exact balance is infeasible or would blow up variance.

Formulation

For ATT estimation with treated group :math:\mathcal{T} and control group :math:\mathcal{C}, solve

.. math::

\min_{w} \; \frac{1}{|\mathcal{C}|}\sum_{i \in \mathcal{C}} w_i^2

\text{s.t.} \quad
\left| \frac{1}{|\mathcal{T}|}\sum_{i \in \mathcal{T}} X_{ij}
      - \sum_{i \in \mathcal{C}} w_i X_{ij} \right|
      \;\le\; \delta_j \sigma_j  \; \forall j,

\sum_{i \in \mathcal{C}} w_i = 1, \; w_i \ge 0.

The variance-minimising objective is equivalent to maximising effective sample size :math:\mathrm{ESS}(w) = (\sum w_i)^2 / \sum w_i^2.

This complements :func:ebalance (exact balance, KL objective) and :func:cbps (covariate-balancing propensity score) — together forming the 2026 triumvirate of modern weighting estimators.

References

Zubizarreta, J.R. (2015). "Stable Weights that Balance Covariates for Estimation with Incomplete Outcome Data." Journal of the American Statistical Association, 110(511), 910-922. [@zubizarreta2015stable]

Wang, Y. and Zubizarreta, J.R. (2020). "Minimal dispersion approximately balancing weights: asymptotic properties and practical considerations." Biometrika, 107(1), 93-105. [@wang2019minimal]

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']

sbw

sbw(data: DataFrame, treat: str, covariates: List[str], y: Optional[str] = None, *, estimand: str = 'att', delta: Union[float, Sequence[float]] = 0.02, objective: str = 'variance', tolerance_scale: str = 'sd', include_squares: bool = False, alpha: float = 0.05, solver_options: Optional[dict] = None) -> SBWResult

Stable Balancing Weights (Zubizarreta 2015) with optional ATT/ATE treatment-effect estimation.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
treat	str	Binary 0/1 treatment indicator column.	required
covariates	list of str	Columns whose means must be balanced.	required
y	str	Outcome column. If provided, a weighted ATT/ATE estimate with HC-robust SE is attached to the returned :class:SBWResult.	None
estimand	('att', 'ate', 'atc')	'att' reweights controls to match treated means (standard); 'atc' reweights treated to match control means; 'ate' reweights each group to match the pooled means.	'att'
delta	float or sequence	Balance tolerance. With tolerance_scale='sd' the constraint is |mean_T(X_j) - weighted mean_C(X_j)| ≤ δ_j · sd(X_j).	0.02
objective	('variance', 'entropy')	Dispersion objective. 'variance' minimises Σ w_i²; 'entropy' minimises Σ w_i log(n · w_i) (KL from uniform).	'variance'
tolerance_scale	('sd', 'raw')	Whether delta is in SD units (standard) or raw units.	'sd'
include_squares	bool	Also balance second-moments (w_j² columns).	False
alpha	float	Significance level for inference on the outcome.	0.05
solver_options	dict	Passed to scipy.optimize.minimize.	None

Returns:

Type	Description
SBWResult

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)
>>> res.estimand
'ATT'
>>> summary_text = res.summary()
>>> res.balance.shape[0]               # one row per covariate
3

sp.overlap_weights(...)

overlap_weights

Overlap weights (Li, Morgan, Zaslavsky 2018).

Overlap weights target the "average treatment effect among the overlap population" (ATO) — those observations with genuine equipoise. The weight function is

w_i = 1 - e(X_i)    for treated  i   (T_i=1)
w_i =     e(X_i)    for control  i   (T_i=0)

which is proportional to the "tilting" that minimises the variance of the resulting weighted treatment-effect estimator subject to exact covariate balance on the moments used to fit e(·) (Li et al. 2018, Theorem 1 & 3). Overlap weights:

• are bounded in [0, 1] — no extreme weights from small propensity scores, so results are stable even with poor overlap;
• exactly balance the log-odds covariates when e(X) is a logit fit;
• target ATO, not ATE, and should be interpreted accordingly.

References

Li, F., Morgan, K.L., Zaslavsky, A.M. (2018). "Balancing Covariates via Propensity Score Weighting." JASA, 113(521), 390-400.

Li, F., Thomas, L.E., Li, F. (2019). "Addressing Extreme Propensity Scores via the Overlap Weights." American Journal of Epidemiology, 188(1), 250-257.

overlap_weights

overlap_weights(data: DataFrame, y: str, treat: str, covariates: List[str], estimand: str = 'ATO', n_bootstrap: int = 500, alpha: float = 0.05, seed: Optional[int] = None, trim: float = 0.0) -> CausalResult

Overlap-weight (ATO) treatment effect estimator.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str	Outcome column.	required
treat	str	Binary 0/1 treatment column.	required
covariates	list of str	Covariates for the logistic propensity-score model.	required
estimand	('ATO', 'ATE', 'ATT', 'ATC', 'matching', 'entropy')	Which generalized-weight scheme to use. All follow Li-Li-Li (2019) Table 1; 'ATO' uses the overlap weights; 'matching' uses the min(e, 1-e) weight; 'entropy' uses -e·log(e) - (1-e)·log(1-e); 'ATE/ATT/ATC' reduce to standard IPW for comparison.	'ATO'
n_bootstrap	int	Paired-sample bootstrap replications for SE.	500
alpha	float		0.05
seed	int		None
trim	float	Optional clip of pscore to [trim, 1-trim]. For overlap weights this is rarely needed — set to 0 by default.	0.0

Returns:

Type	Description
CausalResult	.estimate targets the named estimand; .model_info stores the weight summary, effective sample size, and pscore diagnostics.

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> result = sp.overlap_weights(df, y='log_wage', treat='union',
...                             covariates=['education', 'experience',
...                                         'tenure'],
...                             n_bootstrap=100, seed=42)
>>> result.summary()
>>> result.model_info['effective_sample_size']

>>> # min(e, 1-e) matching weights instead of overlap weights
>>> result = sp.overlap_weights(df, y='log_wage', treat='union',
...                             covariates=['education', 'experience',
...                                         'tenure'],
...                             estimand='matching',
...                             n_bootstrap=100, seed=42)

sp.balance_diagnostics(...)

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) -> 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

Returns:

Type	Description
BalanceDiagnosticsResult	.table has one row per covariate; .summary_stats records max/mean SMDs, imbalance counts, effective sample size, and propensity-score overlap.

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

sp.love_plot(...)

love_plot

love_plot(data: DataFrame, treatment: str, covariates: List[str], 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.	required
covariates	list of str	Covariate columns.	required
weights	array - like	IPW or matching weights. If None, inverse-PS weights are computed.	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')