@@ -18,23 +18,29 @@ def setup(self, *args, **kwargs):
18
18
"c" : xr .DataArray (np .arange (2 * self .n )),
19
19
}
20
20
)
21
- self .ds2d = self .ds1d .expand_dims (z = 10 )
21
+ self .ds2d = self .ds1d .expand_dims (z = 10 ). copy ()
22
22
self .ds1d_mean = self .ds1d .groupby ("b" ).mean ()
23
23
self .ds2d_mean = self .ds2d .groupby ("b" ).mean ()
24
24
25
25
@parameterized (["ndim" ], [(1 , 2 )])
26
26
def time_init (self , ndim ):
27
27
getattr (self , f"ds{ ndim } d" ).groupby ("b" )
28
28
29
- @parameterized (["method" , "ndim" ], [("sum" , "mean" ), (1 , 2 )])
30
- def time_agg_small_num_groups (self , method , ndim ):
29
+ @parameterized (
30
+ ["method" , "ndim" , "use_flox" ], [("sum" , "mean" ), (1 , 2 ), (True , False )]
31
+ )
32
+ def time_agg_small_num_groups (self , method , ndim , use_flox ):
31
33
ds = getattr (self , f"ds{ ndim } d" )
32
- getattr (ds .groupby ("a" ), method )().compute ()
34
+ with xr .set_options (use_flox = use_flox ):
35
+ getattr (ds .groupby ("a" ), method )().compute ()
33
36
34
- @parameterized (["method" , "ndim" ], [("sum" , "mean" ), (1 , 2 )])
35
- def time_agg_large_num_groups (self , method , ndim ):
37
+ @parameterized (
38
+ ["method" , "ndim" , "use_flox" ], [("sum" , "mean" ), (1 , 2 ), (True , False )]
39
+ )
40
+ def time_agg_large_num_groups (self , method , ndim , use_flox ):
36
41
ds = getattr (self , f"ds{ ndim } d" )
37
- getattr (ds .groupby ("b" ), method )().compute ()
42
+ with xr .set_options (use_flox = use_flox ):
43
+ getattr (ds .groupby ("b" ), method )().compute ()
38
44
39
45
def time_binary_op_1d (self ):
40
46
(self .ds1d .groupby ("b" ) - self .ds1d_mean ).compute ()
@@ -115,15 +121,21 @@ def setup(self, *args, **kwargs):
115
121
def time_init (self , ndim ):
116
122
getattr (self , f"ds{ ndim } d" ).resample (time = "D" )
117
123
118
- @parameterized (["method" , "ndim" ], [("sum" , "mean" ), (1 , 2 )])
119
- def time_agg_small_num_groups (self , method , ndim ):
124
+ @parameterized (
125
+ ["method" , "ndim" , "use_flox" ], [("sum" , "mean" ), (1 , 2 ), (True , False )]
126
+ )
127
+ def time_agg_small_num_groups (self , method , ndim , use_flox ):
120
128
ds = getattr (self , f"ds{ ndim } d" )
121
- getattr (ds .resample (time = "3M" ), method )().compute ()
129
+ with xr .set_options (use_flox = use_flox ):
130
+ getattr (ds .resample (time = "3M" ), method )().compute ()
122
131
123
- @parameterized (["method" , "ndim" ], [("sum" , "mean" ), (1 , 2 )])
124
- def time_agg_large_num_groups (self , method , ndim ):
132
+ @parameterized (
133
+ ["method" , "ndim" , "use_flox" ], [("sum" , "mean" ), (1 , 2 ), (True , False )]
134
+ )
135
+ def time_agg_large_num_groups (self , method , ndim , use_flox ):
125
136
ds = getattr (self , f"ds{ ndim } d" )
126
- getattr (ds .resample (time = "48H" ), method )().compute ()
137
+ with xr .set_options (use_flox = use_flox ):
138
+ getattr (ds .resample (time = "48H" ), method )().compute ()
127
139
128
140
129
141
class ResampleDask (Resample ):
@@ -132,3 +144,32 @@ def setup(self, *args, **kwargs):
132
144
super ().setup (** kwargs )
133
145
self .ds1d = self .ds1d .chunk ({"time" : 50 })
134
146
self .ds2d = self .ds2d .chunk ({"time" : 50 , "z" : 4 })
147
+
148
+
149
+ class ResampleCFTime (Resample ):
150
+ def setup (self , * args , ** kwargs ):
151
+ self .ds1d = xr .Dataset (
152
+ {
153
+ "b" : ("time" , np .arange (365.0 * 24 )),
154
+ },
155
+ coords = {
156
+ "time" : xr .date_range (
157
+ "2001-01-01" , freq = "H" , periods = 365 * 24 , calendar = "noleap"
158
+ )
159
+ },
160
+ )
161
+ self .ds2d = self .ds1d .expand_dims (z = 10 )
162
+ self .ds1d_mean = self .ds1d .resample (time = "48H" ).mean ()
163
+ self .ds2d_mean = self .ds2d .resample (time = "48H" ).mean ()
164
+
165
+
166
+ @parameterized (["use_cftime" , "use_flox" ], [[True , False ], [True , False ]])
167
+ class GroupByLongTime :
168
+ def setup (self , use_cftime , use_flox ):
169
+ arr = np .random .randn (10 , 10 , 365 * 30 )
170
+ time = xr .date_range ("2000" , periods = 30 * 365 , use_cftime = use_cftime )
171
+ self .da = xr .DataArray (arr , dims = ("y" , "x" , "time" ), coords = {"time" : time })
172
+
173
+ def time_mean (self , use_cftime , use_flox ):
174
+ with xr .set_options (use_flox = use_flox ):
175
+ self .da .groupby ("time.year" ).mean ()
0 commit comments