From fed10ef15c5e91ebba73e5e047b26b5453b3a7aa Mon Sep 17 00:00:00 2001 From: Kun Jinkao <45487685+Snoopy1866@users.noreply.github.com> Date: Mon, 16 Sep 2024 23:08:05 +0800 Subject: [PATCH] update --- main.py | 65 +-- src/pystatpower/procedures/two_proportion.py | 421 ++++++++----------- 2 files changed, 190 insertions(+), 296 deletions(-) diff --git a/main.py b/main.py index 5fc876a..5012505 100644 --- a/main.py +++ b/main.py @@ -1,60 +1,17 @@ +import pprint from pystatpower.procedures.two_proportion import * -result = ( - TwoProportion(EnumSolvableParameter.N) - .set_alpha(0.05) - .set_power(0.8) - .set_alternative(EnumAlternative.TWO_SIDED) - .set_group_allocation(GroupAllocation(EnumGroupAllocation.EQUAL)) - .set_test_type(EnumTestType.Z_TEST_POOLED) - .set_treatment_proportion(0.95) - .set_reference_proportion(0.80) - .get_solver() - .solve() -) - -print(result) - - -result = ( - TwoProportion(EnumSolvableParameter.N) - .set_alpha(0.05) - .set_power(0.8) - .set_alternative(EnumAlternative.TWO_SIDED) - .set_group_allocation(GroupAllocation(EnumGroupAllocation.FIX_TREATMENT_GROUP).set_treatment_n(100)) - .set_test_type(EnumTestType.Z_TEST_POOLED) - .set_treatment_proportion(0.95) - .set_reference_proportion(0.80) - .get_solver() - .solve() -) - -print(result) - - -result = ( - TwoProportion(EnumSolvableParameter.N) - .set_alpha(0.05) - .set_power(0.8) - .set_alternative(EnumAlternative.TWO_SIDED) - .set_group_allocation(GroupAllocation(EnumGroupAllocation.EQUAL)) - .set_test_type(EnumTestType.Z_TEST_POOLED) - .set_treatment_proportion(0.68) - .set_reference_proportion(0.69) - .get_solver() - .solve() -) - -print(result) - - -result = fun_power( +model = TwoProportionSolveForSize( alpha=0.05, - treatment_n=10000, - reference_n=10000, - treatment_proportion=0.68, - reference_proportion=0.69, + power=0.80, + treatment_proportion=0.95, + reference_proportion=0.80, alternative=EnumAlternative.TWO_SIDED, test_type=EnumTestType.Z_TEST_POOLED, + group_allocation=TwoProportionSolveForSizeGroupAllocator(), ) -print(result) + +model.solve() + + +pprint.pprint(model) diff --git a/src/pystatpower/procedures/two_proportion.py b/src/pystatpower/procedures/two_proportion.py index f5ef0b7..03db2d9 100644 --- a/src/pystatpower/procedures/two_proportion.py +++ b/src/pystatpower/procedures/two_proportion.py @@ -1,5 +1,7 @@ """两独立样本差异性检验""" +from abc import abstractmethod +from dataclasses import dataclass from enum import Enum from math import sqrt @@ -7,7 +9,7 @@ from scipy.optimize import brenth -class EnumSolvableParameter(Enum): +class EnumSolveForParameter(Enum): """求解目标""" N = 1 @@ -33,18 +35,6 @@ class EnumTestType(Enum): Z_TEST_CC_UNPOOLED = 4 -class EnumGroupAllocation(Enum): - """样本量分配方式""" - - EQUAL = 1 - FIX_TREATMENT_GROUP = 2 - FIX_REFERENCE_GROUP = 3 - RATIO_OF_TREATMENT_TO_REFERENCE = 4 - RATIO_OF_REFERENCE_TO_TREATMENT = 5 - PERCENT_OF_TREATMENT = 6 - PERCENT_OF_REFERENCE = 7 - - def fun_power( alpha: float, treatment_n: float, @@ -94,267 +84,214 @@ def fun_power( return power -class GroupAllocation: - def __init__(self, group_allocation_option: EnumGroupAllocation): - self._group_allocation_option = group_allocation_option +class TwoProportion: - def set_treatment_n(self, treatment_n: float): - if self._group_allocation_option != EnumGroupAllocation.FIX_TREATMENT_GROUP: - raise ValueError("treatment_n 只能在 group_allocation 为 FIX_TREATMENT_GROUP 时指定") - self._treatment_n = treatment_n - return self + @abstractmethod + def solve(self): + pass - def set_reference_n(self, reference_n: float): - if self._group_allocation_option != EnumGroupAllocation.FIX_REFERENCE_GROUP: - raise ValueError("reference_n 只能在 group_allocation 为 FIX_REFERENCE_GROUP 时指定") - self._reference_n = reference_n - return self - def set_ratio_of_treatment_to_reference(self, ratio_of_treatment_to_reference: float): - if self._group_allocation_option != EnumGroupAllocation.RATIO_OF_TREATMENT_TO_REFERENCE: - raise ValueError("ratio 只能在 group_allocation 为 RATIO_OF_TREATMENT_TO_REFERENCE 时指定") - self._ratio_of_treatment_to_reference = ratio_of_treatment_to_reference - return self +# solve for sample size +class EnumTwoProportionSolveForSizeGroupAllocation(Enum): + EQUAL = 1 + SIZE_OF_TREATMENT = 2 + SIZE_OF_REFERENCE = 3 + RATIO_OF_TREATMENT_TO_REFERENCE = 4 + RATIO_OF_REFERENCE_TO_TREATMENT = 5 + PERCENT_OF_TREATMENT = 6 + PERCENT_OF_REFERENCE = 7 - def set_ratio_of_reference_to_treatment(self, ratio_of_reference_to_treatment: float): - if self._group_allocation_option != EnumGroupAllocation.RATIO_OF_REFERENCE_TO_TREATMENT: - raise ValueError("ratio 只能在 group_allocation 为 RATIO_OF_REFERENCE_TO_TREATMENT 时指定") - self._ratio_of_reference_to_treatment = ratio_of_reference_to_treatment - return self - def set_percent_of_treatment(self, percent_of_treatment: float): - if self._group_allocation_option != EnumGroupAllocation.PERCENT_OF_TREATMENT: - raise ValueError("percent_of_treatment 只能在 group_allocation 为 PERCENT_OF_TREATMENT 时指定") - self._percent_of_treatment = percent_of_treatment - return self +class TwoProportionSolveForSizeGroupAllocator: + def __init__( + self, + group_allocation_option: EnumTwoProportionSolveForSizeGroupAllocation = EnumTwoProportionSolveForSizeGroupAllocation.EQUAL, + treatment_size: float = None, + reference_size: float = None, + ratio_of_treatment_to_reference: float = None, + ratio_of_reference_to_treatment: float = None, + percent_of_treatment: float = None, + percent_of_reference: float = None, + ): + match group_allocation_option: + case EnumTwoProportionSolveForSizeGroupAllocation.EQUAL: + self.treatment_size_formula = lambda n: n + self.reference_size_formula = lambda n: n + case EnumTwoProportionSolveForSizeGroupAllocation.SIZE_OF_TREATMENT: + self.treatment_size_formula = lambda n: treatment_size + self.reference_size_formula = lambda n: n + case EnumTwoProportionSolveForSizeGroupAllocation.SIZE_OF_REFERENCE: + self.treatment_size_formula = lambda n: n + self.reference_size_formula = lambda n: reference_size + case EnumTwoProportionSolveForSizeGroupAllocation.RATIO_OF_TREATMENT_TO_REFERENCE: + self.treatment_size_formula = lambda n: ratio_of_treatment_to_reference * n + self.reference_size_formula = lambda n: n + case EnumTwoProportionSolveForSizeGroupAllocation.RATIO_OF_REFERENCE_TO_TREATMENT: + self.treatment_size_formula = lambda n: n + self.reference_size_formula = lambda n: ratio_of_reference_to_treatment * n + case EnumTwoProportionSolveForSizeGroupAllocation.PERCENT_OF_TREATMENT: + self.treatment_size_formula = lambda n: n + self.reference_size_formula = lambda n: (1 - percent_of_treatment) / percent_of_treatment * n + case EnumTwoProportionSolveForSizeGroupAllocation.PERCENT_OF_REFERENCE: + self.treatment_size_formula = lambda n: (1 - percent_of_reference) / percent_of_reference * n + self.reference_size_formula = lambda n: n + + +class TwoProportionSolveForSize(TwoProportion): - def set_percent_of_reference(self, percent_of_reference: float): - if self._group_allocation_option != EnumGroupAllocation.PERCENT_OF_REFERENCE: - raise ValueError("percent_of_reference 只能在 group_allocation 为 PERCENT_OF_REFERENCE 时指定") - self._percent_of_reference = percent_of_reference - return self + def __init__( + self, + alpha: float = 0.05, + power: float = 0.80, + alternative: EnumAlternative = EnumAlternative.TWO_SIDED, + test_type: EnumTestType = EnumTestType.Z_TEST_POOLED, + treatment_proportion: float = 0.95, + reference_proportion: float = 0.80, + group_allocation: TwoProportionSolveForSizeGroupAllocator = TwoProportionSolveForSizeGroupAllocator(), + ): + self._alpha = alpha + self._power = power + self._alternative = alternative + self._test_type = test_type + self._treatment_proportion = treatment_proportion + self._reference_proportion = reference_proportion + self._group_allocation = group_allocation + @property + def alpha(self) -> float: + return self._alpha -class TwoProportionSolveForNDesigner: + @alpha.setter + def alpha(self, alpha: float): + self._alpha = alpha + + @property + def power(self) -> float: + return self._power - def __init__(self): - self._config: dict[str, any] = {} + @power.setter + def power(self, power: float): + self._power = power - def set_alpha(self, alpha: float = 0.05): - self._config["alpha"] = alpha - return self + @property + def alternative(self) -> EnumAlternative: + return self._alternative - def set_power(self, power: float = 0.80): - self._config["power"] = power - return self + @alternative.setter + def alternative(self, alternative: EnumAlternative): + self._alternative = alternative - def set_alternative(self, alternative: EnumAlternative = EnumAlternative.TWO_SIDED): - self._config["alternative"] = alternative - return self + @property + def test_type(self) -> EnumTestType: + return self._test_type - def set_test_type(self, test_type: EnumTestType = EnumTestType.Z_TEST_POOLED): - self._config["test_type"] = test_type - return self + @test_type.setter + def test_type(self, test_type: EnumTestType): + self._test_type = test_type - def set_treatment_proportion(self, treatment_proportion: float): - self._config["treatment_proportion"] = treatment_proportion - return self + @property + def treatment_proportion(self) -> float: + return self._treatment_proportion - def set_reference_proportion(self, reference_proportion: float): - self._config["reference_proportion"] = reference_proportion - return self + @treatment_proportion.setter + def treatment_proportion(self, treatment_proportion: float): + self._treatment_proportion = treatment_proportion - def set_group_allocation(self, group_allocation: GroupAllocation = GroupAllocation(EnumGroupAllocation.EQUAL)): - self._config["group_allocation"] = group_allocation - return self + @property + def reference_proportion(self) -> float: + return self._reference_proportion - def set_input_type(self, input_type): - raise NotImplementedError("这个功能还没有实现") + @reference_proportion.setter + def reference_proportion(self, reference_proportion: float): + self._reference_proportion = reference_proportion - def get_solver(self): - return TwoProportionSolveForNSolver(**self._config) + @property + def group_allocation(self) -> TwoProportionSolveForSizeGroupAllocator: + return self._group_allocation + @group_allocation.setter + def group_allocation(self, group_allocation: TwoProportionSolveForSizeGroupAllocator): + self._group_allocation = group_allocation -class TwoProportionSolveForNSolver: + def solve(self): + if None in ( + self._alpha, + self._power, + self._alternative, + self._test_type, + self._treatment_proportion, + self._reference_proportion, + self._group_allocation, + ): + raise ValueError("所有参数必须在调用 solve 之前设置") + self._eval = ( + lambda n: fun_power( + self._alpha, + self._group_allocation.treatment_size_formula(n), + self._group_allocation.reference_size_formula(n), + self._treatment_proportion, + self._reference_proportion, + self._alternative, + self._test_type, + ) + - self._power + ) + try: + n = brenth(self._eval, 1, 1e10) + except ValueError as e: + raise ValueError("无法求解样本量") from e + self._treatment_n = self._group_allocation.treatment_size_formula(n) + self._reference_n = self._group_allocation.reference_size_formula(n) + + +# solve for alpha +class EnumTwoProportionSolveForAlphaGroupAllocation(Enum): + EQUAL = 1 + INDIVIDUAL = 2 + TREATMENT_SIZE_AND_RATIO_OF_TREATMENT_TO_REFERENCE = 3 + TREATMENT_SIZE_AND_RATIO_OF_REFERENCE_TO_TREATMENT = 4 + REFERENCE_SIZE_AND_RATIO_OF_TREATMENT_TO_REFERENCE = 5 + REFERENCE_SIZE_AND_RATIO_OF_REFERENCE_TO_TREATMENT = 6 + TOTAL_SIZE_AND_TREATMENT_PERCENT = 7 + TOTAL_SIZE_AND_REFERENCE_PERCENT = 8 + TOTAL_SIZE_AND_RATIO_OF_TREATMENT_TO_REFERENCE = 9 + TOTAL_SIZE_AND_RATIO_OF_REFERENCE_TO_TREATMENT = 10 + TREATMENT_SIZE_AND_TREATMENT_PERCENT = 11 + TREATMENT_SIZE_AND_REFERENCE_PERCENT = 12 + REFERENCE_SIZE_AND_TREATMENT_PERCENT = 13 + REFERENCE_SIZE_AND_REFERENCE_PERCENT = 14 + + +class TwoProportionSolveForAlphaGroupAllocator: def __init__( self, - alpha: float, - power: float, - alternative: EnumAlternative, - test_type: EnumTestType, - treatment_proportion: float, - reference_proportion: float, - group_allocation: GroupAllocation, + group_allocation_option: EnumTwoProportionSolveForAlphaGroupAllocation = EnumTwoProportionSolveForAlphaGroupAllocation.EQUAL, + total_size: float = None, + treatment_size: float = None, + reference_size: float = None, + ratio_of_treatment_to_reference: float = None, + ratio_of_reference_to_treatment: float = None, + percent_of_treatment: float = None, + percent_of_reference: float = None, ): - self._alpha = alpha - self._power = power - self._alternative = alternative - self._test_type = test_type - self._treatment_proportion = treatment_proportion - self._reference_proportion = reference_proportion - self._group_allocation = group_allocation + match group_allocation_option: + case EnumTwoProportionSolveForAlphaGroupAllocation.EQUAL: + self.treatment_size_formula = lambda n: n + self.reference_size_formula = lambda n: n - def solve(self): - match self._group_allocation._group_allocation_option: - case EnumGroupAllocation.EQUAL: - eval = ( - lambda n: fun_power( - self._alpha, - n, - n, - self._treatment_proportion, - self._reference_proportion, - self._alternative, - self._test_type, - ) - - self._power - ) - try: - n = brenth(eval, 1, 1e10) - except ValueError as e: - raise ValueError("无法求解样本量") from e - case EnumGroupAllocation.FIX_TREATMENT_GROUP: - eval = ( - lambda n: fun_power( - self._alpha, - self._group_allocation._treatment_n, - n, - self._treatment_proportion, - self._reference_proportion, - self._alternative, - self._test_type, - ) - - self._power - ) - try: - n = brenth(eval, 1, 1e10) - except ValueError as e: - raise ValueError("无法求解样本量") from e - case EnumGroupAllocation.FIX_REFERENCE_GROUP: - eval = ( - lambda n: fun_power( - self._alpha, - n, - self._group_allocation._reference_n, - self._treatment_proportion, - self._reference_proportion, - self._alternative, - self._test_type, - ) - - self._power - ) - try: - n = brenth(eval, 1, 1e10) - except ValueError as e: - raise ValueError("无法求解样本量") from e - case EnumGroupAllocation.RATIO_OF_TREATMENT_TO_REFERENCE: - eval = ( - lambda n: fun_power( - self._alpha, - self._group_allocation._ratio_of_treatment_to_reference * n, - n, - self._treatment_proportion, - self._reference_proportion, - self._alternative, - self._test_type, - ) - - self._power - ) - try: - n = brenth(eval, 1, 1e10) - except ValueError as e: - raise ValueError("无法求解样本量") from e - case EnumGroupAllocation.RATIO_OF_REFERENCE_TO_TREATMENT: - eval = ( - lambda n: fun_power( - self._alpha, - n, - self._group_allocation._ratio_of_reference_to_treatment * n, - self._treatment_proportion, - self._reference_proportion, - self._alternative, - self._test_type, - ) - - self._power - ) - try: - n = brenth(eval, 1, 1e10) - except ValueError as e: - raise ValueError("无法求解样本量") from e - case EnumGroupAllocation.PERCENT_OF_TREATMENT: - eval = ( - lambda n: fun_power( - self._alpha, - n, - (1 - self._group_allocation._percent_of_treatment) - / self._group_allocation._percent_of_treatment - * n, - self._treatment_proportion, - self._reference_proportion, - self._alternative, - self._test_type, - ) - - self._power - ) - try: - n = brenth(eval, 1, 1e10) - except ValueError as e: - raise ValueError("无法求解样本量") from e - case EnumGroupAllocation.PERCENT_OF_REFERENCE: - eval = ( - lambda n: fun_power( - self._alpha, - (1 - self._group_allocation._percent_of_reference) - / self._group_allocation._percent_of_reference - * n, - n, - self._treatment_proportion, - self._reference_proportion, - self._alternative, - self._test_type, - ) - - self._power - ) - try: - n = brenth(eval, 1, 1e10) - except ValueError as e: - raise ValueError("无法求解样本量") from e - - return n - - -class TwoProportionSolveForAlphaDesigner: - pass +class TwoProportionSolveForAlpha: -class TwoProportionSolveForPowerDesigner: pass -class TwoProportionSolveForTreatmentProportionDesigner: +class TwoProportionSolveForPower: pass -class TwoProportionSolveForReferenceProportionDesigner: +class TwoProportionSolveForTreatmentProportion: pass -class TwoProportion: - def __new__(cls, solve_for: EnumSolvableParameter): - if not isinstance(solve_for, EnumSolvableParameter): - raise TypeError("solve_for must be an instance of SolvableParameter") - return cls._create_designer(solve_for) - - @staticmethod - def _create_designer(solve_for: EnumSolvableParameter): - match solve_for: - case EnumSolvableParameter.N: - return TwoProportionSolveForNDesigner() - case EnumSolvableParameter.ALPHA: - return TwoProportionSolveForAlphaDesigner() - case EnumSolvableParameter.POWER: - return TwoProportionSolveForPowerDesigner() - case EnumSolvableParameter.TREATMENT_PROPORTION: - return TwoProportionSolveForTreatmentProportionDesigner() - case EnumSolvableParameter.REFERENCE_PROPORTION: - return TwoProportionSolveForReferenceProportionDesigner() +class TwoProportionSolveForReferenceProportion: + pass