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 AxisInt 

20 

21 from pandas.core.generic import NDFrame 

22 

23 

24def preprocess_weights(obj: NDFrame, weights, axis: AxisInt) -> np.ndarray: 

25 """ 

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

27 `.GroupBy.sample`. 

28 

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

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

31 """ 

32 # If a series, align with frame 

33 if isinstance(weights, ABCSeries): 

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

35 

36 # Strings acceptable if a dataframe and axis = 0 

37 if isinstance(weights, str): 

38 if isinstance(obj, ABCDataFrame): 

39 if axis == 0: 

40 try: 

41 weights = obj[weights] 

42 except KeyError as err: 

43 raise KeyError( 

44 "String passed to weights not a valid column" 

45 ) from err 

46 else: 

47 raise ValueError( 

48 "Strings can only be passed to " 

49 "weights when sampling from rows on " 

50 "a DataFrame" 

51 ) 

52 else: 

53 raise ValueError( 

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

55 ) 

56 

57 if isinstance(obj, ABCSeries): 

58 func = obj._constructor 

59 else: 

60 func = obj._constructor_sliced 

61 

62 weights = func(weights, dtype="float64")._values 

63 

64 if len(weights) != obj.shape[axis]: 

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

66 

67 if lib.has_infs(weights): 

68 raise ValueError("weight vector may not include `inf` values") 

69 

70 if (weights < 0).any(): 

71 raise ValueError("weight vector many not include negative values") 

72 

73 missing = np.isnan(weights) 

74 if missing.any(): 

75 # Don't modify weights in place 

76 weights = weights.copy() 

77 weights[missing] = 0 

78 return weights 

79 

80 

81def process_sampling_size( 

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

83) -> int | None: 

84 """ 

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

86 `.GroupBy.sample`. 

87 

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

89 the constant sampling size. 

90 """ 

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

92 if n is None and frac is None: 

93 n = 1 

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

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

96 elif n is not None: 

97 if n < 0: 

98 raise ValueError( 

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

100 ) 

101 if n % 1 != 0: 

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

103 else: 

104 assert frac is not None # for mypy 

105 if frac > 1 and not replace: 

106 raise ValueError( 

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

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

109 ) 

110 if frac < 0: 

111 raise ValueError( 

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

113 ) 

114 

115 return n 

116 

117 

118def sample( 

119 obj_len: int, 

120 size: int, 

121 replace: bool, 

122 weights: np.ndarray | None, 

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

124) -> np.ndarray: 

125 """ 

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

127 

128 Parameters 

129 ---------- 

130 obj_len : int 

131 The length of the indices being considered 

132 size : int 

133 The number of values to choose 

134 replace : bool 

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

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

137 If None, equal probability weighting, otherwise weights according 

138 to the vector normalized 

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

140 State used for the random sampling 

141 

142 Returns 

143 ------- 

144 np.ndarray[np.intp] 

145 """ 

146 if weights is not None: 

147 weight_sum = weights.sum() 

148 if weight_sum != 0: 

149 weights = weights / weight_sum 

150 else: 

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

152 

153 assert weights is not None # for mypy 

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

155 raise ValueError( 

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

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

158 "sample for details." 

159 ) 

160 

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

162 np.intp, copy=False 

163 )