@@ -1209,24 +1209,34 @@ def get_dask_names(ds):
1209
1209
),
1210
1210
)
1211
1211
@pytest .mark .parametrize ("freq" , ["D" , "W" , "5ME" , "YE" ])
1212
- def test_chunk_by_frequency (self , freq , calendar ) -> None :
1212
+ @pytest .mark .parametrize ("add_gap" , [True , False ])
1213
+ def test_chunk_by_frequency (self , freq : str , calendar : str , add_gap : bool ) -> None :
1213
1214
import dask .array
1214
1215
1215
1216
N = 365 * 2
1217
+ ΔN = 28
1218
+ time = xr .date_range (
1219
+ "2001-01-01" , periods = N + ΔN , freq = "D" , calendar = calendar
1220
+ ).to_numpy ()
1221
+ if add_gap :
1222
+ # introduce an empty bin
1223
+ time [31 : 31 + ΔN ] = np .datetime64 ("NaT" )
1224
+ time = time [~ np .isnat (time )]
1225
+ else :
1226
+ time = time [:N ]
1227
+
1216
1228
ds = Dataset (
1217
1229
{
1218
1230
"pr" : ("time" , dask .array .random .random ((N ), chunks = (20 ))),
1219
1231
"pr2d" : (("x" , "time" ), dask .array .random .random ((10 , N ), chunks = (20 ))),
1220
1232
"ones" : ("time" , np .ones ((N ,))),
1221
1233
},
1222
- coords = {
1223
- "time" : xr .date_range (
1224
- "2001-01-01" , periods = N , freq = "D" , calendar = calendar
1225
- )
1226
- },
1234
+ coords = {"time" : time },
1227
1235
)
1228
1236
rechunked = ds .chunk (x = 2 , time = TimeResampler (freq ))
1229
- expected = tuple (ds .ones .resample (time = freq ).sum ().data .tolist ())
1237
+ expected = tuple (
1238
+ ds .ones .resample (time = freq ).sum ().dropna ("time" ).astype (int ).data .tolist ()
1239
+ )
1230
1240
assert rechunked .chunksizes ["time" ] == expected
1231
1241
assert rechunked .chunksizes ["x" ] == (2 ,) * 5
1232
1242
0 commit comments