| 1326 | @pytest.mark.parametrize("freq", ["D", "W", "5ME", "YE"]) |
| 1327 | @pytest.mark.parametrize("add_gap", [True, False]) |
| 1328 | def test_chunk_by_frequency(self, freq: str, calendar: str, add_gap: bool) -> None: |
| 1329 | import dask.array |
| 1330 | |
| 1331 | N = 365 * 2 |
| 1332 | ΔN = 28 # noqa: PLC2401 |
| 1333 | time = xr.date_range( |
| 1334 | "2001-01-01", periods=N + ΔN, freq="D", calendar=calendar |
| 1335 | ).to_numpy(copy=True) |
| 1336 | if add_gap: |
| 1337 | # introduce an empty bin |
| 1338 | time[31 : 31 + ΔN] = np.datetime64("NaT", "us") |
| 1339 | time = time[~np.isnat(time)] |
| 1340 | else: |
| 1341 | time = time[:N] |
| 1342 | |
| 1343 | ds = Dataset( |
| 1344 | { |
| 1345 | "pr": ("time", dask.array.random.random((N), chunks=(20))), |
| 1346 | "pr2d": (("x", "time"), dask.array.random.random((10, N), chunks=(20))), |
| 1347 | "ones": ("time", np.ones((N,))), |
| 1348 | }, |
| 1349 | coords={"time": time}, |
| 1350 | ) |
| 1351 | rechunked = ds.chunk(x=2, time=TimeResampler(freq)) |
| 1352 | expected = tuple( |
| 1353 | ds.ones.resample(time=freq).sum().dropna("time").astype(int).data.tolist() |
| 1354 | ) |
| 1355 | assert rechunked.chunksizes["time"] == expected |
| 1356 | assert rechunked.chunksizes["x"] == (2,) * 5 |
| 1357 | |
| 1358 | rechunked = ds.chunk({"x": 2, "time": TimeResampler(freq)}) |
| 1359 | assert rechunked.chunksizes["time"] == expected |
| 1360 | assert rechunked.chunksizes["x"] == (2,) * 5 |
| 1361 | |
| 1362 | def test_chunk_by_frequency_errors(self): |
| 1363 | ds = Dataset({"foo": ("x", [1, 2, 3])}) |