Coverage for /Users/sebastiana/Documents/Sugarpills/confidence/spotify_confidence/analysis/frequentist/sample_size_calculator.py: 65%

1 from typing import Union, Iterable, Tuple

2

3 from pandas import DataFrame

4

5 from spotify_confidence.analysis.frequentist.confidence_computers.generic_computer import GenericComputer

6 from ..abstract_base_classes.confidence_computer_abc import ConfidenceComputerABC

7 from ..confidence_utils import (

8 listify,

9)

10 from ..constants import BONFERRONI, ZTEST, METHOD_COLUMN_NAME

11

12

13 class SampleSizeCalculator:

14 def __init__(

15 self,

16 data_frame: DataFrame,

17 point_estimate_column: str,

18 var_column: str,

19 is_binary_column: str,

20 categorical_group_columns: Union[None, str, Iterable] = None,

21 interval_size: float = 0.95,

22 correction_method: str = BONFERRONI,

23 confidence_computer: ConfidenceComputerABC = None,

24 metric_column=None,

25 power: float = 0.8,

26 ):

27 if confidence_computer is not None:

28 self._confidence_computer = confidence_computer

29 else:

30 self._confidence_computer = GenericComputer(

31 data_frame=data_frame.assign(**{METHOD_COLUMN_NAME: ZTEST}),

32 numerator_column=None,

33 numerator_sum_squares_column=None,

34 denominator_column=None,

35 categorical_group_columns=listify(categorical_group_columns),

36 ordinal_group_column=None,

37 interval_size=interval_size,

38 correction_method=correction_method.lower(),

39 method_column=METHOD_COLUMN_NAME,

40 bootstrap_samples_column=None,

41 metric_column=metric_column,

42 treatment_column=None,

43 power=power,

44 point_estimate_column=point_estimate_column,

45 var_column=var_column,

46 is_binary_column=is_binary_column,

47 feature_column=None,

48 feature_sum_squares_column=None,

49 feature_cross_sum_column=None,

50 )

51

52 def sample_size(

53 self,

54 treatment_weights: Iterable,

55 mde_column: str,

56 nim_column: str,

57 preferred_direction_column: str,

58 final_expected_sample_size_column: str = None,

59 ) -> DataFrame:

60 """Args:

61 treatment_weights (Iterable): The first weight is treated as control, the rest as treatment groups.

62 mde_column (str): Name of column in source dataframe containing the minimum detectable effect sizes

63 nim_column (str): Name of column in source dataframe containing the non-inferiority margins.

64 preferred_direction_column (str): Name of column in source dataframe containing the preferred direction

65 of the metric, which can be "increase", "decrease" or None, the latter meaning a two sided test will be

66 performed.

67 final_expected_sample_size_column (str): Name of column in source dataframe containing an expected

68 sample size. If this is given, a confidence interval width around the avg will be returned in the

69 "ci_width" column

70 Returns:

71 Dataframe containing a column "required_sample_size_metric" containing the sample size ot the control vs

72 treatment comparison that requires the largest sample size.

73

74 One of mde or nim has to be set for each row of the dataframe. If mde is set, the null hypothesis is no

75 difference between control and treatment, and the sample size returned will be enough to detect a relative

76 difference of mde or and absolute difference of mde*avg. If nim is set the null hypothesis is -nim*avg

77 if preferred_direction is "increase" and +nim*avg if preferred direction is "decrease"

78 """

79 return self._confidence_computer.compute_sample_size(

80 treatment_weights, mde_column, nim_column, preferred_direction_column, final_expected_sample_size_column

81 )

82

83 def optimal_weights_and_sample_size(

84 self, sample_size_df: DataFrame, number_of_groups: int

85 ) -> Tuple[Iterable, int]:

86 """Args:

87 sample_size_df (DataFrame): A data frame returned by the sample_size method of this class

88 number_of_groups (int): Number of groups in the experiment, including control

89 Returns:

90 Tuple (list, int), where the list contains optimal weights and the int is the sample size that would be

91 required if those weights were used

92 """

93 return self._confidence_computer.compute_optimal_weights_and_sample_size(sample_size_df, number_of_groups)

94

95 def powered_effect(

96 self,

97 treatment_weights: Iterable,

98 mde_column: str,

99 nim_column: str,

100 preferred_direction_column: str,

101 sample_size: int,

102 ) -> DataFrame:

103 """Args:

104 treatment_weights (Iterable): The first weight is treated as control, the rest as treatment groups.

105 mde_column (str): Name of column in source dataframe containing the minimum detectable effect sizes

106 nim_column (str): Name of column in source dataframe containing the non-inferiority margins.

107 preferred_direction_column (str): Name of column in source dataframe containing the preferred direction

108 of the metric, which can be "increase", "decrease" or None, the latter meaning a two sided test will be

109 performed.

110 sample_size (int): Total sample size across all groups to base the powered effect calculation on

111 Returns:

112 Dataframe containing a column "powered_effect" containing the powered effect of the control vs

113 treatment comparison that requires the largest powered effect.

114

115 Mde or nims are only needed for some multiple correction methods and are there to give results

116 consistent with the sample size calculations, i.e. if you first calculate a sample size for a

117 specific MDE and then calculate the powered effect for that sample size, the original MDE will be returned.

118 """

119 return self._confidence_computer.compute_powered_effect(

120 treatment_weights, mde_column, nim_column, preferred_direction_column, sample_size

121 )