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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
|
import os
import numpy as np
import pandas as pd
import xarray as xr
from . import parameterized, randint, randn, requires_dask
nx = 2000
ny = 1000
nt = 500
basic_indexes = {
"1scalar": {"x": 0},
"1slice": {"x": slice(0, 3)},
"1slice-1scalar": {"x": 0, "y": slice(None, None, 3)},
"2slicess-1scalar": {"x": slice(3, -3, 3), "y": 1, "t": slice(None, -3, 3)},
}
basic_assignment_values = {
"1scalar": 0,
"1slice": xr.DataArray(randn((3, ny), frac_nan=0.1), dims=["x", "y"]),
"1slice-1scalar": xr.DataArray(randn(int(ny / 3) + 1, frac_nan=0.1), dims=["y"]),
"2slicess-1scalar": xr.DataArray(
randn(np.empty(nx)[slice(3, -3, 3)].size, frac_nan=0.1), dims=["x"]
),
}
outer_indexes = {
"1d": {"x": randint(0, nx, 400)},
"2d": {"x": randint(0, nx, 500), "y": randint(0, ny, 400)},
"2d-1scalar": {"x": randint(0, nx, 100), "y": 1, "t": randint(0, nt, 400)},
}
outer_assignment_values = {
"1d": xr.DataArray(randn((400, ny), frac_nan=0.1), dims=["x", "y"]),
"2d": xr.DataArray(randn((500, 400), frac_nan=0.1), dims=["x", "y"]),
"2d-1scalar": xr.DataArray(randn(100, frac_nan=0.1), dims=["x"]),
}
def make_vectorized_indexes(n_index):
return {
"1-1d": {"x": xr.DataArray(randint(0, nx, n_index), dims="a")},
"2-1d": {
"x": xr.DataArray(randint(0, nx, n_index), dims="a"),
"y": xr.DataArray(randint(0, ny, n_index), dims="a"),
},
"3-2d": {
"x": xr.DataArray(
randint(0, nx, n_index).reshape(n_index // 100, 100), dims=["a", "b"]
),
"y": xr.DataArray(
randint(0, ny, n_index).reshape(n_index // 100, 100), dims=["a", "b"]
),
"t": xr.DataArray(
randint(0, nt, n_index).reshape(n_index // 100, 100), dims=["a", "b"]
),
},
}
vectorized_indexes = make_vectorized_indexes(400)
big_vectorized_indexes = make_vectorized_indexes(400_000)
vectorized_assignment_values = {
"1-1d": xr.DataArray(randn((400, ny)), dims=["a", "y"], coords={"a": randn(400)}),
"2-1d": xr.DataArray(randn(400), dims=["a"], coords={"a": randn(400)}),
"3-2d": xr.DataArray(
randn((4, 100)), dims=["a", "b"], coords={"a": randn(4), "b": randn(100)}
),
}
class Base:
def setup(self, key):
self.ds = xr.Dataset(
{
"var1": (("x", "y"), randn((nx, ny), frac_nan=0.1)),
"var2": (("x", "t"), randn((nx, nt))),
"var3": (("t",), randn(nt)),
},
coords={
"x": np.arange(nx),
"y": np.linspace(0, 1, ny),
"t": pd.date_range("1970-01-01", periods=nt, freq="D"),
"x_coords": ("x", np.linspace(1.1, 2.1, nx)),
},
)
# Benchmark how indexing is slowed down by adding many scalar variable
# to the dataset
# https://github.com/pydata/xarray/pull/9003
self.ds_large = self.ds.merge({f"extra_var{i}": i for i in range(400)})
class Indexing(Base):
@parameterized(["key"], [list(basic_indexes.keys())])
def time_indexing_basic(self, key):
self.ds.isel(**basic_indexes[key]).load()
@parameterized(["key"], [list(outer_indexes.keys())])
def time_indexing_outer(self, key):
self.ds.isel(**outer_indexes[key]).load()
@parameterized(["key"], [list(vectorized_indexes.keys())])
def time_indexing_vectorized(self, key):
self.ds.isel(**vectorized_indexes[key]).load()
@parameterized(["key"], [list(basic_indexes.keys())])
def time_indexing_basic_ds_large(self, key):
# https://github.com/pydata/xarray/pull/9003
self.ds_large.isel(**basic_indexes[key]).load()
class IndexingOnly(Base):
@parameterized(["key"], [list(basic_indexes.keys())])
def time_indexing_basic(self, key):
self.ds.isel(**basic_indexes[key])
@parameterized(["key"], [list(outer_indexes.keys())])
def time_indexing_outer(self, key):
self.ds.isel(**outer_indexes[key])
@parameterized(["key"], [list(big_vectorized_indexes.keys())])
def time_indexing_big_vectorized(self, key):
self.ds.isel(**big_vectorized_indexes[key])
class Assignment(Base):
@parameterized(["key"], [list(basic_indexes.keys())])
def time_assignment_basic(self, key):
ind = basic_indexes[key]
val = basic_assignment_values[key]
self.ds["var1"][ind.get("x", slice(None)), ind.get("y", slice(None))] = val
@parameterized(["key"], [list(outer_indexes.keys())])
def time_assignment_outer(self, key):
ind = outer_indexes[key]
val = outer_assignment_values[key]
self.ds["var1"][ind.get("x", slice(None)), ind.get("y", slice(None))] = val
@parameterized(["key"], [list(vectorized_indexes.keys())])
def time_assignment_vectorized(self, key):
ind = vectorized_indexes[key]
val = vectorized_assignment_values[key]
self.ds["var1"][ind.get("x", slice(None)), ind.get("y", slice(None))] = val
class IndexingDask(Indexing):
def setup(self, key):
requires_dask()
super().setup(key)
self.ds = self.ds.chunk({"x": 100, "y": 50, "t": 50})
class BooleanIndexing:
# https://github.com/pydata/xarray/issues/2227
def setup(self):
self.ds = xr.Dataset(
{"a": ("time", np.arange(10_000_000))},
coords={"time": np.arange(10_000_000)},
)
self.time_filter = self.ds.time > 50_000
def time_indexing(self):
self.ds.isel(time=self.time_filter)
class HugeAxisSmallSliceIndexing:
# https://github.com/pydata/xarray/pull/4560
def setup(self):
self.filepath = "test_indexing_huge_axis_small_slice.nc"
if not os.path.isfile(self.filepath):
xr.Dataset(
{"a": ("x", np.arange(10_000_000))},
coords={"x": np.arange(10_000_000)},
).to_netcdf(self.filepath, format="NETCDF4")
self.ds = xr.open_dataset(self.filepath)
def time_indexing(self):
self.ds.isel(x=slice(100))
def cleanup(self):
self.ds.close()
class AssignmentOptimized:
# https://github.com/pydata/xarray/pull/7382
def setup(self):
self.ds = xr.Dataset(coords={"x": np.arange(500_000)})
self.da = xr.DataArray(np.arange(500_000), dims="x")
def time_assign_no_reindex(self):
# assign with non-indexed DataArray of same dimension size
self.ds.assign(foo=self.da)
def time_assign_identical_indexes(self):
# fastpath index comparison (same index object)
self.ds.assign(foo=self.ds.x)
|
