Skills › Scientific Agent Skills › statistical-power › simulate_power.py
simulate_power.py
314 lines · skills/statistical-power/scripts/simulate_power.py at 5f2bc03
1"""Monte Carlo power when analytical assumptions do not match the design.
3Simulation can represent features omitted by simple analytical approximations
4for logistic/Poisson regression, mixed models, clustering, survival, mediation,
5or interactions:
7 1. simulate a dataset under the assumed truth
8 2. analyze it with the EXACT test/model you will use on real data
9 3. repeat many times; power = fraction of replicates that reach significance
11This module provides the harness (`simulate_power`, `find_sample_size`) plus four
12worked, runnable examples. Copy an example and swap in your own data-generating
13process and analysis -- that is the intended workflow.
15Requires: numpy, scipy, pandas, statsmodels. Survival extensions can use lifelines.
16"""
18from __future__ import annotations
20import math
21from numbers import Integral
22from dataclasses import dataclass, field
23import warnings
25import numpy as np
28@dataclass
29class PowerEstimate:
30 power: float
31 n_sims: int
32 n: int
33 ci_low: float
34 ci_high: float
35 n_failures: int = 0
36 failure_reasons: dict[str, int] = field(default_factory=dict)
37 n_warned: int = 0
38 warning_counts: dict[str, int] = field(default_factory=dict)
40 def __str__(self):
41 return (f"n={self.n}: power={self.power:.3f} "
42 f"(95% MC CI {self.ci_low:.3f}-{self.ci_high:.3f}, {self.n_sims} sims; "
43 f"{self.n_failures} failed, {self.n_warned} warned)")
46class SimulationFitError(RuntimeError):
47 """An anticipated fitting failure, counted separately and as non-rejection."""
50def _fit_decision(fit, coefficient, alpha):
51 """Require convergence and a finite p-value from a statsmodels fit."""
52 from statsmodels.tools.sm_exceptions import (
53 HessianInversionWarning, PerfectSeparationError, PerfectSeparationWarning,
54 )
55 try:
56 with warnings.catch_warnings():
57 warnings.simplefilter("error", PerfectSeparationWarning)
58 warnings.simplefilter("error", HessianInversionWarning)
59 result = fit()
60 except (np.linalg.LinAlgError, PerfectSeparationError, PerfectSeparationWarning,
61 HessianInversionWarning, FloatingPointError) as exc:
62 raise SimulationFitError(type(exc).__name__) from exc
63 converged = (result.converged if hasattr(result, "converged")
64 else result.mle_retvals.get("converged", False))
65 if not converged:
66 raise SimulationFitError("nonconvergence")
67 pvalue = result.pvalues[coefficient]
68 if not np.isfinite(pvalue) or not 0 <= pvalue <= 1:
69 raise SimulationFitError("invalid p-value")
70 return bool(pvalue < alpha)
73def _check_example(alpha, **positive):
74 if not np.isfinite(alpha) or not 0 < alpha < 1:
75 raise ValueError("alpha must lie strictly between 0 and 1")
76 for name, value in positive.items():
77 if not np.isfinite(value) or value <= 0:
78 raise ValueError(f"{name} must be finite and positive")
81def _wilson_ci(k, n, z=1.96):
82 """Wilson interval for a simulated rate conditional on fixed assumptions."""
83 if n == 0:
84 return (0.0, 1.0)
85 p = k / n
86 denom = 1 + z**2 / n
87 center = (p + z**2 / (2 * n)) / denom
88 half = z * math.sqrt(p * (1 - p) / n + z**2 / (4 * n**2)) / denom
89 return (max(0.0, center - half), min(1.0, center + half))
92def simulate_power(gen_and_test, n, n_sims=2000, alpha=0.05, seed=0):
93 """Estimate power at sample size `n`.
95 gen_and_test(n, rng) -> bool : build a dataset of size n under the assumed
96 effect, run the planned analysis, and return True iff it is significant.
97 It receives a numpy Generator `rng` so results are reproducible and each
98 replicate consumes fresh pseudorandom draws. Alpha must be used inside
99 the callback; the harness alpha argument is validated but not forwarded.
100 Raise SimulationFitError for an expected fit failure: it counts as a
101 non-rejection in the full denominator and is reported separately.
102 Other exceptions propagate so programming errors are not hidden.
104 Returns a PowerEstimate including a Wilson confidence interval, so you can
105 tell whether 0.81 vs 0.79 is real or just simulation noise.
106 """
107 if any(isinstance(value, bool) or not isinstance(value, Integral) or value < 1
108 for value in (n, n_sims)):
109 raise ValueError("n and n_sims must be positive integers")
110 if not 0 < alpha < 1:
111 raise ValueError("alpha must lie strictly between 0 and 1")
112 rng = np.random.default_rng(seed)
113 hits = 0
114 failures = {}
115 warning_counts = {}
116 n_warned = 0
117 for _ in range(n_sims):
118 with warnings.catch_warnings(record=True) as caught:
119 warnings.simplefilter("always")
120 try:
121 rejected = gen_and_test(n, rng)
122 except SimulationFitError as exc:
123 reason = str(exc) or "unspecified fit failure"
124 failures[reason] = failures.get(reason, 0) + 1
125 rejected = False
126 if caught:
127 n_warned += 1
128 for warning in caught:
129 name = warning.category.__name__
130 warning_counts[name] = warning_counts.get(name, 0) + 1
131 if not isinstance(rejected, (bool, np.bool_)):
132 raise TypeError("gen_and_test must return a boolean rejection decision, not a p-value")
133 if rejected:
134 hits += 1
135 lo, hi = _wilson_ci(hits, n_sims)
136 return PowerEstimate(power=hits / n_sims, n_sims=n_sims, n=n,
137 ci_low=lo, ci_high=hi, n_failures=sum(failures.values()),
138 failure_reasons=failures, n_warned=n_warned,
139 warning_counts=warning_counts)
142def find_sample_size(gen_and_test, target_power=0.80, n_sims=2000, alpha=0.05,
143 lo=10, hi=2000, seed=0, verbose=True):
144 """Candidate n reaching estimated target power, via noisy bisection over n.
146 Requires a design whose power increases with n; discrete designs, rounding,
147 and fit failures can violate this. A fixed seed makes the search reproducible
148 but not monotone or noise-free. Validate the candidate and neighboring n with
149 more replicates and a fresh seed. Returns (n, PowerEstimate); the reported CI
150 at an adaptively selected n is not a selection-adjusted assurance interval.
151 """
152 if any(isinstance(value, bool) or not isinstance(value, Integral)
153 for value in (lo, hi)) or not 1 <= lo <= hi <= 1_000_000:
154 raise ValueError("bounds must be integers with 1 <= lo <= hi <= 1000000")
155 if not 0 < target_power <= 1:
156 raise ValueError("target_power must lie in (0, 1]")
157 # expand hi until it clears target (guards against too-small upper bound)
158 while True:
159 est_hi = simulate_power(gen_and_test, hi, n_sims, alpha, seed)
160 if est_hi.power >= target_power:
161 break
162 if hi >= 1_000_000:
163 raise ValueError("target power not reached by the sample-size cap of 1000000")
164 lo, hi = hi, min(hi * 2, 1_000_000)
166 best = est_hi
167 while lo < hi:
168 mid = (lo + hi) // 2
169 est = simulate_power(gen_and_test, mid, n_sims, alpha, seed)
170 if verbose:
171 print(est)
172 if est.power >= target_power:
173 hi, best = mid, est
174 else:
175 lo = mid + 1
176 return hi, best
179# ========================================================================== #
180# Worked examples -- run this file directly to see them.
181# Each is a `gen_and_test(n, rng)` you can adapt.
182# ========================================================================== #
184def example_two_group_difference(effect=0.5, sd=1.0, alpha=0.05):
185 """Two-group difference in means (sanity check vs. the closed-form t-test)."""
186 from scipy import stats
187 _check_example(alpha, sd=sd)
188 if not np.isfinite(effect):
189 raise ValueError("effect must be finite")
191 def gen_and_test(n, rng): # n per group
192 if n < 2:
193 raise ValueError("two-group t-test needs at least two observations per group")
194 a = rng.normal(0.0, sd, n)
195 b = rng.normal(effect, sd, n)
196 _, p = stats.ttest_ind(a, b, equal_var=True, alternative="two-sided")
197 return p < alpha
199 return gen_and_test
202def example_logistic_regression(beta=0.8, base_rate=0.2, x_sd=1.0, alpha=0.05):
203 """Power for a single coefficient in logistic regression.
205 beta is the log-odds change per one raw unit of x (SD is x_sd).
206 base_rate is P(y=1 | x=0), not the marginal event rate.
207 """
208 import statsmodels.api as sm
209 from scipy.special import expit
211 _check_example(alpha, x_sd=x_sd)
212 if not np.isfinite(beta) or not np.isfinite(base_rate) or not 0 < base_rate < 1:
213 raise ValueError("beta must be finite and base_rate must lie in (0, 1)")
215 intercept = math.log(base_rate / (1 - base_rate))
217 def gen_and_test(n, rng):
218 if n < 3:
219 raise ValueError("logistic example requires at least three observations")
220 x = rng.normal(0, x_sd, n)
221 logit = intercept + beta * x
222 p = expit(logit)
223 y = rng.binomial(1, p)
224 X = sm.add_constant(x)
225 return _fit_decision(lambda: sm.Logit(y, X).fit(disp=0), 1, alpha)
227 return gen_and_test
230def example_cluster_randomized(effect=0.3, icc=0.05, cluster_size=20,
231 resid_sd=1.0, alpha=0.05):
232 """Cluster-randomized trial: clusters (not individuals) are randomized.
234 Ignoring the clustering (analyzing individuals as independent) would badly
235 overstate power -- that is pseudoreplication. Here `n` is the number of
236 clusters PER ARM; the analysis is a mixed model with a random cluster intercept.
237 """
238 import statsmodels.formula.api as smf
239 import pandas as pd
240 _check_example(alpha, resid_sd=resid_sd)
241 if not np.isfinite(effect) or not np.isfinite(icc) or not 0 <= icc < 1:
242 raise ValueError("effect must be finite and icc must lie in [0, 1)")
243 if isinstance(cluster_size, bool) or not isinstance(cluster_size, Integral) or cluster_size < 2:
244 raise ValueError("cluster_size must be an integer >= 2")
246 # split total variance into between-cluster (tau^2) and residual by the ICC
247 tau = math.sqrt(icc * resid_sd**2 / (1 - icc)) if icc > 0 else 0.0
249 def gen_and_test(n, rng): # n clusters per arm
250 if n < 2:
251 raise ValueError("cluster example requires at least two clusters per arm")
252 rows = []
253 cid = 0
254 for arm in (0, 1):
255 for _ in range(n):
256 u = rng.normal(0, tau) # cluster random effect
257 for _ in range(cluster_size):
258 y = effect * arm + u + rng.normal(0, resid_sd)
259 rows.append((y, arm, cid))
260 cid += 1
261 df = pd.DataFrame(rows, columns=["y", "arm", "cluster"])
262 return _fit_decision(
263 lambda: smf.mixedlm("y ~ arm", df, groups=df["cluster"]).fit(reml=True),
264 "arm", alpha)
266 return gen_and_test
269def example_linear_mixed_repeated(effect=0.4, n_timepoints=3, subj_sd=0.7,
270 resid_sd=1.0, alpha=0.05):
271 """Repeated-measures: a linear time trend within subjects, random intercepts.
273 `n` is the number of subjects; each is measured at n_timepoints occasions.
274 `effect` is the slope per time unit. Tests whether the time slope != 0.
275 """
276 import statsmodels.formula.api as smf
277 import pandas as pd
278 _check_example(alpha, resid_sd=resid_sd)
279 if not np.isfinite(effect) or not np.isfinite(subj_sd) or subj_sd < 0:
280 raise ValueError("effect must be finite and subj_sd finite and nonnegative")
281 if isinstance(n_timepoints, bool) or not isinstance(n_timepoints, Integral) or n_timepoints < 2:
282 raise ValueError("n_timepoints must be an integer >= 2")
284 def gen_and_test(n, rng): # n subjects
285 if n < 2:
286 raise ValueError("repeated-measures example requires at least two subjects")
287 rows = []
288 for s in range(n):
289 b0 = rng.normal(0, subj_sd)
290 for t in range(n_timepoints):
291 y = b0 + effect * t + rng.normal(0, resid_sd)
292 rows.append((y, t, s))
293 df = pd.DataFrame(rows, columns=["y", "time", "subj"])
294 return _fit_decision(
295 lambda: smf.mixedlm("y ~ time", df, groups=df["subj"]).fit(reml=True),
296 "time", alpha)
298 return gen_and_test
301if __name__ == "__main__":
302 print("== two-group difference (compare to closed-form t-test n/group=64) ==")
303 g = example_two_group_difference(effect=0.5)
304 print(simulate_power(g, n=64, n_sims=2000))
306 print("\n== logistic regression, beta=0.8 ==")
307 g = example_logistic_regression(beta=0.8, base_rate=0.2)
308 print(simulate_power(g, n=150, n_sims=1000))
310 print("\n== search: subjects needed for repeated-measures slope=0.4 ==")
311 g = example_linear_mixed_repeated(effect=0.4, n_timepoints=3)
312 n, est = find_sample_size(g, target_power=0.80, n_sims=500, lo=10, hi=60,
313 verbose=False)
314 print(f"-> candidate n={n} subjects; confirm with fresh simulations ({est})")