statspai.utils

utils

Utility functions for StatsPAI.

Provides Stata-style data manipulation tools: - Variable labels (label_var, get_label) - Pairwise correlation with stars (pwcorr) - Winsorization (winsor)

pwcorr

pwcorr(data: DataFrame, vars: Optional[List[str]] = None, stars: bool = True, method: str = 'pearson', decimals: int = 3, output: str = 'text', listwise: bool = False, obs: bool = False) -> Union[str, DataFrame]

Pairwise correlation matrix with significance stars.

Equivalent to Stata's pwcorr var1 var2 var3, star(0.05).

Missing values are handled pairwise by default, exactly as Stata's pwcorr: each correlation uses every row on which that pair is observed, so a missing value in a third variable never changes the correlation between two complete variables. Pass listwise=True for Stata's pwcorr, listwise (one common estimation sample).

Parameters:

Name	Type	Description	Default
data	DataFrame		required
vars	list of str	Variables to correlate. Default: all numeric columns.	None
stars	bool	Show significance stars (* p<0.1, ** p<0.05, *** p<0.01).	True
method	str	'pearson', 'spearman', or 'kendall'.	'pearson'
decimals	int	Decimal places.	3
output	str	'text', 'dataframe', 'latex', 'html'.	'text'
listwise	bool	Drop every row with a missing value in any of vars before computing all correlations (Stata listwise).	False
obs	bool	Print the number of observations used for each pair beneath the coefficient (Stata obs) in text / LaTeX / HTML output.	False

Returns:

Type	Description
str or DataFrame	Formatted correlation matrix. With output='dataframe' the correlation matrix is returned and the p-values and pairwise observation counts are attached as .attrs['pvalues'] and .attrs['nobs'] (both DataFrames). A pair with fewer than three jointly observed rows (or a constant variable) has a NaN correlation rather than a fabricated number.

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> out = sp.pwcorr(df, vars=['log_wage', 'education', 'experience'])
>>> isinstance(out, str)
True
>>> bool('***' in out)   # significant pairwise correlations are starred
True
>>> # Raw correlation matrix instead of formatted text
>>> cm = sp.pwcorr(df, vars=['log_wage', 'education'], output='dataframe')
>>> round(float(cm.loc['log_wage', 'education']), 3)
0.335
>>> int(cm.attrs['nobs'].loc['log_wage', 'education']) == len(df)
True

winsor

winsor(data: DataFrame, vars: Optional[List[str]] = None, cuts: Tuple[float, float] = (1, 99), replace: bool = False, suffix: str = '_w', method: str = 'stata', subset: Optional[Union[str, Series, ndarray]] = None, by: Optional[Union[str, List[str]]] = None, trim: bool = False) -> DataFrame

Winsorize variables at specified percentiles.

Equivalent to Stata's winsor2 var1 var2, cuts(1 99), including its percentile definition, if qualifier (subset=) and by().

Parameters:

Name	Type	Description	Default
data	DataFrame		required
vars	list of str	Variables to winsorize. Default: all numeric columns.	None
cuts	tuple of (float, float)	Lower and upper percentile cutoffs.	(1, 99)
replace	bool	If True, overwrite original columns. If False, create new columns with suffix.	False
suffix	str	Suffix for new winsorized columns (when replace=False).	'_w'
method	str	Percentile definition. 'stata' is Stata's _pctile default, used by winsor2 (the empirical quantile, averaging the two order statistics when n p / 100 is an integer; numpy's 'averaged_inverted_cdf'). Any other value is passed to :func:numpy.percentile as its method; 'linear' reproduces StatsPAI before 1.33.	'stata'
subset	str, boolean Series or array	Rows used to compute the cutoffs and winsorized, as winsor2 ... if <cond>: a :meth:pandas.DataFrame.query string such as "year >= 2007 & year <= 2020" or a boolean mask. Rows outside keep their value when replace=True and get NaN in the new column otherwise (Stata's generate ... if).	None
by	str or list of str	Compute the cutoffs within groups, as winsor2 ..., by(g).	None
trim	bool	Set the values beyond the cutoffs to missing instead of pulling them in (winsor2 ..., trim; truncation). The new columns then take the suffix '_tr' unless suffix is given. Trimming one variable at a time removes a different set of rows for each; a regression on several trimmed variables loses every row trimmed in any of them.	False

Returns:

Type	Description
DataFrame	DataFrame with winsorized variables.

Examples:

>>> import statspai as sp
>>> df = sp.cps_wage()
>>> # Winsorize at 1st and 99th percentile -> adds 'log_wage_w'
>>> out = sp.winsor(df, vars=['log_wage'], cuts=(1, 99))
>>> 'log_wage_w' in out.columns
True
>>> bool(out['log_wage_w'].max() <= df['log_wage'].max())  # tails clipped
True
>>> len(out) == len(df)   # winsorizing does not drop rows
True
>>> # Replace in place instead of adding a suffixed column
>>> repl = sp.winsor(df, vars=['log_wage'], replace=True)
>>> 'log_wage_w' in repl.columns
False

Notes

Winsorization replaces values below the cuts[0]-th percentile with that percentile value, and values above the cuts[1]-th percentile with that percentile value. Unlike trimming, winsorization does not remove observations.

The 'stata' percentile is also R's quantile(type = 2), so sp.winsor(df, [v], cuts=(1, 99)) equals farr::winsorize(v, 0.01) and trim=True equals farr::truncate(v, 0.01).

Before 1.33 the cutoffs used numpy's linear interpolation, which differs from winsor2 on every variable whose n p / 100 is not an integer; on a long-tailed variable the 99th percentile moved by 0.02% (51,003.45 against Stata's 51,015).

dgp_bartik

dgp_bartik(n_regions: int = 50, n_industries: int = 10, effect: float = 1.0, seed: int | None = None, endogenous: bool = False) -> dict

Generate shift-share (Bartik) instrument data.

Parameters:

Name	Type	Description	Default
n_regions	int	Number of regions.	50
n_industries	int	Number of industries.	10
effect	float	True effect of the Bartik instrument on the outcome.	1.0
seed	int or None	Random seed.	None
endogenous	bool	If True, generate a structural shift-share IV instead of the reduced form: an endogenous regressor x = bartik + v and y = effect * x + e with corr(v, e) = 0.6, so OLS of y on x is biased and sp.bartik(data, y='y', endog='x', shares=out['shares'], shocks=out['shocks']) recovers effect. The default (reduced-form) draws are unchanged.	False

Returns:

Type	Description
dict	'data': DataFrame with y, bartik, region (plus x when endogenous=True); 'shares': (n_regions, n_industries) array; 'shocks': (n_industries,) array.

Examples:

>>> import statspai as sp
>>> result = sp.dgp_bartik(n_regions=20, n_industries=5, seed=0)
>>> result['data'].shape
(20, 3)
>>> result['shares'].shape
(20, 5)

dgp_bunching

dgp_bunching(n: int = 10000, kink_point: float = 50000.0, elasticity: float = 0.3, seed: int | None = None, t0: float = 0.0, t1: float = 0.2) -> DataFrame

Generate earnings that bunch at a kink in the marginal tax schedule.

Counterfactual earnings z* are those under the flat rate t0. At the kink the marginal rate rises to t1; with iso-elastic earnings z = z* ((1 - t) / (1 - t0))^elasticity:

• z* <= kink_point -- unaffected;
• z* >= kink_point * ((1 - t0) / (1 - t1))^elasticity -- earn z* ((1 - t1) / (1 - t0))^elasticity;
• in between -- locate exactly at kink_point (the bunchers).

Parameters:

Name	Type	Description	Default
n	int	Sample size.	10000
kink_point	float	Location of the kink (e.g., tax threshold).	50000.0
elasticity	float	Earnings elasticity with respect to the net-of-tax rate.	0.3
seed	int or None	Random seed.	None
t0	float	Marginal tax rate below and above the kink (t0 < t1 < 1).	0.0
t1	float	Marginal tax rate below and above the kink (t0 < t1 < 1).	0.0

Returns:

Type	Description
DataFrame	Columns: income, counterfactual_income. df.attrs holds true_effect (the elasticity), dt (t1 - t0, the argument :func:sp.bunching takes), bunching_upper (the top of the counterfactual interval that bunches) and n_bunchers.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_bunching(n=5000, kink_point=50000, elasticity=0.2, seed=0)
>>> df.attrs['true_effect'], df.attrs['dt']
(0.2, 0.2)
>>> bool((df['income'] <= df['counterfactual_income'] + 1e-9).all())
True
>>> bool((df['income'] == 50000).sum() == df.attrs['n_bunchers'] > 0)
True

dgp_cluster_rct

dgp_cluster_rct(n_clusters: int = 50, cluster_size: int = 20, effect: float = 0.3, icc: float = 0.1, seed: int | None = None) -> DataFrame

Generate cluster-randomised trial data.

Parameters:

Name	Type	Description	Default
n_clusters	int	Number of clusters.	50
cluster_size	int	Number of individuals per cluster.	20
effect	float	Treatment effect.	0.3
icc	float	Intra-cluster correlation coefficient (0-1).	0.1
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: y, treatment, cluster_id, unit_id.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_cluster_rct(n_clusters=10, cluster_size=5, effect=0.5, seed=0)
>>> df.shape
(50, 4)
>>> df.attrs['true_effect']
0.5

dgp_did

dgp_did(n_units: int = 100, n_periods: int = 10, effect: float = 0.5, staggered: bool = False, n_groups: int = 4, heterogeneous: bool = False, seed: int | None = None, ddd: bool = False, ddd_confound: float = 1.0) -> DataFrame

Generate Difference-in-Differences panel data.

Parameters:

Name	Type	Description	Default
n_units	int	Number of cross-sectional units.	100
n_periods	int	Number of time periods.	10
effect	float	Average treatment effect on the treated.	0.5
staggered	bool	If True, treatment adoption is staggered across groups.	False
n_groups	int	Number of treatment-timing groups (used when staggered=True).	4
heterogeneous	bool	If True, the treatment effect varies by unit.	False
seed	int or None	Random seed for reproducibility.	None
ddd	bool	Triple-difference design (non-staggered only). Adds a binary subgroup column (drawn per unit); the treatment effect applies only to subgroup == 1 units of the treated group, and every treated-group unit (both subgroups) also receives a post-period shock of size ddd_confound. Plain DiD on the eligible subgroup is then biased by ddd_confound; DDD recovers effect. The draws for all other columns are identical to ddd=False.	False
ddd_confound	float	Treated-group x post shock shared by both subgroups (ddd=True).	1.0

Returns:

Type	Description
DataFrame	Columns: unit, time, y, treated, first_treat, group (plus subgroup when ddd=True; treated is then the effective treatment group * post * subgroup). df.attrs['true_effect'] stores the average treatment effect.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_did(n_units=50, n_periods=8, effect=1.0, seed=0)
>>> df.shape
(400, 6)
>>> df.attrs['true_effect']
1.0

dgp_iv

dgp_iv(n: int = 500, effect: float = 0.5, first_stage: float = 0.4, n_instruments: int = 1, seed: int | None = None) -> DataFrame

Generate Instrumental Variables data with endogeneity.

Parameters:

Name	Type	Description	Default
n	int	Sample size.	500
effect	float	True causal effect of treatment on outcome.	0.5
first_stage	float	Strength of the instrument in the first stage.	0.4
n_instruments	int	Number of instruments.	1
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: y, treatment, instrument (or instrument_1, ...), x1, x2.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_iv(n=300, effect=0.5, seed=0)
>>> df.attrs['true_effect']
0.5
>>> bool('instrument' in df.columns)
True

dgp_observational

dgp_observational(n: int = 1000, effect: float = 0.5, confounding: float = 0.3, seed: int | None = None) -> DataFrame

Generate observational data with selection on observables.

Parameters:

Name	Type	Description	Default
n	int	Sample size.	1000
effect	float	True treatment effect.	0.5
confounding	float	Strength of confounding (higher = more selection bias).	0.3
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: y, treatment, x1, x2, propensity_score.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_observational(n=500, effect=1.0, seed=0)
>>> df.attrs['true_effect']
1.0
>>> bool('propensity_score' in df.columns)
True

dgp_panel

dgp_panel(n_units: int = 100, n_periods: int = 20, fe_sd: float = 1.0, te_sd: float = 0.5, ar1_coef: float = 0.5, seed: int | None = None) -> DataFrame

Generate panel data with unit/time fixed effects and AR(1) covariate.

Parameters:

Name	Type	Description	Default
n_units	int	Number of units.	100
n_periods	int	Number of time periods.	20
fe_sd	float	Standard deviation of unit fixed effects.	1.0
te_sd	float	Standard deviation of time fixed effects.	0.5
ar1_coef	float	AR(1) coefficient for the covariate process.	0.5
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: unit, time, y, x.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_panel(n_units=20, n_periods=5, seed=0)
>>> df.shape
(100, 4)
>>> df.attrs['true_effect']
1.0

dgp_rct

dgp_rct(n: int = 500, effect: float = 0.3, p_treat: float = 0.5, heterogeneous: bool = False, n_covariates: int = 3, seed: int | None = None) -> DataFrame

Generate Randomised Controlled Trial data.

Parameters:

Name	Type	Description	Default
n	int	Sample size.	500
effect	float	Average treatment effect.	0.3
p_treat	float	Probability of assignment to treatment.	0.5
heterogeneous	bool	If True, the treatment effect varies with the first covariate.	False
n_covariates	int	Number of pre-treatment covariates.	3
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: y, treatment, x1, x2, ...

Examples:

>>> import statspai as sp
>>> df = sp.dgp_rct(n=200, effect=1.0, seed=0)
>>> df.attrs['true_effect']
1.0
>>> bool(0.0 <= df['treatment'].mean() <= 1.0)  # randomized assignment
True

dgp_rd

dgp_rd(n: int = 1000, effect: float = 0.3, cutoff: float = 0.0, fuzzy: bool = False, bandwidth_relevant: float = 0.5, seed: int | None = None) -> DataFrame

Generate Regression Discontinuity data (sharp or fuzzy).

Parameters:

Name	Type	Description	Default
n	int	Sample size.	1000
effect	float	Treatment effect at the cutoff. The effect is constant, so in the fuzzy design it is also the complier effect.	0.3
cutoff	float	RD cutoff value. The running variable is centred on it.	0.0
fuzzy	bool	If True, generate a fuzzy RD design: the probability of treatment is 0.2 + 0.1 (x - cutoff) / s below the cutoff and 0.6 higher above it, where s is the half-width of the support.	False
bandwidth_relevant	float	Spread of the running variable, which is uniform on cutoff +/- 2 * bandwidth_relevant (+/- 1 at the default).	0.5
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: y, x, treatment. attrs['true_effect'] is effect; the fuzzy design also records attrs['first_stage_jump'] = 0.6.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_rd(n=500, effect=0.5, seed=0)
>>> bool('treatment' in df.columns)
True
>>> df.attrs['true_effect']
0.5

dgp_rd_2d

dgp_rd_2d(n: int = 2000, effect: float = 2.0, seed: int | None = None) -> DataFrame

Generate 2D boundary RD data.

Treatment is assigned when x1 >= 0 (vertical boundary). True effect = effect.

Parameters:

Name	Type	Description	Default
n	int	Sample size.	2000
effect	float	Treatment effect at boundary.	2.0
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: y, x1, x2, d (treatment).

Examples:

>>> import statspai as sp
>>> df = sp.dgp_rd_2d(n=2000, effect=2.0, seed=42)
>>> df.attrs['true_effect']
2.0

dgp_rd_hte

dgp_rd_hte(n: int = 3000, ate: float = 2.0, hte_coef: float = 1.5, cutoff: float = 0.0, seed: int | None = None) -> DataFrame

Generate RD data with heterogeneous treatment effects.

CATE(z) = ate + hte_coef * z, where z ~ N(0,1). So ATE = ate and the slope of heterogeneity = hte_coef.

Parameters:

Name	Type	Description	Default
n	int	Sample size.	3000
ate	float	Average treatment effect at cutoff.	2.0
hte_coef	float	Slope of heterogeneity w.r.t. covariate z.	1.5
cutoff	float	RD cutoff.	0.0
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: y, x, z.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_rd_hte(n=3000, ate=2.0, hte_coef=1.5, seed=42)
>>> df.attrs['true_ate']
2.0
>>> df.attrs['true_hte_coef']
1.5

dgp_rd_kink

dgp_rd_kink(n: int = 2000, kink: float = 0.8, cutoff: float = 0.0, seed: int | None = None) -> DataFrame

Generate Regression Kink Design data (Card et al. 2015).

The treatment function has a kink (change in slope) at the cutoff. The true kink effect on the outcome equals kink.

DGP: T = X + kink_T * max(X, 0), Y = 0.4*T + noise True RKD = 0.4 * kink_T / kink_T = kink (simplified)

Parameters:

Name	Type	Description	Default
n	int	Sample size.	2000
kink	float	True kink in slope at cutoff.	0.8
cutoff	float	Kink point location.	0.0
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: y, x, treatment.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_rd_kink(n=2000, kink=0.8, seed=42)
>>> df.attrs['true_kink']
0.8

dgp_rd_multi

dgp_rd_multi(n: int = 3000, cutoffs: list | None = None, effects: list | None = None, seed: int | None = None) -> DataFrame

Generate multi-cutoff RD data.

Parameters:

Name	Type	Description	Default
n	int	Sample size.	3000
cutoffs	list of float	Cutoff values. Default [0.0, 1.0].	None
effects	list of float	Treatment effects at each cutoff. Default [2.0, 3.0].	None
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: y, x, z (covariate).

Examples:

>>> import statspai as sp
>>> df = sp.dgp_rd_multi(n=3000, seed=42)
>>> df.attrs['true_effects']
{0.0: 2.0, 1.0: 3.0}

dgp_rdit

dgp_rdit(n_periods: int = 200, effect: float = 2.0, cutoff_period: int = 100, seasonality: bool = True, seed: int | None = None) -> DataFrame

Generate Regression Discontinuity in Time data (Hausman-Rapson 2018).

Parameters:

Name	Type	Description	Default
n_periods	int	Number of time periods.	200
effect	float	Treatment effect at the policy change date: a permanent level shift of this size.	2.0
cutoff_period	int	Period of the policy change.	100
seasonality	bool	Include monthly seasonality.	True

Notes

y = 0.01 t + effect * 1[t >= cutoff_period] + season + e with AR(1) errors, e_t = 0.3 e_{t-1} + nu_t and nu_t ~ N(0, 0.5^2). The autocorrelation is in the errors only, so the level shift is effect at every horizon. seed : int or None Random seed.

Returns:

Type	Description
DataFrame	Columns: y, time, date.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_rdit(n_periods=200, effect=2.0, seed=42)
>>> df.attrs['true_effect']
2.0

dgp_synth

dgp_synth(n_units: int = 20, n_periods: int = 30, treated_unit: int = 0, treatment_time: int = 20, effect: float = 0.5, seed: int | None = None) -> DataFrame

Generate synthetic control data with a factor model.

Parameters:

Name	Type	Description	Default
n_units	int	Number of units (including the treated unit).	20
n_periods	int	Number of time periods.	30
treated_unit	int	Index of the treated unit.	0
treatment_time	int	Period when treatment begins.	20
effect	float	Treatment effect (additive, constant post-treatment).	0.5
seed	int or None	Random seed.	None

Returns:

Type	Description
DataFrame	Columns: unit, time, y, treated.

Examples:

>>> import statspai as sp
>>> df = sp.dgp_synth(n_units=10, n_periods=30, treatment_time=20, effect=1.0, seed=0)
>>> df.loc[(df['unit'] == 0) & (df['time'] >= 20), 'treated'].unique()
array([1.])

outlier_indicator

outlier_indicator(data: DataFrame, vars: Sequence[str], cuts: tuple = (1, 99), combined: bool = False) -> DataFrame

Flag outlier observations by percentile cutoffs.

Returns 0/1 indicator columns for values outside [cuts[0], cuts[1]] percentile bounds. Mirrors the pywinsor2.outlier_indicator API.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
vars	list of str	Columns to check.	required
cuts	tuple of (float, float)	Lower and upper percentile cutoffs.	(1, 99)
combined	bool	If True, also add a column _outlier_any marking rows that are outliers on any variable.	False

Returns:

Type	Description
DataFrame	Copy of data with added {var}_outlier 0/1 columns.

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({"x": [1.0, 2.0, 3.0, 4.0, 100.0]})
>>> out = sp.outlier_indicator(df, ["x"], cuts=(10, 90))
>>> out["x_outlier"].tolist()
[1, 0, 0, 0, 1]

rank

rank(data: DataFrame, var: str, by: Optional[Union[str, Sequence[str]]] = None, method: str = 'average', ascending: bool = True) -> Series

Rank a variable, optionally within groups.

Equivalent to Stata's egen newvar = rank(var), by(group) with method / ascending mapping to Stata's tie-handling options (checked in Stata 18):

• 'average' → Stata default
• 'min' → track (1 + the number of lower values)
• 'min' with ascending=False → field (1 + the number of higher)
• 'first' → unique (ties in row order)
• 'max' / 'dense' have no Stata counterpart.

sp.egen(df, "rank", var, rank_ties="field") is the same with Stata's if and missing-value handling.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
var	str	Column to rank.	required
by	str or list of str	Group variables. If given, ranks are computed within each group.	None
method	(average, min, max, dense, first)		'average'
ascending	bool		True

Returns:

Type	Description
pd.Series aligned to ``data.index``.

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({"grp": ["a", "a", "b", "b"],
...                    "score": [3.0, 1.0, 4.0, 2.0]})
>>> sp.rank(df, "score").tolist()
[3.0, 1.0, 4.0, 2.0]
>>> sp.rank(df, "score", by="grp").tolist()
[2.0, 1.0, 2.0, 1.0]

rowcount

rowcount(data: DataFrame, vars: Sequence[str]) -> Series

Row-wise count of non-missing values (Stata egen rownonmiss).

Examples:

>>> import statspai as sp
>>> import pandas as pd, numpy as np
>>> df = pd.DataFrame({"math": [80.0, 90.0, np.nan],
...                    "read": [85.0, 95.0, 60.0],
...                    "sci":  [78.0, 88.0, np.nan]})
>>> sp.rowcount(df, ["math", "read", "sci"]).tolist()
[3, 3, 1]

rowmax

rowmax(data: DataFrame, vars: Sequence[str]) -> Series

Row-wise max across vars, ignoring NaN.

Examples:

>>> import statspai as sp
>>> import pandas as pd, numpy as np
>>> df = pd.DataFrame({"math": [80.0, 90.0, np.nan],
...                    "read": [85.0, 95.0, 60.0],
...                    "sci":  [78.0, 88.0, 65.0]})
>>> sp.rowmax(df, ["math", "read", "sci"]).tolist()
[85.0, 95.0, 65.0]

rowmean

rowmean(data: DataFrame, vars: Sequence[str]) -> Series

Row-wise mean across vars, ignoring NaN (Stata egen rowmean).

Examples:

>>> import statspai as sp
>>> import pandas as pd, numpy as np
>>> df = pd.DataFrame({"math": [80.0, 90.0, np.nan],
...                    "read": [85.0, 95.0, 60.0],
...                    "sci":  [78.0, 88.0, 65.0]})
>>> sp.rowmean(df, ["math", "read", "sci"]).tolist()
[81.0, 91.0, 62.5]

rowmin

rowmin(data: DataFrame, vars: Sequence[str]) -> Series

Row-wise min across vars, ignoring NaN.

Examples:

>>> import statspai as sp
>>> import pandas as pd, numpy as np
>>> df = pd.DataFrame({"math": [80.0, 90.0, np.nan],
...                    "read": [85.0, 95.0, 60.0],
...                    "sci":  [78.0, 88.0, 65.0]})
>>> sp.rowmin(df, ["math", "read", "sci"]).tolist()
[78.0, 88.0, 60.0]

rowsd

rowsd(data: DataFrame, vars: Sequence[str]) -> Series

Row-wise standard deviation across vars (Stata egen rowsd).

Uses the sample standard deviation (ddof=1), matching Stata.

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({"q1": [2.0, 5.0],
...                    "q2": [4.0, 5.0],
...                    "q3": [6.0, 5.0]})
>>> sp.rowsd(df, ["q1", "q2", "q3"]).tolist()
[2.0, 0.0]

rowtotal

rowtotal(data: DataFrame, vars: Sequence[str]) -> Series

Row-wise sum across vars, treating NaN as 0 (Stata egen rowtotal).

Examples:

>>> import statspai as sp
>>> import pandas as pd, numpy as np
>>> df = pd.DataFrame({"math": [80.0, 90.0, np.nan],
...                    "read": [85.0, 95.0, 60.0],
...                    "sci":  [78.0, 88.0, 65.0]})
>>> sp.rowtotal(df, ["math", "read", "sci"]).tolist()
[243.0, 273.0, 125.0]

read_data

read_data(path: str, encoding: Optional[str] = None, extended_missing: str = 'nan', **kwargs: Any) -> DataFrame

Read data from any common format, preserving variable labels.

Automatically detects format from file extension and stores Stata/SPSS/SAS variable labels in df.attrs['_labels'].

Parameters:

Name	Type	Description	Default
path	str	File path. Supported: .dta, .csv, .tsv (tab-separated), .xlsx, .xls, .parquet, .sas7bdat, .sav, .feather, .json.	required
encoding	str	Character encoding (for .csv and .tsv).	None
extended_missing	(nan, column)	What to do with Stata's extended missing values .a ... .z in a .dta file (survey codes such as "refused" or "don't know"). 'nan' reads them as NaN, like .; which of them a row held is then gone. 'column' also reads them as NaN and adds, for every variable that has any, a column <var>__miss holding '.a' ... '.z' on those rows and missing elsewhere. Either way the labels attached to these codes are in df.attrs['_missing_labels'].	'nan'
**kwargs	Any	Passed to the underlying pandas reader. A CSV / TSV file over 4 MB read without further keywords (other than sep / usecols) is parsed by pyarrow when it is installed: several times faster, the same column types (checked against pandas on the first rows, with pandas taking over on any difference), dates left as text, and decimals rounded to the nearest double where pandas' default parser can be one unit in the last place off.	{}

Returns:

Type	Description
DataFrame	With variable labels in df.attrs['_labels'] if available.

Warns:

Type	Description
UserWarning	When a .dta file labels an extended missing value and extended_missing='nan': the rows that held it are no longer told apart from ..

Examples:

>>> import statspai as sp
>>> df = sp.read_data('survey.dta')
>>> sp.describe(df)  # shows variable labels from Stata
>>> sp.get_label(df, 'wage')
'Monthly wage in CNY'

>>> df = sp.read_data('data.csv')  # CSV has no labels
>>> sp.label_vars(df, {'wage': 'Monthly wage', 'edu': 'Education'})

write_data

write_data(data: DataFrame, path: Union[str, Path], *, labels: Optional[Dict[str, str]] = None, value_labels: Optional[Dict[str, Dict[Any, str]]] = None, data_label: Optional[str] = None, extended_missing: str = 'column', **kwargs: Any) -> Path

Write data to a file, keeping variable and value labels.

The counterpart of :func:read_data. For Stata .dta the variable labels (data.attrs['_labels']), value labels (data.attrs['_value_labels']) and dataset label (data.attrs['_data_label']) are written into the file, so Stata's describe / label list and :func:read_data see them again. pandas.DataFrame.to_stata writes none of these unless each is passed by hand. So are the rest of what :func:read_data keeps of a .dta file: the names of the value-label sets, the extended missing values, notes and characteristics, and display formats (see Notes).

Parameters:

Name	Type	Description	Default
data	DataFrame	Data to write. The row index is not written.	required
path	str or Path	Output file. The extension picks the format: .dta, .csv, .tsv, .xlsx, .parquet, .feather, .json.	required
labels	dict	{column: label}. Added to (and overriding) the labels already in data.attrs['_labels'].	None
value_labels	dict	{column: {code: text}}. Added to (and overriding) those in data.attrs['_value_labels']. A code may be '.a' ... '.z' to label an extended missing value.	None
data_label	str	Dataset label (Stata's label data). Default: the one in data.attrs['_data_label'], if any.	None
extended_missing	(column, nan)	For .dta only. 'column' undoes read_data(extended_missing='column'): a column <var>__miss holding '.a' ... '.z' beside a numeric <var> is not written as a variable; the rows it marks are written as that missing value in <var>. 'nan' writes such a column as the string variable it is and every missing value as ..	'column'
**kwargs	Any	Passed to the underlying pandas writer. For .dta, version defaults to 118 (Stata 14 and later), the oldest format that stores labels as UTF-8; pass version=114 for older Stata, in which case labels must be Latin-1.	{}

Returns:

Type	Description
Path	The file written.

Raises:

Type	Description
TypeError	data is not a pandas DataFrame.
MethodIncompatibility	(a ValueError.) Unsupported extension; labels / value_labels naming a column that is not in data; a variable or dataset label longer than Stata's 80 characters; value labels on a non-numeric column or on non-integer codes; a <var>__miss code on a row where <var> is not missing.

Warns:

Type	Description
UserWarning	When data carries labels and the format cannot store them (.csv, .tsv, .xlsx, .json). Parquet and Feather store them.

Notes

Labels for columns that are not in data (dropped since the labels were attached) are skipped. Labels of extended missing values (data.attrs['_missing_labels']) are written into the file's label tables.

A value-label set keeps the name in data.attrs['_value_label_names'], and variables that share a name and the same labels share one set in the file, as they did in Stata. A variable whose labels no longer match the others under its name gets a set named after itself. Without a recorded name the set is named after the variable. A name that was attached to a variable without ever being defined is not kept.

Notes (data.attrs['_notes']) and characteristics (data.attrs['_characteristics']), each {variable or '_dta': ...}, are written for the columns still in data.

A display format in data.attrs['_formats'] is written when Stata would accept it for the column; a datetime64 column read from a %td / %tm / %tq variable is written back with that unit unless convert_dates= says otherwise. Each numeric column is stored in the smallest Stata type that holds its values exactly (as Stata's compress does); no value is rounded.

The file is written beside path and moved into place when complete, so a failed write leaves any existing file as it was.

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({'wage': [10.0, 20.0], 'female': [0, 1]})
>>> sp.label_var(df, 'wage', 'Monthly wage (CNY)')
>>> out = sp.write_data(
...     df, 'survey.dta',
...     value_labels={'female': {0: 'male', 1: 'female'}},
... )
>>> sp.get_label(sp.read_data('survey.dta'), 'wage')
'Monthly wage (CNY)'

scalar_iv_projection

scalar_iv_projection(data: DataFrame, treat: str, instruments: List[str], covariates: Optional[List[str]] = None, new_col: Optional[str] = None, return_column: bool = False) -> Union[DataFrame, Series]

Project a vector of instruments onto a scalar first-stage index.

Fits the first-stage OLS regression::

treat ~ intercept + Z1 + Z2 + ... + X1 + X2 + ...

and returns the in-sample fitted values. This scalar "first-stage index" can then be passed as a single-instrument input to estimators that don't accept vector-valued Z (e.g. StatsPAI's PLIV / IIVM implementations).

Parameters:

Name	Type	Description	Default
data	DataFrame		required
treat	str	Endogenous treatment variable name.	required
instruments	list of str	Excluded-instrument column names. Must have length ≥ 1.	required
covariates	list of str	Exogenous controls to include in the first stage.	None
new_col	str	Name of the column for the scalar index. Defaults to f'{treat}_iv_hat'.	None
return_column	bool	If True, return just the fitted pandas Series. If False, return a copy of data with the new column appended.	False

Returns:

Type	Description
DataFrame or Series	By default, data with the fitted-value column appended. With return_column=True, returns the fitted-value Series.

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> age = rng.normal(40, 8, n)
>>> parent_edu = rng.integers(8, 16, n).astype(float)
>>> z1, z2 = rng.normal(0, 1, n), rng.normal(0, 1, n)
>>> schooling = 12 + 0.8 * z1 + 0.5 * z2 + 0.1 * (age - 40) + rng.normal(0, 1, n)
>>> earnings = 5000 + 800 * schooling + 50 * age + rng.normal(0, 2000, n)
>>> df = pd.DataFrame({
...     'earnings': earnings, 'schooling': schooling,
...     'quarter_of_birth': z1, 'distance_to_college': z2,
...     'age': age, 'parent_edu': parent_edu,
... })
>>> # PLIV with two excluded instruments — project onto a scalar index first
>>> df_aug = sp.scalar_iv_projection(
...     df, treat='schooling',
...     instruments=['quarter_of_birth', 'distance_to_college'],
...     covariates=['age', 'parent_edu'],
... )
>>> 'schooling_iv_hat' in df_aug.columns
True
>>> result = sp.dml(df_aug, y='earnings', treat='schooling',
...                 covariates=['age', 'parent_edu'],
...                 model='pliv',
...                 instrument='schooling_iv_hat')

Notes

Projecting instruments onto a scalar index loses power relative to a full multi-instrument 2SLS when the instruments are heterogeneous, but the bias of the causal estimate remains valid under the same exogeneity conditions. A proper multi-instrument DML (vector r(X) and a Cragg-Donald-style first-stage test) is on the roadmap.

decode

decode(df: DataFrame, columns: Union[str, Sequence[str], None] = None, missing: bool = False) -> DataFrame

Replace value-labelled codes by their label texts.

The counterpart of Stata's decode. A .dta file read by sp.read_data keeps its labelled variables as numeric codes (what Stata itself stores and estimates on); this gives the readable version for tables, plots and exports.

Parameters:

Name	Type	Description	Default
df	DataFrame	Data with value labels in df.attrs['_value_labels'].	required
columns	str or list of str	Columns to decode. Default: every column that has value labels.	None
missing	bool	Also show the labels of extended missing values (.a "Refused") instead of leaving those rows missing. Needs the <var>__miss column that sp.read_data(extended_missing='column') adds.	False

Returns:

Type	Description
DataFrame	A copy in which each decoded column is an ordered categorical whose categories follow the codes. A code without a label is shown as the number. Decoded columns no longer appear in attrs['_value_labels']; variable labels are kept.

Raises:

Type	Description
MethodIncompatibility	(a ValueError.) A requested column is not in df or has no value labels.

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({'edu': [3, 1, 2, 1], 'wage': [30., 10., 20., 12.]})
>>> sp.label_values(df, 'edu', {1: 'primary', 2: 'secondary', 3: 'tertiary'})
>>> out = sp.decode(df)
>>> out['edu'].tolist()
['tertiary', 'primary', 'secondary', 'primary']
>>> list(out['edu'].cat.categories)   # code order, not alphabetical
['primary', 'secondary', 'tertiary']
>>> df['edu'].tolist()                # the input is not modified
[3, 1, 2, 1]

describe

describe(df: DataFrame, columns: Optional[list] = None) -> DataFrame

Stata-style describe — variable names, types, labels, and non-missing counts in one table.

Parameters:

Name	Type	Description	Default
df	DataFrame		required
columns	list	Subset of columns. Default: all.	None

Returns:

Type	Description
DataFrame	Columns: variable, type, n, n_missing, label.

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({'wage': [10.0, 20.0, 15.0],
...                    'edu': [12, 16, 14]})
>>> sp.label_var(df, 'wage', 'Monthly wage (CNY)')
>>> tbl = sp.describe(df)
>>> list(tbl.columns)
['variable', 'type', 'n', 'n_missing', 'label', 'value_labels']
>>> tbl['label'].tolist()
['Monthly wage (CNY)', '']
>>> sp.label_values(df, 'edu', {12: 'high school', 16: 'college'})
>>> sp.describe(df)['value_labels'].tolist()
['', '12=high school, 16=college']

get_label

get_label(df: DataFrame, var: str) -> str

Get the label for a variable, falling back to the column name.

Parameters:

Name	Type	Description	Default
df	DataFrame		required
var	str		required

Returns:

Type	Description
str	The label if set, otherwise the column name itself.

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({'wage': [10.0, 20.0], 'edu': [12, 16]})
>>> sp.label_var(df, 'wage', 'Monthly wage (CNY)')
>>> sp.get_label(df, 'wage')
'Monthly wage (CNY)'
>>> sp.get_label(df, 'edu')   # unlabeled -> falls back to column name
'edu'

get_labels

get_labels(df: DataFrame) -> Dict[str, str]

Get all variable labels as a dictionary.

Returns:

Type	Description
dict	{column_name: label} for all labeled columns. Unlabeled columns map to their own name.

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({'wage': [10.0, 20.0], 'edu': [12, 16]})
>>> sp.label_var(df, 'wage', 'Monthly wage (CNY)')
>>> sp.get_labels(df)
{'wage': 'Monthly wage (CNY)', 'edu': 'edu'}

label_values

label_values(df: DataFrame, var: Union[str, Sequence[str]], labels: Optional[Dict[Any, str]]) -> None

Attach value labels to the codes of one or more variables.

Equivalent to Stata's label define yn 0 "No" 1 "Yes" followed by label values female union yn. The data are not changed: the columns keep their numeric codes, and the labels are used by :func:describe, sp.tab, :func:decode, sp.regtable(labels=) and written to .dta / Parquet by sp.write_data.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame to label (modified in place).	required
var	str or list of str	Column(s) the labels apply to.	required
labels	dict or None	{code: text}. Codes are integers; '.a' ... '.z' label Stata's extended missing values. The mapping replaces any labels the variable already has; None (or {}) removes them.	required

Raises:

Type	Description
MethodIncompatibility	(a ValueError.) A column that is not in df, a non-numeric column, or a code that is neither an integer nor '.a' ... '.z'.

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({'female': [0, 1, 1], 'union': [1, 0, 1]})
>>> sp.label_values(df, ['female', 'union'], {0: 'No', 1: 'Yes'})
>>> df.attrs['_value_labels']['union']
{0: 'No', 1: 'Yes'}
>>> sp.decode(df)['female'].tolist()
['No', 'Yes', 'Yes']

label_var

label_var(df: DataFrame, var: str, label: str) -> None

Attach a human-readable label to a variable.

Equivalent to Stata's label variable wage "Monthly wage (CNY)".

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame to label (modified in place).	required
var	str	Column name.	required
label	str	Human-readable label.	required

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({'wage': [10.0, 20.0], 'edu': [12, 16]})
>>> sp.label_var(df, 'wage', 'Monthly wage (CNY)')
>>> sp.label_var(df, 'edu', 'Years of education')
>>> sp.get_label(df, 'wage')
'Monthly wage (CNY)'

label_vars

label_vars(df: DataFrame, labels: Union[Dict[str, str], DataFrame], rename: Optional[Dict[str, str]] = None) -> None

Attach labels to multiple variables at once, or copy them from a frame.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame to label (modified in place).	required
labels	dict or DataFrame	{column_name: label_string}, or a labelled DataFrame to copy from. Copying carries variable labels, value labels, the labels of extended missing values and display formats for every column the two frames share by name, and leaves labels df already has alone. It is the way to get labels back after a pandas operation that drops them: merge, pd.get_dummies, pivot, and pd.concat of frames whose labels differ all return a frame with empty attrs.	required
rename	dict	{old_name: new_name}, only when copying: the columns of labels that go by another name in df. DataFrame.rename keeps attrs but leaves the labels under the old names.	None

Examples:

>>> import statspai as sp
>>> import pandas as pd
>>> df = pd.DataFrame({'wage': [10.0, 20.0],
...                    'edu': [12, 16],
...                    'exp': [5, 8]})
>>> sp.label_vars(df, {
...     'wage': 'Monthly wage (CNY)',
...     'edu': 'Years of education',
...     'exp': 'Work experience (years)',
... })
>>> sp.get_label(df, 'exp')
'Work experience (years)'

A merge drops the labels; copy them back from the frame that had them:

>>> merged = df.merge(pd.DataFrame({'edu': [12, 16], 'z': [0, 1]}), on='edu')
>>> sp.get_label(merged, 'wage')
'wage'
>>> sp.label_vars(merged, df)
>>> sp.get_label(merged, 'wage')
'Monthly wage (CNY)'

>>> short = df.rename(columns={'wage': 'w'})
>>> sp.label_vars(short, df, rename={'wage': 'w'})
>>> sp.get_label(short, 'w')
'Monthly wage (CNY)'

diff

diff(data: DataFrame, var: Union[str, Sequence[str]], n: Union[int, Sequence[int]] = 1, *, id: Optional[str] = None, time: Optional[str] = None, delta: Optional[float] = None) -> Union[Series, DataFrame]

Stata's difference operator Dn.x, read from the time variable.

D.x is x - L.x; D2.x is the difference of the difference (x - 2 L.x + L2.x), not x - L2.x. Missing wherever a period it needs is absent.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
var	str or list of str		required
n	int or list of int	The order of the difference (0: the variable itself).	1
id	str	Default: the declaration of :func:xtset.	None
time	str	Default: the declaration of :func:xtset.	None
delta	float		None

Returns:

Type	Description
Series or DataFrame	Named D.x / D2.x.

Raises:

Type	Description
MethodIncompatibility	As :func:lag.

Examples:

>>> import pandas as pd
>>> import statspai as sp
>>> df = pd.DataFrame({"t": [1, 2, 3, 5], "x": [10.0, 12.0, 15.0, 19.0]})
>>> sp.diff(df, "x", time="t").tolist()
[nan, 2.0, 3.0, nan]

get_xtset

get_xtset(data: DataFrame) -> Optional[Dict[str, Any]]

The declaration sp.xtset stored on data (None if none).

lag

lag(data: DataFrame, var: Union[str, Sequence[str]], n: Union[int, Sequence[int]] = 1, *, id: Optional[str] = None, time: Optional[str] = None, delta: Optional[float] = None) -> Union[Series, DataFrame]

Stata's lag operator Ln.x, read from the time variable.

L.x at time t is x at t - delta in the same panel, and missing when that period is not in the data -- not the row above.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
var	str or list of str	The variable(s) to lag.	required
n	int or list of int	The lag order; 0 returns the variable itself. A list such as [1, 2, 3] is Stata's L(1/3).x.	1
id	str	The panel and time variables. Default: the declaration of :func:xtset.	None
time	str	The panel and time variables. Default: the declaration of :func:xtset.	None
delta	float	The length of one period (default: the declared one, else 1).	None

Returns:

Type	Description
Series or DataFrame	Aligned to data.index, named L.x / L2.x (x for order 0). One variable and one order give a Series.

Raises:

Type	Description
MethodIncompatibility	No time variable given or declared; a time variable with missing, non-integer or repeated values; a non-numeric variable.

Examples:

>>> import pandas as pd
>>> import statspai as sp
>>> df = pd.DataFrame({"t": [1, 2, 3, 5], "x": [10.0, 12.0, 15.0, 19.0]})
>>> sp.lag(df, "x", time="t").tolist()
[nan, 10.0, 12.0, nan]
>>> list(sp.lag(df, "x", n=[1, 2], time="t").columns)
['L.x', 'L2.x']

References

StataCorp. Stata 18 User's Guide, [U] 11.4.4 Time-series varlists.

lead

lead(data: DataFrame, var: Union[str, Sequence[str]], n: Union[int, Sequence[int]] = 1, *, id: Optional[str] = None, time: Optional[str] = None, delta: Optional[float] = None) -> Union[Series, DataFrame]

Stata's lead operator Fn.x, read from the time variable.

F.x at time t is x at t + delta in the same panel, and missing when that period is not in the data.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
var	str or list of str		required
n	int or list of int	The lead order (0: the variable itself).	1
id	str	Default: the declaration of :func:xtset.	None
time	str	Default: the declaration of :func:xtset.	None
delta	float		None

Returns:

Type	Description
Series or DataFrame	Named F.x / F2.x.

Raises:

Type	Description
MethodIncompatibility	As :func:lag.

Examples:

>>> import pandas as pd
>>> import statspai as sp
>>> df = pd.DataFrame({"t": [1, 2, 3, 5], "x": [10.0, 12.0, 15.0, 19.0]})
>>> sp.lead(df, "x", time="t").tolist()
[12.0, 15.0, nan, nan]

xtset

xtset(data: DataFrame, id: Optional[str] = None, time: Optional[str] = None, delta: float = 1.0) -> PanelDeclaration

Declare the panel and time variables of a DataFrame (Stata xtset).

sp.xtset(df, id="firm", time="year") is Stata's xtset firm year; sp.xtset(df, time="t") is tsset t. The declaration is stored in data.attrs (in place, as Stata changes the dataset in memory) and is read by :func:lag / :func:lead / :func:diff and by formulas that use L.x / F.x / D.x. Called with neither id nor time it returns the current declaration.

Parameters:

Name	Type	Description	Default
data	DataFrame	The data. Only data.attrs is changed; the rows keep their order (Stata sorts by panel and time; nothing here depends on the order).	required
id	str	The panel variable.	None
time	str	The time variable: whole numbers of periods (Stata dates are integers) or a pandas Period column.	None
delta	float	The length of one period, Stata's delta(): with delta=5, L.x at t is x at t - 5.	1

Returns:

Type	Description
PanelDeclaration	The variables, and Stata's description of the panel: balance (r(balanced)), whether some panel skips a period (r(gaps)), the number of panels, the first and last period.

Raises:

Type	Description
MethodIncompatibility	When the time variable has missing or non-integer values, repeats a period within a panel (Stata r(451)), is a datetime, or has two periods closer than delta (Stata r(451)); or when nothing has been declared and no variable is given.

Examples:

>>> import pandas as pd
>>> import statspai as sp
>>> df = pd.DataFrame({"id": [1, 1, 1, 2, 2], "t": [1, 2, 4, 1, 2],
...                    "x": [1.0, 4.0, 9.0, 2.0, 3.0]})
>>> decl = sp.xtset(df, id="id", time="t")
>>> decl.balanced, decl.gaps
('unbalanced', True)
>>> sp.lag(df, "x").tolist()   # t = 4 has no t = 3 before it
[nan, 1.0, nan, nan, 2.0]

References

StataCorp. Stata 18 Base Reference Manual, [XT] xtset and [TS] tsset.