1"""
2Missing data handling for arithmetic operations.
3
4In particular, pandas conventions regarding division by zero differ
5from numpy in the following ways:
6 1) np.array([-1, 0, 1], dtype=dtype1) // np.array([0, 0, 0], dtype=dtype2)
7 gives [nan, nan, nan] for most dtype combinations, and [0, 0, 0] for
8 the remaining pairs
9 (the remaining being dtype1==dtype2==intN and dtype==dtype2==uintN).
10
11 pandas convention is to return [-inf, nan, inf] for all dtype
12 combinations.
13
14 Note: the numpy behavior described here is py3-specific.
15
16 2) np.array([-1, 0, 1], dtype=dtype1) % np.array([0, 0, 0], dtype=dtype2)
17 gives precisely the same results as the // operation.
18
19 pandas convention is to return [nan, nan, nan] for all dtype
20 combinations.
21
22 3) divmod behavior consistent with 1) and 2).
23"""
24from __future__ import annotations
25
26import operator
27
28import numpy as np
29
30from pandas.core import roperator
31
32
33def _fill_zeros(result: np.ndarray, x, y):
34 """
35 If this is a reversed op, then flip x,y
36
37 If we have an integer value (or array in y)
38 and we have 0's, fill them with np.nan,
39 return the result.
40
41 Mask the nan's from x.
42 """
43 if result.dtype.kind == "f":
44 return result
45
46 is_variable_type = hasattr(y, "dtype")
47 is_scalar_type = not isinstance(y, np.ndarray)
48
49 if not is_variable_type and not is_scalar_type:
50 # e.g. test_series_ops_name_retention with mod we get here with list/tuple
51 return result
52
53 if is_scalar_type:
54 y = np.array(y)
55
56 if y.dtype.kind in "iu":
57 ymask = y == 0
58 if ymask.any():
59 # GH#7325, mask and nans must be broadcastable
60 mask = ymask & ~np.isnan(result)
61
62 # GH#9308 doing ravel on result and mask can improve putmask perf,
63 # but can also make unwanted copies.
64 result = result.astype("float64", copy=False)
65
66 np.putmask(result, mask, np.nan)
67
68 return result
69
70
71def mask_zero_div_zero(x, y, result: np.ndarray) -> np.ndarray:
72 """
73 Set results of 0 // 0 to np.nan, regardless of the dtypes
74 of the numerator or the denominator.
75
76 Parameters
77 ----------
78 x : ndarray
79 y : ndarray
80 result : ndarray
81
82 Returns
83 -------
84 ndarray
85 The filled result.
86
87 Examples
88 --------
89 >>> x = np.array([1, 0, -1], dtype=np.int64)
90 >>> x
91 array([ 1, 0, -1])
92 >>> y = 0 # int 0; numpy behavior is different with float
93 >>> result = x // y
94 >>> result # raw numpy result does not fill division by zero
95 array([0, 0, 0])
96 >>> mask_zero_div_zero(x, y, result)
97 array([ inf, nan, -inf])
98 """
99
100 if not hasattr(y, "dtype"):
101 # e.g. scalar, tuple
102 y = np.array(y)
103 if not hasattr(x, "dtype"):
104 # e.g scalar, tuple
105 x = np.array(x)
106
107 zmask = y == 0
108
109 if zmask.any():
110 # Flip sign if necessary for -0.0
111 zneg_mask = zmask & np.signbit(y)
112 zpos_mask = zmask & ~zneg_mask
113
114 x_lt0 = x < 0
115 x_gt0 = x > 0
116 nan_mask = zmask & (x == 0)
117 neginf_mask = (zpos_mask & x_lt0) | (zneg_mask & x_gt0)
118 posinf_mask = (zpos_mask & x_gt0) | (zneg_mask & x_lt0)
119
120 if nan_mask.any() or neginf_mask.any() or posinf_mask.any():
121 # Fill negative/0 with -inf, positive/0 with +inf, 0/0 with NaN
122 result = result.astype("float64", copy=False)
123
124 result[nan_mask] = np.nan
125 result[posinf_mask] = np.inf
126 result[neginf_mask] = -np.inf
127
128 return result
129
130
131def dispatch_fill_zeros(op, left, right, result):
132 """
133 Call _fill_zeros with the appropriate fill value depending on the operation,
134 with special logic for divmod and rdivmod.
135
136 Parameters
137 ----------
138 op : function (operator.add, operator.div, ...)
139 left : object (np.ndarray for non-reversed ops)
140 We have excluded ExtensionArrays here
141 right : object (np.ndarray for reversed ops)
142 We have excluded ExtensionArrays here
143 result : ndarray
144
145 Returns
146 -------
147 result : np.ndarray
148
149 Notes
150 -----
151 For divmod and rdivmod, the `result` parameter and returned `result`
152 is a 2-tuple of ndarray objects.
153 """
154 if op is divmod:
155 result = (
156 mask_zero_div_zero(left, right, result[0]),
157 _fill_zeros(result[1], left, right),
158 )
159 elif op is roperator.rdivmod:
160 result = (
161 mask_zero_div_zero(right, left, result[0]),
162 _fill_zeros(result[1], right, left),
163 )
164 elif op is operator.floordiv:
165 # Note: no need to do this for truediv; in py3 numpy behaves the way
166 # we want.
167 result = mask_zero_div_zero(left, right, result)
168 elif op is roperator.rfloordiv:
169 # Note: no need to do this for rtruediv; in py3 numpy behaves the way
170 # we want.
171 result = mask_zero_div_zero(right, left, result)
172 elif op is operator.mod:
173 result = _fill_zeros(result, left, right)
174 elif op is roperator.rmod:
175 result = _fill_zeros(result, right, left)
176 return result