Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/core/sample.py: 16%

1""" 

2Module containing utilities for NDFrame.sample() and .GroupBy.sample() 

3""" 

4 

5from __future__ import annotations 

6 

7from typing import TYPE_CHECKING 

8 

9import numpy as np 

10 

11from pandas._libs import lib 

12 

13from pandas.core.dtypes.generic import ( 

14 ABCDataFrame, 

15 ABCSeries, 

16) 

17 

18if TYPE_CHECKING: 

19 from pandas._typing import ( 

20 AxisInt, 

21 ListLike, 

22 ) 

23 

24 from pandas.core.generic import NDFrame 

25 

26 

27def preprocess_weights( 

28 obj: NDFrame, weights: str | ListLike, axis: AxisInt 

29) -> np.ndarray: 

30 """ 

31 Process and validate the `weights` argument to `NDFrame.sample` and 

32 `.GroupBy.sample`. 

33 

34 Returns `weights` as an ndarray[np.float64], validated except for normalizing 

35 weights (because that must be done groupwise in groupby sampling). 

36 """ 

37 # If a series, align with frame 

38 if isinstance(weights, ABCSeries): 

39 weights = weights.reindex(obj.axes[axis]) 

40 

41 # Strings acceptable if a dataframe and axis = 0 

42 if isinstance(weights, str): 

43 if isinstance(obj, ABCDataFrame): 

44 if axis == 0: 

45 try: 

46 weights = obj[weights] 

47 except KeyError as err: 

48 raise KeyError( 

49 "String passed to weights not a valid column" 

50 ) from err 

51 else: 

52 raise ValueError( 

53 "Strings can only be passed to " 

54 "weights when sampling from rows on " 

55 "a DataFrame" 

56 ) 

57 else: 

58 raise ValueError( 

59 "Strings cannot be passed as weights when sampling from a Series." 

60 ) 

61 

62 if isinstance(obj, ABCSeries): 

63 func = obj._constructor 

64 else: 

65 func = obj._constructor_sliced 

66 

67 weights_arr = np.asarray(func(weights, dtype="float64")._values, dtype=np.float64) 

68 

69 if len(weights_arr) != obj.shape[axis]: 

70 raise ValueError("Weights and axis to be sampled must be of same length") 

71 

72 if lib.has_infs(weights_arr): 

73 raise ValueError("weights vector may not include `inf` values") 

74 

75 if (weights_arr < 0).any(): 

76 raise ValueError("weights vector may not include negative values") 

77 

78 missing = np.isnan(weights_arr) 

79 if missing.any(): 

80 # Don't modify weights in place 

81 weights_arr = weights_arr.copy() 

82 weights_arr[missing] = 0 

83 return weights_arr 

84 

85 

86def process_sampling_size( 

87 n: int | None, frac: float | None, replace: bool 

88) -> int | None: 

89 """ 

90 Process and validate the `n` and `frac` arguments to `NDFrame.sample` and 

91 `.GroupBy.sample`. 

92 

93 Returns None if `frac` should be used (variable sampling sizes), otherwise returns 

94 the constant sampling size. 

95 """ 

96 # If no frac or n, default to n=1. 

97 if n is None and frac is None: 

98 n = 1 

99 elif n is not None and frac is not None: 

100 raise ValueError("Please enter a value for `frac` OR `n`, not both") 

101 elif n is not None: 

102 if n < 0: 

103 raise ValueError( 

104 "A negative number of rows requested. Please provide `n` >= 0." 

105 ) 

106 if n % 1 != 0: 

107 raise ValueError("Only integers accepted as `n` values") 

108 else: 

109 assert frac is not None # for mypy 

110 if frac > 1 and not replace: 

111 raise ValueError( 

112 "Replace has to be set to `True` when " 

113 "upsampling the population `frac` > 1." 

114 ) 

115 if frac < 0: 

116 raise ValueError( 

117 "A negative number of rows requested. Please provide `frac` >= 0." 

118 ) 

119 

120 return n 

121 

122 

123def sample( 

124 obj_len: int, 

125 size: int, 

126 replace: bool, 

127 weights: np.ndarray | None, 

128 random_state: np.random.RandomState | np.random.Generator, 

129) -> np.ndarray: 

130 """ 

131 Randomly sample `size` indices in `np.arange(obj_len)`. 

132 

133 Parameters 

134 ---------- 

135 obj_len : int 

136 The length of the indices being considered 

137 size : int 

138 The number of values to choose 

139 replace : bool 

140 Allow or disallow sampling of the same row more than once. 

141 weights : np.ndarray[np.float64] or None 

142 If None, equal probability weighting, otherwise weights according 

143 to the vector normalized 

144 random_state: np.random.RandomState or np.random.Generator 

145 State used for the random sampling 

146 

147 Returns 

148 ------- 

149 np.ndarray[np.intp] 

150 """ 

151 if weights is not None: 

152 weight_sum = weights.sum() 

153 if weight_sum != 0: 

154 weights = weights / weight_sum 

155 else: 

156 raise ValueError("Invalid weights: weights sum to zero") 

157 

158 assert weights is not None # for mypy 

159 if not replace and size * weights.max() > 1: 

160 raise ValueError( 

161 "Weighted sampling cannot be achieved with replace=False. Either " 

162 "set replace=True or use smaller weights. See the docstring of " 

163 "sample for details." 

164 ) 

165 

166 return random_state.choice(obj_len, size=size, replace=replace, p=weights).astype( 

167 np.intp, copy=False 

168 )