# Sample size for a four-group dose study (vehicle and three doses), from the planned analysis.
# Expected means 100, 95, 90, 85 (vehicle, low, mid, high); SD 12 within groups; two-sided alpha 0.05; power 80%.
# Power from the noncentral t and F distributions. Python 3.13, NumPy, SciPy.
import numpy as np
from scipy import stats

mu = np.array([100, 95, 90, 85]); sd = 12; k = len(mu); alpha = 0.05; target = 0.80

def power_t(diff, n, a):                               # two-group comparison inside the ANOVA: error df = k*n - k
    df = k * n - k; ncp = diff / (sd * np.sqrt(2 / n)); tc = stats.t.ppf(1 - a / 2, df)
    return stats.nct.sf(tc, df, ncp) + stats.nct.cdf(-tc, df, ncp)

def power_f(n):                                        # omnibus one-way ANOVA F-test
    f = np.sqrt(((mu - mu.mean()) ** 2).mean()) / sd
    df1, df2 = k - 1, k * n - k
    return stats.ncf.sf(stats.f.ppf(1 - alpha, df1, df2), df1, df2, f ** 2 * k * n)

def power_trend(n, c=np.array([-3, -1, 1, 3])):        # linear trend contrast
    df = k * n - k; ncp = (c * mu).sum() / (sd * np.sqrt((c ** 2).sum() / n)); tc = stats.t.ppf(1 - alpha / 2, df)
    return stats.nct.sf(tc, df, ncp) + stats.nct.cdf(-tc, df, ncp)

def smallest_n(fun):
    return next(n for n in range(2, 500) if fun(n) >= target)

f = np.sqrt(((mu - mu.mean()) ** 2).mean()) / sd
print(f"Cohen's f from the expected means: {f:.3f}")
n = smallest_n(power_f); print(f"Omnibus ANOVA F-test: n = {n} per group (power {power_f(n):.3f})")
n = smallest_n(power_trend); print(f"Linear trend contrast: n = {n} per group (power {power_trend(n):.3f})")
for diff in (15, 10, 5):
    for a, lab in ((alpha, "alpha 0.05"), (alpha / 3, "alpha 0.05/3 (Bonferroni for 3 comparisons)")):
        n = smallest_n(lambda n: power_t(diff, n, a))
        print(f"Dose vs vehicle, difference {diff}, {lab}: n = {n} per group (power {power_t(diff, n, a):.3f})")
ns = [n for n in range(2, 50) if 10 <= k * n - k <= 20]
print(f"Resource equation (error df 10 to 20): n = {ns[0]} to {ns[-1]} per group")
for n in ns:
    print(f"  n = {n}: power for difference 10 at alpha 0.05/3 = {power_t(10, n, alpha / 3):.3f}; for difference 15 = {power_t(15, n, alpha / 3):.3f}")
