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