statspai.experimental

experimental

Experimental design and analysis tools.

Provides randomization, balance checking, attrition analysis, pre-analysis plan generation, group sequential designs, always-valid (mixture SPRT) inference and CUPED variance reduction for RCTs and A/B tests.

AdaptiveInferenceResult dataclass

Bases: ResultProtocolMixin

Result of :func:adaptive_inference.

Attributes:

Name	Type	Description
estimates	DataFrame	Arm means with standard errors and confidence intervals.
contrasts	DataFrame	Differences between arm means.
method	str

Examples:

>>> import statspai as sp
>>> exp = sp.bandit_experiment(
...     lambda k, rng: [0.0, 0.5][k] + rng.normal(), 200, n_arms=2,
...     sigma=1.0, prob_floor=0.1, seed=0)
>>> r = sp.adaptive_inference(exp.data, "reward", "arm", "prob")
>>> isinstance(r, sp.AdaptiveInferenceResult)
True
>>> r.contrasts["contrast"].tolist()
['1 - 0']

BanditExperimentResult dataclass

Bases: ResultProtocolMixin

Result of :func:bandit_experiment.

Attributes:

Name	Type	Description
data	DataFrame	One row per period: t, arm, reward, prob (the probability with which the chosen arm was assigned) and one prob_<arm> column per arm.
arms	DataFrame	Number of pulls and mean reward by arm.
regret	float or None	Cumulative shortfall of the expected reward relative to always playing the best arm; available when true_means was given or potential outcomes were supplied.

Examples:

>>> import statspai as sp
>>> exp = sp.bandit_experiment(
...     lambda k, rng: [0.0, 0.5][k] + rng.normal(), 100, n_arms=2,
...     sigma=1.0, seed=0)
>>> isinstance(exp, sp.BanditExperimentResult)
True
>>> list(exp.data.columns)
['t', 'arm', 'reward', 'prob', 'prob_0', 'prob_1']

AttritionResult

Bases: ResultProtocolMixin

Results from attrition analysis.

Produced by :func:attrition_test. Exposes overall / arm-specific attrition rates, the differential-attrition chi-squared test, and an optional table of covariate predictors of attrition via .summary().

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> treated = rng.integers(0, 2, n)
>>> observed = (rng.random(n) < np.where(treated == 1, 0.9, 0.75)).astype(int)
>>> age = rng.normal(40, 10, n)
>>> df = pd.DataFrame({"treated": treated,
...                    "endline_observed": observed, "age": age})
>>> res = sp.attrition_test(df, treatment="treated",
...                         observed="endline_observed", covariates=["age"])
>>> type(res).__name__
'AttritionResult'
>>> res.n_total
400
>>> isinstance(res.summary(), str)
True

BalanceResult

Bases: ResultProtocolMixin

Results from balance check.

Produced by :func:balance_check. Exposes the per-covariate balance .table (means, raw and normalized differences, t-test p-values), the omnibus F-test, and .summary() / .plot() (a love plot of normalized differences).

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> treated = rng.integers(0, 2, n)
>>> df = pd.DataFrame({"treated": treated,
...                    "age": rng.normal(40, 10, n),
...                    "income": rng.normal(50, 15, n)})
>>> bal = sp.balance_check(df, treatment="treated",
...                        covariates=["age", "income"])
>>> type(bal).__name__
'BalanceResult'
>>> bal.n_treat + bal.n_control
200
>>> list(bal.table["variable"])
['age', 'income']
>>> isinstance(bal.summary(), str)
True

plot

plot(ax: Any = None, **kwargs: Any) -> Any

Love plot of normalized differences.

RandomizationResult

Bases: ResultProtocolMixin

Results from randomization.

Produced by :func:randomize. Carries the assigned data frame (.data), arm counts and an optional balance summary, with a formatted .summary().

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> df = pd.DataFrame({"age": rng.normal(40, 10, n),
...                    "income": rng.normal(50, 15, n),
...                    "district": rng.integers(0, 4, n)})
>>> res = sp.randomize(df, strata="district",
...                    balance_vars=["age", "income"], seed=1)
>>> type(res).__name__
'RandomizationResult'
>>> int(res.n_treated + res.n_control)
200
>>> isinstance(res.summary(), str)
True

GroupSequentialResult dataclass

Bases: ResultProtocolMixin

Result of :func:group_sequential.

Attributes:

Name	Type	Description
boundaries	DataFrame	One row per look: look, info_frac, the efficacy and futility critical values on the z scale (efficacy_lower / efficacy_upper / futility_lower / futility_upper, NaN where a design has no such bound), their nominal p-values, cumulative alpha_spent (futility stops ignored), alpha_spent_futility (futility stops honoured), cumulative beta_spent under the alternative, and the stopping probabilities under the null and the alternative. With n_fixed given, n is the cumulative sample size at each look.
info_ratio	float	Maximum information over the information of the fixed design with the same alpha and power (the sample-size inflation factor).
power_attained	float	Power of the design at the alternative, which equals the target.
expected_info_h0, expected_info_h1	float	Expected information at stopping, as a fraction of the fixed design's information, under the null and the alternative (futility stops honoured, also for nonbinding bounds; see the Notes of :func:group_sequential for Stata's different convention).

Examples:

>>> import statspai as sp
>>> d = sp.group_sequential(n_looks=3, efficacy="obfleming")
>>> d.boundaries["efficacy_upper"].round(4).tolist()
[3.4711, 2.4544, 2.004]
>>> round(d.info_ratio, 4)
1.0174

decision

decision(z: Sequence[float]) -> DataFrame

Apply the boundaries to observed z-statistics.

Parameters:

Name	Type	Description	Default
z	sequence of float	The z-statistic at each look carried out so far (at most the number of looks of the design), on the scale of the design: a positive z favours the alternative of an 'upper' design and a negative one that of a 'lower' design.	required

Returns:

Type	Description
DataFrame	look, z and action ('continue', 'stop: efficacy' or 'stop: futility'); rows end at the first stop. At the final look a z that does not cross the efficacy bound is a futility (acceptance) decision.

plot

plot(ax: Any = None) -> Any

Plot the boundaries against the information fraction.

MSPRTResult dataclass

Bases: ResultProtocolMixin

Result of :func:msprt.

Attributes:

Name	Type	Description
path	DataFrame	One row per look: n, n_treated, n_control, estimate, se, likelihood_ratio, p_value (always valid, running minimum) and ci_lower / ci_upper (the confidence sequence, running intersection).
p_value	float	Always-valid p-value at the last look.
ci	tuple	Confidence sequence at the last look.
rejected	bool	Whether the null was rejected at some look.
stop_look	int or None	First look (row of path) at which p_value <= alpha.

Examples:

>>> import numpy as np, pandas as pd, statspai as sp
>>> rng = np.random.default_rng(1)
>>> d = rng.integers(0, 2, size=1000)
>>> df = pd.DataFrame({"d": d, "y": 0.5 * d + rng.normal(size=1000)})
>>> r = sp.msprt(df, "y", "d", tau=0.5, sigma=1.0)
>>> r.rejected
True

OptimalDesignResult

Bases: ResultProtocolMixin

Results from optimal design calculation.

Returned by :func:optimal_design. Carries the required total / per-arm sample size, cluster counts and size (for cluster designs), the intra- cluster correlation, the minimum detectable effect, and the target power.

Examples:

>>> import statspai as sp
>>> result = sp.optimal_design(
...     design="cluster", mde=0.2, sigma=1.0, icc=0.05, cluster_size=20
... )
>>> type(result).__name__
'OptimalDesignResult'
>>> result.design_type
'Cluster RCT'
>>> bool(result.n_total > 0 and result.n_clusters > 0)
True

adaptive_inference

adaptive_inference(data: DataFrame, y: str, arm: str, prob: Optional[str] = None, *, probs: Optional[Union[Sequence[str], Dict[Any, str]]] = None, method: str = 'aw', contrasts: Optional[Sequence[Tuple[Any, Any]]] = None, time: Optional[str] = None, alpha: float = 0.05) -> AdaptiveInferenceResult

Confidence intervals for arm means after an adaptive experiment.

Parameters:

Name	Type	Description	Default
data	DataFrame	One row per subject, in the order of assignment (or sortable by time).	required
y	str	Outcome column.	required
arm	str	Column with the assigned arm.	required
prob	str	Column with the probability with which the assigned arm was chosen, given the data available at that time. Sufficient for method="aw", "ipw" and "mean".	None
probs	sequence of str or dict	Columns with the assignment probability of every arm in each period: a dict {arm label: column} or a list aligned with the sorted arm labels. Required for method="aipw". The probabilities of a period must sum to one; if an arm was never assigned, pass a dict that includes its column. When prob is given as well, it must equal the probs entry of the assigned arm.	None
method	('aw', 'aipw', 'ipw', 'mean')	"aw" weights each observation of an arm by the inverse square root of its assignment probability and normalises by the sum of those weights. "aipw" applies square-root-probability weights to augmented inverse-probability scores whose outcome model is the arm's running mean over earlier periods. Both have asymptotically normal studentised statistics under adaptive assignment. "ipw" and "mean" are the inverse-probability weighted mean and the sample mean with their usual standard errors; they are reported for comparison and their intervals are not valid when assignment depended on earlier outcomes.	"aw"
contrasts	sequence of (arm, arm)	Pairs (a, b) for which mean(a) - mean(b) is reported. Defaults to every arm against the first.	None
time	str	Column to sort by before computing "aipw" running means.	None
alpha	float		0.05

Returns:

Type	Description
AdaptiveInferenceResult

Notes

The potential outcomes are assumed independent and identically distributed over time, and assignment may depend on earlier outcomes only. Assignment probabilities must be the ones actually used. Validity needs probabilities that do not fall to zero too fast; a warning is issued when an assigned arm had probability below 1 / T and an error when a probability is zero or the assignment was deterministic.

The variance of a contrast is the sum of the two arm variances: the weighted terms of different arms are uncorrelated because no period contributes an outcome to both.

References

[@hadad2021confidence]

Examples:

>>> import statspai as sp
>>> means = [0.0, 0.5]
>>> exp = sp.bandit_experiment(
...     lambda k, rng: means[k] + rng.normal(), 400, n_arms=2,
...     sigma=1.0, prob_floor=0.05, seed=1)
>>> r = sp.adaptive_inference(exp.data, "reward", "arm", "prob")
>>> r.estimates.shape[0]
2

bandit_allocate

bandit_allocate(data: DataFrame, y: str, arm: str, *, arms: Optional[Sequence[Any]] = None, algorithm: str = 'thompson', model: str = 'gaussian', sigma: Optional[float] = None, horizon: Optional[int] = None, ucb_scale: float = 2.0, epsilon: float = 0.1, prob_floor: float = 0.0, min_pulls: int = 2) -> Series

Assignment probabilities for the next subject of an adaptive experiment.

Parameters:

Name	Type	Description	Default
data	DataFrame	Outcomes observed so far, one row per subject. May be empty.	required
y	str	Outcome (reward) column; larger is better.	required
arm	str	Column with the arm each subject received.	required
arms	sequence	All arm labels, including arms not yet tried. Defaults to the labels present in data.	None
algorithm	('thompson', 'ucb', 'epsilon_greedy', 'uniform')	"thompson" assigns each arm with the posterior probability that it is the best one. "ucb" puts probability one on the arm with the largest upper confidence bound mean + ucb_scale * sd * sqrt(log(horizon) / n). "epsilon_greedy" plays the best arm so far with probability 1 - epsilon and a uniformly drawn arm otherwise.	"thompson"
model	('gaussian', 'bernoulli')	Reward model behind the posterior. "gaussian" uses a flat prior on each mean; "bernoulli" uses uniform Beta priors and needs 0/1 rewards.	"gaussian"
sigma	float	Known reward standard deviation for the Gaussian model, positive. By default the pooled within-arm standard deviation is used. When that cannot be estimated (the rewards do not vary within arms, or there is one observation per arm) the value 1 is used and a warning is issued, because the allocation then depends on the units of the reward.	None
horizon	int	Planned number of subjects, used by "ucb". Defaults to the number observed so far plus one.	None
ucb_scale	float	Width multiplier of the upper confidence bound.	2.0
epsilon	float	Exploration share of "epsilon_greedy".	0.1
prob_floor	float	Lower bound on every assignment probability, between 0 and one over the number of arms. A positive floor keeps later inference well behaved. It holds in every period, the first draws included (see min_pulls).	0.0
min_pulls	int	Arms with fewer observations are drawn first, with equal probability among them. With prob_floor=0 the other arms get probability zero in those periods, so when a single arm is short its draw is deterministic (probability one, which is what is returned and what :func:bandit_experiment records). With a positive prob_floor every arm keeps at least the floor and the short arms share the rest.	2

Returns:

Type	Description
Series	Probabilities indexed by arm label, summing to one.

References

[@thompson1933likelihood], [@lai1985asymptotically]

Examples:

>>> import pandas as pd
>>> import statspai as sp
>>> df = pd.DataFrame({"arm": ["a", "b", "a", "b"], "y": [1.0, 0.2, 0.8, 0.1]})
>>> p = sp.bandit_allocate(df, "y", "arm", sigma=1.0)
>>> bool(p["a"] > p["b"])
True

bandit_experiment

bandit_experiment(reward: Union[Callable[[int, Generator], float], ndarray, DataFrame], n_periods: Optional[int] = None, *, n_arms: Optional[int] = None, algorithm: str = 'thompson', model: str = 'gaussian', sigma: Optional[float] = None, ucb_scale: float = 2.0, epsilon: float = 0.1, prob_floor: float = 0.0, min_pulls: int = 2, batch_size: int = 1, true_means: Optional[Sequence[float]] = None, seed: Optional[int] = None) -> BanditExperimentResult

Run a sequential multi-armed bandit experiment.

Parameters:

Name	Type	Description	Default
reward	(callable, ndarray or DataFrame)	Either a function reward(k, rng) returning the outcome of pulling arm k (an integer from 0), or a T x K table of potential outcomes whose row t holds the outcome each arm would produce in period t.	required
n_periods	int	Number of subjects. Required with a callable; defaults to the number of rows of a potential-outcome table.	None
n_arms	int	Number of arms. Required with a callable.	None
algorithm	str	As in :func:bandit_allocate. "ucb" uses n_periods as its horizon.	'thompson'
model	str	As in :func:bandit_allocate. "ucb" uses n_periods as its horizon.	'thompson'
sigma	str	As in :func:bandit_allocate. "ucb" uses n_periods as its horizon.	'thompson'
ucb_scale	str	As in :func:bandit_allocate. "ucb" uses n_periods as its horizon.	'thompson'
epsilon	str	As in :func:bandit_allocate. "ucb" uses n_periods as its horizon.	'thompson'
prob_floor	str	As in :func:bandit_allocate. "ucb" uses n_periods as its horizon.	'thompson'
min_pulls	str	As in :func:bandit_allocate. "ucb" uses n_periods as its horizon.	'thompson'
batch_size	int	Assignment probabilities are recomputed every batch_size periods.	1
true_means	sequence of float	Mean reward of each arm, used to report regret. With a potential-outcome table the column means are used by default.	None
seed	int		None

Returns:

Type	Description
BanditExperimentResult	.data can be passed to :func:adaptive_inference.

Notes

Rules that drive assignment probabilities to zero quickly collect high rewards during the experiment but leave little information about the arms they abandon. If confidence intervals are wanted afterwards, set prob_floor so that every arm keeps being sampled.

References

[@thompson1933likelihood], [@lai1985asymptotically]

Examples:

>>> import statspai as sp
>>> means = [0.0, 0.5]
>>> exp = sp.bandit_experiment(
...     lambda k, rng: means[k] + rng.normal(), 300, n_arms=2,
...     sigma=1.0, prob_floor=0.05, true_means=means, seed=0)
>>> exp.data.shape[0]
300

contextual_bandit

contextual_bandit(reward: Union[Callable[[int, ndarray, Generator], float], ndarray, DataFrame], contexts: Union[ndarray, DataFrame], *, n_arms: Optional[int] = None, algorithm: str = 'thompson', sigma: Optional[float] = None, prior_precision: float = 1.0, ucb_scale: float = 2.0, epsilon: float = 0.1, prob_floor: float = 0.0, batch_size: int = 1, add_intercept: bool = True, mean_reward: Optional[Callable[[int, ndarray], float]] = None, seed: Optional[int] = None) -> BanditExperimentResult

Run a sequential experiment whose assignment depends on covariates.

Each arm has a linear reward model in the subject's covariates, with a Gaussian posterior that is updated after every batch. Assignment probabilities are computed from the posteriors at the covariates of the arriving subject and are recorded.

Parameters:

Name	Type	Description	Default
reward	(callable, ndarray or DataFrame)	Either a function reward(k, x, rng) returning the outcome of giving arm k to a subject with covariate vector x, or a T x K table of potential outcomes.	required
contexts	(ndarray or DataFrame, shape(T, p))	Covariates of the subjects, in order of arrival. Numeric: a categorical or text column is refused and has to be expanded into 0/1 indicators first. A column may not be called t, arm, reward, prob or prob_<arm>, the names of the columns of the result.	required
n_arms	int	Number of arms. Required with a callable.	None
algorithm	('thompson', 'ucb', 'epsilon_greedy', 'uniform')	"thompson" assigns each arm with the posterior probability that its expected reward at x is the largest. "ucb" plays the arm with the largest mean + ucb_scale * sd at x. "epsilon_greedy" plays the arm with the largest posterior mean with probability 1 - epsilon.	"thompson"
sigma	float	Reward standard deviation. By default it is re-estimated from the residuals of the fitted models as data arrive.	None
prior_precision	float	Precision of the independent normal prior on each coefficient (a ridge penalty).	1.0
ucb_scale	float	As in :func:bandit_experiment.	2.0
epsilon	float	As in :func:bandit_experiment.	2.0
prob_floor	float	As in :func:bandit_experiment.	2.0
batch_size	float	As in :func:bandit_experiment.	2.0
add_intercept	bool	Add a constant to the covariates.	True
mean_reward	callable	mean_reward(k, x), the expected reward of arm k at x, used to report regret against the best arm for each subject. With a potential-outcome table, realised regret against the best potential outcome in each row is reported instead.	None
seed	int		None

Returns:

Type	Description
BanditExperimentResult	.data holds the period, arm, reward, the assignment probability of every arm, and the covariates. .detail["coefficients"] holds the final posterior means, one row per arm.

Notes

The recorded probabilities are exact, not simulated, which is what weighting estimators need afterwards. With them, the value of any fixed rule can be estimated by inverse-probability weighting (:func:statspai.ips with pi_behavior= the prob column). :func:adaptive_inference is for experiments without covariates and should not be applied to these data: once assignment depends on covariates, the weighted arm means are no longer centred on the arm means.

References

[@thompson1933likelihood], [@lai1985asymptotically]

Examples:

>>> import numpy as np
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(300, 1))
>>> exp = sp.contextual_bandit(
...     lambda k, x, g: (x[0] if k == 1 else -x[0]) + g.normal(),
...     X, n_arms=2, sigma=1.0, prob_floor=0.05, seed=0)
>>> exp.data.shape[0]
300

attrition_bounds

attrition_bounds(data: DataFrame, y: str, treatment: str, observed: Optional[str] = None, method: str = 'lee', alpha: float = 0.05, trimming: str = 'quantile') -> Dict[str, Any]

Compute bounds on treatment effects under attrition.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str	Outcome variable.	required
treatment	str	Treatment indicator (0/1).	required
observed	str	Indicator for observed outcome. If None, uses non-missing y.	None
method	str	Bounding method: 'lee' (Lee 2009), 'manski' (worst-case).	'lee'
alpha	float	Not used: no confidence interval is computed. For an interval for the identified set use :func:sp.lee_bounds.	0.05
trimming	(quantile, exact)	Lee trimming rule, as in :func:sp.lee_bounds: Lee's sample-quantile estimator, or fractional trimming of the observations tied at the quantile (what Stata leebounds is written to compute). The Lee bounds equal sp.lee_bounds' exactly.	'quantile'

Returns:

Type	Description
dict	Keys: 'lower_bound', 'upper_bound', 'naive_ate', 'method', 'n_obs'.

Notes

Rows with a missing treatment or observed value, and rows with observed == 1 but a missing outcome, are dropped with a warning (the estimation sample of :func:sp.lee_bounds). treatment and observed must be coded 0/1.

Before 1.30 method='lee' kept n - ceil(q n) observations, one fewer than Lee's sample-quantile rule whenever q n is not an integer, and disagreed with sp.lee_bounds on the same data.

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> treated = rng.integers(0, 2, n)
>>> y = 1.0 * treated + rng.normal(size=n)
>>> observed = (rng.random(n) < np.where(treated == 1, 0.9, 0.75)).astype(int)
>>> y = np.where(observed == 1, y, np.nan)
>>> df = pd.DataFrame({"y": y, "treated": treated, "observed": observed})
>>> res = sp.attrition_bounds(df, y="y", treatment="treated",
...                           observed="observed", method="lee")
>>> sorted(res.keys())
['attrition_rate', 'lower_bound', 'method', 'n_obs', 'n_total',
 'naive_ate', 'upper_bound']
>>> bool(res["lower_bound"] <= res["naive_ate"] <= res["upper_bound"])
True

References

[@lee2009training]

attrition_test

attrition_test(data: DataFrame, treatment: str, observed: str, covariates: Optional[List[str]] = None, correction: bool = True) -> AttritionResult

Test for differential attrition in an RCT.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
treatment	str	Treatment indicator (0/1).	required
observed	str	Indicator for whether outcome is observed (1) or missing (0).	required
covariates	list of str	Baseline covariates to test as predictors of attrition.	None
correction	bool	Yates' continuity correction for the 2x2 differential-attrition test (R chisq.test's default). False gives Pearson's chi2, which is what Stata tabulate treat observed, chi2 reports.	True

Returns:

Type	Description
AttritionResult

Notes

If the attrition-predictor regression for a covariate fails (e.g. constant or degenerate values), its row in covariate_tests is reported as NaN and a StatsPAIWarning names the covariate.

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> treated = rng.integers(0, 2, n)
>>> # attrition slightly higher in the control arm
>>> p_obs = 0.9 - 0.1 * (treated == 0)
>>> df = pd.DataFrame({
...     "treated": treated,
...     "endline_observed": (rng.uniform(size=n) < p_obs).astype(int),
...     "age": rng.normal(40, 10, n),
...     "income": rng.normal(50, 15, n),
...     "education": rng.integers(8, 18, n).astype(float),
... })
>>> result = sp.attrition_test(df, treatment='treated',
...                            observed='endline_observed',
...                            covariates=['age', 'income', 'education'])
>>> _ = result.summary()

balance_check

balance_check(data: DataFrame, treatment: str, covariates: List[str], alpha: float = 0.05, equal_var: bool = False) -> BalanceResult

Check covariate balance between treatment and control.

Computes normalized differences, t-tests, and omnibus F-test.

Equivalent to Stata's iebaltab and R's cobalt::bal.tab().

Parameters:

Name	Type	Description	Default
data	DataFrame		required
treatment	str	Binary treatment variable (0/1).	required
covariates	list of str	Covariates to check balance on.	required
alpha	float	Not used: the table reports p-values, and the summary flags |normalized difference| > 0.25.	0.05
equal_var	bool	Per-covariate difference test. False is Welch's unequal-variance t-test. True is the pooled-variance t-test -- the OLS t-test of the covariate on the treatment dummy -- which is what Stata iebaltab reports (its pair test regresses the variable on the group dummy; its difference is control minus treatment, the opposite sign of diff here).	False

Returns:

Type	Description
BalanceResult

Notes

The omnibus F-test regresses the treatment on the covariates (complete cases); its numerator degrees of freedom are the rank of the covariates, so a covariate that is a linear combination of others is dropped as R lm and Stata regress drop it, with a warning. If the omnibus F-test cannot be computed (e.g. non-numeric or perfectly collinear covariates), omnibus_f and omnibus_p are reported as NaN and a StatsPAIWarning is emitted.

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> df = pd.DataFrame({"treated": rng.integers(0, 2, n),
...                    "age": rng.normal(40, 10, n),
...                    "income": rng.normal(50, 15, n),
...                    "education": rng.normal(12, 3, n)})
>>> bal = sp.balance_check(df, treatment="treated",
...                        covariates=["age", "income", "education"])
>>> type(bal).__name__
'BalanceResult'
>>> list(bal.table["variable"])
['age', 'income', 'education']
>>> bool(isinstance(bal.summary(), str))
True

randomize

randomize(data: DataFrame, n_arms: int = 2, prob: Optional[List[float]] = None, strata: Optional[str] = None, cluster: Optional[str] = None, method: str = 'simple', balance_vars: Optional[List[str]] = None, n_rerand: int = 0, rerand_threshold: float = 0.001, seed: Optional[int] = None, treatment_col: str = 'treatment', rerand_accept: Optional[float] = None) -> RandomizationResult

Randomize units to treatment and control.

Equivalent to R's randomizr::complete_ra() / block_ra() / cluster_ra().

Parameters:

Name	Type	Description	Default
data	DataFrame	Data with units to randomize.	required
n_arms	int	Number of treatment arms.	2
prob	list of float	Probability of each arm. Default: equal.	None
strata	str	Stratification variable for block randomization.	None
cluster	str	Cluster variable for cluster randomization.	None
method	str	'simple' (independent Bernoulli draws; randomizr::simple_ra), 'complete' (fixed arm counts; complete_ra), 'stratified' (complete randomization within each stratum; block_ra) or 'cluster' (complete randomization of clusters; cluster_ra). Passing strata= / cluster= selects the block / cluster design whatever method says. Passing both randomizes whole clusters within each stratum (randomizr::block_and_cluster_ra); every cluster must then lie in one stratum. Missing strata or cluster values raise DataInsufficient, as in randomizr.	'simple'

Notes

Arm counts follow randomizr::complete_ra: floor(N p_j) units per arm, with the misfits allocated at random in proportion to the fractional parts. Before 1.30 (a) method='complete' silently did simple randomization, (b) stratified counts used Python's round-half-to-even and gave every misfit of an odd stratum to a fixed arm (strata of 5 got 2 treated, strata of 7 got 4), (c) cluster randomization was Bernoulli per cluster, and (d) re-randomization redrew by simple randomization, discarding the strata / clusters. balance_vars : list of str, optional Variables to check balance on (for re-randomization). n_rerand : int, default 0 Number of re-randomization iterations (0 = no re-randomization). rerand_threshold : float, default 0.001 Stopping threshold of the n_rerand search, on diff' S^-1 diff with diff the difference in covariate means and S the covariance of the covariates. This distance is not scaled by the arm sizes, so the threshold is not a chi-square quantile; rerand_accept is. rerand_accept : float, optional Rerandomization by Morgan and Rubin's criterion, for two arms under complete randomization. Assignments are redrawn until M = (n1 n0 / n) diff' S^-1 diff is at most the rerand_accept quantile of the chi-square distribution with as many degrees of freedom as balance_vars, so about that share of all assignments is acceptable (0.05 keeps the best-balanced 5%). result.rerandomization records the threshold, the accepted M and the number of draws. Unlike the n_rerand search the accepted set is known, which is what the analysis needs: a randomization test must redraw from it, and regression adjustment for balance_vars (sp.lm_lin) stays valid. When given, n_rerand is the largest number of draws (default ceil(1000 / rerand_accept)) and rerand_threshold is unused. seed : int, optional Random seed for reproducibility. treatment_col : str, default 'treatment' Name of treatment column to create.

Returns:

Type	Description
RandomizationResult

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> df = pd.DataFrame({"district": rng.integers(0, 4, n),
...                    "age": rng.normal(40, 10, n),
...                    "income": rng.normal(50, 15, n)})
>>> result = sp.randomize(df, strata="district",
...                       balance_vars=["age", "income"], seed=1)
>>> type(result).__name__
'RandomizationResult'
>>> int(result.n_treated + result.n_control)
200
>>> bool(isinstance(result.summary(), str))
True
>>> df_randomized = result.data

msprt_lambda

msprt_lambda(estimate, variance, tau: float, theta0: float = 0.0)

Mixture likelihood ratio Lambda for estimates with known variances.

Parameters:

Name	Type	Description	Default
estimate	float or array - like	Estimate of the effect and its variance at each look.	required
variance	float or array - like	Estimate of the effect and its variance at each look.	required
tau	float	Standard deviation of the normal mixing distribution of the effect, in the units of the estimate.	required
theta0	float	Null value.	0

Returns:

Type	Description
ndarray or float	Lambda at each look (computed on the log scale).

Examples:

>>> import statspai as sp
>>> round(float(sp.msprt_lambda(0.0, 1.0, tau=1.0)), 6)
0.707107

optimal_design

optimal_design(design: str = 'individual', sigma: float = 1.0, mde: Optional[float] = None, power: float = 0.8, alpha: float = 0.05, n_arms: int = 2, prop_treat: float = 0.5, icc: float = 0.0, cluster_size: Optional[int] = None, n_clusters: Optional[int] = None, cost_per_cluster: Optional[float] = None, cost_per_unit: Optional[float] = None, r2: float = 0.0, baseline_mean: float = 0.0, n: Optional[int] = None) -> OptimalDesignResult

Compute optimal sample size and design parameters.

Parameters:

Name	Type	Description	Default
design	str	'individual', 'cluster', 'stratified'.	'individual'
sigma	float	Standard deviation of the outcome.	1.0
mde	float	Minimum detectable effect. If None, compute MDE given n.	None
power	float	Statistical power (1 - Type II error).	0.8
alpha	float	Significance level.	0.05
n_arms	int	Number of treatment arms.	2
prop_treat	float	Proportion assigned to treatment.	0.5
icc	float	Intra-cluster correlation (for cluster designs).	0.0
cluster_size	int	Average cluster size.	None
n_clusters	int	Number of clusters (if fixed).	None
cost_per_cluster	float	Cost of adding a cluster (for optimal allocation).	None
cost_per_unit	float	Cost per individual unit.	None
r2	float	R-squared from baseline covariates (variance reduction).	0.0
baseline_mean	float		0.0
n	int	Total number of individuals, used to solve for the MDE when mde is None (individual and stratified designs; cluster designs use n_clusters and cluster_size).	None

Returns:

Type	Description
OptimalDesignResult	n_total is the total sample (individuals); n_per_arm the size of the largest arm (all arms are equal under prop_treat=0.5).

Notes

Two arms with treated share p and total sample N give Var(tau_hat) = sigma^2 (1 - r2) deff / (N p (1 - p)), so N = (z_{1-alpha/2} + z_{power})^2 sigma^2 (1 - r2) deff / (mde^2 p (1 - p)) individuals in total, p N treated and (1 - p) N control (deff = 1 + (cluster_size - 1) icc for cluster designs, 1 otherwise). With n_arms > 2 every arm is sized for its pairwise comparison with control at equal allocation. When cost_per_cluster and cost_per_unit are given, the cost-optimal cluster size sqrt((cost_per_cluster / cost_per_unit) (1 - icc) / icc) is used for the sample-size calculation itself.

Examples:

>>> import statspai as sp
>>> result = sp.optimal_design(
...     design="cluster", mde=0.2, sigma=1.0, icc=0.05, cluster_size=20
... )
>>> result.design_type, result.cluster_size, result.icc
('Cluster RCT', 20, 0.05)
>>> text = result.summary()

switchback_design

switchback_design(n_periods: int, m: int = 0, design: Union[str, int, Sequence[Any], ndarray] = 'optimal', p: float = 0.5, seed: Optional[int] = None) -> DataFrame

Randomization points and an assignment path for a switchback experiment.

Parameters:

Name	Type	Description	Default
n_periods	int	Number of periods T in the experiment.	required
m	int	Order of the carryover effect: a treatment may affect the outcome of its own period and of the next m.	0
design	('optimal', 'every')	Where the coin is flipped. 'optimal' is the minimax design of Bojinov, Simchi-Levi and Zhao (2023, Theorem 2): every period when m = 0; periods 1, 2m+1, 3m+1, ..., (n-2)m+1 when m > 0 and n_periods = n m with n >= 4. 'every' flips in every period, an integer b flips every b periods, and a boolean vector marks the periods directly.	'optimal'
p	float	Probability of treatment at each flip. The optimal design uses a fair coin (their Theorem 1).	0.5
seed	int	Seed for the assignment draw.	None

Returns:

Type	Description
DataFrame	One row per period: period (1-based), randomize (True where the coin is flipped) and treat (the drawn assignment). Run the experiment with treat, add the outcome column, and pass the frame to :func:switchback with design='randomize'.

Examples:

>>> import statspai as sp
>>> plan = sp.switchback_design(12, m=2, seed=0)
>>> plan.loc[plan["randomize"], "period"].tolist()
[1, 5, 7, 9]

References

bojinov2023design