test_sample.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363
  1. import numpy as np
  2. import pytest
  3. from pandas import (
  4. DataFrame,
  5. Index,
  6. Series,
  7. )
  8. import pandas._testing as tm
  9. import pandas.core.common as com
  10. class TestSample:
  11. @pytest.fixture
  12. def obj(self, frame_or_series):
  13. if frame_or_series is Series:
  14. arr = np.random.randn(10)
  15. else:
  16. arr = np.random.randn(10, 10)
  17. return frame_or_series(arr, dtype=None)
  18. @pytest.mark.parametrize("test", list(range(10)))
  19. def test_sample(self, test, obj):
  20. # Fixes issue: 2419
  21. # Check behavior of random_state argument
  22. # Check for stability when receives seed or random state -- run 10
  23. # times.
  24. seed = np.random.randint(0, 100)
  25. tm.assert_equal(
  26. obj.sample(n=4, random_state=seed), obj.sample(n=4, random_state=seed)
  27. )
  28. tm.assert_equal(
  29. obj.sample(frac=0.7, random_state=seed),
  30. obj.sample(frac=0.7, random_state=seed),
  31. )
  32. tm.assert_equal(
  33. obj.sample(n=4, random_state=np.random.RandomState(test)),
  34. obj.sample(n=4, random_state=np.random.RandomState(test)),
  35. )
  36. tm.assert_equal(
  37. obj.sample(frac=0.7, random_state=np.random.RandomState(test)),
  38. obj.sample(frac=0.7, random_state=np.random.RandomState(test)),
  39. )
  40. tm.assert_equal(
  41. obj.sample(frac=2, replace=True, random_state=np.random.RandomState(test)),
  42. obj.sample(frac=2, replace=True, random_state=np.random.RandomState(test)),
  43. )
  44. os1, os2 = [], []
  45. for _ in range(2):
  46. np.random.seed(test)
  47. os1.append(obj.sample(n=4))
  48. os2.append(obj.sample(frac=0.7))
  49. tm.assert_equal(*os1)
  50. tm.assert_equal(*os2)
  51. def test_sample_lengths(self, obj):
  52. # Check lengths are right
  53. assert len(obj.sample(n=4) == 4)
  54. assert len(obj.sample(frac=0.34) == 3)
  55. assert len(obj.sample(frac=0.36) == 4)
  56. def test_sample_invalid_random_state(self, obj):
  57. # Check for error when random_state argument invalid.
  58. msg = (
  59. "random_state must be an integer, array-like, a BitGenerator, Generator, "
  60. "a numpy RandomState, or None"
  61. )
  62. with pytest.raises(ValueError, match=msg):
  63. obj.sample(random_state="a_string")
  64. def test_sample_wont_accept_n_and_frac(self, obj):
  65. # Giving both frac and N throws error
  66. msg = "Please enter a value for `frac` OR `n`, not both"
  67. with pytest.raises(ValueError, match=msg):
  68. obj.sample(n=3, frac=0.3)
  69. def test_sample_requires_positive_n_frac(self, obj):
  70. with pytest.raises(
  71. ValueError,
  72. match="A negative number of rows requested. Please provide `n` >= 0",
  73. ):
  74. obj.sample(n=-3)
  75. with pytest.raises(
  76. ValueError,
  77. match="A negative number of rows requested. Please provide `frac` >= 0",
  78. ):
  79. obj.sample(frac=-0.3)
  80. def test_sample_requires_integer_n(self, obj):
  81. # Make sure float values of `n` give error
  82. with pytest.raises(ValueError, match="Only integers accepted as `n` values"):
  83. obj.sample(n=3.2)
  84. def test_sample_invalid_weight_lengths(self, obj):
  85. # Weight length must be right
  86. msg = "Weights and axis to be sampled must be of same length"
  87. with pytest.raises(ValueError, match=msg):
  88. obj.sample(n=3, weights=[0, 1])
  89. with pytest.raises(ValueError, match=msg):
  90. bad_weights = [0.5] * 11
  91. obj.sample(n=3, weights=bad_weights)
  92. with pytest.raises(ValueError, match="Fewer non-zero entries in p than size"):
  93. bad_weight_series = Series([0, 0, 0.2])
  94. obj.sample(n=4, weights=bad_weight_series)
  95. def test_sample_negative_weights(self, obj):
  96. # Check won't accept negative weights
  97. bad_weights = [-0.1] * 10
  98. msg = "weight vector many not include negative values"
  99. with pytest.raises(ValueError, match=msg):
  100. obj.sample(n=3, weights=bad_weights)
  101. def test_sample_inf_weights(self, obj):
  102. # Check inf and -inf throw errors:
  103. weights_with_inf = [0.1] * 10
  104. weights_with_inf[0] = np.inf
  105. msg = "weight vector may not include `inf` values"
  106. with pytest.raises(ValueError, match=msg):
  107. obj.sample(n=3, weights=weights_with_inf)
  108. weights_with_ninf = [0.1] * 10
  109. weights_with_ninf[0] = -np.inf
  110. with pytest.raises(ValueError, match=msg):
  111. obj.sample(n=3, weights=weights_with_ninf)
  112. def test_sample_zero_weights(self, obj):
  113. # All zeros raises errors
  114. zero_weights = [0] * 10
  115. with pytest.raises(ValueError, match="Invalid weights: weights sum to zero"):
  116. obj.sample(n=3, weights=zero_weights)
  117. def test_sample_missing_weights(self, obj):
  118. # All missing weights
  119. nan_weights = [np.nan] * 10
  120. with pytest.raises(ValueError, match="Invalid weights: weights sum to zero"):
  121. obj.sample(n=3, weights=nan_weights)
  122. def test_sample_none_weights(self, obj):
  123. # Check None are also replaced by zeros.
  124. weights_with_None = [None] * 10
  125. weights_with_None[5] = 0.5
  126. tm.assert_equal(
  127. obj.sample(n=1, axis=0, weights=weights_with_None), obj.iloc[5:6]
  128. )
  129. @pytest.mark.parametrize(
  130. "func_str,arg",
  131. [
  132. ("np.array", [2, 3, 1, 0]),
  133. ("np.random.MT19937", 3),
  134. ("np.random.PCG64", 11),
  135. ],
  136. )
  137. def test_sample_random_state(self, func_str, arg, frame_or_series):
  138. # GH#32503
  139. obj = DataFrame({"col1": range(10, 20), "col2": range(20, 30)})
  140. obj = tm.get_obj(obj, frame_or_series)
  141. result = obj.sample(n=3, random_state=eval(func_str)(arg))
  142. expected = obj.sample(n=3, random_state=com.random_state(eval(func_str)(arg)))
  143. tm.assert_equal(result, expected)
  144. def test_sample_generator(self, frame_or_series):
  145. # GH#38100
  146. obj = frame_or_series(np.arange(100))
  147. rng = np.random.default_rng()
  148. # Consecutive calls should advance the seed
  149. result1 = obj.sample(n=50, random_state=rng)
  150. result2 = obj.sample(n=50, random_state=rng)
  151. assert not (result1.index.values == result2.index.values).all()
  152. # Matching generator initialization must give same result
  153. # Consecutive calls should advance the seed
  154. result1 = obj.sample(n=50, random_state=np.random.default_rng(11))
  155. result2 = obj.sample(n=50, random_state=np.random.default_rng(11))
  156. tm.assert_equal(result1, result2)
  157. def test_sample_upsampling_without_replacement(self, frame_or_series):
  158. # GH#27451
  159. obj = DataFrame({"A": list("abc")})
  160. obj = tm.get_obj(obj, frame_or_series)
  161. msg = (
  162. "Replace has to be set to `True` when "
  163. "upsampling the population `frac` > 1."
  164. )
  165. with pytest.raises(ValueError, match=msg):
  166. obj.sample(frac=2, replace=False)
  167. class TestSampleDataFrame:
  168. # Tests which are relevant only for DataFrame, so these are
  169. # as fully parametrized as they can get.
  170. def test_sample(self):
  171. # GH#2419
  172. # additional specific object based tests
  173. # A few dataframe test with degenerate weights.
  174. easy_weight_list = [0] * 10
  175. easy_weight_list[5] = 1
  176. df = DataFrame(
  177. {
  178. "col1": range(10, 20),
  179. "col2": range(20, 30),
  180. "colString": ["a"] * 10,
  181. "easyweights": easy_weight_list,
  182. }
  183. )
  184. sample1 = df.sample(n=1, weights="easyweights")
  185. tm.assert_frame_equal(sample1, df.iloc[5:6])
  186. # Ensure proper error if string given as weight for Series or
  187. # DataFrame with axis = 1.
  188. ser = Series(range(10))
  189. msg = "Strings cannot be passed as weights when sampling from a Series."
  190. with pytest.raises(ValueError, match=msg):
  191. ser.sample(n=3, weights="weight_column")
  192. msg = (
  193. "Strings can only be passed to weights when sampling from rows on a "
  194. "DataFrame"
  195. )
  196. with pytest.raises(ValueError, match=msg):
  197. df.sample(n=1, weights="weight_column", axis=1)
  198. # Check weighting key error
  199. with pytest.raises(
  200. KeyError, match="'String passed to weights not a valid column'"
  201. ):
  202. df.sample(n=3, weights="not_a_real_column_name")
  203. # Check that re-normalizes weights that don't sum to one.
  204. weights_less_than_1 = [0] * 10
  205. weights_less_than_1[0] = 0.5
  206. tm.assert_frame_equal(df.sample(n=1, weights=weights_less_than_1), df.iloc[:1])
  207. ###
  208. # Test axis argument
  209. ###
  210. # Test axis argument
  211. df = DataFrame({"col1": range(10), "col2": ["a"] * 10})
  212. second_column_weight = [0, 1]
  213. tm.assert_frame_equal(
  214. df.sample(n=1, axis=1, weights=second_column_weight), df[["col2"]]
  215. )
  216. # Different axis arg types
  217. tm.assert_frame_equal(
  218. df.sample(n=1, axis="columns", weights=second_column_weight), df[["col2"]]
  219. )
  220. weight = [0] * 10
  221. weight[5] = 0.5
  222. tm.assert_frame_equal(df.sample(n=1, axis="rows", weights=weight), df.iloc[5:6])
  223. tm.assert_frame_equal(
  224. df.sample(n=1, axis="index", weights=weight), df.iloc[5:6]
  225. )
  226. # Check out of range axis values
  227. msg = "No axis named 2 for object type DataFrame"
  228. with pytest.raises(ValueError, match=msg):
  229. df.sample(n=1, axis=2)
  230. msg = "No axis named not_a_name for object type DataFrame"
  231. with pytest.raises(ValueError, match=msg):
  232. df.sample(n=1, axis="not_a_name")
  233. ser = Series(range(10))
  234. with pytest.raises(ValueError, match="No axis named 1 for object type Series"):
  235. ser.sample(n=1, axis=1)
  236. # Test weight length compared to correct axis
  237. msg = "Weights and axis to be sampled must be of same length"
  238. with pytest.raises(ValueError, match=msg):
  239. df.sample(n=1, axis=1, weights=[0.5] * 10)
  240. def test_sample_axis1(self):
  241. # Check weights with axis = 1
  242. easy_weight_list = [0] * 3
  243. easy_weight_list[2] = 1
  244. df = DataFrame(
  245. {"col1": range(10, 20), "col2": range(20, 30), "colString": ["a"] * 10}
  246. )
  247. sample1 = df.sample(n=1, axis=1, weights=easy_weight_list)
  248. tm.assert_frame_equal(sample1, df[["colString"]])
  249. # Test default axes
  250. tm.assert_frame_equal(
  251. df.sample(n=3, random_state=42), df.sample(n=3, axis=0, random_state=42)
  252. )
  253. def test_sample_aligns_weights_with_frame(self):
  254. # Test that function aligns weights with frame
  255. df = DataFrame({"col1": [5, 6, 7], "col2": ["a", "b", "c"]}, index=[9, 5, 3])
  256. ser = Series([1, 0, 0], index=[3, 5, 9])
  257. tm.assert_frame_equal(df.loc[[3]], df.sample(1, weights=ser))
  258. # Weights have index values to be dropped because not in
  259. # sampled DataFrame
  260. ser2 = Series([0.001, 0, 10000], index=[3, 5, 10])
  261. tm.assert_frame_equal(df.loc[[3]], df.sample(1, weights=ser2))
  262. # Weights have empty values to be filed with zeros
  263. ser3 = Series([0.01, 0], index=[3, 5])
  264. tm.assert_frame_equal(df.loc[[3]], df.sample(1, weights=ser3))
  265. # No overlap in weight and sampled DataFrame indices
  266. ser4 = Series([1, 0], index=[1, 2])
  267. with pytest.raises(ValueError, match="Invalid weights: weights sum to zero"):
  268. df.sample(1, weights=ser4)
  269. def test_sample_is_copy(self):
  270. # GH#27357, GH#30784: ensure the result of sample is an actual copy and
  271. # doesn't track the parent dataframe / doesn't give SettingWithCopy warnings
  272. df = DataFrame(np.random.randn(10, 3), columns=["a", "b", "c"])
  273. df2 = df.sample(3)
  274. with tm.assert_produces_warning(None):
  275. df2["d"] = 1
  276. def test_sample_does_not_modify_weights(self):
  277. # GH-42843
  278. result = np.array([np.nan, 1, np.nan])
  279. expected = result.copy()
  280. ser = Series([1, 2, 3])
  281. # Test numpy array weights won't be modified in place
  282. ser.sample(weights=result)
  283. tm.assert_numpy_array_equal(result, expected)
  284. # Test DataFrame column won't be modified in place
  285. df = DataFrame({"values": [1, 1, 1], "weights": [1, np.nan, np.nan]})
  286. expected = df["weights"].copy()
  287. df.sample(frac=1.0, replace=True, weights="weights")
  288. result = df["weights"]
  289. tm.assert_series_equal(result, expected)
  290. def test_sample_ignore_index(self):
  291. # GH 38581
  292. df = DataFrame(
  293. {"col1": range(10, 20), "col2": range(20, 30), "colString": ["a"] * 10}
  294. )
  295. result = df.sample(3, ignore_index=True)
  296. expected_index = Index(range(3))
  297. tm.assert_index_equal(result.index, expected_index, exact=True)