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 )