Skills › Scientific Agent Skills › statistical-power › power.py
power.py
453 lines · skills/statistical-power/scripts/power.py at 5f2bc03
1"""Unified closed-form power / sample-size interface over statsmodels and scipy.
3One function each for the three things people actually want:
4 - sample_size(...) : solve for n given effect size, alpha, power
5 - power(...) : solve for achieved power given n and effect size
6 - mde(...) : solve for the minimum detectable effect given n and power
7 - power_curve(...) : plot power vs. n for planning figures
9Every call routes to the right statsmodels solver based on `test=`, so callers
10don't need to remember which Power class belongs to which test. All four solver
11quantities (effect_size, nobs, alpha, power) obey the identity "fix three, solve
12the fourth"; these helpers just expose that cleanly.
14Supported tests:
15 t_ind two independent means (Cohen's d)
16 t_paired/t_one paired or one-sample mean (Cohen's d)
17 anova one-way ANOVA, k groups (Cohen's f)
18 two_proportions two independent proportions (give prop1, prop2; uses Cohen's h)
19 one_proportion one proportion vs. a reference (give prop1, prop0)
20 correlation Pearson r (give effect_size=r)
21 chi2 goodness-of-fit / contingency (Cohen's w; give dof)
22 linear_regression added predictors via Cohen's f^2 (give f2 as effect_size, df_num)
24Tested with statsmodels 0.15.0, scipy 1.18.1, numpy 2.5.3, matplotlib 3.11.2.
25"""
27from __future__ import annotations
29import math
30from numbers import Integral, Real
31import warnings
33import numpy as np
34from statsmodels.stats.power import (
35 FTestAnovaPower,
36 GofChisquarePower,
37 NormalIndPower,
38 TTestIndPower,
39 TTestPower,
40)
41from statsmodels.stats.proportion import proportion_effectsize
42from statsmodels.tools.sm_exceptions import ConvergenceWarning
45def _finite(value, name):
46 if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real) or not math.isfinite(value):
47 raise ValueError(f"{name} must be a finite number")
48 return float(value)
51def _positive_integer(value, name, minimum=1):
52 if isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral) or value < minimum:
53 raise ValueError(f"{name} must be an integer >= {minimum}")
54 return int(value)
57def _target(alpha, target):
58 if not alpha < _finite(target, "power") < 1:
59 raise ValueError("target power must be strictly between alpha and 1")
62def _sample_n(test, nobs1, nobs, base):
63 if nobs1 is not None and nobs is not None:
64 raise ValueError("provide only one of nobs1= or nobs=")
65 n = nobs1 if nobs1 is not None else nobs
66 if n is None:
67 raise ValueError("provide nobs1= (sample 1) or nobs= (total observations/pairs)")
68 n = _finite(n, "sample size")
69 minimum = 0
70 if test in ("t_one", "t_paired"):
71 minimum = 1
72 elif test == "t_ind":
73 minimum = max(1, 1 / base["ratio"])
74 elif test == "correlation":
75 minimum = 3
76 elif test == "anova":
77 minimum = base["k_groups"]
78 elif test == "linear_regression":
79 minimum = base["k_total"] + 1
80 if n <= minimum:
81 raise ValueError(f"sample size for {test} must be > {minimum:g}")
82 return n
85# --------------------------------------------------------------------------- #
86# internal: build the statsmodels solver and the kwargs for the given test
87# --------------------------------------------------------------------------- #
88def _resolve(test, effect_size, alpha, alternative, *, solving_effect=False, **kw):
89 """Return (solver, base_kwargs, n_key) for a test.
91 n_key is the keyword the solver uses for "number of observations" so the
92 sample_size/power/mde wrappers can fill in the right argument generically.
93 """
94 test = test.lower()
95 allowed = {
96 "t_ind": {"ratio"}, "t_one": set(), "t_paired": set(),
97 "anova": {"k_groups"}, "two_proportions": {"prop1", "prop2", "ratio"},
98 "one_proportion": {"prop1", "prop0"}, "correlation": set(),
99 "chi2": {"dof"}, "linear_regression": {"df_num", "k_total"},
100 }
101 if test not in allowed:
102 raise ValueError(f"unknown test '{test}'")
103 if set(kw) - allowed[test]:
104 raise ValueError(f"unexpected arguments for {test}: {sorted(set(kw) - allowed[test])}")
105 if not 0 < _finite(alpha, "alpha") < 1:
106 raise ValueError("alpha must lie strictly between 0 and 1")
107 if alternative not in ("two-sided", "larger", "smaller"):
108 raise ValueError("alternative must be 'two-sided', 'larger', or 'smaller'")
109 if test in ("anova", "chi2", "linear_regression") and alternative != "two-sided":
110 raise ValueError(f"{test} is an omnibus upper-tail test; omit alternative")
111 if "ratio" in kw and _finite(kw["ratio"], "ratio") <= 0:
112 raise ValueError("ratio must be positive for a two-sample test")
113 if effect_size is not None:
114 effect_size = _finite(effect_size, "effect_size")
115 if test in ("anova", "chi2", "linear_regression") and effect_size < 0:
116 raise ValueError("f, w, and f-squared effect sizes cannot be negative")
117 if test == "correlation" and not -1 < effect_size < 1:
118 raise ValueError("correlation must lie strictly between -1 and 1")
119 if test in ("one_proportion", "two_proportions") and abs(effect_size) > math.pi:
120 raise ValueError("Cohen's h must lie between -pi and pi")
121 elif not solving_effect and test not in ("one_proportion", "two_proportions"):
122 raise ValueError("provide effect_size=")
124 if test == "t_ind":
125 solver = TTestIndPower()
126 es = effect_size
127 base = dict(effect_size=es, alpha=alpha, alternative=alternative,
128 ratio=kw.get("ratio", 1.0))
129 return solver, base, "nobs1"
131 if test in ("t_paired", "t_one"):
132 solver = TTestPower()
133 base = dict(effect_size=effect_size, alpha=alpha, alternative=alternative)
134 return solver, base, "nobs"
136 if test == "anova":
137 solver = FTestAnovaPower()
138 k = kw.get("k_groups")
139 if k is None:
140 raise ValueError("anova requires k_groups=")
141 k = _positive_integer(k, "k_groups", 2)
142 base = dict(effect_size=effect_size, alpha=alpha, k_groups=k)
143 return solver, base, "nobs" # nobs = TOTAL n across all groups
145 if test in ("two_proportions", "one_proportion"):
146 # convert proportions to Cohen's h, then use the normal approximation
147 reference = "prop2" if test == "two_proportions" else "prop0"
148 p1, p2 = kw.get("prop1"), kw.get(reference)
149 if solving_effect:
150 if p1 is not None or p2 is not None:
151 raise ValueError("MDE returns Cohen's h; omit proportions, then check baseline feasibility")
152 h = None
153 elif effect_size is not None:
154 if p1 is not None or p2 is not None:
155 raise ValueError("provide Cohen's h or proportions, not both")
156 h = effect_size
157 else:
158 if p1 is None or p2 is None:
159 raise ValueError(f"{test} requires prop1= and {reference}= (or effect_size=h)")
160 if not all(0 <= _finite(p, "proportion") <= 1 for p in (p1, p2)):
161 raise ValueError("proportions must lie in [0, 1]")
162 h = float(proportion_effectsize(p1, p2))
163 solver = NormalIndPower()
164 base = dict(effect_size=h, alpha=alpha, alternative=alternative,
165 ratio=kw.get("ratio", 1.0) if test == "two_proportions" else 0.0)
166 # ratio=0 is statsmodels' explicit single-sample branch.
167 return solver, base, "nobs1"
169 if test == "correlation":
170 # Power for Pearson r via Fisher z; handled analytically below, but we
171 # still route through a solver-shaped object for a uniform interface.
172 return "correlation", dict(effect_size=effect_size, alpha=alpha,
173 alternative=alternative), "nobs"
175 if test == "chi2":
176 solver = GofChisquarePower()
177 dof = kw.get("dof")
178 if dof is None:
179 raise ValueError("chi2 requires dof= (n_bins-1, or (r-1)(c-1))")
180 dof = _positive_integer(dof, "dof")
181 base = dict(effect_size=effect_size, alpha=alpha, n_bins=dof + 1)
182 return solver, base, "nobs"
184 if test == "linear_regression":
185 # Explicit total n keeps nuisance predictors in the residual df and
186 # in the assumed noncentrality f^2 * n. See _reg_* helpers below.
187 df_num = kw.get("df_num")
188 if df_num is None:
189 raise ValueError("linear_regression requires df_num= (number of tested predictors)")
190 df_num = _positive_integer(df_num, "df_num")
191 k_total = _positive_integer(kw.get("k_total", df_num), "k_total")
192 if df_num > k_total:
193 raise ValueError("df_num must not exceed k_total (excluding intercept)")
194 return "regression", dict(effect_size=effect_size, alpha=alpha,
195 df_num=df_num, k_total=k_total), "nobs"
197 raise ValueError(f"unknown test '{test}'")
200# --------------------------------------------------------------------------- #
201# correlation power (Fisher z transform) -- normal approximation
202# --------------------------------------------------------------------------- #
203def _corr_power(r, n, alpha, alternative):
204 from scipy import stats
205 z = math.atanh(r)
206 se = 1.0 / math.sqrt(n - 3)
207 if alternative == "two-sided":
208 zc = stats.norm.ppf(1 - alpha / 2)
209 return float(stats.norm.cdf(z / se - zc) + stats.norm.cdf(-z / se - zc))
210 else:
211 zc = stats.norm.ppf(1 - alpha)
212 direction = 1 if alternative == "larger" else -1
213 return float(stats.norm.cdf(direction * z / se - zc))
216def _corr_sample_size(r, alpha, power, alternative):
217 # Invert the same Fisher-z approximation as power(), including both tails.
218 # Solve for n-3, then shift; NormalIndPower(ddof=3)'s default n search
219 # starts outside its domain and can return a spurious n near 3.
220 with warnings.catch_warnings():
221 warnings.simplefilter("error", ConvergenceWarning)
222 n = float(3 + NormalIndPower().solve_power(
223 effect_size=math.atanh(r), alpha=alpha, power=power,
224 ratio=0, alternative=alternative))
225 if (not math.isfinite(n) or n <= 3 or
226 not math.isclose(_corr_power(r, n, alpha, alternative), power, abs_tol=1e-5)):
227 raise RuntimeError("correlation sample-size solution failed the power round-trip check")
228 return n
231# --------------------------------------------------------------------------- #
232# multiple-regression power via the noncentral F (Cohen's f^2)
233# --------------------------------------------------------------------------- #
234def _reg_power(f2, n, df_num, k_total, alpha):
235 """Power of the F-test for df_num tested predictors in a model with k_total
236 total predictors, total sample size n. Noncentrality lambda = f^2 * n."""
237 from scipy import stats
238 df_denom = n - k_total - 1
239 if df_denom <= 0:
240 raise ValueError("regression requires positive residual degrees of freedom")
241 ncp = f2 * n
242 crit = stats.f.ppf(1 - alpha, df_num, df_denom)
243 # SciPy 1.18.1 ncf.sf(..., nc=0) returns a negative result; use central F.
244 if ncp == 0:
245 return float(stats.f.sf(crit, df_num, df_denom))
246 result = float(1 - stats.ncf.cdf(crit, df_num, df_denom, ncp))
247 if not math.isfinite(result) or not 0 <= result <= 1:
248 raise RuntimeError("regression power calculation returned a nonfinite or out-of-range value")
249 return result
252def _reg_sample_size(f2, df_num, k_total, alpha, power):
253 lo, hi = k_total + 2, 1_000_000
254 if lo > hi or _reg_power(f2, hi, df_num, k_total, alpha) < power:
255 raise RuntimeError("sample size did not converge below 1e6")
256 while lo < hi:
257 mid = (lo + hi) // 2
258 if _reg_power(f2, mid, df_num, k_total, alpha) >= power:
259 hi = mid
260 else:
261 lo = mid + 1
262 return lo
265def _reg_mde(n, df_num, k_total, alpha, power):
266 """Smallest detectable f^2 at fixed n, with a checked upper bracket."""
267 from scipy.optimize import brentq
268 hi = 1.0
269 while _reg_power(hi, n, df_num, k_total, alpha) < power:
270 hi *= 2
271 if hi > 1e6:
272 raise RuntimeError("MDE did not converge below f-squared=1e6")
273 return brentq(lambda f2: _reg_power(f2, n, df_num, k_total, alpha) - power, 0, hi)
276# --------------------------------------------------------------------------- #
277# public API
278# --------------------------------------------------------------------------- #
279def sample_size(test, effect_size=None, alpha=0.05, power=0.80,
280 alternative="two-sided", round_up=True, **kw):
281 """Solve for required sample size.
283 For two-sample tests returns n1; n2 = ratio * n1 before integer allocation.
284 ANOVA returns balanced TOTAL n (rounded to a multiple of k_groups).
285 For paired tests n counts pairs; for other tests n counts observations.
286 Regression always searches integer n, even when round_up=False.
287 """
288 solver, base, n_key = _resolve(test, effect_size, alpha, alternative, **kw)
289 _target(alpha, power)
290 es = base["effect_size"]
291 if es == 0 or (alternative == "larger" and es < 0) or (alternative == "smaller" and es > 0):
292 raise ValueError("a nonzero effect in the tested direction is required for target power > alpha")
294 if solver == "correlation":
295 n = _corr_sample_size(effect_size, alpha, power, alternative)
296 return math.ceil(n) if round_up else n
298 if solver == "regression":
299 return _reg_sample_size(base["effect_size"], base["df_num"],
300 base["k_total"], alpha, power)
302 base["power"] = power
303 base[n_key] = None
304 with warnings.catch_warnings():
305 warnings.simplefilter("error", ConvergenceWarning)
306 n = float(solver.solve_power(**base))
307 if not math.isfinite(n) or n <= 0:
308 raise RuntimeError("sample-size solver did not return a finite positive value")
309 _sample_n(test.lower(), n, None, base)
310 check = dict(base, **{n_key: n})
311 check.pop("power")
312 if not math.isclose(float(solver.power(**check)), power, abs_tol=1e-5):
313 raise RuntimeError("sample-size solution failed the power round-trip check")
314 if round_up:
315 multiple = base["k_groups"] if test.lower() == "anova" else 1
316 n = multiple * math.ceil(n / multiple)
317 return n
320def power(test, effect_size=None, nobs1=None, nobs=None, alpha=0.05,
321 alternative="two-sided", **kw):
322 """Solve for achieved power given the sample size.
324 Pass nobs1 for per-group tests (t_ind, two_proportions), nobs for total-n tests.
325 """
326 solver, base, n_key = _resolve(test, effect_size, alpha, alternative, **kw)
328 n = _sample_n(test.lower(), nobs1, nobs, base)
330 if solver == "correlation":
331 return _corr_power(effect_size, n, alpha, alternative)
333 if solver == "regression":
334 return _reg_power(base["effect_size"], n, base["df_num"],
335 base["k_total"], alpha)
337 base[n_key] = n
338 result = float(solver.power(**base))
339 if not math.isfinite(result) or not 0 <= result <= 1:
340 raise RuntimeError("power calculation returned a nonfinite or out-of-range value")
341 return result
344def mde(test, nobs1=None, nobs=None, alpha=0.05, power=0.80,
345 alternative="two-sided", **kw):
346 """Solve for the minimum detectable effect (standardized) at a fixed n.
348 Returns the effect size in the test's native units (d, f, h, w, r, ...).
349 A 'smaller' directional MDE is negative; two-sided returns magnitude.
350 Proportions return Cohen's h without needing prop1/prop2/prop0.
351 """
352 solver, base, n_key = _resolve(test, None, alpha, alternative, solving_effect=True, **kw)
353 _target(alpha, power)
355 n = _sample_n(test.lower(), nobs1, nobs, base)
356 sign = -1 if alternative == "smaller" else 1
358 if solver == "correlation":
359 with warnings.catch_warnings():
360 warnings.simplefilter("error", ConvergenceWarning)
361 z = NormalIndPower().solve_power(nobs1=n - 3, power=power, alpha=alpha,
362 ratio=0, alternative="larger" if sign < 0 else alternative)
363 r = float(math.tanh(z))
364 if not math.isfinite(r) or not 0 <= r < 1:
365 raise RuntimeError("MDE is nonfinite or too close to |r|=1 to represent numerically")
366 if not math.isclose(_corr_power(sign * r, n, alpha, alternative), power, abs_tol=1e-5):
367 raise RuntimeError("correlation MDE solution failed the power round-trip check")
368 return sign * r
370 if solver == "regression":
371 return _reg_mde(n, base["df_num"], base["k_total"], alpha, power)
373 base["power"] = power
374 base["effect_size"] = None
375 base[n_key] = n
376 if sign < 0:
377 base["alternative"] = "larger"
378 with warnings.catch_warnings():
379 warnings.simplefilter("error", ConvergenceWarning)
380 effect = float(solver.solve_power(**base))
381 if not math.isfinite(effect) or effect < 0:
382 raise RuntimeError("MDE solver did not return a finite nonnegative magnitude")
383 if test.lower() in ("one_proportion", "two_proportions") and effect > math.pi:
384 raise ValueError("requested power requires an impossible Cohen's h > pi")
385 check = dict(base, effect_size=effect)
386 check.pop("power")
387 if not math.isclose(float(solver.power(**check)), power, abs_tol=1e-5):
388 raise RuntimeError("MDE solution failed the power round-trip check")
389 return sign * effect
392def power_curve(test, effect_size=None, n_range=None, alpha=0.05, power_target=0.80,
393 alternative="two-sided", save=None, show=False, **kw):
394 """Plot power vs. sample size. Returns (n_array, power_array).
396 n_range iterates sample-1 n for two-sample tests, pairs for paired tests,
397 observations otherwise (ANOVA: balanced total n).
398 """
399 import matplotlib.pyplot as plt
401 test = test.lower()
402 if n_range is None:
403 n_range = range(5, 205, 5)
404 ns = np.array(list(n_range), dtype=float)
406 pwr = []
407 for n in ns:
408 if test in ("t_ind", "two_proportions"):
409 pwr.append(power(test, effect_size=effect_size, nobs1=n, alpha=alpha,
410 alternative=alternative, **kw))
411 else:
412 pwr.append(power(test, effect_size=effect_size, nobs=n, alpha=alpha,
413 alternative=alternative, **kw))
414 pwr = np.array(pwr)
416 fig, ax = plt.subplots(figsize=(7, 4.5))
417 ax.plot(ns, pwr, lw=2)
418 ax.axhline(power_target, ls="--", color="crimson", lw=1,
419 label=f"target power = {power_target:g}")
420 unit = " (sample 1)" if test in ("t_ind", "two_proportions") else " (total)"
421 if test == "t_paired":
422 unit = " (pairs)"
423 ax.set_xlabel("Sample size" + unit)
424 ax.set_ylabel("Power (1 - β)")
425 ax.set_ylim(0, 1.02)
426 es_label = effect_size if effect_size is not None else ""
427 ax.set_title(f"Power curve: {test} (effect = {es_label}, α = {alpha:g})")
428 ax.legend()
429 ax.grid(alpha=0.3)
430 fig.tight_layout()
432 if save:
433 fig.savefig(save, dpi=150)
434 if show:
435 plt.show()
436 plt.close(fig)
437 return ns, pwr
440if __name__ == "__main__":
441 # quick smoke test of the main paths
442 print("t_ind d=0.5, 80% power -> n/group =",
443 sample_size("t_ind", effect_size=0.5, power=0.80))
444 print("anova f=0.25, k=4, 80% -> total n =",
445 sample_size("anova", effect_size=0.25, k_groups=4, power=0.80))
446 print("2 props 0.40 vs 0.55, 80% -> n/group =",
447 sample_size("two_proportions", prop1=0.40, prop2=0.55, power=0.80))
448 print("correlation r=0.30, 80% -> n =",
449 sample_size("correlation", effect_size=0.30, power=0.80))
450 print("MDE for t_ind at n=30/group, 80% power -> d =",
451 round(mde("t_ind", nobs1=30, power=0.80), 3))
452 print("power for t_ind d=0.5 at n=64/group ->",
453 round(power("t_ind", effect_size=0.5, nobs1=64), 3))