|
39 | 39 | "2d-1scalar": xr.DataArray(randn(100, frac_nan=0.1), dims=["x"]), |
40 | 40 | } |
41 | 41 |
|
42 | | -vectorized_indexes = { |
43 | | - "1-1d": {"x": xr.DataArray(randint(0, nx, 400), dims="a")}, |
44 | | - "2-1d": { |
45 | | - "x": xr.DataArray(randint(0, nx, 400), dims="a"), |
46 | | - "y": xr.DataArray(randint(0, ny, 400), dims="a"), |
47 | | - }, |
48 | | - "3-2d": { |
49 | | - "x": xr.DataArray(randint(0, nx, 400).reshape(4, 100), dims=["a", "b"]), |
50 | | - "y": xr.DataArray(randint(0, ny, 400).reshape(4, 100), dims=["a", "b"]), |
51 | | - "t": xr.DataArray(randint(0, nt, 400).reshape(4, 100), dims=["a", "b"]), |
52 | | - }, |
53 | | -} |
| 42 | + |
| 43 | +def make_vectorized_indexes(n_index): |
| 44 | + return { |
| 45 | + "1-1d": {"x": xr.DataArray(randint(0, nx, n_index), dims="a")}, |
| 46 | + "2-1d": { |
| 47 | + "x": xr.DataArray(randint(0, nx, n_index), dims="a"), |
| 48 | + "y": xr.DataArray(randint(0, ny, n_index), dims="a"), |
| 49 | + }, |
| 50 | + "3-2d": { |
| 51 | + "x": xr.DataArray( |
| 52 | + randint(0, nx, n_index).reshape(n_index // 100, 100), dims=["a", "b"] |
| 53 | + ), |
| 54 | + "y": xr.DataArray( |
| 55 | + randint(0, ny, n_index).reshape(n_index // 100, 100), dims=["a", "b"] |
| 56 | + ), |
| 57 | + "t": xr.DataArray( |
| 58 | + randint(0, nt, n_index).reshape(n_index // 100, 100), dims=["a", "b"] |
| 59 | + ), |
| 60 | + }, |
| 61 | + } |
| 62 | + |
| 63 | + |
| 64 | +vectorized_indexes = make_vectorized_indexes(400) |
| 65 | +big_vectorized_indexes = make_vectorized_indexes(400_000) |
54 | 66 |
|
55 | 67 | vectorized_assignment_values = { |
56 | 68 | "1-1d": xr.DataArray(randn((400, ny)), dims=["a", "y"], coords={"a": randn(400)}), |
@@ -101,6 +113,20 @@ def time_indexing_basic_ds_large(self, key): |
101 | 113 | self.ds_large.isel(**basic_indexes[key]).load() |
102 | 114 |
|
103 | 115 |
|
| 116 | +class IndexingOnly(Base): |
| 117 | + @parameterized(["key"], [list(basic_indexes.keys())]) |
| 118 | + def time_indexing_basic(self, key): |
| 119 | + self.ds.isel(**basic_indexes[key]) |
| 120 | + |
| 121 | + @parameterized(["key"], [list(outer_indexes.keys())]) |
| 122 | + def time_indexing_outer(self, key): |
| 123 | + self.ds.isel(**outer_indexes[key]) |
| 124 | + |
| 125 | + @parameterized(["key"], [list(big_vectorized_indexes.keys())]) |
| 126 | + def time_indexing_big_vectorized(self, key): |
| 127 | + self.ds.isel(**big_vectorized_indexes[key]) |
| 128 | + |
| 129 | + |
104 | 130 | class Assignment(Base): |
105 | 131 | @parameterized(["key"], [list(basic_indexes.keys())]) |
106 | 132 | def time_assignment_basic(self, key): |
|
0 commit comments