Skip to content

Commit 0d6a663

Browse files
committed
TST: Additional tests for Series ufuncs
This adds a set of tests for ufuncs on Series. The goal is to establish the correct behavior prior to implementing `Series.__array_ufunc__`. There are two kinds of xfails right now 1. Series[Sparse] fails because `Series.__array_ufunc__` doesn't yet dispatch to `Series.array.__array_ufunc__` 2. `ufunc(series, series)` when the two series are unaligned. It's been determined that these should align, but isn't currently implemented.
1 parent baeb1bf commit 0d6a663

File tree

1 file changed

+152
-0
lines changed

1 file changed

+152
-0
lines changed

pandas/tests/series/test_ufunc.py

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
import string
2+
3+
import numpy as np
4+
import pytest
5+
6+
import pandas as pd
7+
import pandas.util.testing as tm
8+
9+
UNARY_UFUNCS = [np.positive, np.floor, np.exp]
10+
BINARY_UFUNCS = [np.add, np.logaddexp] # -> dunder op
11+
SPARSE = [
12+
pytest.param(True,
13+
marks=pytest.mark.xfail(reason="Series.__array_ufunc__")),
14+
False,
15+
]
16+
SPARSE_IDS = ['sparse', 'dense']
17+
SHUFFLE = [
18+
pytest.param(True, marks=pytest.mark.xfail(reason="GH-26945")),
19+
False
20+
]
21+
22+
23+
@pytest.fixture
24+
def arrays_for_binary_ufunc():
25+
"""
26+
A pair of random, length-100 integer-dtype arrays, that are mostly 0.
27+
"""
28+
a1 = np.random.randint(0, 10, 100)
29+
a2 = np.random.randint(0, 10, 100)
30+
a1[::3] = 0
31+
a2[::4] = 0
32+
return a1, a2
33+
34+
35+
@pytest.mark.parametrize("ufunc", UNARY_UFUNCS)
36+
@pytest.mark.parametrize("sparse", SPARSE, ids=SPARSE_IDS)
37+
def test_unary_ufunc(ufunc, sparse):
38+
array = np.random.randint(0, 10, 10)
39+
array[::2] = 0
40+
if sparse:
41+
array = pd.SparseArray(array, dtype=pd.SparseDtype('int', 0))
42+
43+
index = list(string.ascii_letters[:10])
44+
name = "name"
45+
series = pd.Series(array, index=index, name=name)
46+
47+
result = ufunc(series)
48+
expected = pd.Series(ufunc(array), index=index, name=name)
49+
tm.assert_series_equal(result, expected)
50+
51+
52+
@pytest.mark.parametrize("ufunc", BINARY_UFUNCS)
53+
@pytest.mark.parametrize("sparse", SPARSE, ids=SPARSE_IDS)
54+
@pytest.mark.parametrize("shuffle", SHUFFLE)
55+
@pytest.mark.parametrize("box_other", [True, False])
56+
def test_binary_ufunc(ufunc, sparse, shuffle, box_other,
57+
arrays_for_binary_ufunc):
58+
# Check the invariant that
59+
# ufunc(Series(a), Series(b)) == Series(ufunc(a, b))
60+
# with alignment.
61+
a1, a2 = arrays_for_binary_ufunc
62+
if sparse:
63+
a1 = pd.SparseArray(a1, dtype=pd.SparseDtype('int', 0))
64+
a2 = pd.SparseArray(a2, dtype=pd.SparseDtype('int', 0))
65+
66+
name = "name"
67+
# TODO: verify name when the differ? Take the first? Drop?
68+
s1 = pd.Series(a1, name=name)
69+
s2 = pd.Series(a2, name=name)
70+
71+
# handle shufling / alignment
72+
# If boxing -- ufunc(series, series) -- then we don't need to shuffle
73+
# the other array for the expected, since we align.
74+
# If not boxing -- ufunc(series, array) -- then we do need to shuffle
75+
# the other array, since we *dont'* align
76+
idx = np.random.permutation(len(s1))
77+
if box_other and shuffle:
78+
# ensure we align before applying the ufunc
79+
s2 = s2.take(idx)
80+
elif shuffle:
81+
a2 = a2.take(idx)
82+
83+
result = ufunc(s1, s2)
84+
expected = pd.Series(ufunc(a1, a2), name=name)
85+
tm.assert_series_equal(result, expected)
86+
87+
88+
@pytest.mark.parametrize("ufunc", BINARY_UFUNCS)
89+
@pytest.mark.parametrize("sparse", SPARSE, ids=SPARSE_IDS)
90+
@pytest.mark.parametrize("flip", [True, False])
91+
def test_binary_ufunc_scalar(ufunc, sparse, flip, arrays_for_binary_ufunc):
92+
array, _ = arrays_for_binary_ufunc
93+
if sparse:
94+
array = pd.SparseArray(array)
95+
other = 2
96+
series = pd.Series(array, name="name")
97+
98+
a, b = series, other
99+
c, d = array, other
100+
if flip:
101+
c, d = b, c
102+
a, b = b, a
103+
104+
expected = pd.Series(ufunc(a, b), name="name")
105+
result = pd.Series(ufunc(c, d), name="name")
106+
tm.assert_series_equal(result, expected)
107+
108+
109+
@pytest.mark.parametrize("ufunc", [np.divmod]) # any others?
110+
@pytest.mark.parametrize("sparse", SPARSE, ids=SPARSE_IDS)
111+
@pytest.mark.parametrize("shuffle", SHUFFLE)
112+
@pytest.mark.filterwarnings("ignore:divide by zero:RuntimeWarning")
113+
def test_multiple_ouput_binary_ufuncs(ufunc, sparse, shuffle,
114+
arrays_for_binary_ufunc):
115+
a1, a2 = arrays_for_binary_ufunc
116+
117+
if sparse:
118+
a1 = pd.SparseArray(a1, dtype=pd.SparseDtype('int', 0))
119+
a2 = pd.SparseArray(a2, dtype=pd.SparseDtype('int', 0))
120+
121+
s1 = pd.Series(a1)
122+
s2 = pd.Series(a2)
123+
124+
if shuffle:
125+
# ensure we align before applying the ufunc
126+
s2 = s2.sample(frac=1)
127+
128+
expected = ufunc(a1, a2)
129+
assert isinstance(expected, tuple)
130+
131+
result = ufunc(s1, s2)
132+
assert isinstance(result, tuple)
133+
tm.assert_series_equal(result[0], pd.Series(expected[0]))
134+
tm.assert_series_equal(result[1], pd.Series(expected[1]))
135+
136+
137+
@pytest.mark.parametrize("sparse", SPARSE, ids=SPARSE_IDS)
138+
def test_multiple_ouput_ufunc(sparse, arrays_for_binary_ufunc):
139+
array, _ = arrays_for_binary_ufunc
140+
141+
if sparse:
142+
array = pd.SparseArray(array)
143+
144+
series = pd.Series(array, name="name")
145+
result = np.modf(series)
146+
expected = np.modf(array)
147+
148+
assert isinstance(result, tuple)
149+
assert isinstance(expected, tuple)
150+
151+
tm.assert_series_equal(result[0], pd.Series(expected[0], name="name"))
152+
tm.assert_series_equal(result[1], pd.Series(expected[1], name="name"))

0 commit comments

Comments
 (0)