zhukang
1c7abd81d1
- 310 registry:skill items, 1528 files, 7 categories (hub tags) - registry.json generated by @kortix/registry buildRegistry - nested SKILL.md in sn-da-* renamed SUBSKILL.md (no phantom skills) - provenance + licenses: SOURCES.md
321 lines
12 KiB
Python
321 lines
12 KiB
Python
"""Unified closed-form power / sample-size interface over statsmodels and scipy.
|
||
|
||
One function each for the three things people actually want:
|
||
- sample_size(...) : solve for n given effect size, alpha, power
|
||
- power(...) : solve for achieved power given n and effect size
|
||
- mde(...) : solve for the minimum detectable effect given n and power
|
||
- power_curve(...) : plot power vs. n (or vs. effect size) for planning figures
|
||
|
||
Every call routes to the right statsmodels solver based on `test=`, so callers
|
||
don't need to remember which Power class belongs to which test. All four solver
|
||
quantities (effect_size, nobs, alpha, power) obey the identity "fix three, solve
|
||
the fourth"; these helpers just expose that cleanly.
|
||
|
||
Supported tests:
|
||
t_ind two independent means (Cohen's d)
|
||
t_paired/t_one paired or one-sample mean (Cohen's d)
|
||
anova one-way ANOVA, k groups (Cohen's f)
|
||
two_proportions two independent proportions (give prop1, prop2; uses Cohen's h)
|
||
one_proportion one proportion vs. a reference (give prop1, prop0)
|
||
correlation Pearson r (give effect_size=r)
|
||
chi2 goodness-of-fit / contingency (Cohen's w; give dof)
|
||
linear_regression added predictors via Cohen's f^2 (give f2 as effect_size, df_num)
|
||
|
||
Requires: statsmodels>=0.14.6, scipy>=1.11, numpy, matplotlib.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
from dataclasses import dataclass
|
||
|
||
import numpy as np
|
||
from statsmodels.stats.power import (
|
||
FTestAnovaPower,
|
||
GofChisquarePower,
|
||
NormalIndPower,
|
||
TTestIndPower,
|
||
TTestPower,
|
||
)
|
||
from statsmodels.stats.proportion import proportion_effectsize
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# internal: build the statsmodels solver and the kwargs for the given test
|
||
# --------------------------------------------------------------------------- #
|
||
def _resolve(test, effect_size, alpha, alternative, **kw):
|
||
"""Return (solver, base_kwargs, n_key) for a test.
|
||
|
||
n_key is the keyword the solver uses for "number of observations" so the
|
||
sample_size/power/mde wrappers can fill in the right argument generically.
|
||
"""
|
||
test = test.lower()
|
||
|
||
if test == "t_ind":
|
||
solver = TTestIndPower()
|
||
es = effect_size
|
||
base = dict(effect_size=es, alpha=alpha, alternative=alternative,
|
||
ratio=kw.get("ratio", 1.0))
|
||
return solver, base, "nobs1"
|
||
|
||
if test in ("t_paired", "t_one"):
|
||
solver = TTestPower()
|
||
base = dict(effect_size=effect_size, alpha=alpha, alternative=alternative)
|
||
return solver, base, "nobs"
|
||
|
||
if test == "anova":
|
||
solver = FTestAnovaPower()
|
||
k = kw.get("k_groups")
|
||
if k is None:
|
||
raise ValueError("anova requires k_groups=")
|
||
base = dict(effect_size=effect_size, alpha=alpha, k_groups=k)
|
||
return solver, base, "nobs" # nobs = TOTAL n across all groups
|
||
|
||
if test == "two_proportions":
|
||
# convert proportions to Cohen's h, then use the normal approximation
|
||
p1, p2 = kw.get("prop1"), kw.get("prop2")
|
||
if p1 is None or p2 is None:
|
||
raise ValueError("two_proportions requires prop1= and prop2=")
|
||
h = proportion_effectsize(p1, p2)
|
||
solver = NormalIndPower()
|
||
base = dict(effect_size=h, alpha=alpha, alternative=alternative,
|
||
ratio=kw.get("ratio", 1.0))
|
||
return solver, base, "nobs1"
|
||
|
||
if test == "one_proportion":
|
||
p1, p0 = kw.get("prop1"), kw.get("prop0")
|
||
if p1 is None or p0 is None:
|
||
raise ValueError("one_proportion requires prop1= and prop0=")
|
||
h = proportion_effectsize(p1, p0)
|
||
# one-sample: NormalIndPower with an effectively infinite second group
|
||
solver = NormalIndPower()
|
||
base = dict(effect_size=h, alpha=alpha, alternative=alternative, ratio=0.0)
|
||
return solver, base, "nobs1"
|
||
|
||
if test == "correlation":
|
||
# Power for Pearson r via Fisher z; handled analytically below, but we
|
||
# still route through a solver-shaped object for a uniform interface.
|
||
return "correlation", dict(effect_size=effect_size, alpha=alpha,
|
||
alternative=alternative), "nobs"
|
||
|
||
if test == "chi2":
|
||
solver = GofChisquarePower()
|
||
dof = kw.get("dof")
|
||
if dof is None:
|
||
raise ValueError("chi2 requires dof= (n_bins-1, or (r-1)(c-1))")
|
||
base = dict(effect_size=effect_size, alpha=alpha, n_bins=dof + 1)
|
||
return solver, base, "nobs"
|
||
|
||
if test == "linear_regression":
|
||
# Handled by a dedicated noncentral-F solver (statsmodels' FTestPower is
|
||
# unreliable when solving for sample size). See _reg_* helpers below.
|
||
df_num = kw.get("df_num")
|
||
if df_num is None:
|
||
raise ValueError("linear_regression requires df_num= (number of tested predictors)")
|
||
return "regression", dict(effect_size=effect_size, alpha=alpha,
|
||
df_num=df_num,
|
||
k_total=kw.get("k_total", df_num)), "nobs"
|
||
|
||
raise ValueError(f"unknown test '{test}'")
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# correlation power (Fisher z transform) -- closed form, no statsmodels solver
|
||
# --------------------------------------------------------------------------- #
|
||
def _corr_power(r, n, alpha, alternative):
|
||
from scipy import stats
|
||
z = math.atanh(r)
|
||
se = 1.0 / math.sqrt(n - 3)
|
||
if alternative == "two-sided":
|
||
zc = stats.norm.ppf(1 - alpha / 2)
|
||
else:
|
||
zc = stats.norm.ppf(1 - alpha)
|
||
return float(stats.norm.cdf(z / se - zc) + stats.norm.cdf(-z / se - zc))
|
||
|
||
|
||
def _corr_sample_size(r, alpha, power, alternative):
|
||
from scipy import stats
|
||
z = abs(math.atanh(r))
|
||
if alternative == "two-sided":
|
||
zc = stats.norm.ppf(1 - alpha / 2)
|
||
else:
|
||
zc = stats.norm.ppf(1 - alpha)
|
||
zp = stats.norm.ppf(power)
|
||
return ((zc + zp) / z) ** 2 + 3
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# multiple-regression power via the noncentral F (Cohen's f^2)
|
||
# --------------------------------------------------------------------------- #
|
||
def _reg_power(f2, n, df_num, k_total, alpha):
|
||
"""Power of the F-test for df_num tested predictors in a model with k_total
|
||
total predictors, total sample size n. Noncentrality lambda = f^2 * n."""
|
||
from scipy import stats
|
||
df_denom = n - k_total - 1
|
||
if df_denom < 1:
|
||
return 0.0
|
||
ncp = f2 * n
|
||
crit = stats.f.ppf(1 - alpha, df_num, df_denom)
|
||
return float(1 - stats.ncf.cdf(crit, df_num, df_denom, ncp))
|
||
|
||
|
||
def _reg_sample_size(f2, df_num, k_total, alpha, power):
|
||
n = k_total + 2 # smallest n with df_denom >= 1
|
||
while _reg_power(f2, n, df_num, k_total, alpha) < power:
|
||
n += 1
|
||
if n > 1_000_000:
|
||
raise RuntimeError("sample size did not converge below 1e6")
|
||
return n
|
||
|
||
|
||
def _reg_mde(n, df_num, k_total, alpha, power):
|
||
"""Smallest detectable f^2 at fixed n (bisection)."""
|
||
lo, hi = 1e-6, 100.0
|
||
for _ in range(200):
|
||
mid = (lo + hi) / 2
|
||
if _reg_power(mid, n, df_num, k_total, alpha) < power:
|
||
lo = mid
|
||
else:
|
||
hi = mid
|
||
return (lo + hi) / 2
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# public API
|
||
# --------------------------------------------------------------------------- #
|
||
def sample_size(test, effect_size=None, alpha=0.05, power=0.80,
|
||
alternative="two-sided", round_up=True, **kw):
|
||
"""Solve for required sample size.
|
||
|
||
For per-group tests (t_ind, two_proportions) returns n PER GROUP.
|
||
For anova returns TOTAL n across all groups. For chi2/regression returns total n.
|
||
"""
|
||
solver, base, n_key = _resolve(test, effect_size, alpha, alternative, **kw)
|
||
|
||
if solver == "correlation":
|
||
n = _corr_sample_size(effect_size, alpha, power, alternative)
|
||
return math.ceil(n) if round_up else n
|
||
|
||
if solver == "regression":
|
||
return _reg_sample_size(base["effect_size"], base["df_num"],
|
||
base["k_total"], alpha, power)
|
||
|
||
base["power"] = power
|
||
base[n_key] = None
|
||
n = solver.solve_power(**base)
|
||
if round_up:
|
||
n = math.ceil(n)
|
||
return n
|
||
|
||
|
||
def power(test, effect_size=None, nobs1=None, nobs=None, alpha=0.05,
|
||
alternative="two-sided", **kw):
|
||
"""Solve for achieved power given the sample size.
|
||
|
||
Pass nobs1 for per-group tests (t_ind, two_proportions), nobs for total-n tests.
|
||
"""
|
||
solver, base, n_key = _resolve(test, effect_size, alpha, alternative, **kw)
|
||
|
||
n = nobs1 if nobs1 is not None else nobs
|
||
if n is None:
|
||
raise ValueError("provide nobs1= (per-group tests) or nobs= (total-n tests)")
|
||
|
||
if solver == "correlation":
|
||
return _corr_power(effect_size, n, alpha, alternative)
|
||
|
||
if solver == "regression":
|
||
return _reg_power(base["effect_size"], n, base["df_num"],
|
||
base["k_total"], alpha)
|
||
|
||
base[n_key] = n
|
||
return float(solver.solve_power(**base))
|
||
|
||
|
||
def mde(test, nobs1=None, nobs=None, alpha=0.05, power=0.80,
|
||
alternative="two-sided", **kw):
|
||
"""Solve for the minimum detectable effect (standardized) at a fixed n.
|
||
|
||
Returns the effect size in the test's native units (d, f, h, w, r, ...).
|
||
"""
|
||
solver, base, n_key = _resolve(test, None, alpha, alternative, **kw)
|
||
|
||
n = nobs1 if nobs1 is not None else nobs
|
||
if n is None:
|
||
raise ValueError("provide nobs1= or nobs=")
|
||
|
||
if solver == "correlation":
|
||
# invert the Fisher-z sample-size formula
|
||
from scipy import stats
|
||
zc = stats.norm.ppf(1 - alpha / 2) if alternative == "two-sided" \
|
||
else stats.norm.ppf(1 - alpha)
|
||
zp = stats.norm.ppf(power)
|
||
z = (zc + zp) / math.sqrt(n - 3)
|
||
return float(math.tanh(z))
|
||
|
||
if solver == "regression":
|
||
return _reg_mde(n, base["df_num"], base["k_total"], alpha, power)
|
||
|
||
base["power"] = power
|
||
base["effect_size"] = None
|
||
base[n_key] = n
|
||
return float(solver.solve_power(**base))
|
||
|
||
|
||
def power_curve(test, effect_size=None, n_range=None, alpha=0.05, power_target=0.80,
|
||
alternative="two-sided", save=None, show=False, **kw):
|
||
"""Plot power vs. sample size. Returns (n_array, power_array).
|
||
|
||
n_range iterates the per-group n for per-group tests, total n otherwise.
|
||
"""
|
||
import matplotlib.pyplot as plt
|
||
|
||
if n_range is None:
|
||
n_range = range(5, 205, 5)
|
||
ns = np.array(list(n_range), dtype=float)
|
||
|
||
pwr = []
|
||
for n in ns:
|
||
if test in ("t_ind", "two_proportions"):
|
||
pwr.append(power(test, effect_size=effect_size, nobs1=n, alpha=alpha,
|
||
alternative=alternative, **kw))
|
||
else:
|
||
pwr.append(power(test, effect_size=effect_size, nobs=n, alpha=alpha,
|
||
alternative=alternative, **kw))
|
||
pwr = np.array(pwr)
|
||
|
||
fig, ax = plt.subplots(figsize=(7, 4.5))
|
||
ax.plot(ns, pwr, lw=2)
|
||
ax.axhline(power_target, ls="--", color="crimson", lw=1,
|
||
label=f"target power = {power_target:g}")
|
||
ax.set_xlabel("Sample size" + (" per group" if test in ("t_ind", "two_proportions") else " (total)"))
|
||
ax.set_ylabel("Power (1 - β)")
|
||
ax.set_ylim(0, 1.02)
|
||
es_label = effect_size if effect_size is not None else ""
|
||
ax.set_title(f"Power curve: {test} (effect = {es_label}, α = {alpha:g})")
|
||
ax.legend()
|
||
ax.grid(alpha=0.3)
|
||
fig.tight_layout()
|
||
|
||
if save:
|
||
fig.savefig(save, dpi=150)
|
||
if show:
|
||
plt.show()
|
||
plt.close(fig)
|
||
return ns, pwr
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# quick smoke test of the main paths
|
||
print("t_ind d=0.5, 80% power -> n/group =",
|
||
sample_size("t_ind", effect_size=0.5, power=0.80))
|
||
print("anova f=0.25, k=4, 80% -> total n =",
|
||
sample_size("anova", effect_size=0.25, k_groups=4, power=0.80))
|
||
print("2 props 0.40 vs 0.55, 80% -> n/group =",
|
||
sample_size("two_proportions", prop1=0.40, prop2=0.55, power=0.80))
|
||
print("correlation r=0.30, 80% -> n =",
|
||
sample_size("correlation", effect_size=0.30, power=0.80))
|
||
print("MDE for t_ind at n=30/group, 80% power -> d =",
|
||
round(mde("t_ind", nobs1=30, power=0.80), 3))
|
||
print("power for t_ind d=0.5 at n=64/group ->",
|
||
round(power("t_ind", effect_size=0.5, nobs1=64), 3))
|