1from functools import partial
2
3import numpy as np
4import pytest
5
6import pandas.util._test_decorators as td
7
8import pandas as pd
9import pandas._testing as tm
10
11dtypes = [
12    "int64",
13    "Int64",
14    {"A": "int64", "B": "Int64"},
15]
16
17
18@pytest.mark.parametrize("dtype", dtypes)
19def test_unary_unary(dtype):
20    # unary input, unary output
21    values = np.array([[-1, -1], [1, 1]], dtype="int64")
22    df = pd.DataFrame(values, columns=["A", "B"], index=["a", "b"]).astype(dtype=dtype)
23    result = np.positive(df)
24    expected = pd.DataFrame(
25        np.positive(values), index=df.index, columns=df.columns
26    ).astype(dtype)
27    tm.assert_frame_equal(result, expected)
28
29
30@pytest.mark.parametrize("dtype", dtypes)
31def test_unary_binary(dtype):
32    # unary input, binary output
33    if pd.api.types.is_extension_array_dtype(dtype) or isinstance(dtype, dict):
34        pytest.xfail(reason="Extension / mixed with multiple outuputs not implemented.")
35
36    values = np.array([[-1, -1], [1, 1]], dtype="int64")
37    df = pd.DataFrame(values, columns=["A", "B"], index=["a", "b"]).astype(dtype=dtype)
38    result_pandas = np.modf(df)
39    assert isinstance(result_pandas, tuple)
40    assert len(result_pandas) == 2
41    expected_numpy = np.modf(values)
42
43    for result, b in zip(result_pandas, expected_numpy):
44        expected = pd.DataFrame(b, index=df.index, columns=df.columns)
45        tm.assert_frame_equal(result, expected)
46
47
48@pytest.mark.parametrize("dtype", dtypes)
49def test_binary_input_dispatch_binop(dtype):
50    # binop ufuncs are dispatched to our dunder methods.
51    values = np.array([[-1, -1], [1, 1]], dtype="int64")
52    df = pd.DataFrame(values, columns=["A", "B"], index=["a", "b"]).astype(dtype=dtype)
53    result = np.add(df, df)
54    expected = pd.DataFrame(
55        np.add(values, values), index=df.index, columns=df.columns
56    ).astype(dtype)
57    tm.assert_frame_equal(result, expected)
58
59
60@pytest.mark.parametrize(
61    "func,arg,expected",
62    [
63        (np.add, 1, [2, 3, 4, 5]),
64        (
65            partial(np.add, where=[[False, True], [True, False]]),
66            np.array([[1, 1], [1, 1]]),
67            [0, 3, 4, 0],
68        ),
69        (np.power, np.array([[1, 1], [2, 2]]), [1, 2, 9, 16]),
70        (np.subtract, 2, [-1, 0, 1, 2]),
71        (
72            partial(np.negative, where=np.array([[False, True], [True, False]])),
73            None,
74            [0, -2, -3, 0],
75        ),
76    ],
77)
78def test_ufunc_passes_args(func, arg, expected, request):
79    # GH#40662
80    arr = np.array([[1, 2], [3, 4]])
81    df = pd.DataFrame(arr)
82    result_inplace = np.zeros_like(arr)
83    # 1-argument ufunc
84    if arg is None:
85        result = func(df, out=result_inplace)
86    else:
87        result = func(df, arg, out=result_inplace)
88
89    expected = np.array(expected).reshape(2, 2)
90    tm.assert_numpy_array_equal(result_inplace, expected)
91
92    expected = pd.DataFrame(expected)
93    tm.assert_frame_equal(result, expected)
94
95
96@pytest.mark.parametrize("dtype_a", dtypes)
97@pytest.mark.parametrize("dtype_b", dtypes)
98def test_binary_input_aligns_columns(dtype_a, dtype_b):
99    if (
100        pd.api.types.is_extension_array_dtype(dtype_a)
101        or isinstance(dtype_a, dict)
102        or pd.api.types.is_extension_array_dtype(dtype_b)
103        or isinstance(dtype_b, dict)
104    ):
105        pytest.xfail(reason="Extension / mixed with multiple inputs not implemented.")
106
107    df1 = pd.DataFrame({"A": [1, 2], "B": [3, 4]}).astype(dtype_a)
108
109    if isinstance(dtype_a, dict) and isinstance(dtype_b, dict):
110        dtype_b["C"] = dtype_b.pop("B")
111
112    df2 = pd.DataFrame({"A": [1, 2], "C": [3, 4]}).astype(dtype_b)
113    with tm.assert_produces_warning(FutureWarning):
114        result = np.heaviside(df1, df2)
115    # Expected future behaviour:
116    # expected = np.heaviside(
117    #     np.array([[1, 3, np.nan], [2, 4, np.nan]]),
118    #     np.array([[1, np.nan, 3], [2, np.nan, 4]]),
119    # )
120    # expected = pd.DataFrame(expected, index=[0, 1], columns=["A", "B", "C"])
121    expected = pd.DataFrame([[1.0, 1.0], [1.0, 1.0]], columns=["A", "B"])
122    tm.assert_frame_equal(result, expected)
123
124    # ensure the expected is the same when applying with numpy array
125    result = np.heaviside(df1, df2.values)
126    tm.assert_frame_equal(result, expected)
127
128
129@pytest.mark.parametrize("dtype", dtypes)
130def test_binary_input_aligns_index(dtype):
131    if pd.api.types.is_extension_array_dtype(dtype) or isinstance(dtype, dict):
132        pytest.xfail(reason="Extension / mixed with multiple inputs not implemented.")
133    df1 = pd.DataFrame({"A": [1, 2], "B": [3, 4]}, index=["a", "b"]).astype(dtype)
134    df2 = pd.DataFrame({"A": [1, 2], "B": [3, 4]}, index=["a", "c"]).astype(dtype)
135    with tm.assert_produces_warning(FutureWarning):
136        result = np.heaviside(df1, df2)
137    # Expected future behaviour:
138    # expected = np.heaviside(
139    #     np.array([[1, 3], [3, 4], [np.nan, np.nan]]),
140    #     np.array([[1, 3], [np.nan, np.nan], [3, 4]]),
141    # )
142    # # TODO(FloatArray): this will be Float64Dtype.
143    # expected = pd.DataFrame(expected, index=["a", "b", "c"], columns=["A", "B"])
144    expected = pd.DataFrame(
145        [[1.0, 1.0], [1.0, 1.0]], columns=["A", "B"], index=["a", "b"]
146    )
147    tm.assert_frame_equal(result, expected)
148
149    # ensure the expected is the same when applying with numpy array
150    result = np.heaviside(df1, df2.values)
151    tm.assert_frame_equal(result, expected)
152
153
154@pytest.mark.filterwarnings("ignore:Calling a ufunc on non-aligned:FutureWarning")
155def test_binary_frame_series_raises():
156    # We don't currently implement
157    df = pd.DataFrame({"A": [1, 2]})
158    # with pytest.raises(NotImplementedError, match="logaddexp"):
159    with pytest.raises(ValueError, match=""):
160        np.logaddexp(df, df["A"])
161
162    # with pytest.raises(NotImplementedError, match="logaddexp"):
163    with pytest.raises(ValueError, match=""):
164        np.logaddexp(df["A"], df)
165
166
167def test_unary_accumulate_axis():
168    # https://github.com/pandas-dev/pandas/issues/39259
169    df = pd.DataFrame({"a": [1, 3, 2, 4]})
170    result = np.maximum.accumulate(df)
171    expected = pd.DataFrame({"a": [1, 3, 3, 4]})
172    tm.assert_frame_equal(result, expected)
173
174    df = pd.DataFrame({"a": [1, 3, 2, 4], "b": [0.1, 4.0, 3.0, 2.0]})
175    result = np.maximum.accumulate(df)
176    # in theory could preserve int dtype for default axis=0
177    expected = pd.DataFrame({"a": [1.0, 3.0, 3.0, 4.0], "b": [0.1, 4.0, 4.0, 4.0]})
178    tm.assert_frame_equal(result, expected)
179
180    result = np.maximum.accumulate(df, axis=0)
181    tm.assert_frame_equal(result, expected)
182
183    result = np.maximum.accumulate(df, axis=1)
184    expected = pd.DataFrame({"a": [1.0, 3.0, 2.0, 4.0], "b": [1.0, 4.0, 3.0, 4.0]})
185    tm.assert_frame_equal(result, expected)
186
187
188def test_frame_outer_deprecated():
189    df = pd.DataFrame({"A": [1, 2]})
190    with tm.assert_produces_warning(FutureWarning):
191        np.subtract.outer(df, df)
192
193
194def test_alignment_deprecation():
195    # https://github.com/pandas-dev/pandas/issues/39184
196    df1 = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
197    df2 = pd.DataFrame({"b": [1, 2, 3], "c": [4, 5, 6]})
198    s1 = pd.Series([1, 2], index=["a", "b"])
199    s2 = pd.Series([1, 2], index=["b", "c"])
200
201    # binary dataframe / dataframe
202    expected = pd.DataFrame({"a": [2, 4, 6], "b": [8, 10, 12]})
203
204    with tm.assert_produces_warning(None):
205        # aligned -> no warning!
206        result = np.add(df1, df1)
207    tm.assert_frame_equal(result, expected)
208
209    with tm.assert_produces_warning(FutureWarning):
210        # non-aligned -> warns
211        result = np.add(df1, df2)
212    tm.assert_frame_equal(result, expected)
213
214    result = np.add(df1, df2.values)
215    tm.assert_frame_equal(result, expected)
216
217    result = np.add(df1.values, df2)
218    expected = pd.DataFrame({"b": [2, 4, 6], "c": [8, 10, 12]})
219    tm.assert_frame_equal(result, expected)
220
221    # binary dataframe / series
222    expected = pd.DataFrame({"a": [2, 3, 4], "b": [6, 7, 8]})
223
224    with tm.assert_produces_warning(None):
225        # aligned -> no warning!
226        result = np.add(df1, s1)
227    tm.assert_frame_equal(result, expected)
228
229    with tm.assert_produces_warning(FutureWarning):
230        result = np.add(df1, s2)
231    tm.assert_frame_equal(result, expected)
232
233    with tm.assert_produces_warning(FutureWarning):
234        result = np.add(s2, df1)
235    tm.assert_frame_equal(result, expected)
236
237    result = np.add(df1, s2.values)
238    tm.assert_frame_equal(result, expected)
239
240
241@td.skip_if_no("numba", "0.46.0")
242def test_alignment_deprecation_many_inputs():
243    # https://github.com/pandas-dev/pandas/issues/39184
244    # test that the deprecation also works with > 2 inputs -> using a numba
245    # written ufunc for this because numpy itself doesn't have such ufuncs
246    from numba import float64, vectorize
247
248    @vectorize([float64(float64, float64, float64)])
249    def my_ufunc(x, y, z):
250        return x + y + z
251
252    df1 = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
253    df2 = pd.DataFrame({"b": [1, 2, 3], "c": [4, 5, 6]})
254    df3 = pd.DataFrame({"a": [1, 2, 3], "c": [4, 5, 6]})
255
256    with tm.assert_produces_warning(FutureWarning):
257        result = my_ufunc(df1, df2, df3)
258    expected = pd.DataFrame([[3.0, 12.0], [6.0, 15.0], [9.0, 18.0]], columns=["a", "b"])
259    tm.assert_frame_equal(result, expected)
260
261    # all aligned -> no warning
262    with tm.assert_produces_warning(None):
263        result = my_ufunc(df1, df1, df1)
264    tm.assert_frame_equal(result, expected)
265
266    # mixed frame / arrays
267    with tm.assert_produces_warning(FutureWarning):
268        result = my_ufunc(df1, df2, df3.values)
269    tm.assert_frame_equal(result, expected)
270
271    # single frame -> no warning
272    with tm.assert_produces_warning(None):
273        result = my_ufunc(df1, df2.values, df3.values)
274    tm.assert_frame_equal(result, expected)
275
276    # takes indices of first frame
277    with tm.assert_produces_warning(FutureWarning):
278        result = my_ufunc(df1.values, df2, df3)
279    expected = expected.set_axis(["b", "c"], axis=1)
280    tm.assert_frame_equal(result, expected)
281