test_numpy_compat.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189
  1. import numpy as np
  2. import pytest
  3. from pandas import (
  4. CategoricalIndex,
  5. DatetimeIndex,
  6. Index,
  7. PeriodIndex,
  8. TimedeltaIndex,
  9. isna,
  10. )
  11. import pandas._testing as tm
  12. from pandas.api.types import (
  13. is_complex_dtype,
  14. is_numeric_dtype,
  15. )
  16. from pandas.core.arrays import BooleanArray
  17. from pandas.core.indexes.datetimelike import DatetimeIndexOpsMixin
  18. def test_numpy_ufuncs_out(index):
  19. result = index == index
  20. out = np.empty(index.shape, dtype=bool)
  21. np.equal(index, index, out=out)
  22. tm.assert_numpy_array_equal(out, result)
  23. if not index._is_multi:
  24. # same thing on the ExtensionArray
  25. out = np.empty(index.shape, dtype=bool)
  26. np.equal(index.array, index.array, out=out)
  27. tm.assert_numpy_array_equal(out, result)
  28. @pytest.mark.parametrize(
  29. "func",
  30. [
  31. np.exp,
  32. np.exp2,
  33. np.expm1,
  34. np.log,
  35. np.log2,
  36. np.log10,
  37. np.log1p,
  38. np.sqrt,
  39. np.sin,
  40. np.cos,
  41. np.tan,
  42. np.arcsin,
  43. np.arccos,
  44. np.arctan,
  45. np.sinh,
  46. np.cosh,
  47. np.tanh,
  48. np.arcsinh,
  49. np.arccosh,
  50. np.arctanh,
  51. np.deg2rad,
  52. np.rad2deg,
  53. ],
  54. ids=lambda x: x.__name__,
  55. )
  56. def test_numpy_ufuncs_basic(index, func):
  57. # test ufuncs of numpy, see:
  58. # https://numpy.org/doc/stable/reference/ufuncs.html
  59. if isinstance(index, DatetimeIndexOpsMixin):
  60. with tm.external_error_raised((TypeError, AttributeError)):
  61. with np.errstate(all="ignore"):
  62. func(index)
  63. elif is_numeric_dtype(index) and not (
  64. is_complex_dtype(index) and func in [np.deg2rad, np.rad2deg]
  65. ):
  66. # coerces to float (e.g. np.sin)
  67. with np.errstate(all="ignore"):
  68. result = func(index)
  69. arr_result = func(index.values)
  70. if arr_result.dtype == np.float16:
  71. arr_result = arr_result.astype(np.float32)
  72. exp = Index(arr_result, name=index.name)
  73. tm.assert_index_equal(result, exp)
  74. if isinstance(index.dtype, np.dtype) and is_numeric_dtype(index):
  75. if is_complex_dtype(index):
  76. assert result.dtype == index.dtype
  77. elif index.dtype in ["bool", "int8", "uint8"]:
  78. assert result.dtype in ["float16", "float32"]
  79. elif index.dtype in ["int16", "uint16", "float32"]:
  80. assert result.dtype == "float32"
  81. else:
  82. assert result.dtype == "float64"
  83. else:
  84. # e.g. np.exp with Int64 -> Float64
  85. assert type(result) is Index
  86. # raise AttributeError or TypeError
  87. elif len(index) == 0:
  88. pass
  89. else:
  90. with tm.external_error_raised((TypeError, AttributeError)):
  91. with np.errstate(all="ignore"):
  92. func(index)
  93. @pytest.mark.parametrize(
  94. "func", [np.isfinite, np.isinf, np.isnan, np.signbit], ids=lambda x: x.__name__
  95. )
  96. def test_numpy_ufuncs_other(index, func):
  97. # test ufuncs of numpy, see:
  98. # https://numpy.org/doc/stable/reference/ufuncs.html
  99. if isinstance(index, (DatetimeIndex, TimedeltaIndex)):
  100. if func in (np.isfinite, np.isinf, np.isnan):
  101. # numpy 1.18 changed isinf and isnan to not raise on dt64/td64
  102. result = func(index)
  103. assert isinstance(result, np.ndarray)
  104. out = np.empty(index.shape, dtype=bool)
  105. func(index, out=out)
  106. tm.assert_numpy_array_equal(out, result)
  107. else:
  108. with tm.external_error_raised(TypeError):
  109. func(index)
  110. elif isinstance(index, PeriodIndex):
  111. with tm.external_error_raised(TypeError):
  112. func(index)
  113. elif is_numeric_dtype(index) and not (
  114. is_complex_dtype(index) and func is np.signbit
  115. ):
  116. # Results in bool array
  117. result = func(index)
  118. if not isinstance(index.dtype, np.dtype):
  119. # e.g. Int64 we expect to get BooleanArray back
  120. assert isinstance(result, BooleanArray)
  121. else:
  122. assert isinstance(result, np.ndarray)
  123. out = np.empty(index.shape, dtype=bool)
  124. func(index, out=out)
  125. if not isinstance(index.dtype, np.dtype):
  126. tm.assert_numpy_array_equal(out, result._data)
  127. else:
  128. tm.assert_numpy_array_equal(out, result)
  129. elif len(index) == 0:
  130. pass
  131. else:
  132. with tm.external_error_raised(TypeError):
  133. func(index)
  134. @pytest.mark.parametrize("func", [np.maximum, np.minimum])
  135. def test_numpy_ufuncs_reductions(index, func, request):
  136. # TODO: overlap with tests.series.test_ufunc.test_reductions
  137. if len(index) == 0:
  138. pytest.skip("Test doesn't make sense for empty index.")
  139. if isinstance(index, CategoricalIndex) and index.dtype.ordered is False:
  140. with pytest.raises(TypeError, match="is not ordered for"):
  141. func.reduce(index)
  142. return
  143. else:
  144. result = func.reduce(index)
  145. if func is np.maximum:
  146. expected = index.max(skipna=False)
  147. else:
  148. expected = index.min(skipna=False)
  149. # TODO: do we have cases both with and without NAs?
  150. assert type(result) is type(expected)
  151. if isna(result):
  152. assert isna(expected)
  153. else:
  154. assert result == expected
  155. @pytest.mark.parametrize("func", [np.bitwise_and, np.bitwise_or, np.bitwise_xor])
  156. def test_numpy_ufuncs_bitwise(func):
  157. # https://github.com/pandas-dev/pandas/issues/46769
  158. idx1 = Index([1, 2, 3, 4], dtype="int64")
  159. idx2 = Index([3, 4, 5, 6], dtype="int64")
  160. with tm.assert_produces_warning(None):
  161. result = func(idx1, idx2)
  162. expected = Index(func(idx1.values, idx2.values))
  163. tm.assert_index_equal(result, expected)