diff --git a/src/pystatpower/procedures/two_proportion.py b/src/pystatpower/procedures/two_proportion.py index 95aacde..61ce69f 100644 --- a/src/pystatpower/procedures/two_proportion.py +++ b/src/pystatpower/procedures/two_proportion.py @@ -1,6 +1,10 @@ """两独立样本差异性检验""" from enum import Enum +from math import sqrt + +from scipy.stats import norm +from scipy.optimize import brenth class SolvableParameter(Enum): @@ -41,6 +45,55 @@ class GroupAllocation(Enum): PERCENT_OF_REFERENCE = 7 +def _power( + alpha: float, + treatment_n: float, + reference_n: float, + treatment_proportion: float, + reference_proportion: float, + alternative: Alternative, + test_type: TestType, +): + n1 = treatment_n + n2 = reference_n + p1 = treatment_proportion + p2 = reference_proportion + + # 计算标准误 + match test_type: + case TestType.Z_TEST_POOLED | TestType.Z_TEST_CC_POOLED: + p_hat = (n1 * p1 + n2 * p2) / (n1 + n2) + se = sqrt(p_hat * (1 - p_hat) * (1 / n1 + 1 / n2)) + case TestType.Z_TEST_UNPOOLED | TestType.Z_TEST_CC_UNPOOLED: + se = sqrt(p1 * (1 - p1) / n1 + p2 * (1 - p2) / n2) + case _: + assert False, "未知的检验类型" + + # 连续性校正 + c = 0 + if test_type in [TestType.Z_TEST_CC_POOLED, TestType.Z_TEST_CC_UNPOOLED]: + c = (1 / 2) * (1 / n1 + 1 / n2) + + # 计算检验效能 + match alternative: + case Alternative.TWO_SIDED: + z_alpha = norm.ppf(1 - alpha / 2) + z_stat = [(p1 - p2 - c) / se, (p1 - p2 + c) / se] + power = norm.cdf(-z_alpha - z_stat[0]) + 1 - norm.cdf(z_alpha - z_stat[1]) + case Alternative.ONE_SIDED: + z_alpha = norm.ppf(1 - alpha) + if p1 > p2: + z_stat = (p1 - p2 + c) / se + power = 1 - norm.cdf(z_alpha - z_stat) + elif p1 <= p2: + z_stat = (p1 - p2 - c) / se + power = norm.cdf(-z_alpha - z_stat) + case _: + assert False, "未知的备择假设类型" + + return power + + class TwoProportionDesigner: def __init__(self, solve_for: SolvableParameter): if not isinstance(solve_for, SolvableParameter): @@ -59,10 +112,10 @@ def __init__(self, solve_for: SolvableParameter): return TwoProportionSolveForReferenceProportionDesigner() -class TwoProportionSolveForNDesigner(TwoProportionDesigner): +class TwoProportionSolveForNDesigner: def __init__(self): - self._config = {} + self._config: dict[str, any] = {} def set_alpha(self, alpha: float = 0.05): self._config["alpha"] = alpha