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, treatment="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, caliper_scale: str = 'raw', replace: bool = True, ties: str = 'all', tie_tolerance: float = 0.0, m_order: str = 'smallest_min_dist', mahalanobis_cov: str = 'pooled', exact: Optional[Union[str, Sequence[str]]] = None, covariate_calipers: Optional[Mapping[str, float]] = None, mahvars: Optional[Union[str, Sequence[str]]] = None, bias_correction: bool = False, ps_poly: int = 1, ps_model: str = 'logit', pscore: Optional[str] = None, strata: Optional[str] = None, common_support: str = 'none', kernel: str = 'epan', bwidth: float = 0.06, se_method: str = 'auto', ai_matches: int = 1, bootstrap_reps: int = 200, bootstrap_seed: Optional[int] = None, llr_stata_compat: bool = False, n_strata: int = 5, n_bins: Optional[Union[int, Sequence[int], Mapping[str, Any]]] = None, alpha: float = 0.05, weights: Optional[str] = None, cluster: Optional[str] = None, sampling_weights: Optional[str] = None, k2k: bool = False, design: Any = None) -> 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', 'kernel', 'radius', 'llr'. Legacy values 'psm', 'mahalanobis' also accepted. 'llr' is local linear regression matching (Heckman, Ichimura & Todd 1997): each treated unit's counterfactual is the intercept of a kernel-weighted degree-1 regression of the control outcome on the propensity gap. Note that Stata psmatch2 ..., llr with its default Epanechnikov kernel does not run this estimator — it substitutes nearest-neighbour matching on an lpoly-smoothed outcome. Pass kernel='tricube' to reproduce psmatch2's own LLR routine; kernel='epan' runs genuine LLR here and warns about the divergence.	'nearest'
estimand	str	Target estimand: 'ATT' or 'ATE'. The ATE is available for nearest-neighbour matching (both arms are matched), for method='stratify' and for method='cem' (cells weighted by their size rather than by their number of treated units).	'ATT'
n_matches	int	Number of nearest-neighbor matches per unit.	1
caliper	float	Maximum distance for a valid match. Units left without a match are excluded with a warning; if no unit at all finds one the call raises DataInsufficient. Not read by method='kernel' or method='llr' (the kernel bandwidth is bwidth), as in Stata psmatch2; passing it there warns and sets model_info['caliper_ignored'].	None
caliper_scale	(raw, sd)	'raw' reads caliper on the scale of the distance (Stata psmatch2 , caliper()); 'sd' in standard deviations of the propensity score over the full sample (MatchIt's std.caliper = TRUE), which needs distance='propensity'. It applies to nearest-neighbour matching and to the radius of method='radius' (Stata has no SD caliper; the result equals psmatch2 , radius caliper() at the width caliper * sd).	'raw'
replace	bool	Match with replacement (nearest-neighbor only).	True
ties	(all, first)	What to do, under matching with replacement, when more units are exactly as close as the n_matches-th nearest. 'all' keeps every one of them and splits the weight, as Stata teffects and R Matching::Match do, so the estimate does not depend on the order of the rows. 'first' keeps the first in data order, the rule of Stata psmatch2 without its ties option; a warning then says how many units were affected. With continuous covariates there are no exact ties and the two agree. Before 1.39 the default was 'first'.	'all'
tie_tolerance	float	With ties='all', how far apart two scaled squared distances may be and still count as tied.	0
exact	str or list of str	Nearest-neighbour matching only: a unit may be matched only to units with the same values of these columns (R MatchIt::matchit(exact = ~ ...)). The distance (propensity score, Mahalanobis, ...) then picks the nearest unit inside each cell. A treated unit whose cell holds no control goes unmatched and is left out of the estimate, with the warning for incomplete coverage; model_info['n_treated_without_exact_control'] counts them. The columns need not be among covariates; the propensity score is fitted on covariates only.	None
covariate_calipers	dict	Nearest-neighbour matching only: {column: width}. A pair is allowed only when the two units are within width of each other on every listed column -- the named entries of MatchIt's caliper = c(age = 0.25). caliper_scale='sd' reads the widths in standard deviations of each column over the matching sample (MatchIt's default std.caliper = TRUE); 'raw' on the column's own scale. caliper still applies to the distance.	None
mahvars	str or list of str	Nearest-neighbour propensity-score matching only: match on the Mahalanobis distance of these columns (pooled within-arm covariance), and use the propensity score fitted on covariates only through caliper, which then applies to the score -- R matchit(..., mahvars = ~ ...). m_order='largest' / 'smallest' still order the treated units by the score. The Abadie-Imbens conditional variances keep using the score to find same-arm neighbours, and se_method='abadie_imbens_2016' (which assumes matching on the score alone) is refused with any of exact / covariate_calipers / mahvars. With any of these three options the distance matrix is searched row by row instead of through the sorted-score or k-d tree search, which is slower on large samples but gives the same matches as the unconstrained search would on the allowed pairs.	None
bias_correction	bool	Apply Abadie-Imbens (2011) bias correction via regression adjustment on the matching discrepancy. The regression of the outcome on covariates is fitted on the units used as matches, weighted by the number of times each is used, as in Matching::Match(BiasAdjust = TRUE); through 1.39.6 it was an unweighted fit on the whole pool. The standard error is the conditional one of se_method and does not change.	False
ps_model	(logit, probit)	The binary model for the propensity score. 'probit' is the treatment model of Stata teffects psmatch (y) (d x, probit) and the default of Stata psmatch2.	'logit'
pscore	str	A column that already holds the propensity score (or any index the matching should run on, such as the linear predictor): Stata psmatch2 d, pscore(varname). No treatment model is fitted; covariates are used for the balance table only, and ps_model / ps_poly are ignored. The standard error treats the score as known, so se_method='abadie_imbens_2016', which charges for the estimation of the score, is refused.	None
strata	str	With method='stratify': a column that holds the stratum of every row (the blocks of :func:sp.pscore), used in place of the n_strata quantiles of the score. Rows where it is missing are left out. This is Becker and Ichino's atts.	None
ps_poly	int	Polynomial degree for the propensity score 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. The rule trims treated units only, so it is refused with estimand='ATE'.	'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	str	Standard-error estimator: one of 'auto', 'ai', 'psmatch2', 'abadie_imbens', 'abadie_imbens_2016', 'bootstrap'. Measured over 36 designs x 1000 replications by benchmarks/matching_se_coverage.py (the ratios below are reported SE / true sampling SD; 1.00 is correctly sized): ======================== =========== ================== option SE ratio coverage (nom. .95) ======================== =========== ================== 'ai' 0.56 - 0.91 0.71 - 0.92 'psmatch2' 1.50 - 1.69 0.994 - 1.000 'abadie_imbens' 0.95 - 1.04 0.905 - 0.956 'abadie_imbens_2016' not on the grid (see below) 'bootstrap' 0.95 - 1.23 0.933 - 1.000 ======================== =========== ================== '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, and never reaches nominal coverage on the grid above. 'psmatch2' is Stata psmatch2's homoskedastic analytic ATT SE sqrt(var1/N1 + var0*Σw²/N1²); it errs the other way, so use it to reproduce Stata rather than to size a test. 'abadie_imbens' is the Abadie-Imbens (2006) heteroskedasticity-robust SE (Stata psmatch2 , ai(J)), with J = ai_matches within-arm matches. It is the only option measured to be correctly sized across the grid and is the recommended choice for nearest-neighbour inference. 'abadie_imbens_2016' is the Abadie-Imbens (2016) variance for matching on an estimated propensity score: the Abadie-Imbens (2006) population-ATT variance corrected for the sampling variability of the logit coefficients, sigma^2 - c'V_gamma c + d'V_gamma d. It is what Stata teffects psmatch (y) (t x, logit), atet vce(robust, nn(J+1)) reports (StatsPAI's ai_matches=J same-arm neighbours plus the unit itself is Stata's nn(J+1); the default J=1 is Stata's default nn(2)) and reproduces it at rel 7e-8 on sp.datasets.nsw_dw(): 621.79 against 643.35 for 'abadie_imbens'. The two differ for two documented reasons -- this one targets the population rather than the sample ATT, and only this one charges for the estimated score (on NSW-DW the score term -c'Vc + d'Vd is negative, which is why the corrected SE is the smaller one). It requires distance='propensity', nearest-neighbour matching and bias_correction=False; the formula's terms are returned in model_info['ai2016_components']. With estimand='ATE' it is the ATE counterpart sigma^2 - c'V_gamma c that teffects psmatch, ate reports (estimating the score can only lower the variance of the ATE); that needs ties='all' and matching with replacement. For the ATT and the ATE alike a caliper is a condition every unit of either arm has to meet -- max(n_matches, ai_matches + 1) opposite-arm units and ai_matches + 1 other same-arm units within it -- and, as in Stata, the call stops when one does not. Not on the coverage grid above. 'bootstrap' is an arm-stratified nonparametric bootstrap that re-estimates the propensity score in every replication, so unlike the analytic options it accounts for the sampling variability of the fitted score. It is the default (and the only valid choice) for method='llr', where local linear weights can be negative and Stata psmatch2 itself reports no standard error. Abadie & Imbens (2008) show the bootstrap is not valid for nearest-neighbour matching with a fixed number of matches; requesting it there warns. 'auto' (the default) resolves to 'abadie_imbens' for nearest-neighbour matching, 'psmatch2' for kernel / radius, and 'bootstrap' for 'llr' -- in each case the estimator measured to be correctly sized for that matching scheme. With estimand='ATE' (nearest-neighbour matching) the options are the ATE counterparts, all symmetric in the two arms: 'abadie_imbens' (and 'auto') is the Abadie-Imbens (2006) conditional variance sum_i c_i^2 sigma^2(X_i) of the estimator written as sum_i c_i Y_i (c_i = (2 D_i - 1)(1 + K_M(i)) / N when every unit is matched), with sigma^2(X_i) estimated from the ai_matches nearest same-arm units in the metric the matching runs on. With ai_matches=1 it is what Matching::Match(estimand = "ATE", sample = TRUE, Var.calc = 1) reports (reproduced to 12 digits on the score and on a Mahalanobis distance). It takes the matching variable as given: on an estimated propensity score it is conservative, because estimating the score lowers the variance of the ATE (Abadie and Imbens 2016), and 'abadie_imbens_2016' is the variance that accounts for it. 'abadie_imbens_pop' is the population variance of Matching::Match's default (sample = FALSE, Var.calc = 0; needs every unit matched); 'ai' combines the matched-pair variances of the two legs; 'bootstrap' is as above. 'psmatch2' is defined for the ATT only. method='stratify' and method='cem' accept 'auto' (the within-cell variance, recorded as 'analytic') and 'bootstrap'; distance='exact' accepts 'auto' / 'ai' (the matched-pair SE) and 'bootstrap'. Other values are refused there rather than ignored. A standard error that cannot be estimated (a single matched unit, a stratum with one unit in an arm) is NaN, with NaN p-value and interval and a warning that names the reason. Through 1.39.3 the standard error of the ATE was the matched-pair SE of the treated-to-control leg alone (the ATT's 'ai'), whatever se_method said, and it changed when the arms were relabelled. The default treats the matching variable as given. When that variable is a propensity score estimated on the same data, this is conservative for the ATE: in one simulated design the default standard error averaged 1.3 times the sampling standard deviation (95% intervals covered 99%), while 'abadie_imbens_2016', which adds the correction for the estimated score and is what Stata teffects psmatch reports, averaged 1.02 times (coverage 94.7%). Pass se_method='abadie_imbens_2016' for propensity score matching with replacement when its conditions hold (ties='all', no bias correction, no user-supplied pscore, no binding caliper). The default is kept the same for the ATT and the ATE, and for every distance, so that one rule applies. The default estimates each unit's conditional variance from its J nearest same-arm units. Same-arm units tied with the J-th nearest follow ties, as the matches do: under ties='all' (the default) every tied unit counts, which is psmatch2 , ai(J) ties, and the standard error does not depend on row order (1178.209 on MatchIt::lalonde in any order, Stata's number under set type double; under Stata's default float storage of the conditional variance it differs by 1e-8). Under ties='first' the earliest rows are kept, which is psmatch2 , ai(J) without ties. .. versionchanged:: 1.39.7 Through 1.39.6 the same-arm units were always taken in row order, even under ties='all', so the default standard error changed with the order of the rows (1174 to 1198 on Lalonde). .. versionchanged:: 1.22 'auto' resolved to 'ai' before 1.22, so default nearest-neighbour standard errors change (they get larger, by roughly 1/0.56 to 1/0.91). Point estimates are untouched. Pass se_method='ai' to recover the old numbers.	'auto'
ai_matches	int	Number of within-arm matches J used by the se_method='abadie_imbens' conditional-variance estimate (Stata's ai(J)). Under se_method='abadie_imbens_2016' the same-arm set is the unit plus its J nearest same-arm units, ties included -- Stata teffects ..., vce(robust, nn(J+1)).	1
bootstrap_reps	int	Number of bootstrap replications for se_method='bootstrap'.	200
bootstrap_seed	int	Seed for the bootstrap resampler. Pass one for reproducible standard errors.	None
n_strata	int	Number of strata for method='stratify'.	5
n_bins	(int, sequence or dict)	Coarsening for method='cem'. An integer is the number of equal-width bins used for every covariate. The default follows the cem packages for R and Stata: Sturges' rule gives ceil(log2(n) + 1) cut points from the minimum to the maximum, hence one bin fewer, and a bin is closed on the right. A sequence gives the number per covariate, in the order of covariates; a dict maps a covariate to its number of bins or to its increasing cut edges ({'age': [25, 40, 65], 'female': 2}), and covariates it does not name use Sturges' rule. A 0/1 covariate needs two bins to be matched exactly. cem itself counts cut points, not bins: its x(#k) is n_bins=k - 1, and x(#2) does not split x at all.	None
alpha	float	Significance level.	0.05
weights	str	Frequency weights (Stata [fw=]): positive integers, each row standing for that many identical observations.	None
cluster	str	Not supported; raises with pointers to estimators that are.	None
sampling_weights	str	Sampling (survey) weights, R MatchIt's s.weights, for nearest-neighbour matching (ATT / ATC). The propensity score is the logit weighted by them (unless pscore= is given), matching is unchanged, and the effect is the weighted difference in means of the matched sample with weights matching weight * sampling weight. Its standard error is the one MatchIt's documentation prescribes: cluster-robust on matched sets (sandwich::vcovCL) without replacement, HC3 with replacement; it does not account for estimating the score. Equals MatchIt 4.7.2 + sandwich to 1e-9 (tests/reference_parity/test_match_sweights_parity.py). To reproduce MatchIt's matches pass m_order='largest'.	None
design	SurveyDesign	A complex survey design from :func:sp.svydesign declared on data (same rows), for nearest-neighbour matching (ATT / ATC). The design weights play the role of sampling_weights= (weighted propensity score, effect = weighted difference in means of the matched sample with weights matching weight * design weight); the standard error is the survey linearisation of that regression over strata / PSUs / fpc, with the matched rows as a domain of the full design and the matching weights treated as fixed -- R svyglm(y ~ d, subset(design, matched)) after MatchIt with s.weights, the recipe of MatchIt's documentation (equal to it, tests/reference_parity/test_match_design_parity.py). t intervals on the design df. Like sampling_weights=, it does not account for estimating the score or for the matching. Other methods have no influence function here and are refused; use sp.ipw / sp.aipw / sp.tmle with design=. .. versionadded:: 1.40.0	None
k2k	bool	method='cem' only: after the coarsening, match 1:1 without replacement inside each cell on the Mahalanobis distance of the original covariates, in data order, and drop unmatched units (R MatchIt's method = "cem", k2k = TRUE). The pairs equal MatchIt's given the same cells (tests/reference_parity/test_cem_k2k_parity.py). Balance on the coarsened distribution is measured by :func:sp.l1_imbalance.	False

Returns:

Type	Description
CausalResult

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.integers(20, 60, n),
...                    'edu': rng.integers(8, 18, n),
...                    'exp': rng.normal(10, 5, n)})
>>> logit = -3 + 0.03 * df['age'] + 0.1 * df['edu']
>>> df['training'] = rng.binomial(1, 1 / (1 + np.exp(-logit)))
>>> df['wage'] = (5 + 1.5 * df['training'] + 0.05 * df['age']
...               + 0.3 * df['edu'] + 0.1 * df['exp'] + rng.normal(size=n))

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

>>> # 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() — 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 [aweight=_weight], with weight='fweight' available for the [fweight=] degrees of freedom).

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.

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.

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

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, ps_model: str = 'logit', pscore: Optional[str] = None, radius_weights: str = 'treated', distance: Optional[str] = None, bootstrap_reps: int = 200, bootstrap_seed: Optional[int] = None, llr_stata_compat: bool = False, alpha: float = 0.05, ties: bool = False, ate: bool = False) -> 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. Optional when pscore names a fitted score.	None
pscore	str	A column that already holds the propensity score: Stata psmatch2 d, pscore(varname) outcome(y). The matching runs on that column and no treatment model is fitted. covariates, when also given, are used for the balance table only.	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. 'treated' is the region of Becker and Ichino's pscore / attnd, comsup: every unit whose score lies outside the range of the treated scores is set aside, so controls outside it cannot be matches. It is available with ties=True (or ate=True), and with any method when pscore names a fitted score (the rows outside the region are then left out of the matched frame).	'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'
radius_weights	('treated', 'pairs')	How radius matching weighs the controls. 'treated' averages the controls within the radius of each treated unit, so every matched treated unit counts once (Stata psmatch2, radius). 'pairs' counts every treated-control pair within the radius once, which weighs a control by the number of treated units near it and a treated unit by the number of controls near it: Becker and Ichino's attr, kept for reproducing results obtained with that command. The two differ whenever treated units have different numbers of controls within the radius. With 'pairs' a pair is within the radius when the scores differ by strictly less than caliper, and the standard error is attr's.	'treated'
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'
llr_stata_compat	bool	Reproduce Stata's psmatch2 ..., llr exactly, including its substitution. With its default Epanechnikov kernel psmatch2 does not run local linear regression matching: it smooths the outcome with lpoly, deg(1) fitted on the on-support controls and then runs nearest-neighbour matching, so a treated unit's raw outcome is compared with its match's smoothed one. Set this only to reproduce a published psmatch2 number; leave it False for the estimator method='llr' actually names.	False
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 propensity-score model.	1
ps_model	('logit', 'probit')	The binary model for the propensity score. Stata psmatch2 fits a probit unless its logit option is given, so ps_model='probit' is the counterpart of a Stata call without logit; the default here stays the logit of earlier releases.	'logit'
distance	str	Matching metric; 'propensity' reproduces psmatch2.	'propensity'
alpha	float	Significance level for the ATT confidence interval.	0.05
ties	bool	Stata ties: every control at the minimal propensity-score distance is a match, each receiving 1/m of the treated unit (instead of one arbitrary control). One nearest neighbour only.	False
ate	bool	Stata ate: also match controls to treated units and report the ATU and ATE (model_info['atu'] / ['ate']). Common support becomes two-sided, and a treated row's _weight counts only its use as a control's match (it starts at 0, not 1), so matched_data['_weight'].notna() is the sample psmatch2's _weight != . selects for PSM-DID. The headline att stays the ATT. One nearest neighbour only.	False

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: Union[int, Sequence[int]] = 1, alpha: float = 0.05, vce: str = 'mestimation', dof_adjust: bool = False, estimand: str = 'ATT', base_weights: Optional[str] = None, weights: Optional[str] = None) -> CausalResult

Entropy Balancing treatment effect estimator.

Reweights control units to exactly match treated covariate moments, then estimates ATT via weighted difference in means. estimand='ATE' reweights both arms to the moments of the whole sample and 'ATC' the treated to the moments of the controls (R WeightIt::weightit(method = "ebal", estimand = ...)).

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 or sequence of int	Number of moments to balance: - 1: means only - 2: means and variances - 3: means, variances, and skewness A sequence gives the order covariate by covariate, in the order of covariates (Stata ebalance ..., targets(2 2 1)). A moment that repeats another one in the sample is balanced by it and is left out of the problem: the square of a 0/1 indicator is the indicator, so asking for its variance adds nothing. Such moments are listed in model_info['redundant_moments'].	1
alpha	float		0.05
vce	(mestimation, naive)	Standard error of the ATT. 'mestimation': the sandwich variance of the stacked estimating equations -- treated moment means, the entropy-balancing dual conditions, and the two outcome means -- so the uncertainty from estimating the weights is propagated. This is the variance R WeightIt::lm_weightit(..., vcov = "asympt") reports for method = "ebal", estimand = "ATT", and it agrees with it to ~1e-9 (tests/reference_parity/test_ebalance_weightit_parity.py). Because the weights balance the covariates exactly, outcome variation that the covariates explain cancels out of the ATT, and this variance reflects that. 'naive': the pre-1.31 formula, a weighted two-sample variance that treats the weights as fixed. It ignores the balancing, so it overstates the standard error whenever the covariates predict the outcome (by about 2x on the Track B design: coverage 1.000, size 0.000). Retained only to reproduce earlier numbers.	'mestimation'
dof_adjust	bool	How a second or third moment is matched. By default the raw moments of the controls are set to those of the treated, which is what R ebal and WeightIt do when handed x**2. Stata's ebalance instead matches the sample variance: it scales the centred moment (x - mean_treated)**p by (n / (n - 1))**(p/2) within each group, so the reweighted controls reproduce the treated variance with its n - 1 divisor. dof_adjust=True follows Stata; with moments=1 the two agree. The factor differs from one by O(1/n) and does not enter the first-order variance, so the standard error is computed the same way. Defined for the ATT only (Stata ebalance estimates nothing else).	False
estimand	(ATT, ATE, ATC)	'ATT' balances the controls to the treated moments (Hainmueller 2012; Stata ebalance); 'ATC' the treated to the control moments; 'ATE' each arm to the moments of the pooled sample. Point estimates and the M-estimation standard errors reproduce WeightIt::lm_weightit(vcov = "asympt") for all three (tests/reference_parity/test_ebalance_estimand_parity.py).	'ATT'
base_weights	str	Column of positive base weights q: the reweighted units' weights are kept as close as possible to q in the entropy sense, w_i ∝ q_i exp(lambda'c(X_i)) (Stata ebalance, basewt(), WeightIt base.weights). The balancing targets are the unweighted moments of the target group, as in WeightIt, and the base weights are treated as fixed in the standard error.	None
weights	str	Column of sampling (survey) weights s, WeightIt's s.weights. They are not base weights: the balancing targets become the s-weighted moments of the target group, the reweighted arm is balanced in the s-weighted sense (sum s w c(X) hits the target, with w ∝ q exp(lambda'c(X)) and q the base weights, 1 by default), and both outcome means carry s. The weights used for the effect are s * w. base_weights keeps its meaning (the entropy reference). The M-estimation variance weights every estimating function by s (treated as fixed), as WeightIt::lm_weightit(vcov = "asympt") does; estimate and standard error equal it to 1e-9 (tests/reference_parity/test_ebalance_sweights_parity.py). model_info['weights_full'] holds w without s, like WeightIt's weights. Not combined with dof_adjust or vce='naive'.	None

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

ebalance_weights

ebalance_weights(data: DataFrame, treat: str, covariates: List[str], moments: Union[int, Sequence[int]] = 1, dof_adjust: bool = False) -> ndarray

The entropy-balancing weights alone, one per row of data.

Treated units carry weight 1 and the control weights sum to the number of treated units (the layout of model_info['weights_full'] of :func:ebalance, and of the variable Stata ebalance, generate() creates). data must be free of missing values in the columns used. Raises NumericalInstability when no balancing weights exist.

sp.cbps(...)

cbps

Covariate-balancing propensity score (CBPS) for a binary treatment.

The propensity score is a logistic model pi(X; beta) = expit(X'beta) whose coefficients are chosen so that the inverse-probability weights it implies balance the covariate means, optionally together with the logit score equation [@imai2014covariate]. With n units, n1 of them treated, the sample moments of Imai and Ratkovic (2014) are

score   (eq. 7):   s(beta) = X'(T - pi) / n
balance (eq. 10):  b(beta) = X'w / n,   w = T/pi - (1 - T)/(1 - pi)
balance (eq. 11):  b(beta) = X'w / n,   w = (n/n1) (T - pi)/(1 - pi)

(eq. 10 for the ATE, eq. 11 for the ATT). variant='exact' is the just-identified CBPS, b(beta) = 0; variant='over' stacks g = [s; b] and minimises the GMM criterion g' Sigma^{-1} g with the model-implied moment covariance Sigma of eq. (16) (ATE) or (17) (ATT).

Two observations organise the computation:

• The just-identified moments are the gradient of a strictly concave function of beta (for the ATE sum_T (eta - e^{-eta}) - sum_C (eta + e^{eta}), for the ATT sum_T eta - sum_C e^{eta}, with eta = X'beta), so the solution is unique when it exists and Newton's method with a line search on that function finds it.
• The over-identified fit is the two-step estimator (CBPS(twostep = TRUE), the reference's default): Sigma is evaluated once, at a * beta_MLE with a the minimiser over [0.8, 1.1] of the continuously-updated criterion along that ray, and then held fixed. This convention was established from the reference's outputs (its reported J and mle.J are reproduced by exactly this Sigma, and the scale sits on 0.8 or 1.1 when the unconstrained minimiser lies outside); see docs/dev/cleanroom/cbps.md.

Both criteria are invariant to a non-singular linear reparametrisation of the covariates, so they are solved on centred and scaled covariates and the coefficients mapped back.

The effect is the Hajek (self-normalised) inverse-probability-weighted contrast; its standard error is the standard deviation of the estimate over nonparametric bootstrap refits.

Clean-room implementation written from Imai and Ratkovic (2014); the reference (R CBPS 0.24, GPL >= 2) was used only as a black box.

References

[@imai2014covariate]

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

Fits a logistic propensity score whose inverse-probability weights balance the covariate means and reports the Hajek-weighted treatment effect.

Parameters:

Name	Type	Description	Default
data	DataFrame	Analysis data. Rows with a missing outcome, treatment or covariate are dropped.	required
y	str	Outcome column.	required
treat	str	Treatment column, coded 0/1.	required
covariates	list of str	Covariates of the propensity model; they are also the balanced moments. Categorical columns are expanded to indicators.	required
estimand	('ATE', 'ATT')	'ATE' balances both arms towards the pooled sample (eq. 10 of Imai and Ratkovic 2014); 'ATT' reweights the controls towards the treated (eq. 11).	'ATE'
variant	('exact', 'over')	'exact': just-identified CBPS, the balance conditions alone (exact mean balance when a solution exists). 'over': over-identified CBPS, balance conditions plus the logit score, two-step GMM with the model-implied moment covariance.	'exact'
n_bootstrap	int	Nonparametric bootstrap replications for the standard error (the whole CBPS fit is repeated on each resample). With fewer than two replications the standard error, p-value and interval are NaN.	500
alpha	float	Level of the normal-approximation confidence interval.	0.05
seed	int	Seed of the bootstrap resampling.	None
add_intercept	bool	Add a constant to the propensity model. With False the covariates must contain exactly one constant column.	True
trim	float	When positive, propensity scores are clipped to [trim, 1 - trim] before the weights are formed (units are not dropped). Scores are always kept inside [1e-6, 1 - 1e-6].	0.0

Returns:

Type	Description
CausalResult	estimate is the Hajek IPW contrast. model_info holds the propensity coefficients (beta, in the order intercept then covariates), std_mean_diff_after (weighted mean difference over the pooled within-arm standard deviation, per design column), gmm_objective, converged, the score range and the bootstrap counts.

Raises:

Type	Description
MethodIncompatibility	Unknown estimand / variant, a treatment not coded 0/1, a design with no or several constant columns, or a covariate that is a linear combination of the others.
DataInsufficient	One treatment arm is empty after dropping missing rows.

Examples:

>>> import numpy as np, 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(-0.5 * x)))
>>> df = pd.DataFrame({"y": 1 + 2 * d + x + rng.normal(size=400),
...                    "d": d, "x": x})
>>> res = sp.cbps(df, y="y", treat="d", covariates=["x"],
...               variant="exact", n_bootstrap=50, seed=1)
>>> abs(res.model_info["std_mean_diff_after"]["x"]) < 1e-8
True

References

[@imai2014covariate]

Clean-room implementation; the reference (R CBPS 0.24, GPL >= 2) was used only as a black box.

sp.genmatch(...)

genmatch

Genetic Matching (Diamond & Sekhon 2013).

The generalised distance is

.. math::

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

where :math:S_d is the diagonal matrix of full-sample covariate variances 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).

:math:S_d is diagonal, not the full covariance: that is the metric Matching::Match(Weight = 3, Weight.matrix = W) implements, and it is verified pair-for-pair against that function in tests/reference_parity/test_matching_r_parity.py. The genetic search itself is stochastic and is not reproducible across languages, so what is pinned is the deterministic matching kernel given a supplied W.

Outputs

• the optimal weight vector,
• matched treated-control pair indices,
• a balance table of standardised mean differences pre/post match,
• the ATT estimate + its matched-pair standard error (see the warning under :class:GenMatchResult — it is not a bootstrap).

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

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.73
>>> 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. [@wang2020minimal]

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', 'target', 'group')	Standard deviation delta is quoted in. 'sd' is the full-sample sd; 'target' the sd of the group being matched to and 'group' the sd of the group being reweighted, matching sbw::sbw's bal_std="target" / "group"; 'raw' is unstandardised.	'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.estimand, round(result.estimate, 3)
('ATO', 0.188)
>>> report = result.summary()
>>> ess = result.model_info['effective_sample_size']
>>> bool(0 < ess <= len(df))
True

>>> # 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, 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

sp.love_plot(...)

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