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