1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
|
"""Base classes implementing arithmetic for xarray objects."""
from __future__ import annotations
import numbers
import numpy as np
from xarray.computation.ops import IncludeNumpySameMethods, IncludeReduceMethods
# _typed_ops.py is a generated file
from xarray.core._typed_ops import (
DataArrayGroupByOpsMixin,
DataArrayOpsMixin,
DatasetGroupByOpsMixin,
DatasetOpsMixin,
VariableOpsMixin,
)
from xarray.core.common import ImplementsArrayReduce, ImplementsDatasetReduce
from xarray.core.options import OPTIONS, _get_keep_attrs
from xarray.namedarray.utils import is_duck_array
class SupportsArithmetic:
"""Base class for xarray types that support arithmetic.
Used by Dataset, DataArray, Variable and GroupBy.
"""
__slots__ = ()
# TODO: implement special methods for arithmetic here rather than injecting
# them in xarray/computation/ops.py. Ideally, do so by inheriting from
# numpy.lib.mixins.NDArrayOperatorsMixin.
# TODO: allow extending this with some sort of registration system
_HANDLED_TYPES = (
np.generic,
numbers.Number,
bytes,
str,
)
def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
from xarray.computation.apply_ufunc import apply_ufunc
# See the docstring example for numpy.lib.mixins.NDArrayOperatorsMixin.
out = kwargs.get("out", ())
for x in inputs + out:
if not is_duck_array(x) and not isinstance(
x, self._HANDLED_TYPES + (SupportsArithmetic,)
):
return NotImplemented
if ufunc.signature is not None:
raise NotImplementedError(
f"{ufunc} not supported: xarray objects do not directly implement "
"generalized ufuncs. Instead, use xarray.apply_ufunc or "
"explicitly convert to xarray objects to NumPy arrays "
"(e.g., with `.values`)."
)
if method != "__call__":
# TODO: support other methods, e.g., reduce and accumulate.
raise NotImplementedError(
f"{method} method for ufunc {ufunc} is not implemented on xarray objects, "
"which currently only support the __call__ method. As an "
"alternative, consider explicitly converting xarray objects "
"to NumPy arrays (e.g., with `.values`)."
)
if any(isinstance(o, SupportsArithmetic) for o in out):
# TODO: implement this with logic like _inplace_binary_op. This
# will be necessary to use NDArrayOperatorsMixin.
raise NotImplementedError(
"xarray objects are not yet supported in the `out` argument "
"for ufuncs. As an alternative, consider explicitly "
"converting xarray objects to NumPy arrays (e.g., with "
"`.values`)."
)
join = dataset_join = OPTIONS["arithmetic_join"]
return apply_ufunc(
ufunc,
*inputs,
input_core_dims=((),) * ufunc.nin,
output_core_dims=((),) * ufunc.nout,
join=join,
dataset_join=dataset_join,
dataset_fill_value=np.nan,
kwargs=kwargs,
dask="allowed",
keep_attrs=_get_keep_attrs(default=True),
)
class VariableArithmetic(
ImplementsArrayReduce,
IncludeNumpySameMethods,
SupportsArithmetic,
VariableOpsMixin,
):
__slots__ = ()
# prioritize our operations over those of numpy.ndarray (priority=0)
__array_priority__ = 50
class DatasetArithmetic(
ImplementsDatasetReduce,
SupportsArithmetic,
DatasetOpsMixin,
):
__slots__ = ()
__array_priority__ = 50
class DataArrayArithmetic(
ImplementsArrayReduce,
IncludeNumpySameMethods,
SupportsArithmetic,
DataArrayOpsMixin,
):
__slots__ = ()
# priority must be higher than Variable to properly work with binary ufuncs
__array_priority__ = 60
class DataArrayGroupbyArithmetic(
SupportsArithmetic,
DataArrayGroupByOpsMixin,
):
__slots__ = ()
class DatasetGroupbyArithmetic(
SupportsArithmetic,
DatasetGroupByOpsMixin,
):
__slots__ = ()
class CoarsenArithmetic(IncludeReduceMethods):
__slots__ = ()
|
