Coverage for /Users/sebastiana/Documents/Sugarpills/confidence/spotify_confidence/examples.py: 24%

1# Copyright 2017-2020 Spotify AB

2#

3# Licensed under the Apache License, Version 2.0 (the "License");

4# you may not use this file except in compliance with the License.

5# You may obtain a copy of the License at

6#

7# http://www.apache.org/licenses/LICENSE-2.0

8#

9# Unless required by applicable law or agreed to in writing, software

10# distributed under the License is distributed on an "AS IS" BASIS,

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

12# See the License for the specific language governing permissions and

13# limitations under the License.

14

15 import pandas as pd

16 import numpy as np

17 from itertools import product

18

19

20 def example_data_binomial():

21 """

22 Returns an output dataframe with categorical

23 features (country and test variation), and orginal features (date),

24 as well as number of successes and total observations for each combination

25 """

26 countries = ["ca", "us"]

27 dates = pd.date_range("2018-01-01", "2018-02-01")

28 variation_names = ["test", "control", "test2"]

29

30 # test ca, test us, control ca, control us, test2 ca, test2 us

31 success_rates = [0.3, 0.32, 0.24, 0.22, 0.25, 0.42]

32 n_observations = [50, 80, 30, 50, 40, 50]

33

34 return_df = pd.DataFrame()

35

36 for i, (country, variation) in enumerate(product(countries, variation_names)):

37 df = pd.DataFrame({"date": dates})

38 df["country"] = country

39 df["variation_name"] = variation

40 df["total"] = np.random.poisson(n_observations[i], size=len(dates))

41 df["success"] = df["total"].apply(lambda x: np.random.binomial(x, success_rates[i]))

42 return_df = pd.concat([return_df, df], axis=0)

43

44 return return_df

45

46

47 def example_data_gaussian():

48 df = pd.DataFrame(

49 {

50 "variation_name": [

51 "test",

52 "control",

53 "test2",

54 "test",

55 "control",

56 "test2",

57 "test",

58 "control",

59 "test2",

60 "test",

61 "control",

62 "test2",

63 "test",

64 "control",

65 "test2",

66 ],

67 "nr_of_items": [

68 500,

69 8,

70 100,

71 510,

72 8,

73 100,

74 520,

75 9,

76 104,

77 530,

78 7,

79 100,

80 530,

81 8,

82 103,

83 ],

84 "nr_of_items_sumsq": [

85 2500,

86 12,

87 150,

88 2510,

89 13,

90 140,

91 2520,

92 14,

93 154,

94 2530,

95 15,

96 160,

97 2530,

98 16,

99 103,

100 ],

101 "users": [

102 1010,

103 22,

104 150,

105 1000,

106 20,

107 153,

108 1030,

109 23,

110 154,

111 1000,

112 20,

113 150,

114 1040,

115 21,

116 155,

117 ],

118 "days_since_reg": [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5],

119 }

120 )

121

122 return df