Source code for pystatsbio.power._means

"""Power calculations for t-tests (two-sample, one-sample, paired).

Validates against: R pwr::pwr.t.test()
"""

from __future__ import annotations

import math

from scipy.stats import nct, norm
from scipy.stats import t as t_dist

from pystatsbio.power._common import PowerResult, _check_power_args, _solve_parameter

_VALID_TYPES = ("two.sample", "one.sample", "paired")
_VALID_ALTERNATIVES = ("two.sided", "less", "greater")


# ---------------------------------------------------------------------------
# Internal power computation — also used by _cluster.py
# ---------------------------------------------------------------------------

def _normal_approx_power(
    ncp: float,
    alpha: float,
    alternative: str,
) -> float:
    """Normal approximation to noncentral t power (exact as df -> inf)."""
    if alternative == "two.sided":
        z_crit = norm.ppf(1.0 - alpha / 2.0)
        return float(norm.sf(z_crit - ncp) + norm.cdf(-z_crit - ncp))
    elif alternative == "greater":
        z_crit = norm.ppf(1.0 - alpha)
        return float(norm.sf(z_crit - ncp))
    else:  # less
        z_crit = norm.ppf(alpha)
        return float(norm.cdf(z_crit - ncp))


def _t_test_power(
    n: float,
    d: float,
    alpha: float,
    alternative: str,
    type: str,
) -> float:
    """Compute t-test power for given n, d, alpha.

    Parameters
    ----------
    n : float
        Sample size (per group for two-sample, total for one-sample/paired).
        May be non-integer during root-finding.
    d : float
        Cohen's d (positive = treatment > control).
    alpha : float
        Significance level.
    alternative : str
        'two.sided', 'less', or 'greater'.
    type : str
        'two.sample', 'one.sample', or 'paired'.

    Returns
    -------
    float
        Statistical power in [0, 1].
    """
    # Noncentrality parameter and degrees of freedom
    if type == "two.sample":
        ncp = d * math.sqrt(n / 2.0)
        df = 2.0 * n - 2.0
    else:  # one.sample or paired
        ncp = d * math.sqrt(n)
        df = n - 1.0

    if df < 1.0:
        return 0.0

    # For very large df, go straight to normal approximation (exact in limit).
    if df > 1e5:
        return _normal_approx_power(ncp, alpha, alternative)

    # Power via noncentral t distribution
    if alternative == "two.sided":
        t_crit = t_dist.ppf(1.0 - alpha / 2.0, df)
        # P(|T| > t_crit) under H1
        pwr = float(nct.sf(t_crit, df, ncp) + nct.cdf(-t_crit, df, ncp))
    elif alternative == "greater":
        t_crit = t_dist.ppf(1.0 - alpha, df)
        pwr = float(nct.sf(t_crit, df, ncp))
    else:  # less
        t_crit = t_dist.ppf(alpha, df)
        pwr = float(nct.cdf(t_crit, df, ncp))

    # scipy's nct can return NaN for moderate-to-large noncentrality params
    # (known issue with ncp ~10+ even at moderate df ~1000).  Fall back to
    # the normal approximation which is very accurate for df > ~30.
    if math.isnan(pwr):
        pwr = _normal_approx_power(ncp, alpha, alternative)

    return pwr


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

[docs] def power_t_test( n: int | None = None, d: float | None = None, alpha: float = 0.05, power: float | None = None, alternative: str = "two.sided", type: str = "two.sample", ) -> PowerResult: """Power calculation for t-tests. Exactly one of ``n``, ``d``, ``power`` must be ``None`` — that parameter is solved for given the others. Parameters ---------- n : int or None Sample size per group (two-sample) or total (one-sample/paired). d : float or None Cohen's d effect size. alpha : float Significance level (default 0.05). power : float or None Desired statistical power (1 - beta). alternative : str ``'two.sided'``, ``'less'``, or ``'greater'``. type : str ``'two.sample'``, ``'one.sample'``, or ``'paired'``. Returns ------- PowerResult Examples -------- >>> r = power_t_test(d=0.5, alpha=0.05, power=0.80) >>> r.n 64 >>> r = power_t_test(n=50, d=0.5, alpha=0.05) >>> round(r.power, 3) 0.697 Validates against: R pwr::pwr.t.test() """ # --- Validate --- if alternative not in _VALID_ALTERNATIVES: raise ValueError( f"alternative must be one of {_VALID_ALTERNATIVES}, got {alternative!r}" ) if type not in _VALID_TYPES: raise ValueError(f"type must be one of {_VALID_TYPES}, got {type!r}") solve_for = _check_power_args(n=n, effect=d, power=power, alpha=alpha, effect_name="d") # For two-sided tests, R's pwr works with |d| d_internal = d if d is not None and alternative == "two.sided": d_internal = abs(d) # --- Solve --- if solve_for == "power": assert n is not None and d_internal is not None result_power = _t_test_power(float(n), d_internal, alpha, alternative, type) result_n = n result_d = d elif solve_for == "n": assert d_internal is not None and power is not None if d_internal == 0.0: raise ValueError("Cannot solve for n when d = 0 (no effect)") raw_n = _solve_parameter( func=lambda x: _t_test_power(x, d_internal, alpha, alternative, type), target=power, bracket=(2.0, 1e7), ) result_n = math.ceil(raw_n) result_power = power result_d = d else: # solve_for == "effect" assert n is not None and power is not None if alternative == "two.sided": # Solve for positive d (symmetric) result_d = _solve_parameter( func=lambda x: _t_test_power(float(n), x, alpha, alternative, type), target=power, bracket=(1e-10, 100.0), ) elif alternative == "greater": result_d = _solve_parameter( func=lambda x: _t_test_power(float(n), x, alpha, alternative, type), target=power, bracket=(1e-10, 100.0), ) else: # less result_d = _solve_parameter( func=lambda x: _t_test_power(float(n), x, alpha, alternative, type), target=power, bracket=(-100.0, -1e-10), ) result_n = n result_power = power # --- Method label --- type_labels = { "two.sample": "Two-sample", "one.sample": "One-sample", "paired": "Paired", } method = f"{type_labels[type]} t test power calculation" note = "n is number in *each* group" if type == "two.sample" else "" return PowerResult( n=result_n, power=result_power, effect_size=result_d, alpha=alpha, alternative=alternative, method=method, note=note, )
[docs] def power_paired_t_test( n: int | None = None, d: float | None = None, alpha: float = 0.05, power: float | None = None, alternative: str = "two.sided", ) -> PowerResult: """Convenience wrapper: ``power_t_test`` with ``type='paired'``. Validates against: R pwr::pwr.t.test(type='paired') """ return power_t_test( n=n, d=d, alpha=alpha, power=power, alternative=alternative, type="paired", )