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

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

17 statements  

1from typing import Union, Iterable, Tuple 

2 

3from pandas import DataFrame 

4 

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

6from ..abstract_base_classes.confidence_computer_abc import ConfidenceComputerABC 

7from ..confidence_utils import ( 

8 listify, 

9) 

10from ..constants import BONFERRONI, ZTEST, METHOD_COLUMN_NAME 

11 

12 

13class 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 )