PageRenderTime 38ms CodeModel.GetById 9ms RepoModel.GetById 0ms app.codeStats 0ms

/pandas/sparse/tests/test_arithmetics.py

http://github.com/wesm/pandas
Python | 347 lines | 246 code | 88 blank | 13 comment | 13 complexity | 5113a2a5a7067c02161e7eb70180cae3 MD5 | raw file
Possible License(s): BSD-3-Clause, Apache-2.0 
    

  1. import numpy as np
    2. 

  2. import pandas as pd
    3. 

  3. import pandas.util.testing as tm
    4. 

  4. 

  5. 

  6. class TestSparseArrayArithmetics(tm.TestCase):
    7. 

  7. 

  8.     _multiprocess_can_split_ = True
    9. 

  9. 

  10.     _base = np.array
    11. 

  11.     _klass = pd.SparseArray
    12. 

  12. 

  13.     def _assert(self, a, b):
    14. 

  14.         tm.assert_numpy_array_equal(a, b)
    15. 

  15. 

  16.     def _check_numeric_ops(self, a, b, a_dense, b_dense):
    17. 

  17.         # sparse & sparse
    18. 

  18.         self._assert((a + b).to_dense(), a_dense + b_dense)
    19. 

  19.         self._assert((b + a).to_dense(), b_dense + a_dense)
    20. 

  20. 

  21.         self._assert((a - b).to_dense(), a_dense - b_dense)
    22. 

  22.         self._assert((b - a).to_dense(), b_dense - a_dense)
    23. 

  23. 

  24.         self._assert((a * b).to_dense(), a_dense * b_dense)
    25. 

  25.         self._assert((b * a).to_dense(), b_dense * a_dense)
    26. 

  26. 

  27.         # pandas uses future division
    28. 

  28.         self._assert((a / b).to_dense(), a_dense * 1.0 / b_dense)
    29. 

  29.         self._assert((b / a).to_dense(), b_dense * 1.0 / a_dense)
    30. 

  30. 

  31.         # ToDo: FIXME in GH 13843
    32. 

  32.         if not (self._base == pd.Series and a.dtype == 'int64'):
    33. 

  33.             self._assert((a // b).to_dense(), a_dense // b_dense)
    34. 

  34.             self._assert((b // a).to_dense(), b_dense // a_dense)
    35. 

  35. 

  36.         self._assert((a % b).to_dense(), a_dense % b_dense)
    37. 

  37.         self._assert((b % a).to_dense(), b_dense % a_dense)
    38. 

  38. 

  39.         self._assert((a ** b).to_dense(), a_dense ** b_dense)
    40. 

  40.         self._assert((b ** a).to_dense(), b_dense ** a_dense)
    41. 

  41. 

  42.         # sparse & dense
    43. 

  43.         self._assert((a + b_dense).to_dense(), a_dense + b_dense)
    44. 

  44.         self._assert((b_dense + a).to_dense(), b_dense + a_dense)
    45. 

  45. 

  46.         self._assert((a - b_dense).to_dense(), a_dense - b_dense)
    47. 

  47.         self._assert((b_dense - a).to_dense(), b_dense - a_dense)
    48. 

  48. 

  49.         self._assert((a * b_dense).to_dense(), a_dense * b_dense)
    50. 

  50.         self._assert((b_dense * a).to_dense(), b_dense * a_dense)
    51. 

  51. 

  52.         # pandas uses future division
    53. 

  53.         self._assert((a / b_dense).to_dense(), a_dense * 1.0 / b_dense)
    54. 

  54.         self._assert((b_dense / a).to_dense(), b_dense * 1.0 / a_dense)
    55. 

  55. 

  56.         # ToDo: FIXME in GH 13843
    57. 

  57.         if not (self._base == pd.Series and a.dtype == 'int64'):
    58. 

  58.             self._assert((a // b_dense).to_dense(), a_dense // b_dense)
    59. 

  59.             self._assert((b_dense // a).to_dense(), b_dense // a_dense)
    60. 

  60. 

  61.         self._assert((a % b_dense).to_dense(), a_dense % b_dense)
    62. 

  62.         self._assert((b_dense % a).to_dense(), b_dense % a_dense)
    63. 

  63. 

  64.         self._assert((a ** b_dense).to_dense(), a_dense ** b_dense)
    65. 

  65.         self._assert((b_dense ** a).to_dense(), b_dense ** a_dense)
    66. 

  66. 

  67.     def _check_bool_result(self, res):
    68. 

  68.         tm.assertIsInstance(res, self._klass)
    69. 

  69.         self.assertEqual(res.dtype, np.bool)
    70. 

  70.         self.assertIsInstance(res.fill_value, bool)
    71. 

  71. 

  72.     def _check_comparison_ops(self, a, b, a_dense, b_dense):
    73. 

  73.         # sparse & sparse
    74. 

  74.         self._check_bool_result(a == b)
    75. 

  75.         self._assert((a == b).to_dense(), a_dense == b_dense)
    76. 

  76. 

  77.         self._check_bool_result(a != b)
    78. 

  78.         self._assert((a != b).to_dense(), a_dense != b_dense)
    79. 

  79. 

  80.         self._check_bool_result(a >= b)
    81. 

  81.         self._assert((a >= b).to_dense(), a_dense >= b_dense)
    82. 

  82. 

  83.         self._check_bool_result(a <= b)
    84. 

  84.         self._assert((a <= b).to_dense(), a_dense <= b_dense)
    85. 

  85. 

  86.         self._check_bool_result(a > b)
    87. 

  87.         self._assert((a > b).to_dense(), a_dense > b_dense)
    88. 

  88. 

  89.         self._check_bool_result(a < b)
    90. 

  90.         self._assert((a < b).to_dense(), a_dense < b_dense)
    91. 

  91. 

  92.         # sparse & dense
    93. 

  93.         self._check_bool_result(a == b_dense)
    94. 

  94.         self._assert((a == b_dense).to_dense(), a_dense == b_dense)
    95. 

  95. 

  96.         self._check_bool_result(a != b_dense)
    97. 

  97.         self._assert((a != b_dense).to_dense(), a_dense != b_dense)
    98. 

  98. 

  99.         self._check_bool_result(a >= b_dense)
    100. 

  100.         self._assert((a >= b_dense).to_dense(), a_dense >= b_dense)
    101. 

  101. 

  102.         self._check_bool_result(a <= b_dense)
    103. 

  103.         self._assert((a <= b_dense).to_dense(), a_dense <= b_dense)
    104. 

  104. 

  105.         self._check_bool_result(a > b_dense)
    106. 

  106.         self._assert((a > b_dense).to_dense(), a_dense > b_dense)
    107. 

  107. 

  108.         self._check_bool_result(a < b_dense)
    109. 

  109.         self._assert((a < b_dense).to_dense(), a_dense < b_dense)
    110. 

  110. 

  111.     def test_float_scalar(self):
    112. 

  112.         values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
    113. 

  113. 

  114.         for kind in ['integer', 'block']:
    115. 

  115.             a = self._klass(values, kind=kind)
    116. 

  116.             self._check_numeric_ops(a, 1, values, 1)
    117. 

  117.             self._check_numeric_ops(a, 0, values, 0)
    118. 

  118.             self._check_numeric_ops(a, 3, values, 3)
    119. 

  119. 

  120.             a = self._klass(values, kind=kind, fill_value=0)
    121. 

  121.             self._check_numeric_ops(a, 1, values, 1)
    122. 

  122.             self._check_numeric_ops(a, 0, values, 0)
    123. 

  123.             self._check_numeric_ops(a, 3, values, 3)
    124. 

  124. 

  125.             a = self._klass(values, kind=kind, fill_value=2)
    126. 

  126.             self._check_numeric_ops(a, 1, values, 1)
    127. 

  127.             self._check_numeric_ops(a, 0, values, 0)
    128. 

  128.             self._check_numeric_ops(a, 3, values, 3)
    129. 

  129. 

  130.     def test_float_scalar_comparison(self):
    131. 

  131.         values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
    132. 

  132. 

  133.         for kind in ['integer', 'block']:
    134. 

  134.             a = self._klass(values, kind=kind)
    135. 

  135.             self._check_comparison_ops(a, 1, values, 1)
    136. 

  136.             self._check_comparison_ops(a, 0, values, 0)
    137. 

  137.             self._check_comparison_ops(a, 3, values, 3)
    138. 

  138. 

  139.             a = self._klass(values, kind=kind, fill_value=0)
    140. 

  140.             self._check_comparison_ops(a, 1, values, 1)
    141. 

  141.             self._check_comparison_ops(a, 0, values, 0)
    142. 

  142.             self._check_comparison_ops(a, 3, values, 3)
    143. 

  143. 

  144.             a = self._klass(values, kind=kind, fill_value=2)
    145. 

  145.             self._check_comparison_ops(a, 1, values, 1)
    146. 

  146.             self._check_comparison_ops(a, 0, values, 0)
    147. 

  147.             self._check_comparison_ops(a, 3, values, 3)
    148. 

  148. 

  149.     def test_float_same_index(self):
    150. 

  150.         # when sp_index are the same
    151. 

  151.         for kind in ['integer', 'block']:
    152. 

  152.             values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
    153. 

  153.             rvalues = self._base([np.nan, 2, 3, 4, np.nan, 0, 1, 3, 2, np.nan])
    154. 

  154. 

  155.             a = self._klass(values, kind=kind)
    156. 

  156.             b = self._klass(rvalues, kind=kind)
    157. 

  157.             self._check_numeric_ops(a, b, values, rvalues)
    158. 

  158. 

  159.             values = self._base([0., 1., 2., 6., 0., 0., 1., 2., 1., 0.])
    160. 

  160.             rvalues = self._base([0., 2., 3., 4., 0., 0., 1., 3., 2., 0.])
    161. 

  161. 

  162.             a = self._klass(values, kind=kind, fill_value=0)
    163. 

  163.             b = self._klass(rvalues, kind=kind, fill_value=0)
    164. 

  164.             self._check_numeric_ops(a, b, values, rvalues)
    165. 

  165. 

  166.     def test_float_same_index_comparison(self):
    167. 

  167.         # when sp_index are the same
    168. 

  168.         for kind in ['integer', 'block']:
    169. 

  169.             values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
    170. 

  170.             rvalues = self._base([np.nan, 2, 3, 4, np.nan, 0, 1, 3, 2, np.nan])
    171. 

  171. 

  172.             a = self._klass(values, kind=kind)
    173. 

  173.             b = self._klass(rvalues, kind=kind)
    174. 

  174.             self._check_comparison_ops(a, b, values, rvalues)
    175. 

  175. 

  176.             values = self._base([0., 1., 2., 6., 0., 0., 1., 2., 1., 0.])
    177. 

  177.             rvalues = self._base([0., 2., 3., 4., 0., 0., 1., 3., 2., 0.])
    178. 

  178. 

  179.             a = self._klass(values, kind=kind, fill_value=0)
    180. 

  180.             b = self._klass(rvalues, kind=kind, fill_value=0)
    181. 

  181.             self._check_comparison_ops(a, b, values, rvalues)
    182. 

  182. 

  183.     def test_float_array(self):
    184. 

  184.         values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
    185. 

  185.         rvalues = self._base([2, np.nan, 2, 3, np.nan, 0, 1, 5, 2, np.nan])
    186. 

  186. 

  187.         for kind in ['integer', 'block']:
    188. 

  188.             a = self._klass(values, kind=kind)
    189. 

  189.             b = self._klass(rvalues, kind=kind)
    190. 

  190.             self._check_numeric_ops(a, b, values, rvalues)
    191. 

  191.             self._check_numeric_ops(a, b * 0, values, rvalues * 0)
    192. 

  192. 

  193.             a = self._klass(values, kind=kind, fill_value=0)
    194. 

  194.             b = self._klass(rvalues, kind=kind)
    195. 

  195.             self._check_numeric_ops(a, b, values, rvalues)
    196. 

  196. 

  197.             a = self._klass(values, kind=kind, fill_value=0)
    198. 

  198.             b = self._klass(rvalues, kind=kind, fill_value=0)
    199. 

  199.             self._check_numeric_ops(a, b, values, rvalues)
    200. 

  200. 

  201.             a = self._klass(values, kind=kind, fill_value=1)
    202. 

  202.             b = self._klass(rvalues, kind=kind, fill_value=2)
    203. 

  203.             self._check_numeric_ops(a, b, values, rvalues)
    204. 

  204. 

  205.     def test_float_array_different_kind(self):
    206. 

  206.         values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
    207. 

  207.         rvalues = self._base([2, np.nan, 2, 3, np.nan, 0, 1, 5, 2, np.nan])
    208. 

  208. 

  209.         a = self._klass(values, kind='integer')
    210. 

  210.         b = self._klass(rvalues, kind='block')
    211. 

  211.         self._check_numeric_ops(a, b, values, rvalues)
    212. 

  212.         self._check_numeric_ops(a, b * 0, values, rvalues * 0)
    213. 

  213. 

  214.         a = self._klass(values, kind='integer', fill_value=0)
    215. 

  215.         b = self._klass(rvalues, kind='block')
    216. 

  216.         self._check_numeric_ops(a, b, values, rvalues)
    217. 

  217. 

  218.         a = self._klass(values, kind='integer', fill_value=0)
    219. 

  219.         b = self._klass(rvalues, kind='block', fill_value=0)
    220. 

  220.         self._check_numeric_ops(a, b, values, rvalues)
    221. 

  221. 

  222.         a = self._klass(values, kind='integer', fill_value=1)
    223. 

  223.         b = self._klass(rvalues, kind='block', fill_value=2)
    224. 

  224.         self._check_numeric_ops(a, b, values, rvalues)
    225. 

  225. 

  226.     def test_float_array_comparison(self):
    227. 

  227.         values = self._base([np.nan, 1, 2, 0, np.nan, 0, 1, 2, 1, np.nan])
    228. 

  228.         rvalues = self._base([2, np.nan, 2, 3, np.nan, 0, 1, 5, 2, np.nan])
    229. 

  229. 

  230.         for kind in ['integer', 'block']:
    231. 

  231.             a = self._klass(values, kind=kind)
    232. 

  232.             b = self._klass(rvalues, kind=kind)
    233. 

  233.             self._check_comparison_ops(a, b, values, rvalues)
    234. 

  234.             self._check_comparison_ops(a, b * 0, values, rvalues * 0)
    235. 

  235. 

  236.             a = self._klass(values, kind=kind, fill_value=0)
    237. 

  237.             b = self._klass(rvalues, kind=kind)
    238. 

  238.             self._check_comparison_ops(a, b, values, rvalues)
    239. 

  239. 

  240.             a = self._klass(values, kind=kind, fill_value=0)
    241. 

  241.             b = self._klass(rvalues, kind=kind, fill_value=0)
    242. 

  242.             self._check_comparison_ops(a, b, values, rvalues)
    243. 

  243. 

  244.             a = self._klass(values, kind=kind, fill_value=1)
    245. 

  245.             b = self._klass(rvalues, kind=kind, fill_value=2)
    246. 

  246.             self._check_comparison_ops(a, b, values, rvalues)
    247. 

  247. 

  248.     def test_int_array(self):
    249. 

  249.         # have to specify dtype explicitly until fixing GH 667
    250. 

  250.         dtype = np.int64
    251. 

  251. 

  252.         values = self._base([0, 1, 2, 0, 0, 0, 1, 2, 1, 0], dtype=dtype)
    253. 

  253.         rvalues = self._base([2, 0, 2, 3, 0, 0, 1, 5, 2, 0], dtype=dtype)
    254. 

  254. 

  255.         for kind in ['integer', 'block']:
    256. 

  256.             a = self._klass(values, dtype=dtype, kind=kind)
    257. 

  257.             self.assertEqual(a.dtype, dtype)
    258. 

  258.             b = self._klass(rvalues, dtype=dtype, kind=kind)
    259. 

  259.             self.assertEqual(b.dtype, dtype)
    260. 

  260. 

  261.             self._check_numeric_ops(a, b, values, rvalues)
    262. 

  262.             self._check_numeric_ops(a, b * 0, values, rvalues * 0)
    263. 

  263. 

  264.             a = self._klass(values, fill_value=0, dtype=dtype, kind=kind)
    265. 

  265.             self.assertEqual(a.dtype, dtype)
    266. 

  266.             b = self._klass(rvalues, dtype=dtype, kind=kind)
    267. 

  267.             self.assertEqual(b.dtype, dtype)
    268. 

  268. 

  269.             self._check_numeric_ops(a, b, values, rvalues)
    270. 

  270. 

  271.             a = self._klass(values, fill_value=0, dtype=dtype, kind=kind)
    272. 

  272.             self.assertEqual(a.dtype, dtype)
    273. 

  273.             b = self._klass(rvalues, fill_value=0, dtype=dtype, kind=kind)
    274. 

  274.             self.assertEqual(b.dtype, dtype)
    275. 

  275.             self._check_numeric_ops(a, b, values, rvalues)
    276. 

  276. 

  277.             a = self._klass(values, fill_value=1, dtype=dtype, kind=kind)
    278. 

  278.             self.assertEqual(a.dtype, dtype)
    279. 

  279.             b = self._klass(rvalues, fill_value=2, dtype=dtype, kind=kind)
    280. 

  280.             self.assertEqual(b.dtype, dtype)
    281. 

  281.             self._check_numeric_ops(a, b, values, rvalues)
    282. 

  282. 

  283.     def test_int_array_comparison(self):
    284. 

  284. 

  285.         # int32 NI ATM
    286. 

  286.         for dtype in ['int64']:
    287. 

  287.             values = self._base([0, 1, 2, 0, 0, 0, 1, 2, 1, 0], dtype=dtype)
    288. 

  288.             rvalues = self._base([2, 0, 2, 3, 0, 0, 1, 5, 2, 0], dtype=dtype)
    289. 

  289. 

  290.             for kind in ['integer', 'block']:
    291. 

  291.                 a = self._klass(values, dtype=dtype, kind=kind)
    292. 

  292.                 b = self._klass(rvalues, dtype=dtype, kind=kind)
    293. 

  293.                 self._check_comparison_ops(a, b, values, rvalues)
    294. 

  294.                 self._check_comparison_ops(a, b * 0, values, rvalues * 0)
    295. 

  295. 

  296.                 a = self._klass(values, dtype=dtype, kind=kind, fill_value=0)
    297. 

  297.                 b = self._klass(rvalues, dtype=dtype, kind=kind)
    298. 

  298.                 self._check_comparison_ops(a, b, values, rvalues)
    299. 

  299. 

  300.                 a = self._klass(values, dtype=dtype, kind=kind, fill_value=0)
    301. 

  301.                 b = self._klass(rvalues, dtype=dtype, kind=kind, fill_value=0)
    302. 

  302.                 self._check_comparison_ops(a, b, values, rvalues)
    303. 

  303. 

  304.                 a = self._klass(values, dtype=dtype, kind=kind, fill_value=1)
    305. 

  305.                 b = self._klass(rvalues, dtype=dtype, kind=kind, fill_value=2)
    306. 

  306.                 self._check_comparison_ops(a, b, values, rvalues)
    307. 

  307. 

  308. 

  309. class TestSparseSeriesArithmetic(TestSparseArrayArithmetics):
    310. 

  310. 

  311.     _base = pd.Series
    312. 

  312.     _klass = pd.SparseSeries
    313. 

  313. 

  314.     def _assert(self, a, b):
    315. 

  315.         tm.assert_series_equal(a, b)
    316. 

  316. 

  317.     def _check_bool_result(self, res):
    318. 

  318.         # ToDo: Must return SparseSeries after GH 667
    319. 

  319.         tm.assertIsInstance(res, self._base)
    320. 

  320.         self.assertEqual(res.dtype, np.bool)
    321. 

  321. 

  322.     def test_alignment(self):
    323. 

  323.         da = pd.Series(np.arange(4))
    324. 

  324.         db = pd.Series(np.arange(4), index=[1, 2, 3, 4])
    325. 

  325. 

  326.         sa = pd.SparseSeries(np.arange(4), dtype=np.int64, fill_value=0)
    327. 

  327.         sb = pd.SparseSeries(np.arange(4), index=[1, 2, 3, 4],
    328. 

  328.                              dtype=np.int64, fill_value=0)
    329. 

  329.         self._check_numeric_ops(sa, sb, da, db)
    330. 

  330. 

  331.         sa = pd.SparseSeries(np.arange(4), dtype=np.int64, fill_value=np.nan)
    332. 

  332.         sb = pd.SparseSeries(np.arange(4), index=[1, 2, 3, 4],
    333. 

  333.                              dtype=np.int64, fill_value=np.nan)
    334. 

  334.         self._check_numeric_ops(sa, sb, da, db)
    335. 

  335. 

  336.         da = pd.Series(np.arange(4))
    337. 

  337.         db = pd.Series(np.arange(4), index=[10, 11, 12, 13])
    338. 

  338. 

  339.         sa = pd.SparseSeries(np.arange(4), dtype=np.int64, fill_value=0)
    340. 

  340.         sb = pd.SparseSeries(np.arange(4), index=[10, 11, 12, 13],
    341. 

  341.                              dtype=np.int64, fill_value=0)
    342. 

  342.         self._check_numeric_ops(sa, sb, da, db)
    343. 

  343. 

  344.         sa = pd.SparseSeries(np.arange(4), dtype=np.int64, fill_value=np.nan)
    345. 

  345.         sb = pd.SparseSeries(np.arange(4), index=[10, 11, 12, 13],
    346. 

  346.                              dtype=np.int64, fill_value=np.nan)
    347. 

  347.         self._check_numeric_ops(sa, sb, da, db)
    348. 

  348.