|
5 | 5 | ExtensionArray |
6 | 6 | """ |
7 | 7 | import operator |
| 8 | +from typing import Any, Callable |
| 9 | +import warnings |
8 | 10 |
|
9 | | -from pandas.core.ops import roperator |
| 11 | +import numpy as np |
| 12 | + |
| 13 | +from pandas._libs import lib |
| 14 | + |
| 15 | +from pandas.core.construction import extract_array |
| 16 | +from pandas.core.ops import maybe_dispatch_ufunc_to_dunder_op, roperator |
10 | 17 | from pandas.core.ops.common import unpack_zerodim_and_defer |
11 | 18 |
|
12 | 19 |
|
@@ -140,3 +147,138 @@ def __pow__(self, other): |
140 | 147 | @unpack_zerodim_and_defer("__rpow__") |
141 | 148 | def __rpow__(self, other): |
142 | 149 | return self._arith_method(other, roperator.rpow) |
| 150 | + |
| 151 | + |
| 152 | +def array_ufunc(self, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any): |
| 153 | + """ |
| 154 | + Compatibility with numpy ufuncs. |
| 155 | +
|
| 156 | + See also |
| 157 | + -------- |
| 158 | + numpy.org/doc/stable/reference/arrays.classes.html#numpy.class.__array_ufunc__ |
| 159 | + """ |
| 160 | + from pandas.core.generic import NDFrame |
| 161 | + from pandas.core.internals import BlockManager |
| 162 | + |
| 163 | + cls = type(self) |
| 164 | + |
| 165 | + # for binary ops, use our custom dunder methods |
| 166 | + result = maybe_dispatch_ufunc_to_dunder_op(self, ufunc, method, *inputs, **kwargs) |
| 167 | + if result is not NotImplemented: |
| 168 | + return result |
| 169 | + |
| 170 | + # Determine if we should defer. |
| 171 | + no_defer = (np.ndarray.__array_ufunc__, cls.__array_ufunc__) |
| 172 | + |
| 173 | + for item in inputs: |
| 174 | + higher_priority = ( |
| 175 | + hasattr(item, "__array_priority__") |
| 176 | + and item.__array_priority__ > self.__array_priority__ |
| 177 | + ) |
| 178 | + has_array_ufunc = ( |
| 179 | + hasattr(item, "__array_ufunc__") |
| 180 | + and type(item).__array_ufunc__ not in no_defer |
| 181 | + and not isinstance(item, self._HANDLED_TYPES) |
| 182 | + ) |
| 183 | + if higher_priority or has_array_ufunc: |
| 184 | + return NotImplemented |
| 185 | + |
| 186 | + # align all the inputs. |
| 187 | + types = tuple(type(x) for x in inputs) |
| 188 | + alignable = [x for x, t in zip(inputs, types) if issubclass(t, NDFrame)] |
| 189 | + |
| 190 | + if len(alignable) > 1: |
| 191 | + # This triggers alignment. |
| 192 | + # At the moment, there aren't any ufuncs with more than two inputs |
| 193 | + # so this ends up just being x1.index | x2.index, but we write |
| 194 | + # it to handle *args. |
| 195 | + |
| 196 | + if len(set(types)) > 1: |
| 197 | + # We currently don't handle ufunc(DataFrame, Series) |
| 198 | + # well. Previously this raised an internal ValueError. We might |
| 199 | + # support it someday, so raise a NotImplementedError. |
| 200 | + raise NotImplementedError( |
| 201 | + "Cannot apply ufunc {} to mixed DataFrame and Series " |
| 202 | + "inputs.".format(ufunc) |
| 203 | + ) |
| 204 | + axes = self.axes |
| 205 | + for obj in alignable[1:]: |
| 206 | + # this relies on the fact that we aren't handling mixed |
| 207 | + # series / frame ufuncs. |
| 208 | + for i, (ax1, ax2) in enumerate(zip(axes, obj.axes)): |
| 209 | + axes[i] = ax1.union(ax2) |
| 210 | + |
| 211 | + reconstruct_axes = dict(zip(self._AXIS_ORDERS, axes)) |
| 212 | + inputs = tuple( |
| 213 | + x.reindex(**reconstruct_axes) if issubclass(t, NDFrame) else x |
| 214 | + for x, t in zip(inputs, types) |
| 215 | + ) |
| 216 | + else: |
| 217 | + reconstruct_axes = dict(zip(self._AXIS_ORDERS, self.axes)) |
| 218 | + |
| 219 | + if self.ndim == 1: |
| 220 | + names = [getattr(x, "name") for x in inputs if hasattr(x, "name")] |
| 221 | + name = names[0] if len(set(names)) == 1 else None |
| 222 | + reconstruct_kwargs = {"name": name} |
| 223 | + else: |
| 224 | + reconstruct_kwargs = {} |
| 225 | + |
| 226 | + def reconstruct(result): |
| 227 | + if lib.is_scalar(result): |
| 228 | + return result |
| 229 | + if result.ndim != self.ndim: |
| 230 | + if method == "outer": |
| 231 | + if self.ndim == 2: |
| 232 | + # we already deprecated for Series |
| 233 | + msg = ( |
| 234 | + "outer method for ufunc {} is not implemented on " |
| 235 | + "pandas objects. Returning an ndarray, but in the " |
| 236 | + "future this will raise a 'NotImplementedError'. " |
| 237 | + "Consider explicitly converting the DataFrame " |
| 238 | + "to an array with '.to_numpy()' first." |
| 239 | + ) |
| 240 | + warnings.warn(msg.format(ufunc), FutureWarning, stacklevel=4) |
| 241 | + return result |
| 242 | + raise NotImplementedError |
| 243 | + return result |
| 244 | + if isinstance(result, BlockManager): |
| 245 | + # we went through BlockManager.apply |
| 246 | + result = self._constructor(result, **reconstruct_kwargs, copy=False) |
| 247 | + else: |
| 248 | + # we converted an array, lost our axes |
| 249 | + result = self._constructor( |
| 250 | + result, **reconstruct_axes, **reconstruct_kwargs, copy=False |
| 251 | + ) |
| 252 | + # TODO: When we support multiple values in __finalize__, this |
| 253 | + # should pass alignable to `__fianlize__` instead of self. |
| 254 | + # Then `np.add(a, b)` would consider attrs from both a and b |
| 255 | + # when a and b are NDFrames. |
| 256 | + if len(alignable) == 1: |
| 257 | + result = result.__finalize__(self) |
| 258 | + return result |
| 259 | + |
| 260 | + if self.ndim > 1 and ( |
| 261 | + len(inputs) > 1 or ufunc.nout > 1 # type: ignore[attr-defined] |
| 262 | + ): |
| 263 | + # Just give up on preserving types in the complex case. |
| 264 | + # In theory we could preserve them for them. |
| 265 | + # * nout>1 is doable if BlockManager.apply took nout and |
| 266 | + # returned a Tuple[BlockManager]. |
| 267 | + # * len(inputs) > 1 is doable when we know that we have |
| 268 | + # aligned blocks / dtypes. |
| 269 | + inputs = tuple(np.asarray(x) for x in inputs) |
| 270 | + result = getattr(ufunc, method)(*inputs) |
| 271 | + elif self.ndim == 1: |
| 272 | + # ufunc(series, ...) |
| 273 | + inputs = tuple(extract_array(x, extract_numpy=True) for x in inputs) |
| 274 | + result = getattr(ufunc, method)(*inputs, **kwargs) |
| 275 | + else: |
| 276 | + # ufunc(dataframe) |
| 277 | + mgr = inputs[0]._mgr |
| 278 | + result = mgr.apply(getattr(ufunc, method)) |
| 279 | + |
| 280 | + if ufunc.nout > 1: # type: ignore[attr-defined] |
| 281 | + result = tuple(reconstruct(x) for x in result) |
| 282 | + else: |
| 283 | + result = reconstruct(result) |
| 284 | + return result |
0 commit comments